Skip to main content

winch_codegen/codegen/
env.rs

1use crate::{
2    Result,
3    abi::{ABI, ABISig, wasm_sig},
4    codegen::{BlockSig, BuiltinFunction, BuiltinFunctions, OperandSize, control},
5    isa::TargetIsa,
6};
7use cranelift_codegen::ir::{UserExternalName, UserExternalNameRef};
8use std::collections::{
9    HashMap,
10    hash_map::Entry::{Occupied, Vacant},
11};
12use std::mem;
13use wasmparser::BlockType;
14use wasmtime_environ::{
15    BuiltinFunctionIndex, DefinedFuncIndex, FuncIndex, FuncKey, GlobalIndex, IndexType, Memory,
16    MemoryIndex, ModuleInternedTypeIndex, ModuleTranslation, ModuleTypesBuilder, PrimaryMap,
17    PtrSize, Table, TableIndex, TypeConvert, TypeIndex, VMOffsets, WasmHeapType, WasmValType,
18    collections::TryClone as _,
19};
20
21#[derive(Debug, Clone, Copy)]
22pub struct GlobalData {
23    /// The offset of the global.
24    pub offset: u32,
25    /// True if the global is imported.
26    pub imported: bool,
27    /// The WebAssembly type of the global.
28    pub ty: WasmValType,
29}
30
31/// Table metadata.
32#[derive(Debug, Copy, Clone)]
33pub struct TableData {
34    /// The offset to the base of the table.
35    pub offset: u32,
36    /// The offset to the current elements field.
37    pub current_elems_offset: u32,
38    /// If the table is imported, this field contains the offset to locate the
39    /// base of the table data.
40    pub import_from: Option<u32>,
41    /// The size of the table elements.
42    pub(crate) element_size: OperandSize,
43    /// The size of the current elements field.
44    pub(crate) current_elements_size: OperandSize,
45    /// The type of this table.
46    pub ty: Table,
47}
48
49impl TableData {
50    pub fn index_type(&self) -> WasmValType {
51        match self.ty.idx_type {
52            IndexType::I32 => WasmValType::I32,
53            IndexType::I64 => WasmValType::I64,
54        }
55    }
56}
57
58/// Heap metadata.
59///
60/// Heaps represent a WebAssembly linear memory.
61#[derive(Debug, Copy, Clone)]
62pub struct HeapData {
63    /// The offset to the base of the heap.
64    /// Relative to the `VMContext` pointer if the WebAssembly memory is locally
65    /// defined. Else this is relative to the location of the imported WebAssembly
66    /// memory location.
67    pub offset: u32,
68    /// The offset to the current length field.
69    pub current_length_offset: u32,
70    /// If the WebAssembly memory is imported or shared, this field contains the offset to locate the
71    /// base of the heap.
72    pub import_from: Option<u32>,
73    /// The memory type this heap is associated with.
74    pub memory: Memory,
75}
76
77impl HeapData {
78    pub fn index_type(&self) -> WasmValType {
79        match self.memory.idx_type {
80            IndexType::I32 => WasmValType::I32,
81            IndexType::I64 => WasmValType::I64,
82        }
83    }
84}
85
86/// A function callee.
87/// It categorizes how the callee should be treated
88/// when performing the call.
89#[derive(Clone)]
90pub(crate) enum Callee {
91    /// Locally defined function.
92    Local(FuncIndex),
93    /// Imported function.
94    Import(FuncIndex),
95    /// Function reference.
96    FuncRef(TypeIndex),
97    /// A built-in function.
98    Builtin(BuiltinFunction),
99    /// A built-in function, but the vmctx argument is located at the static
100    /// offset provided from the current function's vmctx.
101    BuiltinWithDifferentVmctx(BuiltinFunction, u32),
102}
103
104/// The function environment.
105///
106/// Contains all information about the module and runtime that is accessible to
107/// to a particular function during code generation.
108pub struct FuncEnv<'a, 'translation: 'a, 'data: 'translation, P: PtrSize> {
109    /// Offsets to the fields within the `VMContext` ptr.
110    pub vmoffsets: &'a VMOffsets<P>,
111    /// Metadata about the translation process of a WebAssembly module.
112    pub translation: &'translation ModuleTranslation<'data>,
113    /// The module's function types.
114    pub types: &'translation ModuleTypesBuilder,
115    /// The built-in functions available to the JIT code.
116    pub builtins: &'translation mut BuiltinFunctions,
117    /// Track resolved table information.
118    resolved_tables: HashMap<TableIndex, TableData>,
119    /// Track resolved heap information.
120    resolved_heaps: HashMap<MemoryIndex, HeapData>,
121    /// A map from [FunctionIndex] to [ABISig], to keep track of the resolved
122    /// function callees.
123    resolved_callees: HashMap<FuncIndex, ABISig>,
124    /// A map from [TypeIndex] to [ABISig], to keep track of the resolved
125    /// indirect function signatures.
126    resolved_sigs: HashMap<TypeIndex, ABISig>,
127    /// A map from [GlobalIndex] to [GlobalData].
128    resolved_globals: HashMap<GlobalIndex, GlobalData>,
129    /// Pointer size represented as a WebAssembly type.
130    ptr_type: WasmValType,
131    /// Whether or not to enable Spectre mitigation on heap bounds checks.
132    heap_access_spectre_mitigation: bool,
133    /// Whether or not to enable Spectre mitigation on table element accesses.
134    table_access_spectre_mitigation: bool,
135    /// Size of pages on the compilation target.
136    pub page_size_log2: u8,
137    name_map: PrimaryMap<UserExternalNameRef, UserExternalName>,
138    name_intern: HashMap<UserExternalName, UserExternalNameRef>,
139}
140
141pub fn ptr_type_from_ptr_size(size: u8) -> WasmValType {
142    (size == 8)
143        .then(|| WasmValType::I64)
144        .unwrap_or_else(|| unimplemented!("Support for non-64-bit architectures"))
145}
146
147impl<'a, 'translation, 'data, P: PtrSize> FuncEnv<'a, 'translation, 'data, P> {
148    /// Create a new function environment.
149    pub fn new(
150        vmoffsets: &'a VMOffsets<P>,
151        translation: &'translation ModuleTranslation<'data>,
152        types: &'translation ModuleTypesBuilder,
153        builtins: &'translation mut BuiltinFunctions,
154        isa: &dyn TargetIsa,
155        ptr_type: WasmValType,
156    ) -> Self {
157        Self {
158            vmoffsets,
159            translation,
160            types,
161            resolved_tables: HashMap::new(),
162            resolved_heaps: HashMap::new(),
163            resolved_callees: HashMap::new(),
164            resolved_sigs: HashMap::new(),
165            resolved_globals: HashMap::new(),
166            ptr_type,
167            heap_access_spectre_mitigation: isa.flags().enable_heap_access_spectre_mitigation(),
168            table_access_spectre_mitigation: isa.flags().enable_table_access_spectre_mitigation(),
169            page_size_log2: isa.page_size_align_log2(),
170            builtins,
171            name_map: Default::default(),
172            name_intern: Default::default(),
173        }
174    }
175
176    /// Derive the [`WasmType`] from the pointer size.
177    pub(crate) fn ptr_type(&self) -> WasmValType {
178        self.ptr_type
179    }
180
181    /// Returns the byte offset of `index` in the module's shared type-ID array.
182    pub(crate) fn shared_type_index_offset(&self, index: ModuleInternedTypeIndex) -> u32 {
183        index
184            .as_u32()
185            .checked_mul(u32::from(self.vmoffsets.size_of_vmshared_type_index()))
186            .unwrap()
187    }
188
189    /// Resolves a [`Callee::FuncRef`] from a type index.
190    pub(crate) fn funcref(&mut self, idx: TypeIndex) -> Callee {
191        Callee::FuncRef(idx)
192    }
193
194    /// Resolves a function [`Callee`] from an index.
195    pub(crate) fn callee_from_index(&mut self, idx: FuncIndex) -> Callee {
196        let import = self.translation.module.is_imported_function(idx);
197        if import {
198            Callee::Import(idx)
199        } else {
200            Callee::Local(idx)
201        }
202    }
203
204    /// Converts a [wasmparser::BlockType] into a [BlockSig].
205    pub(crate) fn resolve_block_sig(&self, ty: BlockType) -> Result<BlockSig> {
206        use BlockType::*;
207        Ok(match ty {
208            Empty => BlockSig::new(control::BlockType::void()),
209            Type(ty) => {
210                let ty = TypeConverter::new(self.translation, self.types).convert_valtype(ty)?;
211                BlockSig::new(control::BlockType::single(ty))
212            }
213            FuncType(idx) => {
214                let sig_index = self.translation.module.types[TypeIndex::from_u32(idx)]
215                    .unwrap_module_type_index();
216                let sig = self.types[sig_index].unwrap_func();
217                BlockSig::new(control::BlockType::func(sig.clone_panic_on_oom()))
218            }
219        })
220    }
221
222    /// Converts a parser heap type into its canonicalized Wasmtime type.
223    pub(crate) fn convert_heap_type(&self, ty: wasmparser::HeapType) -> Result<WasmHeapType> {
224        Ok(TypeConverter::new(self.translation, self.types).convert_heap_type(ty)?)
225    }
226
227    /// Resolves `GlobalData` of a global at the given index.
228    pub fn resolve_global(&mut self, index: GlobalIndex) -> GlobalData {
229        let ty = self.translation.module.globals[index].wasm_ty;
230        let val = || match self.translation.module.defined_global_index(index) {
231            Some(defined_index) => GlobalData {
232                offset: self.vmoffsets.globals().at(defined_index),
233                imported: false,
234                ty,
235            },
236            None => GlobalData {
237                offset: self.vmoffsets.imported_globals().at(index)
238                    + u32::from(self.vmoffsets.ptr.vm_global_import().from()),
239                imported: true,
240                ty,
241            },
242        };
243
244        *self.resolved_globals.entry(index).or_insert_with(val)
245    }
246
247    /// Returns the table information for the given table index.
248    pub fn resolve_table_data(&mut self, index: TableIndex) -> TableData {
249        match self.resolved_tables.entry(index) {
250            Occupied(entry) => *entry.get(),
251            Vacant(entry) => {
252                let (from_offset, base_offset, current_elems_offset) =
253                    match self.translation.module.defined_table_index(index) {
254                        Some(defined) => (
255                            None,
256                            self.vmoffsets.tables().at(defined)
257                                + u32::from(self.vmoffsets.ptr.vm_table_definition().base()),
258                            self.vmoffsets.tables().at(defined)
259                                + u32::from(
260                                    self.vmoffsets.ptr.vm_table_definition().current_elements(),
261                                ),
262                        ),
263                        None => (
264                            Some(
265                                self.vmoffsets.imported_tables().at(index)
266                                    + u32::from(self.vmoffsets.ptr.vm_table_import().from()),
267                            ),
268                            self.vmoffsets.ptr.vm_table_definition().base().into(),
269                            self.vmoffsets
270                                .ptr
271                                .vm_table_definition()
272                                .current_elements()
273                                .into(),
274                        ),
275                    };
276
277                *entry.insert(TableData {
278                    import_from: from_offset,
279                    offset: base_offset,
280                    current_elems_offset,
281                    element_size: OperandSize::from_bytes(self.vmoffsets.ptr.size()),
282                    current_elements_size: OperandSize::from_bytes(
283                        self.vmoffsets.size_of_vmtable_definition_current_elements(),
284                    ),
285                    ty: self.translation.module.tables[index],
286                })
287            }
288        }
289    }
290
291    /// Resolve a `HeapData` from a [MemoryIndex].
292    pub fn resolve_heap(&mut self, index: MemoryIndex) -> HeapData {
293        let mem = self.translation.module.memories[index];
294        let is_shared = mem.shared;
295        match self.resolved_heaps.entry(index) {
296            Occupied(entry) => *entry.get(),
297            Vacant(entry) => {
298                let (import_from, base_offset, current_length_offset) = match self
299                    .translation
300                    .module
301                    .defined_memory_index(index)
302                {
303                    Some(defined) => {
304                        if is_shared {
305                            (
306                                Some(self.vmoffsets.memories().at(defined)),
307                                self.vmoffsets.ptr.vm_memory_definition().base().into(),
308                                self.vmoffsets
309                                    .ptr
310                                    .vm_memory_definition()
311                                    .current_length()
312                                    .into(),
313                            )
314                        } else {
315                            let owned = self.translation.module.owned_memory_index(defined);
316                            (
317                                None,
318                                self.vmoffsets.owned_memories().at(owned)
319                                    + u32::from(self.vmoffsets.ptr.vm_memory_definition().base()),
320                                self.vmoffsets.owned_memories().at(owned)
321                                    + u32::from(
322                                        self.vmoffsets.ptr.vm_memory_definition().current_length(),
323                                    ),
324                            )
325                        }
326                    }
327                    None => (
328                        Some(
329                            self.vmoffsets.imported_memories().at(index)
330                                + u32::from(self.vmoffsets.ptr.vm_memory_import().from()),
331                        ),
332                        self.vmoffsets.ptr.vm_memory_definition().base().into(),
333                        self.vmoffsets
334                            .ptr
335                            .vm_memory_definition()
336                            .current_length()
337                            .into(),
338                    ),
339                };
340
341                let memory = &self.translation.module.memories[index];
342
343                *entry.insert(HeapData {
344                    offset: base_offset,
345                    import_from,
346                    current_length_offset,
347                    memory: *memory,
348                })
349            }
350        }
351    }
352
353    /// Get a [`Table`] from a [`TableIndex`].
354    pub fn table(&mut self, index: TableIndex) -> &Table {
355        &self.translation.module.tables[index]
356    }
357
358    /// Returns true if Spectre mitigations are enabled for heap bounds check.
359    pub fn heap_access_spectre_mitigation(&self) -> bool {
360        self.heap_access_spectre_mitigation
361    }
362
363    /// Returns true if Spectre mitigations are enabled for table element
364    /// accesses.
365    pub fn table_access_spectre_mitigation(&self) -> bool {
366        self.table_access_spectre_mitigation
367    }
368
369    pub(crate) fn callee_sig<'b, A>(&'b mut self, callee: &'b Callee) -> Result<&'b ABISig>
370    where
371        A: ABI,
372    {
373        match callee {
374            Callee::Local(idx) | Callee::Import(idx) => {
375                if self.resolved_callees.contains_key(idx) {
376                    Ok(self.resolved_callees.get(idx).unwrap())
377                } else {
378                    let types = self.translation.get_types();
379                    let types = types.as_ref();
380                    let ty = types[types.core_function_at(idx.as_u32())].unwrap_func();
381                    let converter = TypeConverter::new(self.translation, self.types);
382                    let ty = converter.convert_func_type(&ty)?;
383                    let sig = wasm_sig::<A>(&ty)?;
384                    self.resolved_callees.insert(*idx, sig);
385                    Ok(self.resolved_callees.get(idx).unwrap())
386                }
387            }
388            Callee::FuncRef(idx) => {
389                if self.resolved_sigs.contains_key(idx) {
390                    Ok(self.resolved_sigs.get(idx).unwrap())
391                } else {
392                    let sig_index = self.translation.module.types[*idx].unwrap_module_type_index();
393                    let ty = self.types[sig_index].unwrap_func();
394                    let sig = wasm_sig::<A>(ty)?;
395                    self.resolved_sigs.insert(*idx, sig);
396                    Ok(self.resolved_sigs.get(idx).unwrap())
397                }
398            }
399            Callee::Builtin(b) | Callee::BuiltinWithDifferentVmctx(b, _) => Ok(b.sig()),
400        }
401    }
402
403    /// Creates a name to reference the `builtin` provided.
404    pub fn name_builtin(&mut self, builtin: BuiltinFunctionIndex) -> UserExternalNameRef {
405        let key = FuncKey::WasmToBuiltinTrampoline(builtin);
406        let (namespace, index) = key.into_raw_parts();
407        self.intern_name(UserExternalName { namespace, index })
408    }
409
410    /// Creates a name to reference the wasm function `index` provided.
411    pub fn name_wasm(&mut self, def_func: DefinedFuncIndex) -> UserExternalNameRef {
412        let key = FuncKey::DefinedWasmFunction(self.translation.module_index(), def_func);
413        let (namespace, index) = key.into_raw_parts();
414        self.intern_name(UserExternalName { namespace, index })
415    }
416
417    /// Interns `name` into a `UserExternalNameRef` and ensures that duplicate
418    /// instances of `name` are given a unique name ref index.
419    fn intern_name(&mut self, name: UserExternalName) -> UserExternalNameRef {
420        *self
421            .name_intern
422            .entry(name.clone())
423            .or_insert_with(|| self.name_map.push(name))
424    }
425
426    /// Extracts the name map that was created while translating this function.
427    pub fn take_name_map(&mut self) -> PrimaryMap<UserExternalNameRef, UserExternalName> {
428        self.name_intern.clear();
429        mem::take(&mut self.name_map)
430    }
431}
432
433/// A wrapper struct over a reference to a [ModuleTranslation] and
434/// [ModuleTypesBuilder].
435pub(crate) struct TypeConverter<'a, 'data: 'a> {
436    translation: &'a ModuleTranslation<'data>,
437    types: &'a ModuleTypesBuilder,
438}
439
440impl TypeConvert for TypeConverter<'_, '_> {
441    fn lookup_heap_type(&self, idx: wasmparser::UnpackedIndex) -> WasmHeapType {
442        wasmtime_environ::WasmparserTypeConverter::new(self.types, |idx| {
443            self.translation.module.types[idx].unwrap_module_type_index()
444        })
445        .lookup_heap_type(idx)
446    }
447
448    fn lookup_type_index(
449        &self,
450        index: wasmparser::UnpackedIndex,
451    ) -> wasmtime_environ::EngineOrModuleTypeIndex {
452        wasmtime_environ::WasmparserTypeConverter::new(self.types, |idx| {
453            self.translation.module.types[idx].unwrap_module_type_index()
454        })
455        .lookup_type_index(index)
456    }
457}
458
459impl<'a, 'data> TypeConverter<'a, 'data> {
460    pub fn new(translation: &'a ModuleTranslation<'data>, types: &'a ModuleTypesBuilder) -> Self {
461        Self { translation, types }
462    }
463}