Resolve debug_function_index in O(1) instead of scanning every function

Instance::debug_function_index previously scanned every function in the
instance, comparing VMFuncRef pointers, to find the one matching a given
Func. That makes a whole-store debug snapshot (which calls this once per
captured funcref) quadratic overall, which was raised in review.

Both an instance's imported-function table and its defined-function
funcref table are contiguous, fixed-stride arrays at statically known
VMContext offsets, so a VMFuncRef pointer's position in either array is
computable directly via pointer arithmetic instead of a linear scan.
Imported functions are indexed by FuncIndex directly. Defined ("escaped")
functions are indexed by a compact FuncRefIndex that is *not* assigned in
FuncIndex order (slots are handed out in the order functions are
discovered to escape during module translation, e.g. export declaration
order), so resolving a FuncRefIndex back to a FuncIndex needs an explicit
reverse table rather than arithmetic. That table only depends on compiled
module metadata, not on any particular Instance, so it's built once,
lazily, and cached on Module -- shared by every Instance and every debug
snapshot of that module, rather than rebuilt per lookup.

Adds a regression test with functions placed into a table out of index
order, to pin down that the reverse table is used correctly rather than
assuming (incorrectly) that escape order tracks function-index order.

Author: smarcd

crates/wasmtime/src/runtime/debug.rs

@@ -25,7 +25,7 @@ pub use wasmtime_environ::ModulePC;
 use wasmtime_environ::{
     DefinedFuncIndex, EntityIndex, FrameInstPos, FrameStackShape, FrameStateSlot,
     FrameStateSlotOffset, FrameTableBreakpointData, FrameTableDescriptorIndex, FrameValType,
-    FuncIndex, FuncKey, GlobalIndex, MemoryIndex, TableIndex, TagIndex, Trap,
+    FuncIndex, FuncKey, GlobalIndex, MemoryIndex, PtrSize, TableIndex, TagIndex, Trap,
 };
 use wasmtime_unwinder::{Frame, FrameCursor};
 
