winch_codegen/frame/
mod.rs

1use crate::{
2    Result,
3    abi::{ABI, ABIOperand, ABISig, LocalSlot, align_to},
4    codegen::{CodeGenPhase, Emission, Prologue},
5    masm::{MacroAssembler, SPOffset},
6    stack::needs_stack_map,
7};
8use smallvec::SmallVec;
9use std::marker::PhantomData;
10use std::ops::Range;
11use wasmparser::{BinaryReader, FuncValidator, ValidatorResources};
12use wasmtime_environ::{TypeConvert, WasmValType};
13
14/// WebAssembly locals.
15// TODO:
16// SpiderMonkey's implementation uses 16;
17// (ref: https://searchfox.org/mozilla-central/source/js/src/wasm/WasmBCFrame.h#585)
18// during instrumentation we should measure to verify if this is a good default.
19pub(crate) type WasmLocals = SmallVec<[LocalSlot; 16]>;
20/// Special local slots used by the compiler.
21// Winch's ABI uses two extra parameters to store the callee and caller
22// VMContext pointers.
23// These arguments are spilled and treated as frame locals, but not
24// WebAssembly locals.
25pub(crate) type SpecialLocals = [LocalSlot; 2];
26
27/// Function defined locals start and end in the frame.
28pub(crate) struct DefinedLocalsRange(Range<u32>);
29
30impl DefinedLocalsRange {
31    /// Get a reference to the inner range.
32    pub fn as_range(&self) -> &Range<u32> {
33        &self.0
34    }
35}
36
37/// An abstraction to read the defined locals from the Wasm binary for a function.
38#[derive(Default)]
39pub(crate) struct DefinedLocals {
40    /// The defined locals for a function.
41    pub defined_locals: WasmLocals,
42    /// The size of the defined locals.
43    pub stack_size: u32,
44}
45
46impl DefinedLocals {
47    /// Compute the local slots for a Wasm function.
48    pub fn new<A: ABI>(
49        types: &impl TypeConvert,
50        reader: &mut BinaryReader<'_>,
51        validator: &mut FuncValidator<ValidatorResources>,
52    ) -> Result<Self> {
53        let mut next_stack: u32 = 0;
54        // The first 32 bits of a Wasm binary function describe the number of locals.
55        let local_count = reader.read_var_u32()?;
56        let mut slots: WasmLocals = Default::default();
57
58        for _ in 0..local_count {
59            let position = reader.original_position();
60            let count = reader.read_var_u32()?;
61            let ty = reader.read()?;
62            validator.define_locals(position, count, ty)?;
63
64            let ty = types.convert_valtype(ty)?;
65            for _ in 0..count {
66                let ty_size = <A as ABI>::sizeof(&ty);
67                next_stack = align_to(next_stack, ty_size as u32) + (ty_size as u32);
68                slots.push(LocalSlot::new(ty, next_stack));
69            }
70        }
71
72        Ok(Self {
73            defined_locals: slots,
74            stack_size: next_stack,
75        })
76    }
77}
78
79/// Frame handler abstraction.
80pub(crate) struct Frame<P: CodeGenPhase> {
81    /// The size of the entire local area; the arguments plus the function defined locals.
82    pub locals_size: u32,
83
84    /// The range in the frame corresponding to the defined locals range.
85    pub defined_locals_range: DefinedLocalsRange,
86
87    /// The local slots for the current function.
88    ///
89    /// Locals get calculated when allocating a frame and are readonly
90    /// through the function compilation lifetime.
91    wasm_locals: WasmLocals,
92    /// Special locals used by the internal ABI. See [`SpecialLocals`].
93    special_locals: SpecialLocals,
94
95    /// The slot holding the address of the results area.
96    pub results_base_slot: Option<LocalSlot>,
97    marker: PhantomData<P>,
98
99    /// Frame offsets of SP-addressed locals that stack maps must cover,
100    /// precomputed so that call sites don't scan every local.
101    gc_ref_local_offsets: SmallVec<[SPOffset; 4]>,
102}
103
104impl Frame<Prologue> {
105    /// Allocate a new [`Frame`].
106    pub fn new<A: ABI>(sig: &ABISig, defined_locals: &DefinedLocals) -> Result<Frame<Prologue>> {
107        let (special_locals, mut wasm_locals, mut gc_ref_local_offsets, defined_locals_start) =
108            Self::compute_arg_slots::<A>(sig)?;
109
110        // The defined locals have a zero-based offset by default
111        // so we need to add the defined locals start to the offset.
112        wasm_locals.extend(defined_locals.defined_locals.iter().map(|l| {
113            let slot = LocalSlot::new(l.ty, l.offset + defined_locals_start);
114            if needs_stack_map(&slot.ty) {
115                gc_ref_local_offsets.push(SPOffset::from_u32(slot.offset));
116            }
117            slot
118        }));
119
120        let stack_align = <A as ABI>::stack_align();
121        let defined_locals_end = align_to(
122            defined_locals_start + defined_locals.stack_size,
123            stack_align as u32,
124        );
125
126        // Handle the results base slot for multi value returns.
127        let (results_base_slot, locals_size) = if sig.params.has_retptr() {
128            match sig.params.unwrap_results_area_operand() {
129                // If the results operand is a stack argument, ensure the
130                // offset is correctly calculated, that is, that it includes the
131                // argument base offset.
132                // In this case, the locals size, remains untouched as we don't
133                // need to create an extra slot for it.
134                ABIOperand::Stack { ty, offset, .. } => (
135                    Some(LocalSlot::stack_arg(
136                        *ty,
137                        *offset + (<A as ABI>::arg_base_offset() as u32),
138                    )),
139                    defined_locals_end,
140                ),
141                // If the results operand is a register, we give this register
142                // the same treatment as all the other argument registers and
143                // spill it, therefore, we need to increase the locals size by
144                // one slot.
145                ABIOperand::Reg { ty, size, .. } => {
146                    let offs = align_to(defined_locals_end, *size) + *size;
147                    (
148                        Some(LocalSlot::new(*ty, offs)),
149                        align_to(offs, <A as ABI>::stack_align().into()),
150                    )
151                }
152            }
153        } else {
154            (None, defined_locals_end)
155        };
156
157        Ok(Self {
158            wasm_locals,
159            special_locals,
160            locals_size,
161            defined_locals_range: DefinedLocalsRange(
162                defined_locals_start..(defined_locals_start + defined_locals.stack_size),
163            ),
164            results_base_slot,
165            marker: PhantomData,
166            gc_ref_local_offsets,
167        })
168    }
169
170    /// Returns an iterator over all the [`LocalSlot`]s in the frame, including
171    /// the [`SpecialLocals`].
172    pub fn locals(&self) -> impl Iterator<Item = &LocalSlot> {
173        self.special_locals.iter().chain(self.wasm_locals.iter())
174    }
175
176    /// Prepares the frame for the [`Emission`] code generation phase.
177    pub fn for_emission(self) -> Frame<Emission> {
178        Frame {
179            wasm_locals: self.wasm_locals,
180            special_locals: self.special_locals,
181            locals_size: self.locals_size,
182            defined_locals_range: self.defined_locals_range,
183            results_base_slot: self.results_base_slot,
184            marker: PhantomData,
185            gc_ref_local_offsets: self.gc_ref_local_offsets,
186        }
187    }
188
189    fn compute_arg_slots<A: ABI>(
190        sig: &ABISig,
191    ) -> Result<(SpecialLocals, WasmLocals, SmallVec<[SPOffset; 4]>, u32)> {
192        // Go over the function ABI-signature and
193        // calculate the stack slots.
194        //
195        //  for each parameter p; when p
196        //
197        //  Stack =>
198        //      The slot offset is calculated from the ABIOperand offset
199        //      relative the to the frame pointer (and its inclusions, e.g.
200        //      return address).
201        //
202        //  Register =>
203        //     The slot is calculated by accumulating into the `next_frame_size`
204        //     the size + alignment of the type that the register is holding.
205        //
206        //  NOTE
207        //      This implementation takes inspiration from SpiderMonkey's implementation
208        //      to calculate local slots for function arguments
209        //      (https://searchfox.org/mozilla-central/source/js/src/wasm/WasmBCFrame.cpp#83).
210        //      The main difference is that SpiderMonkey's implementation
211        //      doesn't append any sort of metadata to the locals regarding stack
212        //      addressing mode (stack pointer or frame pointer), the offset is
213        //      declared negative if the local belongs to a stack argument;
214        //      that's enough to later calculate address of the local later on.
215        //
216        //      Winch appends an addressing mode to each slot, in the end
217        //      we want positive addressing from the stack pointer
218        //      for both locals and stack arguments.
219
220        let arg_base_offset = <A as ABI>::arg_base_offset().into();
221        let mut next_stack = 0u32;
222
223        // Skip the results base param; if present, the [Frame] will create
224        // a dedicated slot for it.
225        let mut params_iter = sig.params_without_retptr().into_iter();
226
227        // Handle special local slots.
228        let callee_vmctx = params_iter
229            .next()
230            .map(|arg| Self::abi_arg_slot(&arg, &mut next_stack, arg_base_offset))
231            .expect("Slot for VMContext");
232
233        let caller_vmctx = params_iter
234            .next()
235            .map(|arg| Self::abi_arg_slot(&arg, &mut next_stack, arg_base_offset))
236            .expect("Slot for VMContext");
237
238        let mut gc_ref_local_offsets = SmallVec::new();
239        let slots: WasmLocals = params_iter
240            .map(|arg| {
241                let slot = Self::abi_arg_slot(&arg, &mut next_stack, arg_base_offset);
242                if slot.addressed_from_sp() && needs_stack_map(&slot.ty) {
243                    gc_ref_local_offsets.push(SPOffset::from_u32(slot.offset));
244                }
245                slot
246            })
247            .collect();
248
249        Ok((
250            [callee_vmctx, caller_vmctx],
251            slots,
252            gc_ref_local_offsets,
253            next_stack,
254        ))
255    }
256
257    fn abi_arg_slot(arg: &ABIOperand, next_stack: &mut u32, arg_base_offset: u32) -> LocalSlot {
258        match arg {
259            // Create a local slot, for input register spilling,
260            // with type-size aligned access.
261            ABIOperand::Reg { ty, size, .. } => {
262                *next_stack = align_to(*next_stack, *size) + *size;
263                LocalSlot::new(*ty, *next_stack)
264            }
265            // Create a local slot, with an offset from the arguments base in
266            // the stack; which is the frame pointer + return address. GC
267            // references are re-homed into a frame slot by the prologue so
268            // that stack maps cover them: the caller's outgoing argument
269            // area is not visited by the collector, so a reference left
270            // there goes stale across a collection.
271            ABIOperand::Stack { ty, size, .. } if ty.is_vmgcref_type_and_not_i31() => {
272                *next_stack = align_to(*next_stack, *size) + *size;
273                LocalSlot::new(*ty, *next_stack)
274            }
275            ABIOperand::Stack { ty, offset, .. } => {
276                LocalSlot::stack_arg(*ty, offset + arg_base_offset)
277            }
278        }
279    }
280}
281
282impl Frame<Emission> {
283    /// Get the [`LocalSlot`] for a WebAssembly local.
284    /// This method assumes that the index is bound to u32::MAX, representing
285    /// the index space for WebAssembly locals.
286    ///
287    /// # Panics
288    /// This method panics if the index is not associated to a valid WebAssembly
289    /// local.
290    pub fn get_wasm_local(&self, index: u32) -> &LocalSlot {
291        self.wasm_locals
292            .get(index as usize)
293            .unwrap_or_else(|| panic!(" Expected WebAssembly local at slot: {index}"))
294    }
295
296    /// Get the [`LocalSlot`] for a special local.
297    ///
298    /// # Panics
299    /// This method panics if the index is not associated to a valid special
300    /// local.
301    pub fn get_special_local(&self, index: usize) -> &LocalSlot {
302        self.special_locals
303            .get(index)
304            .unwrap_or_else(|| panic!(" Expected special local at slot: {index}"))
305    }
306
307    /// Get the special [`LocalSlot`] for the `VMContext`.
308    pub fn vmctx_slot(&self) -> &LocalSlot {
309        self.get_special_local(0)
310    }
311
312    /// Frame offsets of the SP-addressed locals that stack maps must cover.
313    pub fn gc_ref_local_offsets(&self) -> &[SPOffset] {
314        &self.gc_ref_local_offsets
315    }
316
317    /// Returns the address of the local at the given index.
318    ///
319    /// # Panics
320    /// This function panics if the index is not associated to a local.
321    pub fn get_local_address<M: MacroAssembler>(
322        &self,
323        index: u32,
324        masm: &mut M,
325    ) -> Result<(WasmValType, M::Address)> {
326        let slot = self.get_wasm_local(index);
327        Ok((slot.ty, masm.local_address(&slot)?))
328    }
329}
winch_codegen/frame/mod.rs

winch_codegen/frame/
mod.rs
