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    pub fn add_legacy_catch(&mut self, catch: crate::ir::LegacyCatch) -> Result<()> {
160        // Find the most recent Try instruction in the parent control block
161        let frame = self.control(1)?; // Parent block, not the try block itself
162        let block = frame.block;
163        let seq = self.func.block_mut(block);
164
165        // The Try instruction should be the last instruction in the parent block
166        if let Some((Instr::Try(ref mut try_instr), _)) = seq.instrs.last_mut() {
167            try_instr.catches.push(catch);
168            return Ok(());
169        }
170
171        anyhow::bail!("No Try instruction found to add catch clause to");
172    }
173}
174
175fn impl_push_control(
176    types: &ModuleTypes,
177    kind: BlockKind,
178    func: &mut LocalFunction,
179    controls: &mut ControlStack,
180    start_types: Box<[ValType]>,
181    end_types: Box<[ValType]>,
182) -> Result<InstrSeqId> {
183    let ty = InstrSeqType::existing(types, &start_types, &end_types).ok_or_else(|| {
184        anyhow::anyhow!(
185            "attempted to push a control frame for an instruction \
186             sequence with a type that does not exist"
187        )
188        .context(format!("type: {:?} -> {:?}", &start_types, &end_types))
189    })?;
190
191    Ok(impl_push_control_with_ty(
192        types,
193        kind,
194        func,
195        controls,
196        ty,
197        start_types,
198        end_types,
199    ))
200}
201
202fn impl_push_control_with_ty(
203    types: &ModuleTypes,
204    kind: BlockKind,
205    func: &mut LocalFunction,
206    controls: &mut ControlStack,
207    ty: InstrSeqType,
208    start_types: Box<[ValType]>,
209    end_types: Box<[ValType]>,
210) -> InstrSeqId {
211    if let InstrSeqType::MultiValue(ty) = ty {
212        debug_assert_eq!(types.params(ty), &start_types[..]);
213        debug_assert_eq!(types.results(ty), &end_types[..]);
214    }
215
216    let block = func.add_block(|id| InstrSeq::new(id, ty));
217
218    controls.push(ControlFrame {
219        start_types,
220        end_types,
221        unreachable: false,
222        block,
223        kind,
224    });
225
226    block
227}
228
229fn impl_pop_control(controls: &mut ControlStack) -> Result<ControlFrame> {
230    controls
231        .last()
232        .ok_or(ErrorKind::InvalidWasm)
233        .context("attempted to pop a frame from an empty control stack")?;
234    let frame = controls.pop().unwrap();
235    Ok(frame)
236}
walrus/module/functions/local_function/context.rs

walrus/module/functions/local_function/
context.rs
