Simplify Wasmi executor call infrastructure (#2000)

* simplify Wasmi executor call infrastructure

* revert changes to replace

these caused a performance regression for call-intense workloads.

* re-introduce reverted change without perf noise

this time benchmarks indicate noise but no regressions

Author: Robbepop

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

@@ -872,6 +872,13 @@ impl Stack {
     }
 
     /// Adjusts `self` after returning from a function.
+    ///
+    /// # Note
+    ///
+    /// The cached `(memory 0)` is only re-extracted if the returned-to frame uses a
+    /// different [`Inst`]. This is sound because every operation that can grow a memory
+    /// refreshes the cache at the growth site instead: the `memory.grow` handler and
+    /// each of the host call paths.
     pub fn pop_frame(
         &mut self,
         store: &mut PrunedStore,
@@ -1066,12 +1073,7 @@ impl ValueStack {
         let callee_end = callee_start.add(callee_size)?;
         self.grow_if_needed(callee_end)?;
         let caller_sp = self.sp(caller_start);
-        let Some(cells) = self.cells_from_to(callee_start, callee_end) else {
-            unsafe { unreachable_unchecked!("must fit slice after `grow_if_needed` operation") }
-        };
-        let Ok(inout) = InOutParams::new(cells, params_len, results_len) else {
-            panic!("todo")
-        };
+        let inout = self.host_inout(callee_start, callee_end, params_len, results_len);
         let control = match caller {
             Some((ip, _, instance)) => ReturnCallHost::Continue((ip, caller_sp, instance)),
             None => ReturnCallHost::Break(caller_sp),
@@ -1094,13 +1096,36 @@ impl ValueStack {
         let callee_end = callee_start.add(callee_size)?;
         self.grow_if_needed(callee_end)?;
         let sp = self.sp(caller_start);
-        let Some(cells) = self.cells_from_to(callee_start, callee_end) else {
+        let inout = self.host_inout(callee_start, callee_end, params_len, results_len);
+        Ok((sp, inout))
+    }
+
+    /// Returns the [`InOutParams`] window `cells[start..end]` for a host function call.
+    ///
+    /// # Note
+    ///
+    /// Both `end - start` and `max(len_params, len_results)` denote the callee size,
+    /// and [`ValueStack::grow_if_needed`] has already been called for `end`.
+    fn host_inout(
+        &mut self,
+        start: SpOffset,
+        end: SpOffset,
+        len_params: usize,
+        len_results: usize,
+    ) -> InOutParams<'_> {
+        debug_assert_eq!(
+            end.into_inner() - start.into_inner(),
+            len_params.max(len_results)
+        );
+        let Some(cells) = self.cells_from_to(start, end) else {
             unsafe { unreachable_unchecked!("must fit slice after `grow_if_needed` operation") }
         };
-        let Ok(inout) = InOutParams::new(cells, params_len, results_len) else {
-            panic!("todo")
+        let Ok(inout) = InOutParams::new(cells, len_params, len_results) else {
+            unsafe {
+                unreachable_unchecked!("host frame cells are sized as max(len_params, len_results)")
+            }
         };
-        Ok((sp, inout))
+        inout
     }
 
     /// Adjusts `self` for a normal function call.
@@ -1369,9 +1394,11 @@ impl CallStack {
     ///
     /// # Note
     ///
-    /// - If `instance` is `Some` it refers to the _callee_ instance of the tail call
-    ///   which is required to be different from the currently used instance.
+    /// - If `instance` is `Some` it refers to the _callee_ instance of the tail call.
+    ///   Only an actual instance change creates a restoration obligation.
     /// - If `instance` is `None` the callee shares the currently used instance.
+    ///
+    /// This mirrors the contract of [`CallStack::push`].
     #[inline(always)]
     fn replace(&mut self, callee_ip: Ip, instance: Option<Inst>) -> Result<SpOffset, TrapCode> {
         let Some(caller_frame) = self.frames.last_mut() else {
@@ -1380,12 +1407,11 @@ impl CallStack {
         let start = caller_frame.start;
         caller_frame.ip = callee_ip;
         if let Some(callee_instance) = instance {
-            debug_assert!(self.instance != Some(callee_instance));
             // The replaced frame's restoration obligation is carried over unchanged.
             // However, if the replaced frame has no such obligation yet, its caller
             // runs in the currently used instance which must be restored when the
             // new frame returns since the callee continues in a different instance.
-            if caller_frame.instance.is_none() {
+            if caller_frame.instance.is_none() && self.instance != Some(callee_instance) {
                 caller_frame.instance = self.instance;
             }
             self.instance = Some(callee_instance);

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

@@ -21,16 +21,18 @@ use crate::{
         ShiftAmount,
     },
     engine::{
+        CodeView,
         DedupFuncType,
         FuncEntry,
+        InOutParams,
         executor::{
             LoadFromCellsByValue,
             StoreToCells,
             handler::{Break, Control, DoneReason, args::Args},
         },
         utils::unreachable_unchecked,
     },
-    func::{FuncEntity, HostFuncEntity},
+    func::{FuncEntity, HostFuncEntity, Trampoline},
     instance::{DataAddr, ElemAddr, FuncAddr, GlobalAddr, InstanceEntity, MemoryAddr, TableAddr},
     ir::{self, Address, BoundedSlotSpan, Local, Offset, Offset16, Slot, SlotAndReg, SlotSpan},
     memory::DataSegmentEntity,
@@ -816,6 +818,38 @@ pub fn return_call_func_entry(
     Control::Continue((callee_ip, callee_sp))
 }
 
+/// Invokes the host function behind `trampoline`.
+///
+/// Returns the host provided error if the host function trapped.
+///
+/// # Note
+///
+/// Takes `store` and `code` instead of the whole [`VmState`] since `inout` already
+/// borrows its [`Stack`].
+///
+/// [`Stack`]: crate::engine::executor::Stack
+#[inline]
+fn invoke_host(
+    store: &mut PrunedStore,
+    code: &mut CodeView,
+    trampoline: Trampoline,
+    instance: Option<Inst>,
+    inout: InOutParams<'_>,
+    call_hooks: CallHooks,
+) -> Result<(), Error> {
+    match store.call_host_func(trampoline, instance, inout, call_hooks) {
+        Ok(()) => {}
+        Err(StoreError::External(error)) => return Err(error),
+        Err(StoreError::Internal(error)) => unsafe {
+            unreachable_unchecked!(
+                "internal interpreter error while executing host function: {error}"
+            )
+        },
+    }
+    code.refresh();
+    Ok(())
+}
+
 pub fn call_host(
     state: &mut VmState,
     func: Func,
@@ -831,21 +865,16 @@ pub fn call_host(
         .stack
         .prepare_host_frame(caller_ip, params, host_func.len_result_cells())
         .into_control()?;
-    match state
-        .store
-        .call_host_func(trampoline, instance, inout, call_hooks)
-    {
-        Ok(()) => {}
-        Err(StoreError::External(error)) => {
-            done!(state, DoneReason::host_error(error, func, params.span()))
-        }
-        Err(StoreError::Internal(error)) => unsafe {
-            unreachable_unchecked!(
-                "internal interpreter error while executing host function: {error}"
-            )
-        },
+    if let Err(error) = invoke_host(
+        state.store,
+        &mut state.code,
+        trampoline,
+        instance,
+        inout,
+        call_hooks,
+    ) {
+        done!(state, DoneReason::host_error(error, func, params.span()))
     }
-    state.code.refresh();
     Control::Continue(sp)
 }
 
@@ -862,27 +891,22 @@ pub fn return_call_host(
         .stack
         .return_prepare_host_frame(params, host_func.len_result_cells(), instance)
         .into_control()?;
-    match state
-        .store
-        .call_host_func(trampoline, Some(instance), inout, CallHooks::Call)
-    {
-        Ok(()) => {}
-        Err(StoreError::External(error)) => {
-            // Note: we won't allow resumption in case the execution would
-            //       have returned with this the host function tail call.
-            let reason = match control {
-                Control::Continue(_) => DoneReason::host_error(error, func, params.span()),
-                Control::Break(_) => DoneReason::error(error),
-            };
-            done!(state, reason)
-        }
-        Err(StoreError::Internal(error)) => unsafe {
-            unreachable_unchecked!(
-                "internal interpreter error while executing host function: {error}"
-            )
-        },
+    if let Err(error) = invoke_host(
+        state.store,
+        &mut state.code,
+        trampoline,
+        Some(instance),
+        inout,
+        CallHooks::Call,
+    ) {
+        // Note: we won't allow resumption in case the execution would
+        //       have returned with this the host function tail call.
+        let reason = match control {
+            Control::Continue(_) => DoneReason::host_error(error, func, params.span()),
+            Control::Break(_) => DoneReason::error(error),
+        };
+        done!(state, reason)
     }
-    state.code.refresh();
     match control {
         Control::Continue((ip, sp, instance)) => Control::Continue((ip, sp, instance)),
         Control::Break(sp) => done!(state, DoneReason::Return(sp)),
@@ -965,9 +989,8 @@ pub fn return_call_wasm_or_host(
     };
     // Hot path: tail-calling a Wasm function. See `call_wasm_or_host` for the shape.
     let callee_instance: Inst = resolve_instance(state.store, wasm_func.instance()).into();
-    let changed_instance = (callee_instance != instance).then_some(callee_instance);
     let (callee_ip, callee_sp) =
-        return_call_func_entry(state, params, wasm_func.func_entry(), changed_instance)?;
+        return_call_func_entry(state, params, wasm_func.func_entry(), Some(callee_instance))?;
     let (instance, mem0, mem0_len) =
         update_instance(state.store, instance, callee_instance, mem0, mem0_len);
     Control::Continue((callee_ip, callee_sp, mem0, mem0_len, instance))