Improve host call performance (#1506)

* try to drop T:'static from wasmi::Store

* turn RestorePrunedWrapper into v-table

This resolves T:'static issues with call_host_func but unfortunately not with store_inner_and_resource_limiter getter.

* add Copy for RestorePrunedWrapper

* use &self in RestorePrunedWrapper methods

* add more explicit lifetimes to PrunedStore::store_inner_and_resource_limiter_ref

* remove commented out code

* use infered lifetimes in store_inner_and_resource_limiter_ref

* remove unnecesary for<'a>

* add some #[inline] annotations

* remove non-static typeid impl

Author: Robbepop

crates/wasmi/src/store/mod.rs

@@ -6,7 +6,7 @@ use self::pruned::RestorePrunedWrapper;
 pub use self::{
     context::{AsContext, AsContextMut, StoreContext, StoreContextMut},
     inner::{StoreInner, Stored},
-    pruned::{PrunedStore, TypedStore},
+    pruned::PrunedStore,
 };
 use crate::{
     collections::arena::Arena,

crates/wasmi/src/store/pruned.rs

@@ -1,6 +1,5 @@
 use super::{CallHooks, FuncInOut, HostFuncEntity, ResourceLimiterRef, StoreInner};
 use crate::{core::hint, CallHook, Error, Instance, Store};
-use alloc::sync::Arc;
 use core::{
     any::{type_name, TypeId},
     fmt::{self, Debug},
@@ -15,25 +14,70 @@ use crate::Engine;
 /// This wrapper exists to provide a `Debug` impl so that `#[derive(Debug)]`
 /// works for [`Store`].
 #[allow(clippy::type_complexity)]
-#[derive(Clone)]
-pub struct RestorePrunedWrapper(Arc<dyn Send + Sync + Fn(&mut PrunedStore) -> &mut dyn TypedStore>);
+#[derive(Copy, Clone)]
+pub struct RestorePrunedWrapper {
+    /// Calls the given [`HostFuncEntity`] with the `params` and `results` on `instance`.
+    ///
+    /// # Errors
+    ///
+    /// If the called host function returned an error.
+    call_host_func: fn(
+        /* store:         */ &mut PrunedStore,
+        /* func:          */ &HostFuncEntity,
+        /* instance:      */ Option<&Instance>,
+        /* params_results:*/ FuncInOut,
+        /* call_hooks:    */ CallHooks,
+    ) -> Result<(), Error>,
+    /// Returns an exclusive reference to [`StoreInner`] and a [`ResourceLimiterRef`].
+    store_inner_and_resource_limiter_ref:
+        fn(&mut PrunedStore) -> (&mut StoreInner, ResourceLimiterRef<'_>),
+}
 impl RestorePrunedWrapper {
     pub fn new<T: 'static>() -> Self {
-        Self(Arc::new(|pruned| -> &mut dyn TypedStore {
-            let Ok(store) = PrunedStore::restore::<T>(pruned) else {
-                panic!(
-                    "failed to convert `PrunedStore` back into `Store<{}>`",
-                    type_name::<T>(),
-                );
-            };
-            store
-        }))
+        Self {
+            call_host_func: |pruned: &mut PrunedStore,
+                             func: &HostFuncEntity,
+                             instance: Option<&Instance>,
+                             params_results: FuncInOut,
+                             call_hooks: CallHooks|
+             -> Result<(), Error> {
+                let store: &mut Store<T> = pruned.restore_or_panic();
+                if matches!(call_hooks, CallHooks::Call) {
+                    store.invoke_call_hook(CallHook::CallingHost)?;
+                }
+                store.call_host_func(func, instance, params_results)?;
+                if matches!(call_hooks, CallHooks::Call) {
+                    store.invoke_call_hook(CallHook::ReturningFromHost)?;
+                }
+                Ok(())
+            },
+            store_inner_and_resource_limiter_ref: |pruned: &mut PrunedStore| {
+                pruned
+                    .restore_or_panic::<T>()
+                    .store_inner_and_resource_limiter_ref()
+            },
+        }
+    }
+}
+impl RestorePrunedWrapper {
+    #[inline]
+    fn call_host_func(
+        &self,
+        pruned: &mut PrunedStore,
+        func: &HostFuncEntity,
+        instance: Option<&Instance>,
+        params_results: FuncInOut,
+        call_hooks: CallHooks,
+    ) -> Result<(), Error> {
+        (self.call_host_func)(pruned, func, instance, params_results, call_hooks)
     }
 
-    /// Restores the [`PrunedStore`] and returns a reference to it via [`TypedStore`].
     #[inline]
-    fn restore<'a>(&self, pruned: &'a mut PrunedStore) -> &'a mut dyn TypedStore {
-        (self.0)(pruned)
+    fn store_inner_and_resource_limiter_ref<'a>(
+        &self,
+        pruned: &'a mut PrunedStore,
+    ) -> (&'a mut StoreInner, ResourceLimiterRef<'a>) {
+        (self.store_inner_and_resource_limiter_ref)(pruned)
     }
 }
 impl Debug for RestorePrunedWrapper {
@@ -94,27 +138,47 @@ impl PrunedStore {
     /// # Errors
     ///
     /// If the host function returns an error.
+    #[inline]
     pub fn call_host_func(
         &mut self,
         func: &HostFuncEntity,
         instance: Option<&Instance>,
         params_results: FuncInOut,
         call_hooks: CallHooks,
     ) -> Result<(), Error> {
-        self.typed_store()
-            .call_host_func(func, instance, params_results, call_hooks)
+        self.pruned.restore_pruned.clone().call_host_func(
+            self,
+            func,
+            instance,
+            params_results,
+            call_hooks,
+        )
     }
 
     /// Returns an exclusive reference to [`StoreInner`] and a [`ResourceLimiterRef`].
+    #[inline]
     pub fn store_inner_and_resource_limiter_ref(
         &mut self,
     ) -> (&mut StoreInner, ResourceLimiterRef) {
-        self.typed_store().store_inner_and_resource_limiter_ref()
+        self.pruned
+            .restore_pruned
+            .clone()
+            .store_inner_and_resource_limiter_ref(self)
     }
 
-    /// Returns the associated [`TypedStore`] of `self`.
-    fn typed_store(&mut self) -> &mut dyn TypedStore {
-        self.pruned.restore_pruned.clone().restore(self)
+    /// Restores `self` to a proper [`Store<T>`] if possible.
+    ///
+    /// # Panics
+    ///
+    /// If the `T` of the resulting [`Store<T>`] does not match the given `T`.
+    fn restore_or_panic<T: 'static>(&mut self) -> &mut Store<T> {
+        let Ok(store) = self.restore() else {
+            panic!(
+                "failed to convert `PrunedStore` back into `Store<{}>`",
+                type_name::<T>(),
+            );
+        };
+        store
     }
 
     /// Restores `self` to a proper [`Store<T>`] if possible.
@@ -145,49 +209,6 @@ impl PrunedStore {
 #[derive(Debug)]
 pub struct PrunedStoreError;
 
-/// Methods available from [`PrunedStore`] that have been restored dynamically.
-pub trait TypedStore {
-    /// Calls the given [`HostFuncEntity`] with the `params` and `results` on `instance`.
-    ///
-    /// # Errors
-    ///
-    /// If the called host function returned an error.
-    fn call_host_func(
-        &mut self,
-        func: &HostFuncEntity,
-        instance: Option<&Instance>,
-        params_results: FuncInOut,
-        call_hooks: CallHooks,
-    ) -> Result<(), Error>;
-
-    /// Returns an exclusive reference to [`StoreInner`] and a [`ResourceLimiterRef`].
-    fn store_inner_and_resource_limiter_ref(&mut self) -> (&mut StoreInner, ResourceLimiterRef);
-}
-
-impl<T> TypedStore for Store<T> {
-    fn call_host_func(
-        &mut self,
-        func: &HostFuncEntity,
-        instance: Option<&Instance>,
-        params_results: FuncInOut,
-        call_hooks: CallHooks,
-    ) -> Result<(), Error> {
-        if matches!(call_hooks, CallHooks::Call) {
-            <Store<T>>::invoke_call_hook(self, CallHook::CallingHost)?;
-        }
-        <Store<T>>::call_host_func(self, func, instance, params_results)?;
-        if matches!(call_hooks, CallHooks::Call) {
-            <Store<T>>::invoke_call_hook(self, CallHook::ReturningFromHost)?;
-        }
-        Ok(())
-    }
-
-    #[inline]
-    fn store_inner_and_resource_limiter_ref(&mut self) -> (&mut StoreInner, ResourceLimiterRef) {
-        <Store<T>>::store_inner_and_resource_limiter_ref(self)
-    }
-}
-
 #[test]
 fn pruning_works() {
     let engine = Engine::default();