@@ -321,6 +321,10 @@ impl Instance {
     /// This is the inverse of [`Instance::debug_function`]. It is only
     /// available when guest debugging is enabled for the store's engine, and
     /// returns `None` if `func` is not one of this instance's functions.
+    ///
+    /// This runs in constant time: it locates `func`'s `VMFuncRef` directly
+    /// via pointer arithmetic against this instance's `VMContext` layout,
+    /// rather than scanning every function in the instance.
     pub fn debug_function_index(
         &self,
         mut store: impl AsContextMut,
@@ -331,17 +335,36 @@ impl Instance {
             return None;
         }
 
-        let func_ref = func.vm_func_ref(store);
-        let module = self._module(store).env_module();
-        for index in 0..u32::try_from(module.functions.len()).ok()? {
-            let candidate = self
-                .debug_export(store, FuncIndex::from_bits(index).into())?
-                .into_func()?;
-            if candidate.vm_func_ref(store) == func_ref {
-                return Some(index);
-            }
+        let func_ref = func.vm_func_ref(store).as_ptr() as usize;
+        let module = self._module(store);
+        let offsets = module.offsets();
+        let stride = usize::from(offsets.ptr.vm_func_ref().size());
+        let vmctx = store.instance(self.id()).vmctx().as_ptr() as usize;
+
+        // Imported functions live in this instance's `imported_functions`
+        // array, indexed directly by `FuncIndex`.
+        let imports_begin =
+            vmctx + usize::try_from(offsets.vmctx_imported_functions_begin()).ok()?;
+        if let Some(index) = array_index_of(
+            func_ref,
+            imports_begin,
+            offsets.num_imported_functions,
+            stride,
+        ) {
+            return Some(index);
         }
-        None
+
+        // Defined ("escaped") functions live in a separate `func_refs`
+        // array, indexed by a compact `FuncRefIndex` that isn't in
+        // `FuncIndex` order; translate it back via the module's cached
+        // reverse table.
+        let func_refs_begin = vmctx + usize::try_from(offsets.vmctx_func_refs_begin()).ok()?;
+        let position =
+            array_index_of(func_ref, func_refs_begin, offsets.num_escaped_funcs, stride)?;
+        module
+            .debug_func_ref_to_func_index()
+            .get(position as usize)
+            .map(|index| index.as_u32())
     }
 
     /// Get access to a tag within this instance's tag index space.
@@ -384,6 +407,17 @@ impl Instance {
     }
 }
 
+/// If `ptr` falls within the `len`-element array of `stride`-byte entries
+/// starting at `begin`, returns its index within that array.
+fn array_index_of(ptr: usize, begin: usize, len: u32, stride: usize) -> Option<u32> {
+    let offset = ptr.checked_sub(begin)?;
+    if stride == 0 || offset % stride != 0 {
+        return None;
+    }
+    let index = u32::try_from(offset / stride).ok()?;
+    (index < len).then_some(index)
+}
+
 impl<'a, T> StoreContext<'a, T> {
     /// Return all breakpoints.
     pub fn breakpoints(self) -> Option<impl Iterator<Item = Breakpoint> + 'a> {

crates/wasmtime/src/runtime/module.rs

@@ -19,8 +19,9 @@ use core::ptr::NonNull;
 use std::{fs::File, path::Path};
 use wasmparser::{Parser, ValidPayload, Validator};
 use wasmtime_environ::{
-    CompiledFunctionsTable, CompiledModuleInfo, EntityIndex, FuncKey, HostPtr, ModuleTypes,
-    ObjectKind, StaticModuleIndex, TypeTrace, VMOffsets, VMSharedTypeIndex, WasmChecksum,
+    CompiledFunctionsTable, CompiledModuleInfo, EntityIndex, EntityRef, FuncIndex, FuncKey,
+    HostPtr, ModuleTypes, ObjectKind, StaticModuleIndex, TypeTrace, VMOffsets, VMSharedTypeIndex,
+    WasmChecksum, packed_option::ReservedValue,
 };
 mod registry;
 
@@ -151,6 +152,23 @@ struct ModuleInner {
     /// instantiated.
     memory_images: OnceLock<Option<ModuleMemoryImages>>,
 
+    /// Lazily-built reverse map from a function's `FuncRefIndex` (its
+    /// position within this module's `func_refs` array) back to its
+    /// `FuncIndex`.
+    ///
+    /// This exists solely to support `Instance::debug_function_index`,
+    /// which needs to invert a `VMFuncRef` pointer back to the function
+    /// index it came from. `FuncRefIndex`s are *not* assigned in `FuncIndex`
+    /// order (they're assigned in the order functions are discovered to
+    /// escape while translating the module, e.g. export declaration order),
+    /// so that inversion can't be done with arithmetic or a binary search;
+    /// it needs an explicit table. Building the table is a purely
+    /// module-level computation (it only depends on compiled module
+    /// metadata, not on any particular `Instance`), so it's cached here and
+    /// shared by every `Instance` created from this `Module`, rather than
+    /// being rebuilt per-instance or per-lookup.
+    debug_func_ref_to_func_index: OnceLock<Vec<FuncIndex>>,
+
     /// Flag indicating whether this module can be serialized or not.
     #[cfg(any(feature = "cranelift", feature = "winch"))]
     serializable: bool,
@@ -553,6 +571,7 @@ impl Module {
                 engine: engine.clone(),
                 code,
                 memory_images: OnceLock::new(),
+                debug_func_ref_to_func_index: OnceLock::new(),
                 module,
                 #[cfg(any(feature = "cranelift", feature = "winch"))]
                 serializable,
@@ -1118,6 +1137,25 @@ impl Module {
         &self.inner.offsets
     }
 
+    /// Returns a table mapping each escaping function's `FuncRefIndex`
+    /// (i.e. its position in the `func_refs` array within an instance's
+    /// `VMContext`) back to its `FuncIndex`.
+    ///
+    /// Built lazily on first use and cached for the lifetime of this
+    /// `Module`; see the doc comment on `ModuleInner::debug_func_ref_to_func_index`.
+    pub(crate) fn debug_func_ref_to_func_index(&self) -> &[FuncIndex] {
+        self.inner.debug_func_ref_to_func_index.get_or_init(|| {
+            let env_module = self.env_module();
+            let mut map = vec![FuncIndex::reserved_value(); env_module.num_escaped_funcs];
+            for (func_index, ty) in env_module.functions.iter() {
+                if ty.is_escaping() {
+                    map[ty.func_ref.index()] = func_index;
+                }
+            }
+            map
+        })
+    }
+
     /// Return the unique-within-Engine ID for this module.
     ///
     /// Allows distinguishing module identities when introspecting

tests/all/debug.rs

@@ -545,6 +545,46 @@ fn debug_function_index_requires_guest_debug() -> wasmtime::Result<()> {
     Ok(())
 }
 
+// `debug_function_index` resolves a function's position in its module by
+// inverting the pointer arithmetic used to locate a `VMFuncRef` (rather than
+// scanning every function looking for a pointer match). Functions are only
+// assigned a slot in that array once they "escape" the module (e.g. via
+// export, `ref.func`, or table placement), and slots are handed out in the
+// order escaping is *discovered* during translation -- not in function-index
+// order. This module deliberately declares four private functions but only
+// places them into the table in reverse order, so their escape order (and
+// thus their position in the funcref array) is the opposite of their
+// function index. If the lookup assumed escape order tracked function-index
+// order, this would resolve every index but the endpoints incorrectly.
+#[test]
+fn debug_function_index_with_non_monotonic_escape_order() -> wasmtime::Result<()> {
+    let (module, mut store) = get_module_and_store(
+        |_| {},
+        r#"
+        (module
+          (table 4 funcref)
+          (elem (i32.const 0) $f3 $f2 $f1 $f0)
+          (func $f0 (result i32) i32.const 0)
+          (func $f1 (result i32) i32.const 1)
+          (func $f2 (result i32) i32.const 2)
+          (func $f3 (result i32) i32.const 3))
+        "#,
+    )?;
+    let instance = Instance::new(&mut store, &module, &[])?;
+
+    for index in 0..4u32 {
+        let f = instance.debug_function(&mut store, index).unwrap();
+        assert_eq!(
+            instance.debug_function_index(&mut store, &f),
+            Some(index),
+            "function {index} must round-trip through debug_function_index \
+             even though its funcref slot was assigned out of index order"
+        );
+    }
+
+    Ok(())
+}
+
 #[test]
 #[cfg_attr(miri, ignore)]
 #[cfg(target_pointer_width = "64")] // Threads not supported on 32-bit systems.