Make InstanceEntity dynamically sized to drop one indirection in the executor (#2005)

* add HandleAndCache::cast_* methods

* use new HandleAndCache::cast_* methods

* make HandleAndCache::handle private

* turn InstanceEntity into dynamically sized type

* adjust StoreInner to store Box<InstanceEntity>

* update Handle trait to allow ?Sized Entity

* relax const assert

* add ThinPtr utility type

* rename InstanceHeader -> InstanceEntityHeader

also make it private

* make all InstanceEntityHeader methods private

* add ThinPtr<InstanceEntity> API

* add missing HANDLES_OFFSET const for new API

* use ThinPtr<InstanceEntity> in the executor

* remove no longer needed InstanceEntity APIs

* introduce typed CacheAndEntity<T> type

* adjust executor usage

* remove unnecessary InstanceHandle trait again

* rename HandleAndCache -> AnyHandleAndEntity

* rename ThinPtr<InstanceEntity> APIs

* return Option from ThinPtr<InstanceEntity>::entry APIs

* use names from std for ThinPtr API

* add validity doc section to Inst

* re-design ThinPtr::as_byte_ptr -> cast<U>

* fix some minor doc issues

* improve InstanceEntity[Header] docs

* add AnyHandle::assert_kind

* re-design AnyHandleAndEntity conversions to typed wrapper

* elaborate comment in InstanceEntity::alloc

* use new typed conversion APIs in

* update API impl of ThinPtr<InstanceEntity>

* apply improvements to InstanceEntity::alloc from code review

* remove unused imports

* add new instantiation tests

* add new test to guard dangling Inst pointer

* remove superflous comma

* move ptr.rs -> instance/thin_ptr.rs

* elaborate on the safety comment for a Sync impl

* use crate imports

* add cfg-guards for release builds

* use public APIs in tests

* move InstanceEntity defs/impls into entity.rs submodule

* clean-ups for StoreInner

* improve HANDLES_OFFSET computation

* add another const assert

* use HANDLES_OFFSET to verify Layout usage

* carve out manual handles_offset checker

* remove ExactSizeIterator bounds from alloc

* clean-up alloc method

Author: Robbepop

crates/wasmi/src/engine/executor/handler/state.rs

@@ -23,19 +23,12 @@ use crate::{
         },
         utils::unreachable_unchecked,
     },
-    instance::InstanceEntity,
+    instance::{InstanceEntity, ThinPtr},
     ir::{self, BoundedSlotSpan, Slot, SlotSpan},
     store::PrunedStore,
 };
 use alloc::vec::Vec;
-use core::{
-    cmp,
-    marker::PhantomData,
-    mem,
-    ops,
-    ptr::{self, NonNull},
-    slice,
-};
+use core::{cmp, marker::PhantomData, mem, ops, ptr, slice};
 
 pub struct VmState<'vm> {
     pub store: &'vm mut PrunedStore,
@@ -151,6 +144,22 @@ impl DoneReason {
 }
 
 /// A thin-wrapper around a non-owned [`InstanceEntity`].
