Fix signatures registered with modules-in-components (#6694)

* Fix signatures registered with modules-in-components

This commit fixes a minor issue in
`FunctionIndices::link_and_append_code` which previously ended up only
filling out the `wasm_to_native_trampolines` field for the first module
rather than all the modules. Additionally the first module might have
too many entries that encompass all modules instead of just its own
entries. The fix in this commit is to refactor this logic to ensure that
the necessary maps are present for all translations.

While technically a bug that can be surfaced through the embedder API
it's pretty obscure. The given test here panics beforehand but succeeds
afterwards, but this is moreso prep for some future resource-related
work where this map will need persisting into the component metadata
side of things.

* Fix build warning

crates/wasmtime/src/compiler.rs

@@ -24,12 +24,11 @@
 
 use crate::Engine;
 use anyhow::Result;
-use std::collections::{btree_map, BTreeMap};
+use std::collections::{btree_map, BTreeMap, BTreeSet};
 use std::{any::Any, collections::HashMap};
 use wasmtime_environ::{
- Compiler, DefinedFuncIndex, FuncIndex, FunctionBodyData, FunctionLoc, ModuleTranslation,
- ModuleType, ModuleTypes, PrimaryMap, SignatureIndex, StaticModuleIndex, Tunables,
- WasmFunctionInfo,
+ Compiler, DefinedFuncIndex, FuncIndex, FunctionBodyData, ModuleTranslation, ModuleType,
+ ModuleTypes, PrimaryMap, SignatureIndex, StaticModuleIndex, Tunables, WasmFunctionInfo,
 };
 use wasmtime_jit::{CompiledFunctionInfo, CompiledModuleInfo};
 
@@ -586,6 +585,13 @@ impl FunctionIndices {
  .remove(&CompileKey::NATIVE_TO_WASM_TRAMPOLINE_KIND)
  .unwrap_or_default();
 
+ // NB: unlike the above maps this is not emptied out during iteration
+ // since each module may reach into different portions of this map.
+ let wasm_to_native_trampolines = self
+ .indices
+ .remove(&CompileKey::WASM_TO_NATIVE_TRAMPOLINE_KIND)
+ .unwrap_or_default();
+
  artifacts.modules = translations
  .into_iter()
  .map(|(module, translation)| {
@@ -619,16 +625,20 @@ impl FunctionIndices {
  })
  .collect();
 
- let wasm_to_native_trampolines: Vec<(SignatureIndex, FunctionLoc)> = self
- .indices
- .remove(&CompileKey::WASM_TO_NATIVE_TRAMPOLINE_KIND)
- .into_iter()
- .flat_map(|x| x)
- .map(|(key, i)| {
- (
- SignatureIndex::from_u32(key.index),
- symbol_ids_and_locs[i.unwrap_function()].1,
- )
+ let unique_and_sorted_sigs = translation
+ .module
+ .types
+ .iter()
+ .map(|(_, ty)| match ty {
+ ModuleType::Function(ty) => *ty,
+ })
+ .collect::<BTreeSet<_>>();
+ let wasm_to_native_trampolines = unique_and_sorted_sigs
+ .iter()
+ .map(|idx| {
+ let key = CompileKey::wasm_to_native_trampoline(*idx);
+ let compiled = wasm_to_native_trampolines[&key];
+ (*idx, symbol_ids_and_locs[compiled.unwrap_function()].1)
  })
  .collect();
 
@@ -678,17 +688,17 @@ pub struct Artifacts {
  #[cfg(feature = "component-model")]
  pub lowerings: PrimaryMap<
  wasmtime_environ::component::LoweredIndex,
- wasmtime_environ::component::AllCallFunc<FunctionLoc>,
+ wasmtime_environ::component::AllCallFunc<wasmtime_environ::FunctionLoc>,
  >,
  #[cfg(feature = "component-model")]
  pub always_traps: PrimaryMap<
  wasmtime_environ::component::RuntimeAlwaysTrapIndex,
- wasmtime_environ::component::AllCallFunc<FunctionLoc>,
+ wasmtime_environ::component::AllCallFunc<wasmtime_environ::FunctionLoc>,
  >,
  #[cfg(feature = "component-model")]
  pub transcoders: PrimaryMap<
  wasmtime_environ::component::RuntimeTranscoderIndex,
- wasmtime_environ::component::AllCallFunc<FunctionLoc>,
+ wasmtime_environ::component::AllCallFunc<wasmtime_environ::FunctionLoc>,
  >,
 }
 

tests/all/component_model/aot.rs

@@ -98,3 +98,24 @@ fn cannot_serialize_exported_module() -> Result<()> {
  assert!(module.serialize().is_err());
  Ok(())
 }
+
+#[test]
+fn usable_exported_modules() -> Result<()> {
+ let engine = super::engine();
+ let component = Component::new(
+ &engine,
+ r#"(component
+ (core module $m)
+ (core module $m1 (export "a")
+ (import "" "" (func (param i32)))
+ )
+ )"#,
+ )?;
+ let mut store = Store::new(&engine, ());
+ let instance = Linker::new(&engine).instantiate(&mut store, &component)?;
+ let module = instance.get_module(&mut store, "a").unwrap();
+ let mut core_linker = wasmtime::Linker::new(&engine);
+ core_linker.func_wrap("", "", |_: u32| {})?;
+ core_linker.instantiate(&mut store, &module)?;
+ Ok(())
+}