perf: borrow the executing function from its module instance

Every call and return between two functions cloned the callee's
`Shared<WasmFunction>` and dropped the previous one, two refcount updates
each way, and looked the function up in the store.

A module instance now keeps its own functions, which the store allocates
contiguously. The interpreter holds the instance for the whole run, so
the executor borrows the executing function and its module from it: a
call or return within the instance switches a reference, and a direct
call to one of the module's own functions skips the address table, the
host check and the owner check.

Execution that continues in another instance's function (a call through
an import, table or reference, a return, or an exception unwinding into
it) ends the run, and `InterpreterRuntime` resumes that frame with an
executor for its instance. Fuel and time budgets carry over, so a run
suspends at the same points as before.

Author: matthargett

crates/tinywasm/src/instance.rs

@@ -60,6 +60,11 @@ struct ModuleInstanceInner {
     id: ModuleInstanceId,
     type_addrs: Box<[TypeAddr]>,
     func_addrs: Box<[FuncAddr]>,
+    /// The module's own (non-imported) functions, in index order. The store allocates them
+    /// contiguously from `local_func_base`, so an executor can borrow any of them for a whole run.
+    local_funcs: Box<[Shared<WasmFunction>]>,
+    local_func_base: FuncAddr,
+    imported_funcs: u32,
     table_addrs: Box<[TableAddr]>,
     mem_addrs: Box<[MemAddr]>,
     #[cfg(feature = "std")]
@@ -84,6 +89,20 @@ impl ModuleInstance {
         self.0.func_addrs[addr as usize]
     }
 
+    /// The store address and body of the module's own function `idx`, or `None` for an import.
+    #[inline(always)]
+    pub(crate) fn local_func_by_index(&self, idx: FuncAddr) -> Option<(FuncAddr, &WasmFunction)> {
+        let local = idx.wrapping_sub(self.0.imported_funcs);
+        let func = self.0.local_funcs.get(local as usize)?;
+        Some((self.0.local_func_base + local, func))
+    }
+
+    /// The body of the function at store address `addr`, if this instance owns it.
+    #[inline(always)]
+    pub(crate) fn local_func(&self, addr: FuncAddr) -> Option<&WasmFunction> {
+        self.0.local_funcs.get(addr.wrapping_sub(self.0.local_func_base) as usize).map(|func| &**func)
+    }
+
     /// resolve a table address to the global store address
     #[inline]
     pub(crate) fn resolve_table_addr(&self, addr: TableAddr) -> TableAddr {
@@ -211,7 +230,9 @@ impl ModuleInstance {
         let id = store.next_module_instance_id();
         let mut addrs = crate::imports::ResolvedImports::new(store, module, &type_addrs, imports)?;
         let imported_funcs = addrs.funcs.len();
-        addrs.funcs.extend(store.init_funcs(&module.funcs, id, &module.func_type_idxs[imported_funcs..], &type_addrs));
+        let local_funcs = store.init_funcs(&module.funcs, id, &module.func_type_idxs[imported_funcs..], &type_addrs);
+        let local_func_base = local_funcs.start;
+        addrs.funcs.extend(local_funcs);
         addrs.tags.extend(store.init_tags(&module.tags, &type_addrs));
         let limiter = store.engine.config().resource_limiter.clone();
         if !module.skip_local_memory_allocation {
@@ -238,6 +259,9 @@ impl ModuleInstance {
             id,
             type_addrs,
             func_addrs: addrs.funcs.into_boxed_slice(),
+            local_funcs: module.funcs.clone(),
+            local_func_base,
+            imported_funcs: imported_funcs as u32,
             table_addrs: addrs.tables.into_boxed_slice(),
             mem_addrs: addrs.memories.into_boxed_slice(),
             #[cfg(feature = "std")]

crates/tinywasm/src/interpreter/executor/dispatch.rs

@@ -24,67 +24,76 @@ macro_rules! define_stable_dispatch {
     };
 }
 
-impl Executor<'_> {
+impl Executor<'_, '_> {
     instruction_handlers!(define_stable_dispatch);
 
+    /// Runs until the call completes (`None`) or continues in another module instance's frame.
     #[inline(always)]
-    pub(crate) fn run_to_completion(mut self) -> Result<()> {
+    pub(crate) fn run_to_completion(mut self) -> Result<Option<CallFrame>> {
         let mut instr_ptr = self.cf.instr_ptr;
         loop {
             match Self::exec_step(&mut self, instr_ptr)?.next_instr_ptr() {
                 Some(next_instr_ptr) => instr_ptr = next_instr_ptr,
-                None => return cold!(Ok(())),
+                None => return cold!(Ok(self.left())),
             }
         }
     }
 
+    /// Runs `chunk_left` instructions, then checkpoint by checkpoint until `time_budget` has
+    /// elapsed since `start`.
     #[cfg(feature = "std")]
     #[inline(always)]
-    pub(crate) fn run_with_time_budget(mut self, time_budget: core::time::Duration) -> Result<ExecState> {
-        use crate::std::time::Instant;
-
-        if time_budget.is_zero() {
-            return Ok(ExecState::Suspended(self.cf));
-        }
-        let start = Instant::now();
+    pub(crate) fn run_with_time_budget(
+        mut self,
+        start: crate::std::time::Instant,
+        time_budget: core::time::Duration,
+        mut chunk_left: u32,
+    ) -> Result<RunEnd> {
         let mut instr_ptr = self.cf.instr_ptr;
         loop {
-            for _ in 0..CHECKPOINT_INTERVAL {
+            while chunk_left != 0 {
+                chunk_left -= 1;
                 match Self::exec_step(&mut self, instr_ptr)?.next_instr_ptr() {
                     Some(next_instr_ptr) => instr_ptr = next_instr_ptr,
-                    None => return Ok(ExecState::Completed),
+                    None => return Ok(self.run_end(chunk_left)),
                 }
             }
+            chunk_left = CHECKPOINT_INTERVAL;
 
             if start.elapsed() >= time_budget {
                 self.cf.instr_ptr = instr_ptr;
-                return Ok(ExecState::Suspended(self.cf));
+                return Ok(RunEnd::State(ExecState::Suspended(self.cf)));
             }
         }
     }
 
+    /// Runs `chunk_left` instructions, then checkpoint by checkpoint until the store's fuel is out.
     #[inline(always)]
-    pub(crate) fn run_with_fuel(mut self, fuel: u32) -> Result<ExecState> {
+    pub(crate) fn run_with_fuel(mut self, mut chunk_left: u32) -> Result<RunEnd> {
         self.fuel_metered = true;
-        self.store.execution_fuel = fuel;
-        if self.store.execution_fuel == 0 {
-            return Ok(ExecState::Suspended(self.cf));
-        }
-
         let mut instr_ptr = self.cf.instr_ptr;
         loop {
-            for _ in 0..CHECKPOINT_INTERVAL {
+            while chunk_left != 0 {
+                chunk_left -= 1;
                 match Self::exec_step(&mut self, instr_ptr)?.next_instr_ptr() {
                     Some(next_instr_ptr) => instr_ptr = next_instr_ptr,
-                    None => return Ok(ExecState::Completed),
+                    None => return Ok(self.run_end(chunk_left)),
                 }
             }
+            chunk_left = CHECKPOINT_INTERVAL;
 
             self.store.execution_fuel = self.store.execution_fuel.saturating_sub(CHECKPOINT_INTERVAL);
             if self.store.execution_fuel == 0 {
                 self.cf.instr_ptr = instr_ptr;
-                return Ok(ExecState::Suspended(self.cf));
+                return Ok(RunEnd::State(ExecState::Suspended(self.cf)));
             }
         }
     }
+
+    fn run_end(&self, chunk_left: u32) -> RunEnd {
+        match self.left() {
+            Some(frame) => RunEnd::Left(frame, chunk_left),
+            None => RunEnd::State(ExecState::Completed),
+        }
+    }
 }

crates/tinywasm/src/interpreter/executor/dispatch_become.rs

@@ -3,8 +3,9 @@ use super::*;
 struct Unbudgeted;
 struct Bounded;
 
-type UnbudgetedHandler = for<'store> fn(&mut Executor<'store>, usize, Instruction) -> ExecResult<()>;
-type BoundedHandler = for<'store> fn(&mut Executor<'store>, usize, Instruction, u32) -> ExecResult<()>;
+type UnbudgetedHandler = for<'store, 'module> fn(&mut Executor<'store, 'module>, usize, Instruction) -> ExecResult<()>;
+type BoundedHandler =
+    for<'store, 'module> fn(&mut Executor<'store, 'module>, usize, Instruction, u32) -> ExecResult<()>;
 
 #[cold]
 #[inline(never)]
@@ -28,7 +29,7 @@ macro_rules! define_unbudgeted_tail_dispatch {
         $(
             #[allow(non_snake_case, unreachable_code, unused_imports, unused_macros, unused_variables)]
             fn $variant(
-                $executor: &mut Executor<'_>,
+                $executor: &mut Executor<'_, '_>,
                 $instr_ptr: usize,
                 instruction: Instruction,
             ) -> ExecResult<()> {
@@ -78,7 +79,7 @@ macro_rules! define_bounded_tail_dispatch {
         $(
             #[allow(non_snake_case, unreachable_code, unused_imports, unused_macros, unused_variables)]
             fn $variant(
-                $executor: &mut Executor<'_>,
+                $executor: &mut Executor<'_, '_>,
                 $instr_ptr: usize,
                 instruction: Instruction,
                 instructions_until_checkpoint: u32,
@@ -103,7 +104,11 @@ macro_rules! define_bounded_tail_dispatch {
                         match $flow.next_instr_ptr() {
                             Some(next_instr_ptr) => $dispatch_next!(next_instr_ptr),
                             None => return cold!({
-                                $executor.completed = true;
+                                if $executor.left {
+                                    $executor.chunk_left = instructions_until_checkpoint;
+                                } else {
+                                    $executor.completed = true;
+                                }
                                 Ok(())
                             }),
                         }
@@ -130,62 +135,76 @@ impl Unbudgeted {
 impl Bounded {
     instruction_handlers!(define_bounded_tail_dispatch);
 
+    /// Runs up to `chunk_left` (at least 1) instructions from `executor.cf`.
     #[inline(always)]
-    fn run(executor: &mut Executor<'_>) -> ExecResult<()> {
+    fn run(executor: &mut Executor<'_, '_>, chunk_left: u32) -> ExecResult<()> {
         let instr_ptr = executor.cf.instr_ptr;
         let instruction = executor.func.instructions[instr_ptr];
         let handler = Self::handler_for(instruction.opcode());
-        handler(executor, instr_ptr, instruction, CHECKPOINT_INTERVAL - 1)
+        handler(executor, instr_ptr, instruction, chunk_left - 1)
     }
 }
 
-impl<'store> Executor<'store> {
+impl Executor<'_, '_> {
+    /// Runs until the call completes (`None`) or continues in another module instance's frame.
     #[inline(always)]
-    pub(crate) fn run_to_completion(mut self) -> Result<()> {
+    pub(crate) fn run_to_completion(mut self) -> Result<Option<CallFrame>> {
         let instr_ptr = self.cf.instr_ptr;
         let instruction = self.func.instructions[instr_ptr];
         let handler = Unbudgeted::handler_for(instruction.opcode());
-        Ok(handler(&mut self, instr_ptr, instruction)?)
+        handler(&mut self, instr_ptr, instruction)?;
+        Ok(self.left())
     }
 
+    /// Runs `chunk_left` instructions, then checkpoint by checkpoint until `time_budget` has
+    /// elapsed since `start`.
     #[cfg(feature = "std")]
     #[inline(always)]
-    pub(crate) fn run_with_time_budget(mut self, time_budget: core::time::Duration) -> Result<ExecState> {
-        use crate::std::time::Instant;
-
-        if time_budget.is_zero() {
-            return Ok(ExecState::Suspended(self.cf));
-        }
-        let start = Instant::now();
-
+    pub(crate) fn run_with_time_budget(
+        mut self,
+        start: crate::std::time::Instant,
+        time_budget: core::time::Duration,
+        mut chunk_left: u32,
+    ) -> Result<RunEnd> {
         loop {
-            Bounded::run(&mut self)?;
-            if self.completed {
-                return cold!(Ok(ExecState::Completed));
+            if chunk_left != 0 {
+                Bounded::run(&mut self, chunk_left)?;
+                if let Some(end) = self.run_end() {
+                    return Ok(end);
+                }
             }
+            chunk_left = CHECKPOINT_INTERVAL;
             if start.elapsed() >= time_budget {
-                return cold!(Ok(ExecState::Suspended(self.cf)));
+                return cold!(Ok(RunEnd::State(ExecState::Suspended(self.cf))));
             }
         }
     }
 
+    /// Runs `chunk_left` instructions, then checkpoint by checkpoint until the store's fuel is out.
     #[inline(always)]
-    pub(crate) fn run_with_fuel(mut self, fuel: u32) -> Result<ExecState> {
+    pub(crate) fn run_with_fuel(mut self, mut chunk_left: u32) -> Result<RunEnd> {
         self.fuel_metered = true;
-        self.store.execution_fuel = fuel;
-        if self.store.execution_fuel == 0 {
-            return Ok(ExecState::Suspended(self.cf));
-        }
-
         loop {
-            Bounded::run(&mut self)?;
-            if self.completed {
-                return cold!(Ok(ExecState::Completed));
+            if chunk_left != 0 {
+                Bounded::run(&mut self, chunk_left)?;
+                if let Some(end) = self.run_end() {
+                    return Ok(end);
+                }
             }
+            chunk_left = CHECKPOINT_INTERVAL;
             self.store.execution_fuel = self.store.execution_fuel.saturating_sub(CHECKPOINT_INTERVAL);
             if self.store.execution_fuel == 0 {
-                return cold!(Ok(ExecState::Suspended(self.cf)));
+                return cold!(Ok(RunEnd::State(ExecState::Suspended(self.cf))));
             }
         }
     }
+
+    /// How a bounded chain that stopped ended, unless it stopped at a checkpoint.
+    #[inline(always)]
+    fn run_end(&self) -> Option<RunEnd> {
+        if self.completed {
+            return cold!(Some(RunEnd::State(ExecState::Completed)));
+        }
+        self.left().map(|frame| RunEnd::Left(frame, self.chunk_left))
+    }
 }