+///
+/// # Validity
+///
+/// Every `Inst` is created from a live [`InstanceEntity`] that the [`Store`] currently owns,
+/// and the executor keeps that instance alive and warmed up for as long as the `Inst` is
+/// reachable — it is stored in [`Args`] and in the [`CallStack`] frames of the very execution
+/// that owns the [`Store`]. Since the entity is boxed, allocating further instances (e.g. from
+/// a host call) does not move it.
+///
+/// The `unsafe` methods of [`ThinPtr<InstanceEntity>`] therefore hold for any `Inst` reached
+/// through [`Inst::as_ptr`], and call sites need only justify the *kind* of the address they
+/// pass, not the liveness of the instance.
+///
+/// [`Store`]: crate::Store
+/// [`Args`]: super::Args
+/// [`ThinPtr<InstanceEntity>`]: ThinPtr
 #[derive(Debug, Copy, Clone)]
 #[repr(transparent)]
 pub struct Inst {
@@ -166,7 +175,7 @@ pub struct Inst {
     ///   integer and float won't be a terrible trade-off.
     value: InstRepr,
     /// Marks `Inst` as logically containing a shared raw pointer.
-    marker: PhantomData<*const InstanceEntity>,
+    marker: PhantomData<ThinPtr<InstanceEntity>>,
 }
 
 /// The underlying float representation of `Inst` for 64-bit platforms.
@@ -185,60 +194,44 @@ const _: () = {
 impl From<&'_ InstanceEntity> for Inst {
     #[inline]
     fn from(entity: &'_ InstanceEntity) -> Self {
-        let addr = (entity as *const InstanceEntity).expose_provenance();
-        Self::from_addr(addr)
+        Self::from(ThinPtr::from_ref(entity))
     }
 }
 
-impl From<NonNull<InstanceEntity>> for Inst {
+impl From<ThinPtr<InstanceEntity>> for Inst {
     #[inline]
-    fn from(entity: NonNull<InstanceEntity>) -> Self {
-        let addr = entity.as_ptr().expose_provenance();
-        Self::from_addr(addr)
+    fn from(entity: ThinPtr<InstanceEntity>) -> Self {
+        let value = InstRepr::from_ne_bytes(entity.expose_provenance().to_ne_bytes());
+        Self {
+            value,
+            marker: PhantomData,
+        }
     }
 }
 
 impl PartialEq for Inst {
     #[inline]
     fn eq(&self, other: &Self) -> bool {
-        self.as_ptr().addr() == other.as_ptr().addr()
+        self.as_ptr() == other.as_ptr()
     }
 }
 impl Eq for Inst {}
 
 impl Inst {
-    /// Creates a new [`Inst`] from the given `addr` value.
-    #[inline]
-    fn from_addr(addr: usize) -> Self {
-        let value = InstRepr::from_ne_bytes(addr.to_ne_bytes());
-        Self {
-            value,
-            marker: PhantomData,
-        }
-    }
-
-    /// Converts the underlying representation back into its original pointer value.
-    #[inline]
-    pub fn as_ptr(&self) -> NonNull<InstanceEntity> {
-        let bits = usize::from_ne_bytes(self.value.to_ne_bytes());
-        let ptr = ptr::with_exposed_provenance::<InstanceEntity>(bits);
-        unsafe { NonNull::new_unchecked(ptr as *mut InstanceEntity) }
-    }
-
-    /// Returns a shared reference to the referenced [`InstanceEntity`].
+    /// Returns the thin pointer to the referenced [`InstanceEntity`].
     ///
-    /// # Safety
+    /// # Note
     ///
-    /// The caller must ensure that:
+    /// This cannot be a `NonNull<InstanceEntity>` since [`InstanceEntity`] is dynamically
+    /// sized. Its thin-pointer API is provided by [`ThinPtr<InstanceEntity>`].
     ///
-    /// - The [`Inst`] was constructed from a valid, properly aligned
-    ///   `InstanceEntity` pointer.
-    /// - The referenced [`InstanceEntity`] remains alive and is not
-    ///   mutably accessed for the entire duration of the returned
-    ///   reference.
+    /// [`ThinPtr<InstanceEntity>`]: ThinPtr
     #[inline]
-    pub unsafe fn as_ref(&self) -> &InstanceEntity {
-        unsafe { self.as_ptr().as_ref() }
+    pub fn as_ptr(&self) -> ThinPtr<InstanceEntity> {
+        let addr = usize::from_ne_bytes(self.value.to_ne_bytes());
+        // Safety: `Inst` can only ever be created from a `ThinPtr<InstanceEntity>` whose
+        //         provenance has been exposed.
+        unsafe { ThinPtr::with_exposed_provenance(addr) }
     }
 }
 
@@ -269,8 +262,8 @@ mod inst_tests {
     use super::*;
 
     const _: fn() = || {
-        fn assert_send<T: Send>() {}
-        fn assert_sync<T: Sync>() {}
+        fn assert_send<T: ?Sized + Send>() {}
+        fn assert_sync<T: ?Sized + Sync>() {}
 
         assert_send::<InstanceEntity>();
         assert_sync::<InstanceEntity>();

crates/wasmi/src/engine/executor/handler/utils.rs

@@ -539,22 +539,26 @@ pub fn exec_copy_span_des(sp: Sp, dst: SlotSpan, src: SlotSpan, len: u16) {
 /// so the default memory's address is not guaranteed to be zero.
 #[inline]
 pub fn is_default_memory(instance: Inst, memory: ir::MemoryAddr) -> bool {
-    let instance = unsafe { instance.as_ref() };
+    // SAFETY: `instance` refers to a live instance for the duration of the call.
+    let layout = unsafe { instance.as_ptr().layout() };
     matches!(
-        instance.layout().memory_addr(0),
+        layout.memory_addr(0),
         Some(addr) if u32::from(addr) == u32::from(memory)
     )
 }
 
 pub fn extract_mem0(_store: &mut PrunedStore, inst: Inst) -> (Mem0Ptr, Mem0Len) {
-    let instance = unsafe { inst.as_ref() };
-    let Some(addr) = instance.layout().memory_addr(0) else {
+    // SAFETY: `inst` refers to a live instance for the duration of the call.
+    let Some(addr) = (unsafe { inst.as_ptr().layout() }).memory_addr(0) else {
         return (Mem0Ptr::from([].as_mut_ptr()), Mem0Len::from(0));
     };
-    let Some(mem0) = instance.get_memory_ptr(addr) else {
-        unsafe { unreachable_unchecked!() }
+    // SAFETY: `addr` stems from the instance layout and thus addresses a memory entry
+    //         whose cache was warmed at instantiation.
+    let Some(mem0) = (unsafe { inst.as_ptr().get_memory(addr) }) else {
+        unsafe { unreachable_unchecked!("missing memory at: {addr:?}") }
     };
     // SAFETY: warmed at instantiation; the `_store` borrow scopes exclusive memory access.
+    let mem0 = mem0.entity();
     let mem0 = unsafe { &mut *mem0.as_ptr() }.data_mut();
     let mem0_ptr = mem0.as_mut_ptr();
     let mem0_len = mem0.len();
@@ -583,7 +587,7 @@ pub fn memory_slice_mut(
 
 macro_rules! impl_resolve_from_instance {
     (
-        $( fn $fn:ident(inst: Inst, store: &mut StoreInner, $param:ident: ir::$param_ty:ident) -> &mut $ret:ty = $getter:expr );* $(;)?
+        $( fn $fn:ident(inst: Inst, store: &mut StoreInner, $param:ident: ir::$param_ty:ident) -> &mut $ret:ty = $entry:ident );* $(;)?
     ) => {
         $(
             /// Resolves the entity from the warmed up instance cache.
@@ -592,10 +596,9 @@ macro_rules! impl_resolve_from_instance {
             /// cannot alias a concurrent store mutation (this is the safe wrapper).
             #[inline]
             pub fn $fn(inst: Inst, _store: &mut StoreInner, $param: ir::$param_ty) -> &mut $ret {
-                let inst = unsafe { inst.as_ref() };
-                let raw_addr = ::core::primitive::u32::from($param);
-                let addr = <$param_ty>::from(raw_addr);
-                let Some(ptr) = $getter(inst, addr) else {
+                let addr = <$param_ty>::from(::core::primitive::u32::from($param));
+                // SAFETY: `addr` addresses an entry of this kind by translation invariant.
+                let Some(entry) = (unsafe { inst.as_ptr().$entry(addr) }) else {
                     unsafe {
                         $crate::engine::utils::unreachable_unchecked!(
                             ::core::concat!("missing ", ::core::stringify!($param), " at: {:?}"),
@@ -604,22 +607,22 @@ macro_rules! impl_resolve_from_instance {
                     }
                 };
                 // SAFETY: warmed at instantiation; the `_store` borrow scopes the reference.
-                unsafe { &mut *ptr.as_ptr() }
+                unsafe { &mut *entry.entity().as_ptr() }
             }
         )*
     };
 }
 impl_resolve_from_instance! {
-    fn load_data(inst: Inst, store: &mut StoreInner, data: ir::DataAddr) -> &mut DataSegmentEntity = InstanceEntity::get_data_ptr;
-    fn load_elem(inst: Inst, store: &mut StoreInner, elem: ir::ElemAddr) -> &mut ElementSegmentEntity = InstanceEntity::get_elem_ptr;
-    fn load_global(inst: Inst, store: &mut StoreInner, global: ir::GlobalAddr) -> &mut GlobalEntity = InstanceEntity::get_global_ptr;
-    fn load_memory(inst: Inst, store: &mut StoreInner, memory: ir::MemoryAddr) -> &mut MemoryEntity = InstanceEntity::get_memory_ptr;
-    fn load_table(inst: Inst, store: &mut StoreInner, table: ir::TableAddr) -> &mut TableEntity = InstanceEntity::get_table_ptr;
+    fn load_data(inst: Inst, store: &mut StoreInner, data: ir::DataAddr) -> &mut DataSegmentEntity = get_data;
+    fn load_elem(inst: Inst, store: &mut StoreInner, elem: ir::ElemAddr) -> &mut ElementSegmentEntity = get_elem;
+    fn load_global(inst: Inst, store: &mut StoreInner, global: ir::GlobalAddr) -> &mut GlobalEntity = get_global;
+    fn load_memory(inst: Inst, store: &mut StoreInner, memory: ir::MemoryAddr) -> &mut MemoryEntity = get_memory;
+    fn load_table(inst: Inst, store: &mut StoreInner, table: ir::TableAddr) -> &mut TableEntity = get_table;
 }
 
 macro_rules! impl_resolve_ptr_from_instance {
     (
-        $( fn $fn:ident(inst: Inst, $param:ident: ir::$param_ty:ident) -> $ret:ty = $getter:expr );* $(;)?
+        $( fn $fn:ident(inst: Inst, $param:ident: ir::$param_ty:ident) -> $ret:ty = $entry:ident );* $(;)?
     ) => {
         $(
             /// Resolves a pointer to the entity from the warmed up instance cache.
@@ -637,27 +640,26 @@ macro_rules! impl_resolve_ptr_from_instance {
             /// ensure the resulting reference does not alias any other live reference.
             #[inline]
             pub unsafe fn $fn(inst: Inst, $param: ir::$param_ty) -> NonNull<$ret> {
-                let inst = unsafe { inst.as_ref() };
-                let raw_addr = ::core::primitive::u32::from($param);
-                let addr = <$param_ty>::from(raw_addr);
-                let Some(ptr) = $getter(inst, addr) else {
+                let addr = <$param_ty>::from(::core::primitive::u32::from($param));
+                // Safety: guaranteed by the caller.
+                let Some(entry) = (unsafe { inst.as_ptr().$entry(addr) }) else {
                     unsafe {
                         $crate::engine::utils::unreachable_unchecked!(
                             ::core::concat!("missing ", ::core::stringify!($param), " at: {:?}"),
                             addr,
                         )
                     }
                 };
-                ptr
+                entry.entity()
             }
         )*
     };
 }
 impl_resolve_ptr_from_instance! {
-    fn load_memory_ptr(inst: Inst, memory: ir::MemoryAddr) -> MemoryEntity = InstanceEntity::get_memory_ptr;
-    fn load_table_ptr(inst: Inst, table: ir::TableAddr) -> TableEntity = InstanceEntity::get_table_ptr;
-    fn load_data_ptr(inst: Inst, data: ir::DataAddr) -> DataSegmentEntity = InstanceEntity::get_data_ptr;
-    fn load_elem_ptr(inst: Inst, elem: ir::ElemAddr) -> ElementSegmentEntity = InstanceEntity::get_elem_ptr;
+    fn load_memory_ptr(inst: Inst, memory: ir::MemoryAddr) -> MemoryEntity = get_memory;
+    fn load_table_ptr(inst: Inst, table: ir::TableAddr) -> TableEntity = get_table;
+    fn load_data_ptr(inst: Inst, data: ir::DataAddr) -> DataSegmentEntity = get_data;
+    fn load_elem_ptr(inst: Inst, elem: ir::ElemAddr) -> ElementSegmentEntity = get_elem;
 }
 
 /// Resolves the [`Func`] handle and its warmed up cached [`FuncEntity`] pointer from `inst`.
@@ -666,43 +668,51 @@ impl_resolve_ptr_from_instance! {
 /// and must not be turned into a reference that aliases a concurrent store mutation.
 #[inline]
 pub fn load_func_entry(inst: Inst, func: ir::FuncAddr) -> (Func, NonNull<FuncEntity>) {
-    let instance = unsafe { inst.as_ref() };
     let addr = FuncAddr::from(u32::from(func));
-    let Some(func_entry) = instance.get_func_entry(addr) else {
-        unsafe { unreachable_unchecked!("missing func at: {:?}", addr) }
+    // SAFETY: `addr` addresses a func entry by translation invariant and its cache was
+    //         warmed at instantiation.
+    let Some(entry) = (unsafe { inst.as_ptr().get_func(addr) }) else {
+        unsafe { unreachable_unchecked!("missing func at: {addr:?}") }
     };
-    func_entry
+    (entry.handle(), entry.entity())
 }
 
 macro_rules! impl_fetch_from_instance {
     (
-        $( fn $fn:ident($param:ident: $ty:ty as $addr:ty) -> $ret:ty = $getter:expr );* $(;)?
+        $( fn $fn:ident($param:ident: ir::$param_ty:ident) -> $ret:ty = $entry:ident );* $(;)?
     ) => {
         $(
-            pub fn $fn(instance: Inst, $param: $ty) -> $ret {
-                let instance = unsafe { instance.as_ref() };
-                let raw_addr = ::core::primitive::u32::from($param);
-                let addr = <$addr>::from(raw_addr);
-                let Some($param) = $getter(instance, addr) else {
+            pub fn $fn(instance: Inst, $param: ir::$param_ty) -> $ret {
+                let addr = <$param_ty>::from(::core::primitive::u32::from($param));
+                // SAFETY: `addr` addresses an entry of this kind by translation invariant.
+                let Some(entry) = (unsafe { instance.as_ptr().$entry(addr) }) else {
                     unsafe {
                         $crate::engine::utils::unreachable_unchecked!(
                             ::core::concat!("missing ", ::core::stringify!($param), " at: {:?}"),
                             addr,
                         )
                     }
                 };
-                $param
+                entry.handle()
             }
         )*
     };
 }
 impl_fetch_from_instance! {
-    fn fetch_func(func: ir::FuncAddr as FuncAddr) -> Func = InstanceEntity::get_func;
-    fn fetch_memory(memory: ir::MemoryAddr as MemoryAddr) -> Memory = InstanceEntity::get_memory;
-    fn fetch_table(table: ir::TableAddr as TableAddr) -> Table = InstanceEntity::get_table;
-    fn fetch_func_type(func_type: ir::FuncType as u32) -> DedupFuncType = {
-        |instance: &InstanceEntity, index: u32| instance.get_signature(index).copied()
+    fn fetch_func(func: ir::FuncAddr) -> Func = get_func;
+    fn fetch_memory(memory: ir::MemoryAddr) -> Memory = get_memory;
+    fn fetch_table(table: ir::TableAddr) -> Table = get_table;
+}
+
+/// Fetches the [`DedupFuncType`] at `func_type` from the instance's function types.
+#[inline]
+pub fn fetch_func_type(instance: Inst, func_type: ir::FuncType) -> DedupFuncType {
+    let index = u32::from(func_type);
+    // SAFETY: `instance` refers to a live instance for the duration of the call.
+    let Some(func_type) = (unsafe { instance.as_ptr().get_signature(index) }) else {
+        unsafe { unreachable_unchecked!("missing func type at: {index}") }
     };
+    *func_type
 }
 
 macro_rules! impl_resolve_from_store {

crates/wasmi/src/func/caller.rs

@@ -37,7 +37,7 @@ impl<'a, T> Caller<'a, T> {
         let Some(instance) = &self.instance else {
             return None;
         };
-        let instance = unsafe { instance.as_ref() };
+        let instance = unsafe { instance.as_ptr().as_ref() };
         instance.get_export(name)
     }
 

crates/wasmi/src/handle.rs

@@ -6,7 +6,10 @@ pub trait Handle: Copy {
     /// The raw representation of the handle.
     type Raw: ArenaKey;
     /// The store owned entity type of the handle.
-    type Entity;
+    ///
+    /// This may be unsized: [`InstanceEntity`](crate::InstanceEntity) is a dynamically
+    /// sized type and thus owned as a `Box<InstanceEntity>` by its arena.
+    type Entity: ?Sized;
     /// The owner associating reference.
     type Owned<T>;