diff --git a/crates/jit/src/instantiate.rs b/crates/jit/src/instantiate.rs index df6a17fa24..df2cd77c31 100644 --- a/crates/jit/src/instantiate.rs +++ b/crates/jit/src/instantiate.rs @@ -205,7 +205,7 @@ pub struct CompiledModule { artifacts: CompilationArtifacts, code: Arc, finished_functions: FinishedFunctions, - trampolines: PrimaryMap, + trampolines: Vec<(SignatureIndex, VMTrampoline)>, } impl CompiledModule { @@ -292,7 +292,7 @@ impl CompiledModule { } /// Returns the per-signature trampolines for this module. - pub fn trampolines(&self) -> &PrimaryMap { + pub fn trampolines(&self) -> &[(SignatureIndex, VMTrampoline)] { &self.trampolines } @@ -444,7 +444,7 @@ fn build_code_memory( CodeMemory, (*const u8, usize), PrimaryMap, - PrimaryMap, + Vec<(SignatureIndex, VMTrampoline)>, ), String, > { @@ -464,12 +464,15 @@ fn build_code_memory( ); } - let mut trampolines = PrimaryMap::new(); - for (i, fat_ptr) in allocation.trampolines() { - let fat_ptr = - unsafe { std::mem::transmute::<*const VMFunctionBody, VMTrampoline>(fat_ptr.as_ptr()) }; - assert_eq!(trampolines.push(fat_ptr), i); - } + let trampolines = allocation + .trampolines() + .map(|(i, fat_ptr)| { + let fnptr = unsafe { + std::mem::transmute::<*const VMFunctionBody, VMTrampoline>(fat_ptr.as_ptr()) + }; + (i, fnptr) + }) + .collect(); let code_range = allocation.code_range(); diff --git a/crates/jit/src/object.rs b/crates/jit/src/object.rs index 24b431e597..3516be2b13 100644 --- a/crates/jit/src/object.rs +++ b/crates/jit/src/object.rs @@ -4,11 +4,11 @@ use super::trampoline::build_trampoline; use cranelift_frontend::FunctionBuilderContext; use object::write::Object; use serde::{Deserialize, Serialize}; +use std::collections::BTreeSet; 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, ModuleTranslation, TypeTables}; +use wasmtime_environ::{CompiledFunctions, ModuleTranslation, ModuleType, TypeTables}; use wasmtime_obj::{ObjectBuilder, ObjectBuilderTarget}; pub use wasmtime_obj::utils; @@ -39,22 +39,26 @@ pub(crate) fn build_object( .map(|info| ObjectUnwindInfo::Func(translation.module.func_index(index), info.clone())) })); - let mut trampolines = PrimaryMap::with_capacity(types.native_signatures.len()); + // Build trampolines for every signature that can be used by this module. + let signatures = translation + .module + .types + .values() + .filter_map(|t| match t { + ModuleType::Function(f) => Some(*f), + _ => None, + }) + .collect::>(); + let mut trampolines = Vec::with_capacity(signatures.len()); let mut cx = FunctionBuilderContext::new(); - // Build trampolines for every signature. - // - // TODO: for the module linking proposal this builds too many native - // signatures. This builds trampolines for all signatures for all modules - // for each module. That's a lot of trampolines! We should instead figure - // out a way to share trampolines amongst all modules when compiling - // module-linking modules. - for (i, native_sig) in types.native_signatures.iter() { + for i in signatures { + let native_sig = &types.native_signatures[i]; let func = build_trampoline(isa, &mut cx, native_sig, std::mem::size_of::())?; // Preserve trampoline function unwind info. if let Some(info) = &func.unwind_info { unwind_info.push(ObjectUnwindInfo::Trampoline(i, info.clone())) } - trampolines.push(func); + trampolines.push((i, func)); } let target = ObjectBuilderTarget::new(isa.triple().architecture)?; diff --git a/crates/obj/src/builder.rs b/crates/obj/src/builder.rs index 046ae5f784..805adbb806 100644 --- a/crates/obj/src/builder.rs +++ b/crates/obj/src/builder.rs @@ -257,7 +257,7 @@ pub struct ObjectBuilder<'a> { module: &'a Module, code_alignment: u64, compilation: &'a CompiledFunctions, - trampolines: PrimaryMap, + trampolines: Vec<(SignatureIndex, CompiledFunction)>, dwarf_sections: Vec, } @@ -271,7 +271,7 @@ impl<'a> ObjectBuilder<'a> { target, module, code_alignment: 1, - trampolines: PrimaryMap::new(), + trampolines: Vec::new(), dwarf_sections: vec![], compilation, } @@ -284,7 +284,7 @@ impl<'a> ObjectBuilder<'a> { pub fn set_trampolines( &mut self, - trampolines: PrimaryMap, + trampolines: Vec<(SignatureIndex, CompiledFunction)>, ) -> &mut Self { self.trampolines = trampolines; self @@ -359,7 +359,7 @@ impl<'a> ObjectBuilder<'a> { } let mut trampolines = Vec::new(); for (i, func) in self.trampolines.iter() { - let name = utils::trampoline_symbol_name(i).as_bytes().to_vec(); + let name = utils::trampoline_symbol_name(*i).as_bytes().to_vec(); trampolines.push(append_func(name, func)); } @@ -399,7 +399,7 @@ impl<'a> ObjectBuilder<'a> { } } - for (func, symbol) in self.trampolines.values().zip(trampolines) { + for ((_, func), symbol) in self.trampolines.iter().zip(trampolines) { let (_, off) = obj.symbol_section_and_offset(symbol).unwrap(); for r in to_object_relocations( func.relocations.iter(), diff --git a/crates/wasmtime/src/store.rs b/crates/wasmtime/src/store.rs index f36a8897cb..f01c86662f 100644 --- a/crates/wasmtime/src/store.rs +++ b/crates/wasmtime/src/store.rs @@ -320,10 +320,11 @@ impl Store { } fn register_signatures(&self, module: &Module) { - let trampolines = module.compiled_module().trampolines(); let mut signatures = self.signatures().borrow_mut(); - for (index, wasm) in module.types().wasm_signatures.iter() { - signatures.register(wasm, trampolines[index]); + let types = module.types(); + for (index, trampoline) in module.compiled_module().trampolines() { + let wasm = &types.wasm_signatures[*index]; + signatures.register(wasm, *trampoline); } }