Cache FuncEntry in FuncEntity and make use of this caching in the executor (#1996)

* add FuncEntry to FuncEntity and use it in executor

* fix performance call_internal regressions

* fix stackoverflow while retaining performance wins

* fix broken intra doc links

Author: Robbepop

crates/wasmi/src/engine/code_map/mod.rs

@@ -550,6 +550,49 @@ pub struct FuncEntry {
     state: AtomicU8,
 }
 
+/// A stable address of a [`FuncEntry`] in the engine's append-only [`CodeMap`].
+///
+/// Holding this instead of an `EngineFunc` lets callers reach the [`FuncEntry`] directly, skipping
+/// the [`CodeMap`] lookup. It stores the address as an exposed-provenance `usize` (the same
+/// convention the `call_internal` handler uses for its baked pointer) which keeps it `Send + Sync`
+/// without an `unsafe impl`.
+///
+/// [`CodeMap`]: crate::engine::CodeMap
+#[derive(Debug, Copy, Clone)]
+pub struct FuncEntryPtr(usize);
+
+impl From<usize> for FuncEntryPtr {
+    /// Creates a [`FuncEntryPtr`] from a previously exposed [`FuncEntry`] address.
+    ///
+    /// Used by the `call_internal` handlers to recover the address baked into the bytecode.
+    fn from(addr: usize) -> Self {
+        Self(addr)
+    }
+}
+
+impl FuncEntryPtr {
+    /// Creates a [`FuncEntryPtr`] from `entry`, exposing its provenance.
+    pub fn new(entry: &FuncEntry) -> Self {
+        Self(ptr::from_ref(entry).expose_provenance())
+    }
+
+    /// Returns a shared reference to the pointed-to [`FuncEntry`].
+    ///
+    /// # Safety
+    ///
+    /// The engine owning the [`FuncEntry`] must outlive `'a`. This holds for a [`FuncEntryPtr`]
+    /// stored in a [`WasmFuncEntity`] since its [`Func`] cannot outlive the owning engine. The
+    /// [`FuncEntry`] is only ever accessed through shared references (interior mutation is guarded
+    /// by atomics), so no aliasing `&mut` is ever formed.
+    ///
+    /// [`WasmFuncEntity`]: crate::func::WasmFuncEntity
+    /// [`Func`]: crate::Func
+    pub unsafe fn get<'a>(self) -> &'a FuncEntry {
+        // SAFETY: guaranteed by the caller (see above); provenance was exposed in `new`.
+        unsafe { &*ptr::with_exposed_provenance::<FuncEntry>(self.0) }
+    }
+}
+
 impl Drop for FuncEntry {
     fn drop(&mut self) {
         match self.state.load(Ordering::Acquire) {

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

@@ -221,7 +221,7 @@ impl Args {
     }
 
     /// Calls `func` with `params` on `instance` with `state` using `self`.
-    #[inline]
+    #[inline(always)]
     pub fn call_func_entry(
         &mut self,
         state: &mut VmState,
@@ -234,7 +234,7 @@ impl Args {
     }
 
     /// Tail-calls `func` with `params` on `instance` with `state` using `self`.
-    #[inline]
+    #[inline(always)]
     pub fn return_call_func_entry(
         &mut self,
         state: &mut VmState,

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

@@ -19,7 +19,7 @@ use crate::{
     TrapCode,
     core::{CoreTable, RawRef, ReadAs, WriteAs, wasm},
     engine::{
-        FuncEntry,
+        FuncEntryPtr,
         eval,
         executor::handler::{
             Control,
@@ -33,7 +33,7 @@ use crate::{
     ir::{self, BoundedSlotSpan},
     store::StoreError,
 };
-use core::{cmp, ptr};
+use core::cmp;
 
 fn identity<T>(value: T) -> T {
     value
@@ -205,14 +205,9 @@ execution_handler! {
     ) -> Done = {
         let mut args = Args::from_parts(ip, sp, mem0, mem0_len, instance, ireg, freg32, freg64);
         let crate::ir::decode::CallInternal { params, func } = unsafe { args.decode_op() };
-        // Safety:
-        //
-        // `func` is the exposed address of a `FuncEntity` in the engine's append-only `CodeMap`.
-        //
-        //  - the used `with_exposed_provenance` recovers the original provenance.
-        //  - its allocation is never moved or freed while this bytecode runs.
-        //  - the `FuncEntry` type mutates only guarded by lock-free atomics.
-        let func = unsafe { &*ptr::with_exposed_provenance::<FuncEntry>(usize::from(func)) };
+        // SAFETY: `func` is the exposed address of a `FuncEntry` in the engine's append-only
+        //         `CodeMap`, baked into the bytecode; it stays valid while this bytecode runs.
+        let func = unsafe { FuncEntryPtr::from(usize::from(func)).get() };
         args.call_func_entry(state, func, params, None)?;
         dispatch!(state, args)
     }
@@ -287,14 +282,8 @@ execution_handler! {
     ) -> Done = {
         let mut args = Args::from_parts(ip, sp, mem0, mem0_len, instance, ireg, freg32, freg64);
         let crate::ir::decode::ReturnCallInternal { params, func } = unsafe { args.decode_op() };
-        // Safety:
-        //
-        // `func` is the exposed address of a `FuncEntity` in the engine's append-only `CodeMap`.
-        //
-        //  - the used `with_exposed_provenance` recovers the original provenance.
-        //  - its allocation is never moved or freed while this bytecode runs.
-        //  - the `FuncEntry` type mutates only guarded by lock-free atomics.
-        let func = unsafe { &*ptr::with_exposed_provenance::<FuncEntry>(usize::from(func)) };
+        // SAFETY: see `call_internal`.
+        let func = unsafe { FuncEntryPtr::from(usize::from(func)).get() };
         args.return_call_func_entry(state, func, params, None)?;
         dispatch!(state, args)
     }

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

@@ -22,7 +22,6 @@ use crate::{
     },
     engine::{
         DedupFuncType,
-        EngineFunc,
         FuncEntry,
         executor::{
             LoadFromCellsByValue,
@@ -85,22 +84,6 @@ macro_rules! compile_or_get_func_entry {
     }};
 }
 
-pub fn compile_or_get_func(state: &mut VmState, func: EngineFunc) -> Result<(Ip, u16, u16), Error> {
-    let Some(func_entry) = state.code.entry(func) else {
-        unreachable!("missing function entry at: {func:?}")
-    };
-    compile_or_get_func_entry(state, func_entry)
-}
-
-macro_rules! compile_or_get_func {
-    ($state:expr, $func:expr) => {{
-        match $crate::engine::executor::handler::utils::compile_or_get_func($state, $func) {
-            Ok((ip, len_local_slots, len_stack_slots)) => (ip, len_local_slots, len_stack_slots),
-            Err(error) => done!($state, DoneReason::error(error)),
-        }
-    }};
-}
-
 macro_rules! trap {
     ($trap_code:expr) => {{
         return $crate::engine::executor::handler::Control::Break(
@@ -789,7 +772,7 @@ pub fn update_instance(
     (new_instance, mem0, mem0_len)
 }
 
-#[inline]
+#[inline(always)]
 pub fn call_func_entry(
     state: &mut VmState,
     caller_ip: Ip,
@@ -812,30 +795,7 @@ pub fn call_func_entry(
     Control::Continue((callee_ip, callee_sp))
 }
 
-#[inline(never)]
-pub fn call_wasm(
-    state: &mut VmState,
-    caller_ip: Ip,
-    params: BoundedSlotSpan,
-    func: EngineFunc,
-    instance: Option<Inst>,
-) -> Control<(Ip, Sp), Break> {
-    let (callee_ip, len_local_slots, len_stack_slots) = compile_or_get_func!(state, func);
-    let callee_sp = state
-        .stack
-        .push_frame(
-            Some(caller_ip),
-            callee_ip,
-            params,
-            len_local_slots,
-            len_stack_slots,
-            instance,
-        )
-        .into_control()?;
-    Control::Continue((callee_ip, callee_sp))
-}
-
-#[inline]
+#[inline(always)]
 pub fn return_call_func_entry(
     state: &mut VmState,
     params: BoundedSlotSpan,
@@ -856,27 +816,6 @@ pub fn return_call_func_entry(
     Control::Continue((callee_ip, callee_sp))
 }
 
-#[inline(never)]
-pub fn return_call_wasm(
-    state: &mut VmState,
-    params: BoundedSlotSpan,
-    func: EngineFunc,
-    instance: Option<Inst>,
-) -> Control<(Ip, Sp), Break> {
-    let (callee_ip, len_local_slots, len_stack_slots) = compile_or_get_func!(state, func);
-    let callee_sp = state
-        .stack
-        .replace_frame(
-            callee_ip,
-            params,
-            len_local_slots,
-            len_stack_slots,
-            instance,
-        )
-        .into_control()?;
-    Control::Continue((callee_ip, callee_sp))
-}
-
 pub fn call_host(
     state: &mut VmState,
     func: Func,
@@ -950,6 +889,7 @@ pub fn return_call_host(
     }
 }
 
+#[inline]
 #[expect(clippy::too_many_arguments)]
 pub fn call_wasm_or_host(
     state: &mut VmState,
@@ -965,39 +905,41 @@ pub fn call_wasm_or_host(
     //         (indirect calls); the reference is only used to copy out the callee data below,
     //         before any store mutation, so it never aliases a `&mut` into the funcs arena.
     let func_entity = unsafe { func_entity.as_ref() };
-    let next_state = match func_entity {
-        FuncEntity::Wasm(wasm_func) => {
-            let func = wasm_func.func_body();
-            let callee_instance = *wasm_func.instance();
-            let callee_instance: Inst = resolve_instance(state.store, &callee_instance).into();
-            let (callee_ip, callee_sp) =
-                call_wasm(state, caller_ip, params, func, Some(callee_instance))?;
-            let (instance, mem0, mem0_len) =
-                update_instance(state.store, instance, callee_instance, mem0, mem0_len);
-            (callee_ip, callee_sp, mem0, mem0_len, instance)
-        }
+    let wasm_func = match func_entity {
+        FuncEntity::Wasm(wasm_func) => wasm_func,
         FuncEntity::Host(host_func) => {
-            let host_func = *host_func;
             let sp = call_host(
                 state,
                 func,
                 Some(caller_ip),
-                host_func,
+                *host_func,
                 params,
                 Some(instance),
                 CallHooks::Call,
             )?;
-            // Host functions may re-enter WASM (e.g. calling cabi_realloc)
-            // which can trigger memory.grow, invalidating the cached mem0
-            // pointer. Re-extract to avoid stale pointer dereference.
+            // Note: host functions may grow memories, invalidating the cached `(memory 0)`.
+            //       Therefore, it is required to re-extract `(memory 0)` to avoid a stale cache.
             let (mem0, mem0_len) = extract_mem0(state.store, instance);
-            (caller_ip, sp, mem0, mem0_len, instance)
+            return Control::Continue((caller_ip, sp, mem0, mem0_len, instance));
         }
     };
-    Control::Continue(next_state)
+    // Hot path: calling a Wasm function. Uses the cached `FuncEntry` and the same inlined
+    // machinery as `call_internal`, differing only in the possible instance switch.
+    let callee_instance: Inst = resolve_instance(state.store, wasm_func.instance()).into();
+    let (callee_ip, callee_sp) = call_func_entry(
+        state,
+        caller_ip,
+        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))
 }
 
 /// Tail-call (`return_call`) twin of [`call_wasm_or_host`].
+#[inline]
 pub fn return_call_wasm_or_host(
     state: &mut VmState,
     func: Func,
@@ -1009,25 +951,22 @@ pub fn return_call_wasm_or_host(
 ) -> Control<(Ip, Sp, Mem0Ptr, Mem0Len, Inst), Break> {
     // SAFETY: see `call_wasm_or_host`.
     let func_entity = unsafe { func_entity.as_ref() };
-    let (callee_ip, sp, new_instance) = match func_entity {
-        FuncEntity::Wasm(wasm_func) => {
-            let wasm_func_body = wasm_func.func_body();
-            let callee_instance = *wasm_func.instance();
-            let callee_instance: Inst = resolve_instance(state.store, &callee_instance).into();
-            let changed_instance = match callee_instance != instance {
-                true => Some(callee_instance),
-                false => None,
-            };
-            let (callee_ip, callee_sp) =
-                return_call_wasm(state, params, wasm_func_body, changed_instance)?;
-            (callee_ip, callee_sp, callee_instance)
-        }
+    let wasm_func = match func_entity {
+        FuncEntity::Wasm(wasm_func) => wasm_func,
         FuncEntity::Host(host_func) => {
-            let host_func = *host_func;
-            return_call_host(state, func, host_func, params, instance)?
+            let (callee_ip, sp, new_instance) =
+                return_call_host(state, func, *host_func, params, instance)?;
+            let (instance, mem0, mem0_len) =
+                update_instance(state.store, instance, new_instance, mem0, mem0_len);
+            return Control::Continue((callee_ip, sp, mem0, mem0_len, instance));
         }
     };
+    // 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)?;
     let (instance, mem0, mem0_len) =
-        update_instance(state.store, instance, new_instance, mem0, mem0_len);
-    Control::Continue((callee_ip, sp, mem0, mem0_len, instance))
+        update_instance(state.store, instance, callee_instance, mem0, mem0_len);
+    Control::Continue((callee_ip, callee_sp, mem0, mem0_len, instance))
 }

crates/wasmi/src/engine/mod.rs

@@ -15,6 +15,7 @@ mod utils;
 pub(crate) use self::translator::ValidatingFuncTranslator;
 pub(crate) use self::{
     block_type::BlockType,
+    code_map::{FuncEntry, FuncEntryPtr},
     executor::{
         InOutParams,
         InOutResults,
@@ -61,7 +62,6 @@ use crate::{
     Func,
     FuncType,
     StoreContextMut,
-    engine::code_map::FuncEntry,
     module::{FuncIdx, ModuleHeader},
 };
 use alloc::{