Add InstanceEntity handle cache with eager warm-up (#1991)

* make StoreInner use StableArena for most things

* fix StoreInner internal resolution API

* add HandleAndCache

* add HandleAndCache to InstanceEntity

* add StoreInner::[try_]resolve_instance_mut

* use &mut InstanceEntity for Inst in executor

* add as_ptr_mut and as_mut methods to Inst type

* add missing inline annotations

* make HandleAndCache return reference instead of options

* add new InstanceEntity API to load using its cache

* fix broken intra doc links

* partially use new instance cache in executor

* return Option from HandleAndCache API

* remove commented out line of code

* switch to eager initialization for InstanceEntity cache

* adjust executor for caching changes

* rename methods for better parity

* add missing docs and inline

* use InstanceEntity::entry helper in more places

* add low-level utils::load_x_ptr helpers to executor

these allow using the instance cache in more places.

* use the instance cache in more places in the executor

* remove no longer used StoreInner APIs

* rename InstanceEntity::*_v2 methods to `*_ptr`

* rename Resolve[Mut] assoc methods

* create Inst from &InstanceEntity (again)

no more need to have mutable InstanceEntity with eager caching

* remove StoreInner::[try_]resolve_instance_mut

* remove unused Inst::{as_ptr_mut,as_mut} methods

* fix bug in {memory,table}.copy with aliasing imports

* add regression test for aliased memory imports

* add StoreInner::[try_]resolve_*_ptr methods

* use StoreInner::resolve_*_ptr methods to fix UB

* remove unused StoreInner methods

Author: Robbepop

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

@@ -19,24 +19,7 @@ use crate::{
                 Sp,
                 VmState,
             },
-            utils::{
-                self,
-                GetValue,
-                IntoControl as _,
-                SetValue,
-                fetch_data,
-                fetch_elem,
-                fetch_global,
-                fetch_memory,
-                fetch_table,
-                get_value,
-                resolve_data_mut,
-                resolve_elem_mut,
-                resolve_global_mut,
-                resolve_memory_mut,
-                resolve_table_mut,
-                set_value,
-            },
+            utils::{self, GetValue, IntoControl as _, SetValue, get_value, set_value},
         },
     },
     ir::{self, BoundedSlotSpan, BranchOffset},
@@ -163,20 +146,19 @@ impl Args {
     /// Returns the bytes of the `memory`.
     #[inline]
     pub fn fetch_memory_bytes<'a>(
-        &self,
+        &mut self,
         state: &'a mut VmState,
-        memory: ir::MemoryAddr,
+        addr: ir::MemoryAddr,
     ) -> &'a mut [u8] {
-        if utils::is_default_memory(self.instance, memory) {
+        if utils::is_default_memory(self.instance, addr) {
             return self.fetch_default_memory_bytes();
         }
-        let memory = utils::fetch_memory(self.instance, memory);
-        utils::resolve_memory_mut(state.store, &memory).data_mut()
+        self.fetch_memory(state, addr).data_mut()
     }
 
     /// Returns the bytes of the default memory at index 0.
     #[inline]
-    pub fn fetch_default_memory_bytes<'a>(&self) -> &'a mut [u8] {
+    pub fn fetch_default_memory_bytes<'a>(&mut self) -> &'a mut [u8] {
         state::mem0_bytes::<'a>(self.mem0_ptr, self.mem0_len)
     }
 
@@ -185,54 +167,49 @@ impl Args {
     pub fn fetch_memory<'a>(
         &mut self,
         state: &'a mut VmState,
-        index: ir::MemoryAddr,
+        addr: ir::MemoryAddr,
     ) -> &'a mut CoreMemory {
-        let global = fetch_memory(self.instance, index);
-        resolve_memory_mut(state.store, &global)
+        utils::load_memory(self.instance, state.store.inner_mut(), addr)
     }
 
     /// Returns an exclusive reference to the global at `index`.
     #[inline]
     pub fn fetch_global<'a>(
         &mut self,
         state: &'a mut VmState,
-        index: ir::GlobalAddr,
+        addr: ir::GlobalAddr,
     ) -> &'a mut CoreGlobal {
-        let global = fetch_global(self.instance, index);
-        resolve_global_mut(state.store, &global)
+        utils::load_global(self.instance, state.store.inner_mut(), addr)
     }
 
     /// Returns an exclusive reference to the table at `index`.
     #[inline]
     pub fn fetch_table<'a>(
         &mut self,
         state: &'a mut VmState,
-        index: ir::TableAddr,
+        addr: ir::TableAddr,
     ) -> &'a mut CoreTable {
-        let table = fetch_table(self.instance, index);
-        resolve_table_mut(state.store, &table)
+        utils::load_table(self.instance, state.store.inner_mut(), addr)
     }
 
     /// Returns an exclusive reference to the element segment at `index`.
     #[inline]
     pub fn fetch_elem<'a>(
         &mut self,
         state: &'a mut VmState,
-        index: ir::ElemAddr,
+        addr: ir::ElemAddr,
     ) -> &'a mut CoreElementSegment {
