Drop T: 'static bounds from wasmi::Store again (#1507)

* 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

* redesign RestorePrunedWrapper API

* rename RestorePrunedWrapper -> PrunedStoreVTable

* replace commented out code with just code

* add docs

* drop T:'static bounds from Store

Author: Robbepop

crates/wasmi/src/engine/executor/instrs/memory.rs

@@ -1,6 +1,6 @@
 use super::{Executor, InstructionPtr};
 use crate::{
-    core::{MemoryError, ResourceLimiterRef, TrapCode},
+    core::{MemoryError, TrapCode},
     engine::{utils::unreachable_unchecked, ResumableOutOfFuelError},
     ir::{
         index::{Data, Memory},
@@ -76,8 +76,7 @@ impl Executor<'_> {
         delta: Reg,
     ) -> Result<(), Error> {
         let delta: u64 = self.get_register_as(delta);
-        let (store, mut resource_limiter) = store.store_inner_and_resource_limiter_ref();
-        self.execute_memory_grow_impl(store, result, delta, &mut resource_limiter)
+        self.execute_memory_grow_impl(store, result, delta)
     }
 
     /// Executes an [`Instruction::MemoryGrowImm`].
@@ -87,41 +86,38 @@ impl Executor<'_> {
         result: Reg,
         delta: Const32<u64>,
     ) -> Result<(), Error> {
-        let (store, mut resource_limiter) = store.store_inner_and_resource_limiter_ref();
         let delta = u64::from(delta);
-        self.execute_memory_grow_impl(store, result, delta, &mut resource_limiter)
+        self.execute_memory_grow_impl(store, result, delta)
     }
 
     /// Executes a generic `memory.grow` instruction.
     #[inline(never)]
