Move native signatures out of Module (#2362)

After compilation there's actually no need to hold onto the native
signature for a wasm function type, so this commit moves out the
`ir::Signature` value from a `Module` into a separate field that's
deallocated when compilation is finished. This simplifies the
`SignatureRegistry` because it only needs to track wasm functino types
and it also means less work is done for `Func::wrap`.

crates/cranelift/src/func_environ.rs

@@ -5,7 +5,7 @@ use cranelift_codegen::ir::immediates::{Offset32, Uimm64};
 use cranelift_codegen::ir::types::*;
 use cranelift_codegen::ir::{AbiParam, ArgumentPurpose, Function, InstBuilder, Signature};
 use cranelift_codegen::isa::{self, TargetFrontendConfig};
-use cranelift_entity::EntityRef;
+use cranelift_entity::{EntityRef, PrimaryMap};
 use cranelift_frontend::FunctionBuilder;
 use cranelift_wasm::{
     self, FuncIndex, GlobalIndex, GlobalVariable, MemoryIndex, SignatureIndex, TableIndex,
@@ -99,6 +99,9 @@ pub struct FuncEnvironment<'module_environment> {
     /// The module-level environment which this function-level environment belongs to.
     module: &'module_environment Module,
 
+    /// The native signatures for each type signature in this module
+    native_signatures: &'module_environment PrimaryMap<SignatureIndex, ir::Signature>,
+
     /// The Cranelift global holding the vmctx address.
     vmctx: Option<ir::GlobalValue>,
 
@@ -115,6 +118,7 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
     pub fn new(
         target_config: TargetFrontendConfig,
         module: &'module_environment Module,
+        native_signatures: &'module_environment PrimaryMap<SignatureIndex, ir::Signature>,
         tunables: &'module_environment Tunables,
     ) -> Self {
         let builtin_function_signatures = BuiltinFunctionSignatures::new(
@@ -129,6 +133,7 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
         Self {
             target_config,
             module,
+            native_signatures,
             vmctx: None,
             builtin_function_signatures,
             offsets: VMOffsets::new(target_config.pointer_bytes(), module),
@@ -993,16 +998,17 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
         func: &mut ir::Function,
         index: SignatureIndex,
     ) -> WasmResult<ir::SigRef> {
-        Ok(func.import_signature(self.module.signatures[index].1.clone()))
+        Ok(func.import_signature(self.native_signatures[index].clone()))
     }
 
     fn make_direct_func(
         &mut self,
         func: &mut ir::Function,
         index: FuncIndex,
     ) -> WasmResult<ir::FuncRef> {
-        let sig = self.module.native_func_signature(index);
-        let signature = func.import_signature(sig.clone());
+        let sig_index = self.module.functions[index];
+        let sig = self.native_signatures[sig_index].clone();
+        let signature = func.import_signature(sig);
         let name = get_func_name(index);
         Ok(func.import_function(ir::ExtFuncData {
             name,

crates/cranelift/src/lib.rs

@@ -362,12 +362,18 @@ impl Compiler for Cranelift {
         let func_index = module.func_index(func_index);
         let mut context = Context::new();
         context.func.name = get_func_name(func_index);
-        context.func.signature = module.native_func_signature(func_index).clone();
+        let sig_index = module.functions[func_index];
+        context.func.signature = translation.native_signatures[sig_index].clone();
         if tunables.debug_info {
             context.func.collect_debug_info();
         }
 
-        let mut func_env = FuncEnvironment::new(isa.frontend_config(), module, tunables);
+        let mut func_env = FuncEnvironment::new(
+            isa.frontend_config(),
+            module,
+            &translation.native_signatures,
+            tunables,
+        );
 
         // We use these as constant offsets below in
         // `stack_limit_from_arguments`, so assert their values here. This

crates/environ/src/module.rs

@@ -2,7 +2,6 @@
 
 use crate::tunables::Tunables;
 use crate::WASM_MAX_PAGES;
-use cranelift_codegen::ir;
 use cranelift_entity::{EntityRef, PrimaryMap};
 use cranelift_wasm::{
     DataIndex, DefinedFuncIndex, DefinedGlobalIndex, DefinedMemoryIndex, DefinedTableIndex,
@@ -169,7 +168,7 @@ pub struct Module {
     pub func_names: HashMap<FuncIndex, String>,
 
     /// Unprocessed signatures exactly as provided by `declare_signature()`.
-    pub signatures: PrimaryMap<SignatureIndex, (WasmFuncType, ir::Signature)>,
+    pub signatures: PrimaryMap<SignatureIndex, WasmFuncType>,
 
     /// Number of imported functions in the module.
     pub num_imported_funcs: usize,
@@ -319,16 +318,10 @@ impl Module {
         index.index() < self.num_imported_globals
     }
 
-    /// Convenience method for looking up the native signature of a compiled
-    /// Wasm function.
-    pub fn native_func_signature(&self, func_index: FuncIndex) -> &ir::Signature {
-        &self.signatures[self.functions[func_index]].1
-    }
-
     /// Convenience method for looking up the original Wasm signature of a
     /// function.
     pub fn wasm_func_type(&self, func_index: FuncIndex) -> &WasmFuncType {
-        &self.signatures[self.functions[func_index]].0
+        &self.signatures[self.functions[func_index]]
     }
 }
 

crates/environ/src/module_environ.rs

@@ -35,6 +35,9 @@ pub struct ModuleTranslation<'data> {
     /// Module information.
     pub module: Module,
 
+    /// Map of native signatures
+    pub native_signatures: PrimaryMap<SignatureIndex, ir::Signature>,
+
     /// References to the function bodies.
     pub function_body_inputs: PrimaryMap<DefinedFuncIndex, FunctionBodyData<'data>>,
 
@@ -111,6 +114,7 @@ impl<'data> ModuleEnvironment<'data> {
             result: ModuleTranslation {
                 target_config,
                 module: Module::new(),
+                native_signatures: PrimaryMap::new(),
                 function_body_inputs: PrimaryMap::new(),
                 data_initializers: Vec::new(),
                 tunables: tunables.clone(),
@@ -198,17 +202,17 @@ impl<'data> TargetEnvironment for ModuleEnvironment<'data> {
 /// environment-dependent wasm instructions. These functions should not be called by the user.
 impl<'data> cranelift_wasm::ModuleEnvironment<'data> for ModuleEnvironment<'data> {
     fn reserve_signatures(&mut self, num: u32) -> WasmResult<()> {
-        self.result
-            .module
-            .signatures
-            .reserve_exact(usize::try_from(num).unwrap());
+        let num = usize::try_from(num).unwrap();
+        self.result.module.signatures.reserve_exact(num);
+        self.result.native_signatures.reserve_exact(num);
         Ok(())
     }
 
     fn declare_signature(&mut self, wasm: WasmFuncType, sig: ir::Signature) -> WasmResult<()> {
         let sig = translate_signature(sig, self.pointer_type());
         // TODO: Deduplicate signatures.
-        self.result.module.signatures.push((wasm, sig));
+        self.result.module.signatures.push(wasm);
+        self.result.native_signatures.push(sig);
         Ok(())
     }
 
@@ -461,7 +465,7 @@ impl<'data> cranelift_wasm::ModuleEnvironment<'data> for ModuleEnvironment<'data
             }
             info.wasm_file.funcs.push(FunctionMetadata {
                 locals: locals.into_boxed_slice(),
-                params: sig.0.params.iter().cloned().map(|i| i.into()).collect(),
+                params: sig.params.iter().cloned().map(|i| i.into()).collect(),
             });
         }
         self.result

crates/jit/src/compiler.rs

@@ -158,8 +158,7 @@ impl Compiler {
             vec![]
         };
 
-        let (obj, unwind_info) =
-            build_object(&*self.isa, &translation.module, &funcs, dwarf_sections)?;
+        let (obj, unwind_info) = build_object(&*self.isa, &translation, &funcs, dwarf_sections)?;
 
         Ok(Compilation {
             obj,

crates/jit/src/object.rs

@@ -8,7 +8,7 @@ use wasmtime_debug::DwarfSection;
 use wasmtime_environ::entity::PrimaryMap;
 use wasmtime_environ::isa::{unwind::UnwindInfo, TargetIsa};
 use wasmtime_environ::wasm::{FuncIndex, SignatureIndex};
-use wasmtime_environ::{CompiledFunctions, Module};
+use wasmtime_environ::{CompiledFunctions, ModuleTranslation};
 use wasmtime_obj::{ObjectBuilder, ObjectBuilderTarget};
 
 pub use wasmtime_obj::utils;
@@ -23,7 +23,7 @@ pub enum ObjectUnwindInfo {
 // Builds ELF image from the module `Compilation`.
 pub(crate) fn build_object(
     isa: &dyn TargetIsa,
-    module: &Module,
+    translation: &ModuleTranslation,
     funcs: &CompiledFunctions,
     dwarf_sections: Vec<DwarfSection>,
 ) -> Result<(Object, Vec<ObjectUnwindInfo>), anyhow::Error> {
@@ -39,13 +39,13 @@ pub(crate) fn build_object(
     unwind_info.extend(funcs.iter().filter_map(|(index, func)| {
         func.unwind_info
             .as_ref()
-            .map(|info| ObjectUnwindInfo::Func(module.func_index(index), info.clone()))
+            .map(|info| ObjectUnwindInfo::Func(translation.module.func_index(index), info.clone()))
     }));
 
-    let mut trampolines = PrimaryMap::with_capacity(module.signatures.len());
+    let mut trampolines = PrimaryMap::with_capacity(translation.module.signatures.len());
     let mut cx = FunctionBuilderContext::new();
     // Build trampolines for every signature.
-    for (i, (_, native_sig)) in module.signatures.iter() {
+    for (i, native_sig) in translation.native_signatures.iter() {
         let func = build_trampoline(isa, &mut cx, native_sig, std::mem::size_of::<u128>())?;
         // Preserve trampoline function unwind info.
         if let Some(info) = &func.unwind_info {
@@ -55,7 +55,7 @@ pub(crate) fn build_object(
     }
 
     let target = ObjectBuilderTarget::new(isa.triple().architecture)?;
-    let mut builder = ObjectBuilder::new(target, module, funcs);
+    let mut builder = ObjectBuilder::new(target, &translation.module, funcs);
     builder
         .set_code_alignment(CODE_SECTION_ALIGNMENT)
         .set_trampolines(trampolines)

crates/lightbeam/wasmtime/src/lib.rs

@@ -12,8 +12,9 @@ use wasmtime_environ::wasm::{
     GlobalIndex, MemoryIndex, SignatureIndex, TableIndex,
 };
 use wasmtime_environ::{
-    BuiltinFunctionIndex, CompileError, CompiledFunction, Compiler, FunctionBodyData, Module,
-    ModuleTranslation, Relocation, RelocationTarget, TrapInformation, VMOffsets,
+    entity::PrimaryMap, BuiltinFunctionIndex, CompileError, CompiledFunction, Compiler,
+    FunctionBodyData, Module, ModuleTranslation, Relocation, RelocationTarget, TrapInformation,
+    VMOffsets,
 };
 
 /// A compiler that compiles a WebAssembly module with Lightbeam, directly translating the Wasm file.
@@ -32,7 +33,7 @@ impl Compiler for Lightbeam {
         }
         let func_index = translation.module.func_index(i);
 
-        let env = FuncEnvironment::new(isa.frontend_config().pointer_bytes(), &translation.module);
+        let env = FuncEnvironment::new(isa.frontend_config().pointer_bytes(), translation);
         let mut codegen_session: CodeGenSession<_> = CodeGenSession::new(
             translation.function_body_inputs.len() as u32,
             &env,
@@ -180,15 +181,18 @@ struct FuncEnvironment<'module_environment> {
     /// The module-level environment which this function-level environment belongs to.
     module: &'module_environment Module,
 
+    native_signatures: &'module_environment PrimaryMap<SignatureIndex, ir::Signature>,
+
     /// Offsets to struct fields accessed by JIT code.
     offsets: VMOffsets,
 }
 
 impl<'module_environment> FuncEnvironment<'module_environment> {
-    fn new(pointer_bytes: u8, module: &'module_environment Module) -> Self {
+    fn new(pointer_bytes: u8, translation: &'module_environment ModuleTranslation<'_>) -> Self {
         Self {
-            module,
-            offsets: VMOffsets::new(pointer_bytes, module),
+            module: &translation.module,
+            offsets: VMOffsets::new(pointer_bytes, &translation.module),
+            native_signatures: &translation.native_signatures,
         }
     }
 }
@@ -227,7 +231,7 @@ impl lightbeam::ModuleContext for FuncEnvironment<'_> {
     }
 
     fn signature(&self, index: u32) -> &Self::Signature {
-        &self.module.signatures[SignatureIndex::from_u32(index)].1
+        &self.native_signatures[SignatureIndex::from_u32(index)]
     }
 
     fn defined_table_index(&self, table_index: u32) -> Option<u32> {

crates/wasmtime/src/func.rs

@@ -541,7 +541,7 @@ impl Func {
         // Signatures should always be registered in the store's registry of
         // shared signatures, so we should be able to unwrap safely here.
         let signatures = self.instance.store.signatures().borrow();
-        let (wft, _, _) = signatures
+        let (wft, _) = signatures
             .lookup_shared(self.sig_index())
             .expect("signature should be registered");
 
@@ -554,7 +554,7 @@ impl Func {
     /// Returns the number of parameters that this function takes.
     pub fn param_arity(&self) -> usize {
         let signatures = self.instance.store.signatures().borrow();
-        let (sig, _, _) = signatures
+        let (sig, _) = signatures
             .lookup_shared(self.sig_index())
             .expect("signature should be registered");
         sig.params.len()
@@ -563,7 +563,7 @@ impl Func {
     /// Returns the number of results this function produces.
     pub fn result_arity(&self) -> usize {
         let signatures = self.instance.store.signatures().borrow();
-        let (sig, _, _) = signatures
+        let (sig, _) = signatures
             .lookup_shared(self.sig_index())
             .expect("signature should be registered");
         sig.returns.len()
@@ -657,7 +657,7 @@ impl Func {
             .borrow()
             .lookup_shared(unsafe { export.anyfunc.as_ref().type_index })
             .expect("failed to retrieve trampoline from module")
-            .2;
+            .1;
 
         Func {
             instance,

crates/wasmtime/src/instance.rs

@@ -293,7 +293,7 @@ fn with_imports<R>(
                 let ty = store
                     .signatures()
                     .borrow()
-                    .lookup(&m.signatures[m.functions[i]].0)
+                    .lookup(&m.signatures[m.functions[i]])
                     .ok_or_else(|| anyhow!("function types incompatible"))?;
                 if !func.matches_expected(ty) {
                     bail!("function types incompatible");

crates/wasmtime/src/runtime.rs

@@ -915,10 +915,9 @@ impl Store {
         module: &'a wasmtime_environ::Module,
     ) -> impl Fn(wasm::SignatureIndex) -> VMSharedSignatureIndex + 'a {
         move |index| {
-            let (wasm, _native) = &module.signatures[index];
             self.signatures()
                 .borrow()
-                .lookup(wasm)
+                .lookup(&module.signatures[index])
                 .expect("signature not previously registered")
         }
     }
@@ -993,8 +992,8 @@ impl Store {
         let trampolines = module.compiled_module().trampolines();
         let module = module.compiled_module().module();
         let mut signatures = self.signatures().borrow_mut();
-        for (index, (wasm, native)) in module.signatures.iter() {
-            signatures.register(wasm, native, trampolines[index]);
+        for (index, wasm) in module.signatures.iter() {
+            signatures.register(wasm, trampolines[index]);
         }
     }
 

crates/wasmtime/src/sig_registry.rs

@@ -3,7 +3,7 @@
 
 use std::collections::{hash_map, HashMap};
 use std::convert::TryFrom;
-use wasmtime_environ::{ir, wasm::WasmFuncType};
+use wasmtime_environ::wasm::WasmFuncType;
 use wasmtime_runtime::{VMSharedSignatureIndex, VMTrampoline};
 
 /// WebAssembly requires that the caller and callee signatures in an indirect
@@ -25,8 +25,6 @@ pub struct SignatureRegistry {
 struct Entry {
     // The WebAssembly type signature, using wasm types.
     wasm: WasmFuncType,
-    // The native signature we're using for this wasm type signature.
-    native: ir::Signature,
     // The native trampoline used to invoke this type signature from `Func`.
     // Note that the code memory for this trampoline is not owned by this
     // type, but instead it's expected to be owned by the store that this
@@ -39,7 +37,6 @@ impl SignatureRegistry {
     pub fn register(
         &mut self,
         wasm: &WasmFuncType,
-        native: &ir::Signature,
         trampoline: VMTrampoline,
     ) -> VMSharedSignatureIndex {
         let len = self.wasm2index.len();
@@ -57,7 +54,6 @@ impl SignatureRegistry {
                 let index = VMSharedSignatureIndex::new(u32::try_from(len).unwrap());
                 self.index_map.push(Entry {
                     wasm: wasm.clone(),
-                    native: native.clone(),
                     trampoline,
                 });
                 entry.insert(index);
@@ -78,9 +74,9 @@ impl SignatureRegistry {
     pub fn lookup_shared(
         &self,
         idx: VMSharedSignatureIndex,
-    ) -> Option<(&WasmFuncType, &ir::Signature, VMTrampoline)> {
+    ) -> Option<(&WasmFuncType, VMTrampoline)> {
         self.index_map
             .get(idx.bits() as usize)
-            .map(|e| (&e.wasm, &e.native, e.trampoline))
+            .map(|e| (&e.wasm, e.trampoline))
     }
 }