walrus/module/functions/local_function/
context.rs

1//! Context needed when validating instructions and constructing our `Instr` IR.
2
3use crate::error::{ErrorKind, Result};
4use crate::ir::{BlockKind, Instr, InstrLocId, InstrSeq, InstrSeqId, InstrSeqType};
5use crate::module::functions::{FunctionId, LocalFunction};
6use crate::module::Module;
7use crate::parse::IndicesToIds;
8use crate::ty::ValType;
9use crate::{ModuleTypes, TypeId};
10use anyhow::Context;
11
12#[derive(Debug)]
13pub(crate) struct ControlFrame {
14    /// The parameter types of the block (checked before entering the block).
15    pub start_types: Box<[ValType]>,
16
17    /// The result type of the block (used to check its result).
18    pub end_types: Box<[ValType]>,
19
20    /// If `true`, then this frame is unreachable. This is used to handle
21    /// stack-polymorphic typing after unconditional branches.
22    pub unreachable: bool,
23
24    /// The id of this control frame's block.
25    pub block: InstrSeqId,
26
27    /// This control frame's kind of block, eg loop vs block vs if/else.
28    pub kind: BlockKind,
29}
30
31/// The control frame stack.
32pub(crate) type ControlStack = Vec<ControlFrame>;
33
34#[derive(Debug)]
35pub(crate) struct ValidationContext<'a> {
36    /// The module that we're adding a function for.
37    pub module: &'a Module,
38
39    /// Mapping of indexes back to ids.
40    pub indices: &'a IndicesToIds,
41
42    /// The arena id of `func`.
43    pub func_id: FunctionId,
44
45    /// The function being validated/constructed.
46    pub func: &'a mut LocalFunction,
47
48    /// The control frames stack.
49    pub controls: &'a mut ControlStack,
50
51    /// If we're currently parsing an if/else instruction, where we're at
52    pub if_else: Vec<IfElseState>,
53}
54
55#[derive(Debug)]
56pub struct IfElseState {
57    pub start: InstrLocId,
58    pub consequent: InstrSeqId,
59    pub alternative: Option<InstrSeqId>,
60}
61
62impl<'a> ValidationContext<'a> {
63    /// Create a new function context.
64    pub fn new(
65        module: &'a Module,
66        indices: &'a IndicesToIds,
67        func_id: FunctionId,
68        func: &'a mut LocalFunction,
69        controls: &'a mut ControlStack,
70    ) -> ValidationContext<'a> {
71        ValidationContext {
72            module,
73            indices,
74            func_id,
75            func,
76            controls,
77            if_else: Vec::new(),
78        }
79    }
80
81    pub fn push_control(
82        &mut self,
83        kind: BlockKind,
84        start_types: Box<[ValType]>,
85        end_types: Box<[ValType]>,
86    ) -> Result<InstrSeqId> {
87        impl_push_control(
88            &self.module.types,
89            kind,
90            self.func,
91            self.controls,
92            start_types,
93            end_types,
94        )
95    }
96
97    pub fn push_control_with_ty(&mut self, kind: BlockKind, ty: TypeId) -> InstrSeqId {
98        let (start_types, end_types) = self.module.types.params_results(ty);
99        let start_types: Box<[_]> = start_types.into();
100        let end_types: Box<[_]> = end_types.into();
101        impl_push_control_with_ty(
102            &self.module.types,
103            kind,
104            self.func,
105            self.controls,
106            ty.into(),
107            start_types,
108            end_types,
109        )
110    }
111
112    pub fn pop_control(&mut self) -> Result<(ControlFrame, InstrSeqId)> {
113        let frame = impl_pop_control(self.controls)?;
114        let block = frame.block;
115        Ok((frame, block))
116    }
117
118    pub fn unreachable(&mut self) {
119        let frame = self.controls.last_mut().unwrap();
120        frame.unreachable = true;
121    }
122
123    pub fn control(&self, n: usize) -> Result<&ControlFrame> {
124        if n >= self.controls.len() {
125            anyhow::bail!("jump to nonexistent control block");
126        }
127        let idx = self.controls.len() - n - 1;
128        Ok(&self.controls[idx])
129    }
130
131    pub fn alloc_instr_in_block(
132        &mut self,
133        block: InstrSeqId,
134        instr: impl Into<Instr>,
135        loc: InstrLocId,
136    ) {
137        self.func.block_mut(block).instrs.push((instr.into(), loc));
138    }
139
140    pub fn alloc_instr_in_control(
141        &mut self,
142        control: usize,
143        instr: impl Into<Instr>,
144        loc: InstrLocId,
145    ) -> Result<()> {
146        let frame = self.control(control)?;
147        if frame.unreachable {
148            return Ok(());
149        }
150        let block = frame.block;
151        self.alloc_instr_in_block(block, instr, loc);
152        Ok(())
153    }
154
155    pub fn alloc_instr(&mut self, instr: impl Into<Instr>, loc: InstrLocId) {
156        self.alloc_instr_in_control(0, instr, loc).unwrap();
157    }
158}
159
160fn impl_push_control(
161    types: &ModuleTypes,
162    kind: BlockKind,
163    func: &mut LocalFunction,
164    controls: &mut ControlStack,
165    start_types: Box<[ValType]>,
166    end_types: Box<[ValType]>,
167) -> Result<InstrSeqId> {
168    let ty = InstrSeqType::existing(types, &start_types, &end_types).ok_or_else(|| {
169        anyhow::anyhow!(
170            "attempted to push a control frame for an instruction \
171             sequence with a type that does not exist"
172        )
173        .context(format!("type: {:?} -> {:?}", &start_types, &end_types))
174    })?;
175
176    Ok(impl_push_control_with_ty(
177        types,
178        kind,
179        func,
180        controls,
181        ty,
182        start_types,
183        end_types,
184    ))
185}
186
187fn impl_push_control_with_ty(
188    types: &ModuleTypes,
189    kind: BlockKind,
190    func: &mut LocalFunction,
191    controls: &mut ControlStack,
192    ty: InstrSeqType,
193    start_types: Box<[ValType]>,
194    end_types: Box<[ValType]>,
195) -> InstrSeqId {
196    if let InstrSeqType::MultiValue(ty) = ty {
197        debug_assert_eq!(types.params(ty), &start_types[..]);
198        debug_assert_eq!(types.results(ty), &end_types[..]);
199    }
200
201    let block = func.add_block(|id| InstrSeq::new(id, ty));
202
203    controls.push(ControlFrame {
204        start_types,
205        end_types,
206        unreachable: false,
207        block,
208        kind,
209    });
210
211    block
212}
213
214fn impl_pop_control(controls: &mut ControlStack) -> Result<ControlFrame> {
215    controls
216        .last()
217        .ok_or(ErrorKind::InvalidWasm)
218        .context("attempted to pop a frame from an empty control stack")?;
219    let frame = controls.pop().unwrap();
220    Ok(frame)
221}