-        let elem = fetch_elem(self.instance, index);
-        resolve_elem_mut(state.store, &elem)
+        utils::load_elem(self.instance, state.store.inner_mut(), addr)
     }
 
     /// Returns an exclusive reference to the data segment at `index`.
     #[inline]
     pub fn fetch_data<'a>(
         &mut self,
         state: &'a mut VmState,
-        index: ir::DataAddr,
+        addr: ir::DataAddr,
     ) -> &'a mut DataSegmentEntity {
-        let elem = fetch_data(self.instance, index);
-        resolve_data_mut(state.store, &elem)
+        utils::load_data(self.instance, state.store.inner_mut(), addr)
     }
 
     /// Reloads the data pointer and length of the default memory at index 0 from `state`.
@@ -270,7 +247,7 @@ impl Args {
     /// Resolves the [`Func`] at `table[index]` of type `func_type` using `state`.
     #[inline]
     pub fn resolve_indirect_func<Idx>(
-        &self,
+        &mut self,
         state: &mut VmState<'_>,
         index: Idx,
         table: ir::TableAddr,

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

@@ -487,15 +487,24 @@ execution_handler! {
             memory_copy_within(state, &mut args, ip, dst_memory, dst_index, src_index, len)?;
             dispatch!(state, args)
         }
-        let dst_memory = utils::fetch_memory(instance, dst_memory);
-        let src_memory = utils::fetch_memory(instance, src_memory);
-        let (src_memory, dst_memory, fuel) = state
-            .store
-            .inner_mut()
-            .resolve_memory_pair_and_fuel(&src_memory, &dst_memory);
-        // These accesses just perform the bounds checks required by the Wasm spec.
-        let src_bytes = utils::memory_slice(src_memory, src_index, len).into_control()?;
-        let dst_bytes = utils::memory_slice_mut(dst_memory, dst_index, len).into_control()?;
+        let src_ptr = unsafe { utils::load_memory_ptr(instance, src_memory) };
+        let mut dst_ptr = unsafe { utils::load_memory_ptr(instance, dst_memory) };
+        if src_ptr == dst_ptr {
+            // Distinct memory indices can still resolve to the same store entity (e.g. the same
+            // memory imported under two names), so branch on the resolved entity before forming
+            // the aliasing references below.
+            memory_copy_within(state, &mut args, ip, dst_memory, dst_index, src_index, len)?;
+            dispatch!(state, args)
+        }
+        // SAFETY: `src_ptr != dst_ptr`, so these are distinct entities whose references do not
+        //         alias. These accesses just perform the bounds checks required by the Wasm spec.
+        let src_memory = unsafe { src_ptr.as_ref() };
+        let dst_memory = unsafe { dst_ptr.as_mut() };
+        let fuel = state.store.inner_mut().fuel_mut();
+        let src_bytes = utils::memory_slice(src_memory, src_index, len)
+            .into_control()?;
+        let dst_bytes = utils::memory_slice_mut(dst_memory, dst_index, len)
+            .into_control()?;
         consume_fuel!(
             state,
             ip,
@@ -517,8 +526,9 @@ fn memory_copy_within(
     src_index: usize,
     len: usize,
 ) -> Control<(), Break> {
-    let memory = utils::fetch_memory(args.instance, dst_memory);
-    let (memory, fuel) = state.store.inner_mut().resolve_memory_and_fuel_mut(&memory);
+    // SAFETY: `args.instance` is live and warmed up; `memory` is the only entity accessed here.
+    let memory = unsafe { utils::load_memory_ptr(args.instance, dst_memory).as_mut() };
+    let fuel = state.store.inner_mut().fuel_mut();
     // These accesses just perform the bounds checks required by the Wasm spec.
     utils::memory_slice(memory, src_index, len).into_control()?;
     utils::memory_slice(memory, dst_index, len).into_control()?;
@@ -558,8 +568,9 @@ execution_handler! {
         let Ok(len) = usize::try_from(len) else {
             trap!(TrapCode::MemoryOutOfBounds)
         };
-        let memory = utils::fetch_memory(instance, memory);
-        let (memory, fuel) = state.store.inner_mut().resolve_memory_and_fuel_mut(&memory);
+        // SAFETY: `instance` is live and warmed up; `memory` is the only entity accessed here.
+        let memory = unsafe { utils::load_memory_ptr(instance, memory).as_mut() };
+        let fuel = state.store.inner_mut().fuel_mut();
         let slice = utils::memory_slice_mut(memory, dst, len).into_control()?;
         consume_fuel!(state, ip, args, fuel, |costs| costs.fuel_for_copying_values::<u8>(len as u64));
         slice.fill(value);
@@ -599,13 +610,11 @@ execution_handler! {
         let Ok(len) = usize::try_from(len) else {
             trap!(TrapCode::MemoryOutOfBounds)
         };
-        let (memory, data, fuel) = state
-            .store
-            .inner_mut()
-            .resolve_memory_init_params(
-                &utils::fetch_memory(instance, memory),
-                &utils::fetch_data(instance, data),
-            );
+        // SAFETY: `instance` is live and warmed up. `memory` and `data` live in different
+        // arenas, so their references are to distinct entities and cannot alias.
+        let memory = unsafe { utils::load_memory_ptr(instance, memory).as_mut() };
+        let data = unsafe { utils::load_data_ptr(instance, data).as_ref() };
+        let fuel = state.store.inner_mut().fuel_mut();
         let memory = utils::memory_slice_mut(memory, dst_index, len).into_control()?;
         let Some(data) = data
             .bytes()
@@ -731,32 +740,26 @@ execution_handler! {
         let dst: u64 = args.get(dst);
         let src: u64 = args.get(src);
         let len: u64 = args.get(len);
-        if dst_table == src_table {
-            // Case: copy within the same table
-            let table = utils::fetch_table(instance, dst_table);
-            let (table, fuel) = state.store.inner_mut().resolve_table_and_fuel_mut(&table);
-            if let Err(error) = table.copy_within(dst, src, len, Some(fuel)) {
-                let trap_code = match error {
-                    TableError::CopyOutOfBounds => TrapCode::TableOutOfBounds,
-                    TableError::OutOfSystemMemory => TrapCode::OutOfSystemMemory,
-                    TableError::OutOfFuel { required_fuel } => {
-                        args.set_ip(ip);
-                        out_of_fuel!(state, args, required_fuel)
-                    }
-                    _ => panic!("table.copy: unexpected error: {error}"),
-                };
-                trap!(trap_code)
-            }
-            dispatch!(state, args)
-        }
-        // Case: copy between two different tables
-        let dst_table = utils::fetch_table(instance, dst_table);
-        let src_table = utils::fetch_table(instance, src_table);
-        let (dst_table, src_table, fuel) = state
-            .store
-            .inner_mut()
-            .resolve_table_pair_and_fuel(&dst_table, &src_table);
-        if let Err(error) = CoreTable::copy(dst_table, dst, src_table, src, len, Some(fuel)) {
+        let src_ptr = unsafe { utils::load_table_ptr(instance, src_table) };
+        let mut dst_ptr = unsafe { utils::load_table_ptr(instance, dst_table) };
+        // Distinct table indices can still resolve to the same store entity (e.g. the same table
+        // imported under two names), so branch on the resolved entity, not just the index.
+        let result = if src_ptr == dst_ptr {
+            // Case: copy within the same table.
+            // SAFETY: `dst_ptr` is warmed up and is the only entity accessed here.
+            let table = unsafe { dst_ptr.as_mut() };
+            let fuel = state.store.inner_mut().fuel_mut();
+            table.copy_within(dst, src, len, Some(fuel))
+        } else {
+            // Case: copy between two distinct tables.
+            // SAFETY: `src_ptr != dst_ptr`, so these are distinct entities whose references do
+            //         not alias.
+            let src_table = unsafe { src_ptr.as_ref() };
+            let dst_table = unsafe { dst_ptr.as_mut() };
+            let fuel = state.store.inner_mut().fuel_mut();
+            CoreTable::copy(dst_table, dst, src_table, src, len, Some(fuel))
+        };
+        if let Err(error) = result {
             let trap_code = match error {
                 TableError::CopyOutOfBounds => TrapCode::TableOutOfBounds,
                 TableError::OutOfSystemMemory => TrapCode::OutOfSystemMemory,
@@ -794,8 +797,9 @@ execution_handler! {
         let dst: u64 = args.get(dst);
         let len: u64 = args.get(len);
         let value: RawRef = args.get(value);
-        let table = utils::fetch_table(instance, table);
-        let (table, fuel) = state.store.inner_mut().resolve_table_and_fuel_mut(&table);
+        // SAFETY: `instance` is live and warmed up; `table` is the only entity accessed here.
+        let table = unsafe { utils::load_table_ptr(instance, table).as_mut() };
+        let fuel = state.store.inner_mut().fuel_mut();
         if let Err(error) = table.fill_raw(dst, value, len, Some(fuel)) {
             let trap_code = match error {
                 TableError::OutOfSystemMemory => TrapCode::OutOfSystemMemory,
@@ -835,12 +839,11 @@ execution_handler! {
         let dst: u64 = args.get(dst);
         let src: u32 = args.get(src);
         let len: u32 = args.get(len);
-        let table = utils::fetch_table(args.instance, table);
-        let elem = utils::fetch_elem(args.instance, elem);
-        let (table, element, fuel) = state
-            .store
-            .inner_mut()
-            .resolve_table_init_params(&table, &elem);
+        // SAFETY: `args.instance` is live and warmed up. `table` and `elem` live in different
+        //         arenas, so their references are to distinct entities and cannot alias.
+        let table = unsafe { utils::load_table_ptr(args.instance, table).as_mut() };
+        let element = unsafe { utils::load_elem_ptr(args.instance, elem).as_ref() };
+        let fuel = state.store.inner_mut().fuel_mut();
         if let Err(error) = table.init(element.as_ref(), dst, src, len, Some(fuel)) {
             let trap_code = match error {
                 TableError::OutOfSystemMemory => TrapCode::OutOfSystemMemory,