winch_codegen/codegen/
env.rs

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