Remove some allocations in CodeMemory (#3253)

* Remove some allocations in `CodeMemory`

This commit removes the `FinishedFunctions` type as well as allocations
associated with trampolines when allocating inside of a `CodeMemory`.
The main goal of this commit is to improve the time spent in
`CodeMemory` where currently today a good portion of time is spent
simply parsing symbol names and trying to extract function indices from
them. Instead this commit implements a new strategy (different from #3236)
where compilation records offset/length information for all
functions/trampolines so this doesn't need to be re-learned from the
object file later.

A consequence of this commit is that this offset information will be
decoded/encoded through `bincode` unconditionally, but we can also
optimize that later if necessary as well.

Internally this involved quite a bit of refactoring since the previous
map for `FinishedFunctions` was relatively heavily relied upon.

* comments

crates/cranelift/src/compiler.rs

@@ -21,13 +21,13 @@ use cranelift_wasm::{
 use object::write::Object;
 use std::any::Any;
 use std::cmp;
-use std::collections::{BTreeMap, BTreeSet};
+use std::collections::BTreeMap;
 use std::convert::TryFrom;
 use std::mem;
 use std::sync::Mutex;
 use wasmtime_environ::{
     AddressMapSection, CompileError, FilePos, FlagValue, FunctionBodyData, FunctionInfo,
-    InstructionAddressMap, Module, ModuleTranslation, StackMapInformation, TrapCode,
+    InstructionAddressMap, Module, ModuleTranslation, StackMapInformation, Trampoline, TrapCode,
     TrapEncodingBuilder, TrapInformation, Tunables, TypeTables, VMOffsets,
 };
 
@@ -120,12 +120,12 @@ impl wasmtime_environ::Compiler for Compiler {
         let func_index = module.func_index(func_index);
         let mut context = Context::new();
         context.func.name = get_func_name(func_index);
-        context.func.signature = func_signature(isa, module, types, func_index);
+        context.func.signature = func_signature(isa, translation, types, func_index);
         if tunables.generate_native_debuginfo {
             context.func.collect_debug_info();
         }
 
-        let mut func_env = FuncEnvironment::new(isa, module, types, tunables);
+        let mut func_env = FuncEnvironment::new(isa, translation, types, tunables);
 
         // We use these as constant offsets below in
         // `stack_limit_from_arguments`, so assert their values here. This
@@ -201,6 +201,7 @@ impl wasmtime_environ::Compiler for Compiler {
             None
         };
 
+        let length = u32::try_from(code_buf.len()).unwrap();
         Ok(Box::new(CompiledFunction {
             body: code_buf,
             jt_offsets: context.func.jt_offsets,
@@ -212,6 +213,8 @@ impl wasmtime_environ::Compiler for Compiler {
             info: FunctionInfo {
                 start_srcloc: address_transform.start_srcloc,
                 stack_maps: stack_map_sink.finish(),
+                start: 0,
+                length,
             },
             address_map: address_transform,
         }))
@@ -224,40 +227,38 @@ impl wasmtime_environ::Compiler for Compiler {
         funcs: PrimaryMap<DefinedFuncIndex, Box<dyn Any + Send>>,
         emit_dwarf: bool,
         obj: &mut Object,
-    ) -> Result<PrimaryMap<DefinedFuncIndex, FunctionInfo>> {
+    ) -> Result<(PrimaryMap<DefinedFuncIndex, FunctionInfo>, Vec<Trampoline>)> {
         const CODE_SECTION_ALIGNMENT: u64 = 0x1000;
         let funcs: crate::CompiledFunctions = funcs
             .into_iter()
             .map(|(_i, f)| *f.downcast().unwrap())
             .collect();
 
-        // Build trampolines for every signature that can be used by this module.
-        let signatures = translation
-            .module
-            .functions
-            .iter()
-            .filter_map(|(i, sig)| match translation.module.defined_func_index(i) {
-                Some(i) if !translation.module.possibly_exported_funcs.contains(&i) => None,
-                _ => Some(*sig),
-            })
-            .collect::<BTreeSet<_>>();
-        let mut trampolines = Vec::with_capacity(signatures.len());
-        for i in signatures {
-            let func = self.host_to_wasm_trampoline(&types.wasm_signatures[i])?;
-            trampolines.push((i, func));
-        }
-
         let mut builder = ObjectBuilder::new(obj, &translation.module);
         let mut addrs = AddressMapSection::default();
         let mut traps = TrapEncodingBuilder::default();
+        let compiled_trampolines = translation
+            .exported_signatures
+            .iter()
+            .map(|i| self.host_to_wasm_trampoline(&types.wasm_signatures[*i]))
+            .collect::<Result<Vec<_>, _>>()?;
 
+        let mut func_starts = Vec::with_capacity(funcs.len());
         for (i, func) in funcs.iter() {
             let range = builder.func(i, func);
             addrs.push(range.clone(), &func.address_map.instructions);
-            traps.push(range, &func.traps);
+            traps.push(range.clone(), &func.traps);
+            func_starts.push(range.start);
         }
-        for (i, func) in trampolines.iter() {
+
-            builder.trampoline(*i, func);
+        // Build trampolines for every signature that can be used by this module.
+        let mut trampolines = Vec::with_capacity(translation.exported_signatures.len());
+        for (i, func) in translation
+            .exported_signatures
+            .iter()
+            .zip(&compiled_trampolines)
+        {
+            trampolines.push(builder.trampoline(*i, &func));
         }
         builder.align_text_to(CODE_SECTION_ALIGNMENT);
 
@@ -295,7 +296,17 @@ impl wasmtime_environ::Compiler for Compiler {
         addrs.append_to(obj);
         traps.append_to(obj);
 
-        Ok(funcs.into_iter().map(|(_, f)| f.info).collect())
+        Ok((
+            funcs
+                .into_iter()
+                .zip(func_starts)
+                .map(|((_, mut f), start)| {
+                    f.info.start = start;
+                    f.info
+                })
+                .collect(),
+            trampolines,
+        ))
     }
 
     fn emit_trampoline_obj(
@@ -303,15 +314,15 @@ impl wasmtime_environ::Compiler for Compiler {
         ty: &WasmFuncType,
         host_fn: usize,
         obj: &mut Object,
-    ) -> Result<()> {
+    ) -> Result<(Trampoline, Trampoline)> {
         let host_to_wasm = self.host_to_wasm_trampoline(ty)?;
         let wasm_to_host = self.wasm_to_host_trampoline(ty, host_fn)?;
         let module = Module::new();
         let mut builder = ObjectBuilder::new(obj, &module);
-        builder.trampoline(SignatureIndex::new(0), &host_to_wasm);
+        let a = builder.trampoline(SignatureIndex::new(0), &host_to_wasm);
-        builder.trampoline(SignatureIndex::new(1), &wasm_to_host);
+        let b = builder.trampoline(SignatureIndex::new(1), &wasm_to_host);
         builder.finish(&*self.isa)?;
-        Ok(())
+        Ok((a, b))
     }
 
     fn triple(&self) -> &target_lexicon::Triple {

crates/cranelift/src/func_environ.rs

@@ -16,8 +16,8 @@ use std::convert::TryFrom;
 use std::mem;
 use wasmparser::Operator;
 use wasmtime_environ::{
-    BuiltinFunctionIndex, MemoryPlan, MemoryStyle, Module, TableStyle, Tunables, TypeTables,
+    BuiltinFunctionIndex, MemoryPlan, MemoryStyle, Module, ModuleTranslation, TableStyle, Tunables,
-    VMOffsets, INTERRUPTED, WASM_PAGE_SIZE,
+    TypeTables, VMOffsets, INTERRUPTED, WASM_PAGE_SIZE,
 };
 
 /// Compute an `ir::ExternalName` for a given wasm function index.
@@ -111,6 +111,7 @@ wasmtime_environ::foreach_builtin_function!(declare_function_signatures);
 pub struct FuncEnvironment<'module_environment> {
     isa: &'module_environment (dyn TargetIsa + 'module_environment),
     module: &'module_environment Module,
+    translation: &'module_environment ModuleTranslation<'module_environment>,
     types: &'module_environment TypeTables,
 
     /// The Cranelift global holding the vmctx address.
@@ -142,7 +143,7 @@ pub struct FuncEnvironment<'module_environment> {
 impl<'module_environment> FuncEnvironment<'module_environment> {
     pub fn new(
         isa: &'module_environment (dyn TargetIsa + 'module_environment),
-        module: &'module_environment Module,
+        translation: &'module_environment ModuleTranslation<'module_environment>,
         types: &'module_environment TypeTables,
         tunables: &'module_environment Tunables,
     ) -> Self {
@@ -157,11 +158,12 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
         );
         Self {
             isa,
-            module,
+            module: &translation.module,
+            translation,
             types,
             vmctx: None,
             builtin_function_signatures,
-            offsets: VMOffsets::new(isa.pointer_bytes(), module),
+            offsets: VMOffsets::new(isa.pointer_bytes(), &translation.module),
             tunables,
             fuel_var: Variable::new(0),
             vminterrupts_ptr: Variable::new(0),
@@ -1289,7 +1291,7 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
         func: &mut ir::Function,
         index: FuncIndex,
     ) -> WasmResult<ir::FuncRef> {
-        let sig = crate::func_signature(self.isa, self.module, self.types, index);
+        let sig = crate::func_signature(self.isa, self.translation, self.types, index);
         let signature = func.import_signature(sig);
         let name = get_func_name(index);
         Ok(func.import_function(ir::ExtFuncData {

crates/cranelift/src/lib.rs

@@ -95,7 +95,7 @@ use cranelift_entity::PrimaryMap;
 use cranelift_wasm::{DefinedFuncIndex, FuncIndex, WasmFuncType, WasmType};
 use target_lexicon::CallingConvention;
 use wasmtime_environ::{
-    FilePos, FunctionInfo, InstructionAddressMap, Module, TrapInformation, TypeTables,
+    FilePos, FunctionInfo, InstructionAddressMap, ModuleTranslation, TrapInformation, TypeTables,
 };
 
 pub use builder::builder;
@@ -257,16 +257,16 @@ fn indirect_signature(isa: &dyn TargetIsa, wasm: &WasmFuncType) -> ir::Signature
 /// use a custom theoretically faster calling convention instead of the default.
 fn func_signature(
     isa: &dyn TargetIsa,
-    module: &Module,
+    translation: &ModuleTranslation,
     types: &TypeTables,
     index: FuncIndex,
 ) -> ir::Signature {
-    let call_conv = match module.defined_func_index(index) {
+    let call_conv = match translation.module.defined_func_index(index) {
         // If this is a defined function in the module and it's never possibly
         // exported, then we can optimize this function to use the fastest
         // calling convention since it's purely an internal implementation
         // detail of the module itself.
-        Some(idx) if !module.possibly_exported_funcs.contains(&idx) => CallConv::Fast,
+        Some(idx) if !translation.escaped_funcs.contains(&idx) => CallConv::Fast,
 
         // ... otherwise if it's an imported function or if it's a possibly
         // exported function then we use the default ABI wasmtime would
@@ -277,7 +277,7 @@ fn func_signature(
     push_types(
         isa,
         &mut sig,
-        &types.wasm_signatures[module.functions[index]],
+        &types.wasm_signatures[translation.module.functions[index]],
     );
     return sig;
 }

crates/cranelift/src/obj.rs

@@ -37,7 +37,7 @@ use std::convert::TryFrom;
 use std::ops::Range;
 use wasmtime_environ::obj;
 use wasmtime_environ::{
-    DefinedFuncIndex, EntityRef, FuncIndex, Module, PrimaryMap, SignatureIndex,
+    DefinedFuncIndex, EntityRef, FuncIndex, Module, PrimaryMap, SignatureIndex, Trampoline,
 };
 
 const TEXT_SECTION_NAME: &[u8] = b".text";
@@ -230,9 +230,14 @@ impl<'a> ObjectBuilder<'a> {
         range
     }
 
-    pub fn trampoline(&mut self, sig: SignatureIndex, func: &'a CompiledFunction) {
+    pub fn trampoline(&mut self, sig: SignatureIndex, func: &'a CompiledFunction) -> Trampoline {
         let name = obj::trampoline_symbol_name(sig);
-        self.append_func(name.into_bytes(), func);
+        let (_, range) = self.append_func(name.into_bytes(), func);
+        Trampoline {
+            signature: sig,
+            start: range.start,
+            length: u32::try_from(range.end - range.start).unwrap(),
+        }
     }
 
     pub fn align_text_to(&mut self, align: u64) {

crates/environ/src/compilation.rs

@@ -2,8 +2,8 @@
 //! module.
 
 use crate::{
-    DefinedFuncIndex, FilePos, FunctionBodyData, ModuleTranslation, PrimaryMap, StackMap, Tunables,
+    DefinedFuncIndex, FilePos, FunctionBodyData, ModuleTranslation, PrimaryMap, SignatureIndex,
-    TypeTables, WasmError, WasmFuncType,
+    StackMap, Tunables, TypeTables, WasmError, WasmFuncType,
 };
 use anyhow::Result;
 use object::write::Object;
@@ -22,6 +22,25 @@ use thiserror::Error;
 pub struct FunctionInfo {
     pub start_srcloc: FilePos,
     pub stack_maps: Vec<StackMapInformation>,
+
+    /// Offset in the text section of where this function starts.
+    pub start: u64,
+    /// The size of the compiled function, in bytes.
+    pub length: u32,
+}
+
+/// Information about a compiled trampoline which the host can call to enter
+/// wasm.
+#[derive(Serialize, Deserialize, Clone)]
+#[allow(missing_docs)]
+pub struct Trampoline {
+    /// The signature this trampoline is for
+    pub signature: SignatureIndex,
+
+    /// Offset in the text section of where this function starts.
+    pub start: u64,
+    /// The size of the compiled function, in bytes.
+    pub length: u32,
 }
 
 /// The offset within a function of a GC safepoint, and its associated stack
@@ -154,19 +173,21 @@ pub trait Compiler: Send + Sync {
         funcs: PrimaryMap<DefinedFuncIndex, Box<dyn Any + Send>>,
         emit_dwarf: bool,
         obj: &mut Object,
-    ) -> Result<PrimaryMap<DefinedFuncIndex, FunctionInfo>>;
+    ) -> Result<(PrimaryMap<DefinedFuncIndex, FunctionInfo>, Vec<Trampoline>)>;
 
     /// Inserts two functions for host-to-wasm and wasm-to-host trampolines into
     /// the `obj` provided.
     ///
     /// This will configure the same sections as `emit_obj`, but will likely be
-    /// much smaller.
+    /// much smaller. The two returned `Trampoline` structures describe where to
+    /// find the host-to-wasm and wasm-to-host trampolines in the text section,
+    /// respectively.
     fn emit_trampoline_obj(
         &self,
         ty: &WasmFuncType,
         host_fn: usize,
         obj: &mut Object,
-    ) -> Result<()>;
+    ) -> Result<(Trampoline, Trampoline)>;
 
     /// Creates a new `Object` file which is used to build the results of a
     /// compilation into.

crates/environ/src/lib.rs

@@ -43,6 +43,7 @@ pub use crate::stack_map::StackMap;
 pub use crate::trap_encoding::*;
 pub use crate::tunables::Tunables;
 pub use crate::vmoffsets::*;
+pub use object;
 
 // Reexport all of these type-level since they're quite commonly used and it's
 // much easier to refer to everything through one crate rather than importing

crates/environ/src/module.rs

@@ -3,7 +3,7 @@
 use crate::{EntityRef, ModuleTranslation, PrimaryMap, Tunables, WASM_PAGE_SIZE};
 use indexmap::IndexMap;
 use serde::{Deserialize, Serialize};
-use std::collections::{BTreeMap, BTreeSet};
+use std::collections::BTreeMap;
 use std::convert::TryFrom;
 use std::ops::Range;
 use wasmtime_types::*;
@@ -422,10 +422,6 @@ pub struct Module {
 
     /// The type of each nested wasm module this module contains.
     pub modules: PrimaryMap<ModuleIndex, ModuleTypeIndex>,
-
-    /// The set of defined functions within this module which are located in
-    /// element segments.
-    pub possibly_exported_funcs: BTreeSet<DefinedFuncIndex>,
 }
 
 /// Initialization routines for creating an instance, encompassing imports,

crates/environ/src/module_environ.rs

@@ -10,7 +10,7 @@ use crate::{
 };
 use cranelift_entity::packed_option::ReservedValue;
 use std::borrow::Cow;
-use std::collections::{hash_map::Entry, HashMap};
+use std::collections::{hash_map::Entry, HashMap, HashSet};
 use std::convert::{TryFrom, TryInto};
 use std::mem;
 use std::path::PathBuf;
@@ -61,6 +61,23 @@ pub struct ModuleTranslation<'data> {
     /// References to the function bodies.
     pub function_body_inputs: PrimaryMap<DefinedFuncIndex, FunctionBodyData<'data>>,
 
+    /// The set of defined functions within this module which are "possibly
+    /// exported" which means that the host can possibly call them. This
+    /// includes functions such as:
+    ///
+    /// * Exported functions
+    /// * Functions in element segments
+    /// * Functions via `ref.func` instructions
+    ///
+    /// This set is used to determine the set of type signatures that need
+    /// trampolines for the host to call into.
+    pub escaped_funcs: HashSet<DefinedFuncIndex>,
+
+    /// A list of type signatures which are considered exported from this
+    /// module, or those that can possibly be called. This list is sorted, and
+    /// trampolines for each of these signatures are required.
+    pub exported_signatures: Vec<SignatureIndex>,
+
     /// DWARF debug information, if enabled, parsed from the module.
     pub debuginfo: DebugInfoData<'data>,
 
@@ -221,6 +238,22 @@ impl<'data> ModuleEnvironment<'data> {
             Payload::End => {
                 validator.end()?;
 
+                // With the `escaped_funcs` set of functions finished
+                // we can calculate the set of signatures that are exported as
+                // the set of exported functions' signatures.
+                self.result.exported_signatures = self
+                    .result
+                    .module
+                    .functions
+                    .iter()
+                    .filter_map(|(i, sig)| match self.result.module.defined_func_index(i) {
+                        Some(i) if !self.result.escaped_funcs.contains(&i) => None,
+                        _ => Some(*sig),
+                    })
+                    .collect();
+                self.result.exported_signatures.sort_unstable();
+                self.result.exported_signatures.dedup();
+
                 self.result.creation_artifacts.shrink_to_fit();
                 self.result.creation_modules.shrink_to_fit();
 
@@ -417,7 +450,7 @@ impl<'data> ModuleEnvironment<'data> {
                         Operator::RefNull { ty: _ } => GlobalInit::RefNullConst,
                         Operator::RefFunc { function_index } => {
                             let index = FuncIndex::from_u32(function_index);
-                            self.flag_func_possibly_exported(index);
+                            self.flag_func_escaped(index);
                             GlobalInit::RefFunc(index)
                         }
                         Operator::GlobalGet { global_index } => {
@@ -446,7 +479,7 @@ impl<'data> ModuleEnvironment<'data> {
                     let entity = match kind {
                         ExternalKind::Function => {
                             let index = FuncIndex::from_u32(index);
-                            self.flag_func_possibly_exported(index);
+                            self.flag_func_escaped(index);
                             EntityIndex::Function(index)
                         }
                         ExternalKind::Table => EntityIndex::Table(TableIndex::from_u32(index)),
@@ -471,7 +504,7 @@ impl<'data> ModuleEnvironment<'data> {
                 validator.start_section(func, &range)?;
 
                 let func_index = FuncIndex::from_u32(func);
-                self.flag_func_possibly_exported(func_index);
+                self.flag_func_escaped(func_index);
                 debug_assert!(self.result.module.start_func.is_none());
                 self.result.module.start_func = Some(func_index);
             }
@@ -497,7 +530,7 @@ impl<'data> ModuleEnvironment<'data> {
                         elements.push(match item? {
                             ElementItem::Func(f) => {
                                 let f = FuncIndex::from_u32(f);
-                                self.flag_func_possibly_exported(f);
+                                self.flag_func_escaped(f);
                                 f
                             }
                             ElementItem::Null(_ty) => FuncIndex::reserved_value(),
@@ -1113,9 +1146,9 @@ and for re-adding support for interface types you can see this issue:
             .push(ModuleSignature { imports, exports })
     }
 
-    fn flag_func_possibly_exported(&mut self, func: FuncIndex) {
+    fn flag_func_escaped(&mut self, func: FuncIndex) {
         if let Some(idx) = self.result.module.defined_func_index(func) {
-            self.result.module.possibly_exported_funcs.insert(idx);
+            self.result.escaped_funcs.insert(idx);
         }
     }
 

crates/jit/src/code_memory.rs

@@ -2,12 +2,9 @@
 
 use crate::unwind::UnwindRegistration;
 use anyhow::{Context, Result};
-use object::read::{File as ObjectFile, Object, ObjectSection, ObjectSymbol};
+use object::read::{File as ObjectFile, Object, ObjectSection};
-use std::collections::BTreeMap;
 use std::mem::ManuallyDrop;
-use wasmtime_environ::obj::{try_parse_func_name, try_parse_trampoline_name};
+use wasmtime_runtime::Mmap;
-use wasmtime_environ::{FuncIndex, SignatureIndex};
-use wasmtime_runtime::{Mmap, VMFunctionBody};
 
 struct CodeMemoryEntry {
     mmap: ManuallyDrop<Mmap>,
@@ -38,42 +35,6 @@ impl Drop for CodeMemoryEntry {
     }
 }
 
-pub struct CodeMemoryObjectAllocation<'a> {
-    pub code_range: &'a mut [u8],
-    funcs: BTreeMap<FuncIndex, (usize, usize)>,
-    trampolines: BTreeMap<SignatureIndex, (usize, usize)>,
-}
-
-impl<'a> CodeMemoryObjectAllocation<'a> {
-    pub fn funcs_len(&self) -> usize {
-        self.funcs.len()
-    }
-
-    pub fn trampolines_len(&self) -> usize {
-        self.trampolines.len()
-    }
-
-    pub fn funcs(&'a self) -> impl Iterator<Item = (FuncIndex, &'a mut [VMFunctionBody])> + 'a {
-        let buf = self.code_range as *const _ as *mut [u8];
-        self.funcs.iter().map(move |(i, (start, len))| {
-            (*i, unsafe {
-                CodeMemory::view_as_mut_vmfunc_slice(&mut (*buf)[*start..*start + *len])
-            })
-        })
-    }
-
-    pub fn trampolines(
-        &'a self,
-    ) -> impl Iterator<Item = (SignatureIndex, &'a mut [VMFunctionBody])> + 'a {
-        let buf = self.code_range as *const _ as *mut [u8];
-        self.trampolines.iter().map(move |(i, (start, len))| {
-            (*i, unsafe {
-                CodeMemory::view_as_mut_vmfunc_slice(&mut (*buf)[*start..*start + *len])
-            })
-        })
-    }
-}
-
 /// Memory manager for executable code.
 pub struct CodeMemory {
     entries: Vec<CodeMemoryEntry>,
@@ -132,39 +93,25 @@ impl CodeMemory {
         self.published = self.entries.len();
     }
 
-    /// Convert mut a slice from u8 to VMFunctionBody.
-    fn view_as_mut_vmfunc_slice(slice: &mut [u8]) -> &mut [VMFunctionBody] {
-        let byte_ptr: *mut [u8] = slice;
-        let body_ptr = byte_ptr as *mut [VMFunctionBody];
-        unsafe { &mut *body_ptr }
-    }
-
     /// Alternative to `allocate_for_object`, but when the object file isn't
     /// already parsed.
-    pub fn allocate_for_object_unparsed<'a>(
+    pub fn allocate_for_object_unparsed<'a, 'b>(
         &'a mut self,
-        obj: &[u8],
+        obj: &'b [u8],
-    ) -> Result<CodeMemoryObjectAllocation<'a>> {
+    ) -> Result<(&'a mut [u8], ObjectFile<'b>)> {
         let obj = ObjectFile::parse(obj)?;
-        self.allocate_for_object(&obj)
+        Ok((self.allocate_for_object(&obj)?, obj))
     }
 
     /// Allocates and copies the ELF image code section into CodeMemory.
     /// Returns references to functions and trampolines defined there.
-    pub fn allocate_for_object<'a>(
+    pub fn allocate_for_object(&mut self, obj: &ObjectFile) -> Result<&mut [u8]> {
-        &'a mut self,
-        obj: &ObjectFile,
-    ) -> Result<CodeMemoryObjectAllocation<'a>> {
         let text_section = obj.section_by_name(".text").unwrap();
         let text_section_size = text_section.size() as usize;
 
         if text_section_size == 0 {
             // No code in the image.
-            return Ok(CodeMemoryObjectAllocation {
+            return Ok(&mut []);
-                code_range: &mut [],
-                funcs: BTreeMap::new(),
-                trampolines: BTreeMap::new(),
-            });
         }
 
         // Find the platform-specific unwind section, if present, which contains
@@ -195,29 +142,6 @@ impl CodeMemory {
             );
         }
 
-        // Track locations of all defined functions and trampolines.
+        Ok(&mut entry.mmap.as_mut_slice()[..text_section_size])
-        let mut funcs = BTreeMap::new();
-        let mut trampolines = BTreeMap::new();
-        for sym in obj.symbols() {
-            match sym.name() {
-                Ok(name) => {
-                    if let Some(index) = try_parse_func_name(name) {
-                        let is_import = sym.section_index().is_none();
-                        if !is_import {
-                            funcs.insert(index, (sym.address() as usize, sym.size() as usize));
-                        }
-                    } else if let Some(index) = try_parse_trampoline_name(name) {
-                        trampolines.insert(index, (sym.address() as usize, sym.size() as usize));
-                    }
-                }
-                Err(_) => (),
-            }
-        }
-
-        Ok(CodeMemoryObjectAllocation {
-            code_range: &mut entry.mmap.as_mut_slice()[..text_section_size],
-            funcs,
-            trampolines,
-        })
     }
 }

crates/jit/src/debug.rs

@@ -27,9 +27,6 @@ pub fn create_gdbjit_image(
         }
     }
 
-    // let mut file = ::std::fs::File::create(::std::path::Path::new("test.o")).expect("file");
-    // ::std::io::Write::write_all(&mut file, &bytes).expect("write");
-
     Ok(bytes)
 }
 

crates/jit/src/instantiate.rs

@@ -8,9 +8,8 @@ use crate::debug::create_gdbjit_image;
 use crate::link::link_module;
 use crate::{MmapVec, ProfilingAgent};
 use anyhow::{anyhow, Context, Result};
-use object::read::File;
 use object::write::{Object, StandardSegment};
-use object::{Object as _, ObjectSection, SectionKind};
+use object::{File, Object as _, ObjectSection, SectionKind};
 use serde::{Deserialize, Serialize};
 use std::ops::Range;
 use std::sync::Arc;
@@ -18,7 +17,8 @@ use thiserror::Error;
 use wasmtime_environ::{
     CompileError, DefinedFuncIndex, FunctionInfo, InstanceSignature, InstanceTypeIndex, Module,
     ModuleSignature, ModuleTranslation, ModuleTypeIndex, PrimaryMap, SignatureIndex,
-    StackMapInformation, Tunables, WasmFuncType, ELF_WASMTIME_ADDRMAP, ELF_WASMTIME_TRAPS,
+    StackMapInformation, Trampoline, Tunables, WasmFuncType, ELF_WASMTIME_ADDRMAP,
+    ELF_WASMTIME_TRAPS,
 };
 use wasmtime_runtime::{GdbJitImageRegistration, InstantiationError, VMFunctionBody, VMTrampoline};
 
@@ -80,6 +80,10 @@ pub struct CompiledModuleInfo {
     /// Metadata about each compiled function.
     funcs: PrimaryMap<DefinedFuncIndex, FunctionInfo>,
 
+    /// The trampolines compiled into the text section and their start/length
+    /// relative to the start of the text section.
+    trampolines: Vec<Trampoline>,
+
     /// General compilation metadata.
     meta: Metadata,
 }
@@ -124,6 +128,7 @@ pub fn finish_compile(
     translation: ModuleTranslation<'_>,
     mut obj: Object,
     funcs: PrimaryMap<DefinedFuncIndex, FunctionInfo>,
+    trampolines: Vec<Trampoline>,
     tunables: &Tunables,
 ) -> Result<(MmapVec, CompiledModuleInfo)> {
     let ModuleTranslation {
@@ -192,6 +197,7 @@ pub fn finish_compile(
     let info = CompiledModuleInfo {
         module,
         funcs,
+        trampolines,
         meta: Metadata {
             native_debug_info_present: tunables.generate_native_debuginfo,
             has_unparsed_debuginfo,
@@ -221,10 +227,6 @@ pub fn finish_compile(
     }
 }
 
-struct FinishedFunctions(PrimaryMap<DefinedFuncIndex, *mut [VMFunctionBody]>);
-unsafe impl Send for FinishedFunctions {}
-unsafe impl Sync for FinishedFunctions {}
-
 /// This is intended to mirror the type tables in `wasmtime_environ`, except that
 /// it doesn't store the native signatures which are no longer needed past compilation.
 #[derive(Serialize, Deserialize)]
@@ -259,10 +261,9 @@ pub struct CompiledModule {
     mmap: MmapVec,
     module: Arc<Module>,
     funcs: PrimaryMap<DefinedFuncIndex, FunctionInfo>,
+    trampolines: Vec<Trampoline>,
     meta: Metadata,
     code: Arc<ModuleCode>,
-    finished_functions: FinishedFunctions,
-    trampolines: Vec<(SignatureIndex, VMTrampoline)>,
 }
 
 impl CompiledModule {
@@ -305,27 +306,27 @@ impl CompiledModule {
         };
         let module = Arc::new(info.module);
         let funcs = info.funcs;
+        let trampolines = info.trampolines;
         let wasm_data = subslice_range(section(ELF_WASM_DATA)?, &mmap);
         let address_map_data = subslice_range(section(ELF_WASMTIME_ADDRMAP)?, &mmap);
         let trap_data = subslice_range(section(ELF_WASMTIME_TRAPS)?, &mmap);
 
         // Allocate all of the compiled functions into executable memory,
         // copying over their contents.
-        let (code_memory, code_range, finished_functions, trampolines) =
+        let (code_memory, code_range) = build_code_memory(&obj).map_err(|message| {
-            build_code_memory(&obj, &module).map_err(|message| {
+            SetupError::Instantiate(InstantiationError::Resource(anyhow::anyhow!(
-                SetupError::Instantiate(InstantiationError::Resource(anyhow::anyhow!(
+                "failed to build code memory for functions: {}",
-                    "failed to build code memory for functions: {}",
+                message
-                    message
+            )))
-                )))
+        })?;
-            })?;
 
-        let finished_functions = FinishedFunctions(finished_functions);
         let start = code_range.0 as usize;
         let end = start + code_range.1;
 
         let mut ret = Self {
             meta: info.meta,
             funcs,
+            trampolines,
             module,
             mmap,
             wasm_data,
@@ -336,8 +337,6 @@ impl CompiledModule {
                 code_memory,
                 dbg_jit_registration: None,
             }),
-            finished_functions,
-            trampolines,
         };
         ret.register_debug_and_profiling(profiler)?;
 
@@ -400,6 +399,11 @@ impl CompiledModule {
         &self.module
     }
 
+    /// Returns the `FunctionInfo` map for all defined functions.
+    pub fn functions(&self) -> &PrimaryMap<DefinedFuncIndex, FunctionInfo> {
+        &self.funcs
+    }
+
     /// Return a reference to a mutable module (if possible).
     pub fn module_mut(&mut self) -> Option<&mut Module> {
         Arc::get_mut(&mut self.module)
@@ -407,13 +411,28 @@ impl CompiledModule {
 
     /// Returns the map of all finished JIT functions compiled for this module
     #[inline]
-    pub fn finished_functions(&self) -> &PrimaryMap<DefinedFuncIndex, *mut [VMFunctionBody]> {
+    pub fn finished_functions(
-        &self.finished_functions.0
+        &self,
+    ) -> impl ExactSizeIterator<Item = (DefinedFuncIndex, *mut [VMFunctionBody])> + '_ {
+        self.funcs.iter().map(move |(i, info)| {
+            (
+                i,
+                std::ptr::slice_from_raw_parts_mut(
+                    (self.code.range.0 + info.start as usize) as *mut VMFunctionBody,
+                    info.length as usize,
+                ),
+            )
+        })
     }
 
     /// Returns the per-signature trampolines for this module.
-    pub fn trampolines(&self) -> &[(SignatureIndex, VMTrampoline)] {
+    pub fn trampolines(&self) -> impl Iterator<Item = (SignatureIndex, VMTrampoline)> + '_ {
-        &self.trampolines
+        self.trampolines.iter().map(move |info| {
+            (info.signature, unsafe {
+                let ptr = self.code.range.0 + info.start as usize;
+                std::mem::transmute::<usize, VMTrampoline>(ptr)
+            })
+        })
     }
 
     /// Returns the stack map information for all functions defined in this
@@ -425,25 +444,24 @@ impl CompiledModule {
         &self,
     ) -> impl Iterator<Item = (*mut [VMFunctionBody], &[StackMapInformation])> {
         self.finished_functions()
-            .values()
+            .map(|(_, f)| f)
-            .copied()
             .zip(self.funcs.values().map(|f| f.stack_maps.as_slice()))
     }
 
     /// Lookups a defined function by a program counter value.
     ///
     /// Returns the defined function index and the relative address of
-    /// `text_offfset` within the function itself.
+    /// `text_offset` within the function itself.
     pub fn func_by_text_offset(&self, text_offset: usize) -> Option<(DefinedFuncIndex, u32)> {
-        let functions = self.finished_functions();
+        let text_offset = text_offset as u64;
 
-        let text_section = self.code().range().0;
+        let index = match self
-        let pc = text_section + text_offset;
+            .funcs
-        let index = match functions.binary_search_values_by_key(&pc, |body| unsafe {
+            .binary_search_values_by_key(&text_offset, |info| {
-            debug_assert!(!(**body).is_empty());
+                debug_assert!(info.length > 0);
-            // Return the inclusive "end" of the function
+                // Return the inclusive "end" of the function
-            (**body).as_ptr() as usize + (**body).len() - 1
+                info.start + u64::from(info.length) - 1
-        }) {
+            }) {
             Ok(k) => {
                 // Exact match, pc is at the end of this function
                 k
@@ -456,18 +474,15 @@ impl CompiledModule {
             }
         };
 
-        let body = functions.get(index)?;
+        let body = self.funcs.get(index)?;
-        let (start, end) = unsafe {
+        let start = body.start;
-            let ptr = (**body).as_ptr();
+        let end = body.start + u64::from(body.length);
-            let len = (**body).len();
-            (ptr as usize, ptr as usize + len)
-        };
 
-        if pc < start || end < pc {
+        if text_offset < start || end < text_offset {
             return None;
         }
 
-        Some((index, (text_offset - (start - text_section)) as u32))
+        Some((index, (text_offset - body.start) as u32))
     }
 
     /// Gets the function information for a given function index.
@@ -539,55 +554,19 @@ impl<'a> SymbolizeContext<'a> {
     }
 }
 
-fn build_code_memory(
+fn build_code_memory(obj: &File) -> Result<(CodeMemory, (*const u8, usize))> {
-    obj: &File,
-    module: &Module,
-) -> Result<(
-    CodeMemory,
-    (*const u8, usize),
-    PrimaryMap<DefinedFuncIndex, *mut [VMFunctionBody]>,
-    Vec<(SignatureIndex, VMTrampoline)>,
-)> {
     let mut code_memory = CodeMemory::new();
 
     let allocation = code_memory.allocate_for_object(obj)?;
 
-    // Populate the finished functions from the allocation
+    link_module(obj, allocation);
-    let mut finished_functions = PrimaryMap::with_capacity(allocation.funcs_len());
-    for (i, fat_ptr) in allocation.funcs() {
-        let start = fat_ptr.as_ptr() as usize;
-        let fat_ptr: *mut [VMFunctionBody] = fat_ptr;
-        // Assert that the function bodies are pushed in sort order
-        // This property is relied upon to search for functions by PC values
-        assert!(
-            start
-                > finished_functions
-                    .last()
-                    .map(|f: &*mut [VMFunctionBody]| unsafe { (**f).as_ptr() as usize })
-                    .unwrap_or(0)
-        );
-        assert_eq!(
-            Some(finished_functions.push(fat_ptr)),
-            module.defined_func_index(i)
-        );
-    }
 
-    // Populate the trampolines from the allocation
+    let code_range = (allocation.as_ptr(), allocation.len());
-    let mut trampolines = Vec::with_capacity(allocation.trampolines_len());
-    for (i, fat_ptr) in allocation.trampolines() {
-        let fnptr =
-            unsafe { std::mem::transmute::<*const VMFunctionBody, VMTrampoline>(fat_ptr.as_ptr()) };
-        trampolines.push((i, fnptr));
-    }
-
-    link_module(obj, allocation.code_range);
-
-    let code_range = (allocation.code_range.as_ptr(), allocation.code_range.len());
 
     // Make all code compiled thus far executable.
     code_memory.publish();
 
-    Ok((code_memory, code_range, finished_functions, trampolines))
+    Ok((code_memory, code_range))
 }
 
 /// Returns the range of `inner` within `outer`, such that `outer[range]` is the

crates/jit/src/profiling/jitdump_linux.rs

@@ -290,7 +290,7 @@ impl State {
         let tid = pid; // ThreadId does appear to track underlying thread. Using PID.
 
         for (idx, func) in module.finished_functions() {
-            let (addr, len) = unsafe { ((**func).as_ptr() as *const u8, (**func).len()) };
+            let (addr, len) = unsafe { ((*func).as_ptr() as *const u8, (*func).len()) };
             if let Some(img) = &dbg_image {
                 if let Err(err) = self.dump_from_debug_image(img, "wasm", addr, len, pid, tid) {
                     println!(

crates/jit/src/profiling/vtune_linux.rs

@@ -121,7 +121,7 @@ impl State {
         let global_module_id = MODULE_ID.fetch_add(1, atomic::Ordering::SeqCst);
 
         for (idx, func) in module.finished_functions() {
-            let (addr, len) = unsafe { ((**func).as_ptr() as *const u8, (**func).len()) };
+            let (addr, len) = unsafe { ((*func).as_ptr() as *const u8, (*func).len()) };
             let default_filename = "wasm_file";
             let default_module_name = String::from("wasm_module");
             let module_name = module

crates/lightbeam/wasmtime/src/lib.rs

@@ -17,7 +17,7 @@ use wasmtime_environ::{
 };
 use wasmtime_environ::{
     DefinedFuncIndex, DefinedGlobalIndex, DefinedMemoryIndex, DefinedTableIndex, FuncIndex,
-    GlobalIndex, MemoryIndex, TableIndex, TypeIndex, WasmFuncType,
+    GlobalIndex, MemoryIndex, TableIndex, Trampoline, TypeIndex, WasmFuncType,
 };
 
 /// A compiler that compiles a WebAssembly module with Lightbeam, directly translating the Wasm file.
@@ -86,7 +86,7 @@ impl Compiler for Lightbeam {
         _funcs: PrimaryMap<DefinedFuncIndex, Box<dyn Any + Send>>,
         _emit_dwarf: bool,
         _obj: &mut Object,
-    ) -> Result<PrimaryMap<DefinedFuncIndex, FunctionInfo>> {
+    ) -> Result<(PrimaryMap<DefinedFuncIndex, FunctionInfo>, Vec<Trampoline>)> {
         unimplemented!()
     }
 
@@ -95,7 +95,7 @@ impl Compiler for Lightbeam {
         _ty: &WasmFuncType,
         _host_fn: usize,
         _obj: &mut Object,
-    ) -> Result<()> {
+    ) -> Result<(Trampoline, Trampoline)> {
         unimplemented!()
     }
 

crates/runtime/src/instance/allocator.rs

@@ -4,7 +4,7 @@ use crate::memory::{DefaultMemoryCreator, Memory};
 use crate::table::Table;
 use crate::traphandlers::Trap;
 use crate::vmcontext::{
-    VMBuiltinFunctionsArray, VMCallerCheckedAnyfunc, VMContext, VMFunctionBody, VMGlobalDefinition,
+    VMBuiltinFunctionsArray, VMCallerCheckedAnyfunc, VMContext, VMGlobalDefinition,
     VMSharedSignatureIndex,
 };
 use crate::Store;
@@ -18,8 +18,8 @@ use std::slice;
 use std::sync::Arc;
 use thiserror::Error;
 use wasmtime_environ::{
-    DefinedFuncIndex, DefinedMemoryIndex, DefinedTableIndex, EntityRef, EntitySet, GlobalInit,
+    DefinedFuncIndex, DefinedMemoryIndex, DefinedTableIndex, EntityRef, EntitySet, FunctionInfo,
-    HostPtr, MemoryInitialization, MemoryInitializer, Module, ModuleType, PrimaryMap,
+    GlobalInit, HostPtr, MemoryInitialization, MemoryInitializer, Module, ModuleType, PrimaryMap,
     SignatureIndex, TableInitializer, TrapCode, VMOffsets, WasmType, WASM_PAGE_SIZE,
 };
 
@@ -34,8 +34,12 @@ pub struct InstanceAllocationRequest<'a> {
     /// The module being instantiated.
     pub module: Arc<Module>,
 
-    /// The finished (JIT) functions for the module.
+    /// The base address of where JIT functions are located.
-    pub finished_functions: &'a PrimaryMap<DefinedFuncIndex, *mut [VMFunctionBody]>,
+    pub image_base: usize,
+
+    /// Descriptors about each compiled function, such as the offset from
+    /// `image_base`.
+    pub functions: &'a PrimaryMap<DefinedFuncIndex, FunctionInfo>,
 
     /// The imports to use for the instantiation.
     pub imports: Imports<'a>,
@@ -483,7 +487,8 @@ unsafe fn initialize_vmcontext(instance: &mut Instance, req: InstanceAllocationR
 
         let (func_ptr, vmctx) = if let Some(def_index) = instance.module.defined_func_index(index) {
             (
-                NonNull::new(req.finished_functions[def_index] as *mut _).unwrap(),
+                NonNull::new((req.image_base + req.functions[def_index].start as usize) as *mut _)
+                    .unwrap(),
                 instance.vmctx_ptr(),
             )
         } else {

crates/runtime/src/instance/allocator/pooling.rs

@@ -1393,7 +1393,7 @@ mod test {
 
         let mut handles = Vec::new();
         let module = Arc::new(Module::default());
-        let finished_functions = &PrimaryMap::new();
+        let functions = &PrimaryMap::new();
 
         for _ in (0..3).rev() {
             handles.push(
@@ -1402,7 +1402,8 @@ mod test {
                         PoolingAllocationStrategy::NextAvailable,
                         InstanceAllocationRequest {
                             module: module.clone(),
-                            finished_functions,
+                            image_base: 0,
+                            functions,
                             imports: Imports {
                                 functions: &[],
                                 tables: &[],
@@ -1425,7 +1426,8 @@ mod test {
             PoolingAllocationStrategy::NextAvailable,
             InstanceAllocationRequest {
                 module: module.clone(),
-                finished_functions,
+                functions,
+                image_base: 0,
                 imports: Imports {
                     functions: &[],
                     tables: &[],

crates/runtime/src/instance/allocator/pooling/uffd.rs

@@ -512,7 +512,7 @@ mod test {
 
             let mut handles = Vec::new();
             let module = Arc::new(module);
-            let finished_functions = &PrimaryMap::new();
+            let functions = &PrimaryMap::new();
 
             // Allocate the maximum number of instances with the maximum number of memories
             for _ in 0..instances.max_instances {
@@ -522,7 +522,8 @@ mod test {
                             PoolingAllocationStrategy::Random,
                             InstanceAllocationRequest {
                                 module: module.clone(),
-                                finished_functions,
+                                image_base: 0,
+                                functions,
                                 imports: Imports {
                                     functions: &[],
                                     tables: &[],

crates/wasmtime/src/instance.rs

@@ -732,7 +732,8 @@ impl<'a> Instantiator<'a> {
                     .allocator()
                     .allocate(InstanceAllocationRequest {
                         module: compiled_module.module().clone(),
-                        finished_functions: compiled_module.finished_functions(),
+                        image_base: compiled_module.code().range().0,
+                        functions: compiled_module.functions(),
                         imports: self.cur.build(),
                         shared_signatures: self.cur.module.signatures().as_module_map().into(),
                         host_state: Box::new(Instance(instance_to_be)),

crates/wasmtime/src/module.rs

@@ -398,7 +398,7 @@ impl Module {
                 .collect();
 
             let mut obj = engine.compiler().object()?;
-            let funcs = engine.compiler().emit_obj(
+            let (funcs, trampolines) = engine.compiler().emit_obj(
                 &translation,
                 &types,
                 funcs,
@@ -412,7 +412,8 @@ impl Module {
                 translation.try_paged_init();
             }
 
-            let (mmap, info) = wasmtime_jit::finish_compile(translation, obj, funcs, tunables)?;
+            let (mmap, info) =
+                wasmtime_jit::finish_compile(translation, obj, funcs, trampolines, tunables)?;
             Ok((mmap, Some(info)))
         })?;
 
@@ -486,7 +487,7 @@ impl Module {
         let signatures = Arc::new(SignatureCollection::new_for_module(
             engine.signatures(),
             &types.wasm_signatures,
-            modules.iter().flat_map(|m| m.trampolines().iter().cloned()),
+            modules.iter().flat_map(|m| m.trampolines()),
         ));
 
         let module = modules.remove(main_module);

crates/wasmtime/src/module/registry.rs

@@ -52,7 +52,7 @@ impl ModuleRegistry {
         // and for schemes like uffd this performs lazy initialization which
         // could use the module in the future. For that reason we continue to
         // register empty modules and retain them.
-        if compiled_module.finished_functions().is_empty() {
+        if compiled_module.finished_functions().len() == 0 {
             self.modules_without_code.push(compiled_module.clone());
             return;
         }
@@ -539,13 +539,19 @@ fn test_frame_info() -> Result<(), anyhow::Error> {
     GlobalModuleRegistry::with(|modules| {
         for (i, alloc) in module.compiled_module().finished_functions() {
             let (start, end) = unsafe {
-                let ptr = (**alloc).as_ptr();
+                let ptr = (*alloc).as_ptr();
-                let len = (**alloc).len();
+                let len = (*alloc).len();
                 (ptr as usize, ptr as usize + len)
             };
             for pc in start..end {
                 let (frame, _, _) = modules.lookup_frame_info(pc).unwrap();
-                assert!(frame.func_index() == i.as_u32());
+                assert!(
+                    frame.func_index() == i.as_u32(),
+                    "lookup of {:#x} returned {}, expected {}",
+                    pc,
+                    frame.func_index(),
+                    i.as_u32()
+                );
             }
         }
     });

crates/wasmtime/src/store.rs

@@ -194,7 +194,7 @@ impl<T> Store<T> {
     /// tables created to 10,000. This can be overridden with the
     /// [`Store::limiter`] configuration method.
     pub fn new(engine: &Engine, data: T) -> Self {
-        let finished_functions = &Default::default();
+        let functions = &Default::default();
         // Wasmtime uses the callee argument to host functions to learn about
         // the original pointer to the `Store` itself, allowing it to
         // reconstruct a `StoreContextMut<T>`. When we initially call a `Func`,
@@ -206,7 +206,8 @@ impl<T> Store<T> {
             OnDemandInstanceAllocator::default()
                 .allocate(InstanceAllocationRequest {
                     host_state: Box::new(()),
-                    finished_functions,
+                    image_base: 0,
+                    functions,
                     shared_signatures: None.into(),
                     imports: Default::default(),
                     module: Arc::new(wasmtime_environ::Module::default()),

crates/wasmtime/src/trampoline.rs

@@ -16,24 +16,22 @@ use crate::{GlobalType, MemoryType, TableType, Val};
 use anyhow::Result;
 use std::any::Any;
 use std::sync::Arc;
-use wasmtime_environ::{
+use wasmtime_environ::{EntityIndex, GlobalIndex, MemoryIndex, Module, TableIndex};
-    DefinedFuncIndex, EntityIndex, GlobalIndex, MemoryIndex, Module, PrimaryMap, TableIndex,
-};
 use wasmtime_runtime::{
     Imports, InstanceAllocationRequest, InstanceAllocator, OnDemandInstanceAllocator,
-    VMFunctionBody, VMFunctionImport, VMSharedSignatureIndex,
+    VMFunctionImport, VMSharedSignatureIndex,
 };
 
 fn create_handle(
     module: Module,
     store: &mut StoreOpaque<'_>,
-    finished_functions: PrimaryMap<DefinedFuncIndex, *mut [VMFunctionBody]>,
     host_state: Box<dyn Any + Send + Sync>,
     func_imports: &[VMFunctionImport],
     shared_signature_id: Option<VMSharedSignatureIndex>,
 ) -> Result<InstanceId> {
     let mut imports = Imports::default();
     imports.functions = func_imports;
+    let functions = &Default::default();
 
     unsafe {
         let config = store.engine().config();
@@ -43,7 +41,8 @@ fn create_handle(
         let handle = OnDemandInstanceAllocator::new(config.mem_creator.clone(), 0).allocate(
             InstanceAllocationRequest {
                 module: Arc::new(module),
-                finished_functions: &finished_functions,
+                functions,
+                image_base: 0,
                 imports,
                 shared_signatures: shared_signature_id.into(),
                 host_state,

crates/wasmtime/src/trampoline/func.rs

@@ -77,22 +77,23 @@ pub fn create_function(
     engine: &Engine,
 ) -> Result<(InstanceHandle, VMTrampoline)> {
     let mut obj = engine.compiler().object()?;
-    engine
+    let (t1, t2) = engine.compiler().emit_trampoline_obj(
-        .compiler()
+        ft.as_wasm_func_type(),
-        .emit_trampoline_obj(ft.as_wasm_func_type(), stub_fn as usize, &mut obj)?;
+        stub_fn as usize,
+        &mut obj,
+    )?;
     let obj = obj.write()?;
 
+    // Copy the results of JIT compilation into executable memory, and this will
+    // also take care of unwind table registration.
     let mut code_memory = CodeMemory::new();
-    let alloc = code_memory.allocate_for_object_unparsed(&obj)?;
+    let (alloc, _obj) = code_memory.allocate_for_object_unparsed(&obj)?;
-    let mut trampolines = alloc.trampolines();
+
-    let (host_i, host_trampoline) = trampolines.next().unwrap();
+    // Extract the host/wasm trampolines from the results of compilation since
-    assert_eq!(host_i.as_u32(), 0);
+    // we know their start/length.
-    let (wasm_i, wasm_trampoline) = trampolines.next().unwrap();
+    let host_trampoline = alloc[t1.start as usize..][..t1.length as usize].as_ptr();
-    assert_eq!(wasm_i.as_u32(), 1);
+    let wasm_trampoline = &mut alloc[t2.start as usize..][..t2.length as usize];
-    assert!(trampolines.next().is_none());
+    let wasm_trampoline = wasm_trampoline as *mut [u8] as *mut [VMFunctionBody];
-    let host_trampoline = host_trampoline.as_ptr();
-    let wasm_trampoline = wasm_trampoline as *mut [_];
-    drop(trampolines);
 
     code_memory.publish();
 
@@ -104,8 +105,7 @@ pub fn create_function(
             sig,
             Box::new(TrampolineState { func, code_memory }),
         )?;
-        let host_trampoline =
+        let host_trampoline = std::mem::transmute::<*const u8, VMTrampoline>(host_trampoline);
-            std::mem::transmute::<*const VMFunctionBody, VMTrampoline>(host_trampoline);
         Ok((instance, host_trampoline))
     }
 }
@@ -116,7 +116,8 @@ pub unsafe fn create_raw_function(
     host_state: Box<dyn Any + Send + Sync>,
 ) -> Result<InstanceHandle> {
     let mut module = Module::new();
-    let mut finished_functions = PrimaryMap::new();
+    let mut functions = PrimaryMap::new();
+    functions.push(Default::default());
 
     let sig_id = SignatureIndex::from_u32(u32::max_value() - 1);
     module.types.push(ModuleType::Function(sig_id));
@@ -124,12 +125,12 @@ pub unsafe fn create_raw_function(
     module
         .exports
         .insert(String::new(), EntityIndex::Function(func_id));
-    finished_functions.push(func);
 
     Ok(
         OnDemandInstanceAllocator::default().allocate(InstanceAllocationRequest {
             module: Arc::new(module),
-            finished_functions: &finished_functions,
+            functions: &functions,
+            image_base: (*func).as_ptr() as usize,
             imports: Imports::default(),
             shared_signatures: sig.into(),
             host_state,

crates/wasmtime/src/trampoline/global.rs

@@ -2,9 +2,7 @@ use crate::store::{InstanceId, StoreOpaque};
 use crate::trampoline::create_handle;
 use crate::{GlobalType, Mutability, Val};
 use anyhow::Result;
-use wasmtime_environ::{
+use wasmtime_environ::{EntityIndex, Global, GlobalInit, Module, ModuleType, SignatureIndex};
-    EntityIndex, Global, GlobalInit, Module, ModuleType, PrimaryMap, SignatureIndex,
-};
 use wasmtime_runtime::VMFunctionImport;
 
 pub fn create_global(store: &mut StoreOpaque<'_>, gt: &GlobalType, val: Val) -> Result<InstanceId> {
@@ -69,7 +67,6 @@ pub fn create_global(store: &mut StoreOpaque<'_>, gt: &GlobalType, val: Val) ->
     let id = create_handle(
         module,
         store,
-        PrimaryMap::new(),
         Box::new(()),
         &func_imports,
         shared_signature_id,

crates/wasmtime/src/trampoline/memory.rs

@@ -5,7 +5,7 @@ use crate::MemoryType;
 use anyhow::{anyhow, Result};
 use std::convert::TryFrom;
 use std::sync::Arc;
-use wasmtime_environ::{EntityIndex, MemoryPlan, MemoryStyle, Module, PrimaryMap, WASM_PAGE_SIZE};
+use wasmtime_environ::{EntityIndex, MemoryPlan, MemoryStyle, Module, WASM_PAGE_SIZE};
 use wasmtime_runtime::{RuntimeLinearMemory, RuntimeMemoryCreator, VMMemoryDefinition};
 
 pub fn create_memory(store: &mut StoreOpaque<'_>, memory: &MemoryType) -> Result<InstanceId> {
@@ -20,7 +20,7 @@ pub fn create_memory(store: &mut StoreOpaque<'_>, memory: &MemoryType) -> Result
         .exports
         .insert(String::new(), EntityIndex::Memory(memory_id));
 
-    create_handle(module, store, PrimaryMap::new(), Box::new(()), &[], None)
+    create_handle(module, store, Box::new(()), &[], None)
 }
 
 struct LinearMemoryProxy {

crates/wasmtime/src/trampoline/table.rs

@@ -2,7 +2,7 @@ use crate::store::{InstanceId, StoreOpaque};
 use crate::trampoline::create_handle;
 use crate::TableType;
 use anyhow::Result;
-use wasmtime_environ::{EntityIndex, Module, PrimaryMap};
+use wasmtime_environ::{EntityIndex, Module};
 
 pub fn create_table(store: &mut StoreOpaque<'_>, table: &TableType) -> Result<InstanceId> {
     let mut module = Module::new();
@@ -16,5 +16,5 @@ pub fn create_table(store: &mut StoreOpaque<'_>, table: &TableType) -> Result<In
         .exports
         .insert(String::new(), EntityIndex::Table(table_id));
 
-    create_handle(module, store, PrimaryMap::new(), Box::new(()), &[], None)
+    create_handle(module, store, Box::new(()), &[], None)
 }