1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
//! Data structures for control flow emission.
use crate::{Error, Result};
use smallvec::SmallVec;
use wasmparser::BlockType;

/// The type of the control stack frame.
#[repr(u8)]
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum ControlStackFrameType {
    /// The if control stack frame.
    ///
    /// true is has else block, otherwise false.
    If(bool),
    /// The else control stack frame.
    Else,
    /// The loop control stack frame.
    Loop,
    /// The block control stack frame.
    Block,
}

/// Holds the necessary metadata to support the smission
/// of control flow instructions.
///
/// NOTE: The output of control flow should be placed on
/// the stack, so we don't need to store the result type.
#[derive(Clone)]
pub struct ControlStackFrame {
    /// The type of the control stack frame.
    ///
    /// If loop, break it while popping.
    pub ty: ControlStackFrameType,
    /// The program counter offset at the beginning of if.
    pub original_pc_offset: u16,
    /// The return values of the block.
    ///
    /// Could be useful for validation.
    result: BlockType,

    /// Original stack pointer.
    pub original_sp: u8,
}

impl ControlStackFrame {
    /// Create a new control stack frame.
    pub fn new(
        ty: ControlStackFrameType,
        original_pc_offset: u16,
        original_sp: u8,
        result: BlockType,
    ) -> Self {
        Self {
            ty,
            original_pc_offset,
            original_sp,
            result,
        }
    }

    /// Get the offset of the original program counter.
    pub fn pc_offset(&self) -> u16 {
        self.original_pc_offset
    }

    /// Get the result type of the control stack frame.
    pub fn result(&self) -> BlockType {
        self.result
    }
}

/// The control stack.
#[derive(Default)]
pub struct ControlStack {
    /// Stack frames for control flow.
    ///
    /// The 32 is set arbitrarily, we can adjust it as we see fit.
    stack: SmallVec<[ControlStackFrame; 32]>,
}

impl ControlStack {
    /// The total depth of the control stack.
    pub fn depth(&self) -> usize {
        self.stack.len()
    }

    /// Mark the else block of an if.
    pub fn mark_else(&mut self) -> Result<ControlStackFrame> {
        let last = self
            .stack
            .last_mut()
            .ok_or_else(|| Error::ControlStackUnderflow)?;

        if last.ty != ControlStackFrameType::If(false) {
            return Err(Error::InvalidElseBlock(last.original_pc_offset));
        }

        last.ty = ControlStackFrameType::If(true);
        Ok(last.clone())
    }

    /// Push a control stack frame.
    pub fn push(&mut self, frame: ControlStackFrame) {
        self.stack.push(frame);
    }

    /// Pop a control stack frame.
    pub fn pop(&mut self) -> Result<ControlStackFrame> {
        self.stack.pop().ok_or_else(|| Error::ControlStackUnderflow)
    }

    /// Get the label of the control stack frame at given depth.
    pub fn label_from_depth(&self, mut depth: u32) -> Result<u16> {
        for frame in self.stack.iter().rev() {
            if frame.ty == ControlStackFrameType::Else {
                continue;
            }

            if depth == 0 {
                return Ok(frame.pc_offset());
            }

            depth -= 1;
        }

        Err(Error::InvalidDepth(depth as usize))
    }

    /// Get the return type of the control stack frame at given depth.
    pub fn ret_ty(&self, depth: usize) -> Result<BlockType> {
        if depth == 0 {
            return Err(Error::InvalidDepth(depth));
        }

        self.stack
            .get(self.depth() - depth)
            .map(|f| f.result)
            .ok_or_else(|| Error::InvalidDepth(depth))
    }

    /// Get the type of the control stack frame at given depth.
    pub fn ty(&self, depth: usize) -> Result<ControlStackFrameType> {
        if depth == 0 {
            return Err(Error::InvalidDepth(depth));
        }

        self.stack
            .get(self.depth() - depth)
            .map(|f| f.ty)
            .ok_or_else(|| Error::InvalidDepth(depth))
    }
}