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