winch_codegen/frame/
mod.rs

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

winch_codegen/frame/
mod.rs
