Lazily populate a store's trampoline map

crates/wasmtime/src/func.rs

@@ -718,7 +718,7 @@ impl Func {
     }
 
     pub(crate) fn sig_index(&self, data: &StoreData) -> VMSharedSignatureIndex {
-        unsafe { data[self.0].export().anyfunc.as_ref().type_index }
+        data[self.0].sig_index()
     }
 
     /// Invokes this function with the `params` given and writes returned values
@@ -2083,18 +2083,18 @@ impl HostFunc {
     /// Can only be inserted into stores with a matching `T` relative to when
     /// this `HostFunc` was first created.
     pub unsafe fn to_func(self: &Arc<Self>, store: &mut StoreOpaque) -> Func {
-        self.register_trampoline(store);
+        self.validate_store(store);
         let me = self.clone();
         Func::from_func_kind(FuncKind::SharedHost(me), store)
     }
 
     /// Same as [`HostFunc::to_func`], different ownership.
     unsafe fn into_func(self, store: &mut StoreOpaque) -> Func {
-        self.register_trampoline(store);
+        self.validate_store(store);
         Func::from_func_kind(FuncKind::Host(self), store)
     }
 
-    unsafe fn register_trampoline(&self, store: &mut StoreOpaque) {
+    fn validate_store(&self, store: &mut StoreOpaque) {
         // This assert is required to ensure that we can indeed safely insert
         // `self` into the `store` provided, otherwise the type information we
         // have listed won't be correct. This is possible to hit with the public
@@ -2103,8 +2103,6 @@ impl HostFunc {
             Engine::same(&self.engine, store.engine()),
             "cannot use a store with a different engine than a linker was created with",
         );
-        let idx = self.export.anyfunc.as_ref().type_index;
-        store.register_host_trampoline(idx, self.trampoline);
     }
 
     pub(crate) fn sig_index(&self) -> VMSharedSignatureIndex {
@@ -2128,7 +2126,7 @@ impl Drop for HostFunc {
 
 impl FuncData {
     #[inline]
-    fn trampoline(&self) -> VMTrampoline {
+    pub(crate) fn trampoline(&self) -> VMTrampoline {
         match &self.kind {
             FuncKind::StoreOwned { trampoline, .. } => *trampoline,
             FuncKind::SharedHost(host) => host.trampoline,
@@ -2140,6 +2138,10 @@ impl FuncData {
     fn export(&self) -> &ExportFunction {
         self.kind.export()
     }
+
+    pub(crate) fn sig_index(&self) -> VMSharedSignatureIndex {
+        unsafe { self.export().anyfunc.as_ref().type_index }
+    }
 }
 
 impl FuncKind {

crates/wasmtime/src/store.rs

@@ -259,7 +259,12 @@ pub struct StoreOpaque {
     signal_handler: Option<Box<SignalHandler<'static>>>,
     externref_activations_table: VMExternRefActivationsTable,
     modules: ModuleRegistry,
+
+    // See documentation on `StoreOpaque::lookup_trampoline` for what these
+    // fields are doing.
     host_trampolines: HashMap<VMSharedSignatureIndex, VMTrampoline>,
+    host_func_trampolines_registered: usize,
+
     // Numbers of resources instantiated in this store, and their limits
     instance_count: usize,
     instance_limit: usize,
@@ -436,6 +441,7 @@ impl<T> Store<T> {
                 externref_activations_table: VMExternRefActivationsTable::new(),
                 modules: ModuleRegistry::default(),
                 host_trampolines: HashMap::default(),
+                host_func_trampolines_registered: 0,
                 instance_count: 0,
                 instance_limit: crate::DEFAULT_INSTANCE_LIMIT,
                 memory_count: 0,
@@ -1106,14 +1112,6 @@ impl StoreOpaque {
         &mut self.store_data
     }
 
-    pub fn register_host_trampoline(
-        &mut self,
-        idx: VMSharedSignatureIndex,
-        trampoline: VMTrampoline,
-    ) {
-        self.host_trampolines.insert(idx, trampoline);
-    }
-
     pub fn interrupt_handle(&self) -> Result<InterruptHandle> {
         if self.engine.config().tunables.interruptable {
             Ok(InterruptHandle {
@@ -1166,17 +1164,87 @@ impl StoreOpaque {
         unsafe { wasmtime_runtime::gc(&self.modules, &mut self.externref_activations_table) }
     }
 
-    pub fn lookup_trampoline(&self, anyfunc: &VMCallerCheckedAnyfunc) -> VMTrampoline {
-        // Look up the trampoline with the store's trampolines (from `Func`).
+    /// Looks up the corresponding `VMTrampoline` which can be used to enter
+    /// wasm  given an anyfunc function pointer.
+    ///
+    /// This is a somewhat complicated implementation at this time unfortnately.
+    /// Trampolines are a sort of side-channel of information which is
+    /// specifically juggled by the `wasmtime` crate in a careful fashion. The
+    /// sources for trampolines are:
+    ///
+    /// * Compiled modules - each compiled module has a trampoline for all
+    ///   signatures of functions that escape the module (e.g. exports)
+    /// * `Func::new` - host-defined functions with a dynamic signature get an
+    ///   on-the-fly-compiled trampoline (e.g. JIT-compiled as part of the
+    ///   `Func::new` call).
+    /// * `Func::wrap` - host-defined functions where the trampoline is
+    ///   monomorphized in Rust and compiled by LLVM.
+    ///
+    /// The purpose of this function is that given some wasm function pointer we
+    /// need to find the trampoline for it. For compiled wasm modules this is
+    /// pretty easy, the code pointer of the function pointer will point us
+    /// at a wasm module which has a table of trampolines-by-type that we can
+    /// lookup.
+    ///
+    /// If this lookup fails, however, then we're trying to get the trampoline
+    /// for a wasm function pointer defined by the host. The trampoline isn't
+    /// actually stored in the wasm function pointer itself so we need
+    /// side-channels of information. To achieve this a lazy scheme is
+    /// implemented here based on the assumption that most trampoline lookups
+    /// happen for wasm-defined functions, not host-defined functions.
+    ///
+    /// The `Store` already has a list of all functions in
+    /// `self.store_data().funcs`, it's just not indexed in a nice fashion by
+    /// type index or similar. To solve this there's an internal map in each
+    /// store, `host_trampolines`, which maps from a type index to the
+    /// store-owned trampoline. The actual population of this map, however, is
+    /// deferred to this function itself.
+    ///
+    /// In the theory that wasm functions are the primary lookups then we don't
+    /// want to make insertion of a host function into the store more expensive
+    /// than it otherwise has to be. We could update the `host_trampolines`
+    /// whenever a host function is inserted into the store, but this is a
+    /// relatively expensive hash map insertion. Instead the work is deferred
+    /// here.
+    ///
+    /// This all means that if the lookup of the trampoline fails within
+    /// `self.host_trampolines` we lazily populate `self.host_trampolines` by
+    /// iterating over `self.store_data().funcs`, inserting trampolines as we
+    /// go. If we find the right trampoline then it's returned.
+    pub fn lookup_trampoline(&mut self, anyfunc: &VMCallerCheckedAnyfunc) -> VMTrampoline {
+        // First try to see if the `anyfunc` belongs to any module. Each module
+        // has its own map of trampolines-per-type-index and the code pointer in
+        // the `anyfunc` will enable us to quickly find a module.
+        if let Some(trampoline) = self.modules.lookup_trampoline(anyfunc) {
+            return trampoline;
+        }
+
+        // Next cosult the list of store-local host trampolines. This is
+        // primarily populated by functions created by `Func::new` or similar
+        // creation functions, host-defined functions.
         if let Some(trampoline) = self.host_trampolines.get(&anyfunc.type_index) {
             return *trampoline;
         }
 
-        // Look up the trampoline with the registered modules
-        if let Some(trampoline) = self.modules.lookup_trampoline(anyfunc) {
-            return trampoline;
+        // If no trampoline was found then it means that it hasn't been loaded
+        // into the `host_trampolines`. Skip over all the ones we've looked at
+        // so far and start inserting into `self.host_trampolines`, returning
+        // the actual trampoline once found.
+        for f in self
+            .store_data
+            .funcs()
+            .skip(self.host_func_trampolines_registered)
+        {
+            self.host_func_trampolines_registered += 1;
+            self.host_trampolines.insert(f.sig_index(), f.trampoline());
+            if f.sig_index() == anyfunc.type_index {
+                return f.trampoline();
+            }
         }
 
+        // If reached this is a bug in Wasmtime. Lookup of a trampoline should
+        // only happen for wasm functions or host functions, all of which should
+        // be indexed by the above.
         panic!("trampoline missing")
     }
 

crates/wasmtime/src/store/data.rs

@@ -99,6 +99,10 @@ impl StoreData {
         }
         true
     }
+
+    pub(crate) fn funcs(&self) -> impl Iterator<Item = &crate::func::FuncData> {
+        self.funcs.iter()
+    }
 }
 
 impl<T> Index<Stored<T>> for StoreData