-    fn execute_memory_grow_impl<'store>(
+    fn execute_memory_grow_impl(
         &mut self,
-        store: &'store mut StoreInner,
+        store: &mut PrunedStore,
         result: Reg,
         delta: u64,
-        resource_limiter: &mut ResourceLimiterRef<'store>,
     ) -> Result<(), Error> {
         let memory = self.fetch_memory_index(1);
         if delta == 0 {
             // Case: growing by 0 pages means there is nothing to do
-            self.execute_memory_size_impl(store, result, memory);
+            self.execute_memory_size_impl(store.inner(), result, memory);
             return self.try_next_instr_at(2);
         }
         let memory = self.get_memory(memory);
-        let (memory, fuel) = store.resolve_memory_and_fuel_mut(&memory);
-        let return_value = memory.grow(delta, Some(fuel), resource_limiter);
-        let return_value = match return_value {
+        let return_value = match store.grow_memory(&memory, delta) {
             Ok(return_value) => {
                 // The `memory.grow` operation might have invalidated the cached
                 // linear memory so we need to reset it in order for the cache to
                 // reload in case it is used again.
                 //
                 // Safety: the instance has not changed thus calling this is valid.
-                unsafe { self.cache.update_memory(store) };
+                unsafe { self.cache.update_memory(store.inner_mut()) };
                 return_value
             }
             Err(MemoryError::OutOfBoundsGrowth | MemoryError::OutOfSystemMemory) => {
-                match memory.ty().is_64() {
+                let memory_ty = store.inner().resolve_memory(&memory).ty();
+                match memory_ty.is_64() {
                     true => u64::MAX,
                     false => u64::from(u32::MAX),
                 }

crates/wasmi/src/engine/executor/instrs/table.rs

@@ -1,6 +1,6 @@
 use super::{Executor, InstructionPtr};
 use crate::{
-    core::{ResourceLimiterRef, Table as CoreTable, TableError, TrapCode},
+    core::{Table as CoreTable, TableError, TrapCode},
     engine::{utils::unreachable_unchecked, ResumableOutOfFuelError},
     ir::{
         index::{Elem, Table},
@@ -294,8 +294,7 @@ impl Executor<'_> {
         value: Reg,
     ) -> Result<(), Error> {
         let delta: u64 = self.get_register_as(delta);
-        let (store, mut resource_limiter) = store.store_inner_and_resource_limiter_ref();
-        self.execute_table_grow_impl(store, result, delta, value, &mut resource_limiter)
+        self.execute_table_grow_impl(store, result, delta, value)
     }
 
     /// Executes an [`Instruction::TableGrowImm`].
@@ -307,34 +306,31 @@ impl Executor<'_> {
         value: Reg,
     ) -> Result<(), Error> {
         let delta: u64 = delta.into();
-        let (store, mut resource_limiter) = store.store_inner_and_resource_limiter_ref();
-        self.execute_table_grow_impl(store, result, delta, value, &mut resource_limiter)
+        self.execute_table_grow_impl(store, result, delta, value)
     }
 
     /// Executes a generic `table.grow` instruction.
     #[inline(never)]
-    fn execute_table_grow_impl<'store>(
+    fn execute_table_grow_impl(
         &mut self,
-        store: &'store mut StoreInner,
+        store: &mut PrunedStore,
         result: Reg,
         delta: u64,
         value: Reg,
-        resource_limiter: &mut ResourceLimiterRef<'store>,
     ) -> Result<(), Error> {
         let table_index = self.fetch_table_index(1);
         if delta == 0 {
             // Case: growing by 0 elements means there is nothing to do
-            self.execute_table_size_impl(store, result, table_index);
+            self.execute_table_size_impl(store.inner(), result, table_index);
             return self.try_next_instr_at(2);
         }
         let table = self.get_table(table_index);
         let value = self.get_register(value);
-        let (table, fuel) = store.resolve_table_and_fuel_mut(&table);
-        let return_value = table.grow_untyped(delta, value, Some(fuel), resource_limiter);
-        let return_value = match return_value {
+        let return_value = match store.grow_table(&table, delta, value) {
             Ok(return_value) => return_value,
             Err(TableError::GrowOutOfBounds | TableError::OutOfSystemMemory) => {
-                match table.ty().is_64() {
+                let table_ty = store.inner().resolve_table(&table).ty();
+                match table_ty.is_64() {
                     true => u64::MAX,
                     false => u64::from(u32::MAX),
                 }

crates/wasmi/src/store/mod.rs

@@ -1,8 +1,9 @@
 mod context;
 mod inner;
 mod pruned;
+mod typeid;
 
-use self::pruned::RestorePrunedWrapper;
+use self::pruned::PrunedStoreVTable;
 pub use self::{
     context::{AsContext, AsContextMut, StoreContext, StoreContextMut},
     inner::{StoreInner, Stored},
@@ -41,27 +42,27 @@ pub struct Store<T> {
     /// restored `T` matches the original `T` of the `store`.
     id: TypeId,
     /// Used to restore a [`PrunedStore`] to a [`Store<T>`].
-    restore_pruned: RestorePrunedWrapper,
+    restore_pruned: PrunedStoreVTable,
 }
 
 impl<T> Default for Store<T>
 where
-    T: Default + 'static,
+    T: Default,
 {
     fn default() -> Self {
         let engine = Engine::default();
         Self::new(&engine, T::default())
     }
 }
 
-impl<T: 'static> Store<T> {
+impl<T> Store<T> {
     /// Creates a new store.
     pub fn new(engine: &Engine, data: T) -> Self {
         Self {
             inner: StoreInner::new(engine),
             typed: TypedStoreInner::new(data),
-            id: TypeId::of::<T>(),
-            restore_pruned: RestorePrunedWrapper::new::<T>(),
+            id: typeid::of::<T>(),
+            restore_pruned: PrunedStoreVTable::new::<T>(),
         }
     }
 }

crates/wasmi/src/store/pruned.rs

@@ -1,7 +1,15 @@
-use super::{CallHooks, FuncInOut, HostFuncEntity, ResourceLimiterRef, StoreInner};
-use crate::{core::hint, CallHook, Error, Instance, Store};
+use super::{typeid, CallHooks, FuncInOut, HostFuncEntity, StoreInner};
+use crate::{
+    core::{hint, MemoryError, TableError, UntypedVal},
+    CallHook,
+    Error,
+    Instance,
+    Memory,
+    Store,
+    Table,
+};
 use core::{
-    any::{type_name, TypeId},
+    any::type_name,
     fmt::{self, Debug},
     mem,
 };
@@ -15,25 +23,31 @@ use crate::Engine;
 /// works for [`Store`].
 #[allow(clippy::type_complexity)]
 #[derive(Copy, Clone)]
-pub struct RestorePrunedWrapper {
+pub struct PrunedStoreVTable {
     /// 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,
+        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<'_>),
+    /// Grows `memory` by `delta` pages.
+    grow_memory: fn(&mut PrunedStore, memory: &Memory, delta: u64) -> Result<u64, MemoryError>,
+    /// Grows `table` by `delta` items filling with `init`.
+    grow_table: fn(
+        &mut PrunedStore,
+        table: &Table,
+        delta: u64,
+        init: UntypedVal,
+    ) -> Result<u64, TableError>,
 }
-impl RestorePrunedWrapper {
-    pub fn new<T: 'static>() -> Self {
+impl PrunedStoreVTable {
+    pub fn new<T>() -> Self {
         Self {
             call_host_func: |pruned: &mut PrunedStore,
                              func: &HostFuncEntity,
@@ -51,15 +65,29 @@ impl RestorePrunedWrapper {
                 }
                 Ok(())
             },
-            store_inner_and_resource_limiter_ref: |pruned: &mut PrunedStore| {
-                pruned
-                    .restore_or_panic::<T>()
-                    .store_inner_and_resource_limiter_ref()
+            grow_memory: |pruned: &mut PrunedStore,
+                          memory: &Memory,
+                          delta: u64|
+             -> Result<u64, MemoryError> {
+                let store: &mut Store<T> = pruned.restore_or_panic();
+                let (store, mut resource_limiter) = store.store_inner_and_resource_limiter_ref();
+                let (memory, fuel) = store.resolve_memory_and_fuel_mut(memory);
+                memory.grow(delta, Some(fuel), &mut resource_limiter)
+            },
+            grow_table: |pruned: &mut PrunedStore,
+                         table: &Table,
+                         delta: u64,
+                         init: UntypedVal|
+             -> Result<u64, TableError> {
+                let store: &mut Store<T> = pruned.restore_or_panic();
+                let (store, mut resource_limiter) = store.store_inner_and_resource_limiter_ref();
+                let (table, fuel) = store.resolve_table_and_fuel_mut(table);
+                table.grow_untyped(delta, init, Some(fuel), &mut resource_limiter)
             },
         }
     }
 }
-impl RestorePrunedWrapper {
+impl PrunedStoreVTable {
     #[inline]
     fn call_host_func(
         &self,
@@ -73,14 +101,27 @@ impl RestorePrunedWrapper {
     }
 
     #[inline]
-    fn store_inner_and_resource_limiter_ref<'a>(
+    fn grow_memory(
+        &self,
+        pruned: &mut PrunedStore,
+        memory: &Memory,
+        delta: u64,
+    ) -> Result<u64, MemoryError> {
+        (self.grow_memory)(pruned, memory, delta)
+    }
+
+    #[inline]
+    fn grow_table(
         &self,
-        pruned: &'a mut PrunedStore,
-    ) -> (&'a mut StoreInner, ResourceLimiterRef<'a>) {
-        (self.store_inner_and_resource_limiter_ref)(pruned)
+        pruned: &mut PrunedStore,
+        table: &Table,
+        delta: u64,
+        init: UntypedVal,
+    ) -> Result<u64, TableError> {
+        (self.grow_table)(pruned, table, delta, init)
     }
 }
-impl Debug for RestorePrunedWrapper {
+impl Debug for PrunedStoreVTable {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         write!(f, "RestorePrunedWrapper")
     }
@@ -155,23 +196,35 @@ impl PrunedStore {
         )
     }
 
-    /// Returns an exclusive reference to [`StoreInner`] and a [`ResourceLimiterRef`].
+    /// Grows the [`Memory`] by `delta` pages and returns the result.
+    #[inline]
+    pub fn grow_memory(&mut self, memory: &Memory, delta: u64) -> Result<u64, MemoryError> {
+        self.pruned
+            .restore_pruned
+            .clone()
+            .grow_memory(self, memory, delta)
+    }
+
+    /// Grows the [`Table`] by `delta` items and returns the result.
     #[inline]
-    pub fn store_inner_and_resource_limiter_ref(
+    pub fn grow_table(
         &mut self,
-    ) -> (&mut StoreInner, ResourceLimiterRef) {
+        table: &Table,
+        delta: u64,
+        init: UntypedVal,
+    ) -> Result<u64, TableError> {
         self.pruned
             .restore_pruned
             .clone()
-            .store_inner_and_resource_limiter_ref(self)
+            .grow_table(self, table, delta, init)
     }
 
     /// 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> {
+    fn restore_or_panic<T>(&mut self) -> &mut Store<T> {
         let Ok(store) = self.restore() else {
             panic!(
                 "failed to convert `PrunedStore` back into `Store<{}>`",
@@ -187,8 +240,8 @@ impl PrunedStore {
     ///
     /// If the `T` of the resulting [`Store<T>`] does not match the given `T`.
     #[inline]
-    fn restore<T: 'static>(&mut self) -> Result<&mut Store<T>, PrunedStoreError> {
-        if hint::unlikely(TypeId::of::<T>() != self.pruned.id) {
+    fn restore<T>(&mut self) -> Result<&mut Store<T>, PrunedStoreError> {
+        if hint::unlikely(typeid::of::<T>() != self.pruned.id) {
             return Err(PrunedStoreError);
         }
         let store = {

crates/wasmi/src/store/typeid.rs

@@ -0,0 +1,33 @@
+use core::{any::TypeId, marker::PhantomData, mem};
+
+/// Returns the [`TypeId`] of `T`.
+///
+/// # Note
+///
+/// - `T` does _not_ need to be `'static` for this to work.
+/// - This uses a trick copied from the [`typeid` crate](https://docs.rs/typeid) by [dtolnay](https://github.com/dtolnay).
+#[must_use]
+#[inline(always)]
+pub fn of<T>() -> TypeId
+where
+    T: ?Sized,
+{
+    trait NonStaticAny {
+        fn get_type_id(&self) -> TypeId
+        where
+            Self: 'static;
+    }
+    impl<T: ?Sized> NonStaticAny for PhantomData<T> {
+        #[inline(always)]
+        fn get_type_id(&self) -> TypeId
+        where
+            Self: 'static,
+        {
+            TypeId::of::<T>()
+        }
+    }
+    let phantom_data = PhantomData::<T>;
+    NonStaticAny::get_type_id(unsafe {
+        mem::transmute::<&dyn NonStaticAny, &(dyn NonStaticAny + 'static)>(&phantom_data)
+    })
+}