1 use crate::{ 2 Result, 3 abi::{ABI, ABIOperand, ABISig, LocalSlot, align_to}, 4 codegen::{CodeGenPhase, Emission, Prologue}, 5 masm::MacroAssembler, 6 }; 7 use smallvec::SmallVec; 8 use std::marker::PhantomData; 9 use std::ops::Range; 10 use wasmparser::{BinaryReader, FuncValidator, ValidatorResources}; 11 use 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. 18 pub(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. 24 pub(crate) type SpecialLocals = [LocalSlot; 2]; 25 26 /// Function defined locals start and end in the frame. 27 pub(crate) struct DefinedLocalsRange(Range<u32>); 28 29 impl DefinedLocalsRange { 30 /// Get a reference to the inner range. as_range(&self) -> &Range<u32>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)] 38 pub(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 45 impl DefinedLocals { 46 /// Compute the local slots for a Wasm function. new<A: ABI>( types: &impl TypeConvert, reader: &mut BinaryReader<'_>, validator: &mut FuncValidator<ValidatorResources>, ) -> Result<Self>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. 79 pub(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 99 impl Frame<Prologue> { 100 /// Allocate a new [`Frame`]. new<A: ABI>(sig: &ABISig, defined_locals: &DefinedLocals) -> Result<Frame<Prologue>>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`]. locals(&self) -> impl Iterator<Item = &LocalSlot>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. for_emission(self) -> Frame<Emission>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 compute_arg_slots<A: ABI>(sig: &ABISig) -> Result<(SpecialLocals, WasmLocals, u32)>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 abi_arg_slot(arg: &ABIOperand, next_stack: &mut u32, arg_base_offset: u32) -> LocalSlot235 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 252 impl 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. get_wasm_local(&self, index: u32) -> &LocalSlot260 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. get_special_local(&self, index: usize) -> &LocalSlot271 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`. vmctx_slot(&self) -> &LocalSlot278 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. get_local_address<M: MacroAssembler>( &self, index: u32, masm: &mut M, ) -> Result<(WasmValType, M::Address)>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 } 295