winch_codegen/codegen/
env.rs

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