pref: borrow instructions in tail call dispatch

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

@@ -7,7 +7,7 @@ macro_rules! define_stable_dispatch {
         fn exec_step($executor: &mut Self, $instr_ptr: usize) -> ExecResult<ExecFlow> {
             macro_rules! $dispatch_next {
                 ($next_instr_ptr:expr) => {{
-                    return Ok(ExecFlow::next($next_instr_ptr));
+                    return Ok(ExecFlow::Next($next_instr_ptr));
                 }};
             }
             macro_rules! $dispatch_flow {
@@ -19,7 +19,7 @@ macro_rules! define_stable_dispatch {
             match &$executor.func.instructions[$instr_ptr] {
                 $($variant $(($($arg),*))? $({ $($field),* })? => $body,)*
             }
-            Ok(ExecFlow::next($instr_ptr + 1))
+            Ok(ExecFlow::Next($instr_ptr + 1))
         }
     };
 }

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

@@ -3,7 +3,7 @@ use super::*;
 struct Unbudgeted;
 struct Bounded;
 
-type UnbudgetedHandler = for<'store> fn(&mut Executor<'store>, usize, Instruction) -> ExecResult<()>;
+type UnbudgetedHandler = for<'store> fn(&mut Executor<'store>, &[Instruction], usize, Instruction) -> ExecResult<()>;
 type BoundedHandler = for<'store> fn(&mut Executor<'store>, usize, Instruction, u32) -> ExecResult<()>;
 
 #[cold]
@@ -29,22 +29,27 @@ macro_rules! define_unbudgeted_tail_dispatch {
             #[allow(non_snake_case, unreachable_code, unused_imports, unused_macros, unused_variables)]
             fn $variant(
                 $executor: &mut Executor<'_>,
+                instructions: &[Instruction],
                 $instr_ptr: usize,
                 instruction: Instruction,
             ) -> ExecResult<()> {
                 macro_rules! $dispatch_next {
                     ($next_instr_ptr:expr) => {{
                         let next_instr_ptr = $next_instr_ptr;
-                        let instruction = $executor.func.instructions[next_instr_ptr];
+                        let instruction = instructions[next_instr_ptr];
                         let handler = Self::handler_for(instruction.opcode());
-                        become handler($executor, next_instr_ptr, instruction);
+                        become handler($executor, instructions, next_instr_ptr, instruction);
                     }};
                 }
                 macro_rules! $dispatch_flow {
                     ($flow:expr) => {{
-                        match $flow.next_instr_ptr() {
-                            Some(next_instr_ptr) => $dispatch_next!(next_instr_ptr),
-                            None => return cold!(Ok(())),
+                        match $flow {
+                            ExecFlow::Next(next_instr_ptr) => $dispatch_next!(next_instr_ptr),
+                            ExecFlow::Switch(next_instr_ptr) => {
+                                $executor.cf.instr_ptr = next_instr_ptr;
+                                return Ok(());
+                            },
+                            ExecFlow::Complete => return cold!({ $executor.completed = true; Ok(()) }),
                         }
                     }};
                 }
@@ -142,10 +147,17 @@ impl Bounded {
 impl<'store> Executor<'store> {
     #[inline(always)]
     pub(crate) fn run_to_completion(mut self) -> Result<()> {
-        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)?)
+        loop {
+            let func = self.func.clone();
+            let instructions = &func.instructions;
+            let instr_ptr = self.cf.instr_ptr;
+            let instruction = instructions[instr_ptr];
+            let handler = Unbudgeted::handler_for(instruction.opcode());
+            handler(&mut self, instructions, instr_ptr, instruction)?;
+            if self.completed {
+                return Ok(());
+            }
+        }
     }
 
     #[cfg(feature = "std")]

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

@@ -55,20 +55,19 @@ impl From<ExecError> for Error {
 }
 
 #[derive(Clone, Copy)]
-struct ExecFlow(usize);
+enum ExecFlow {
+    Next(usize),
+    Switch(usize),
+    Complete,
+}
 
 impl ExecFlow {
-    const COMPLETE: Self = Self(usize::MAX);
-
-    #[inline(always)]
-    fn next(instr_ptr: usize) -> Self {
-        debug_assert_ne!(instr_ptr, Self::COMPLETE.0);
-        Self(instr_ptr)
-    }
-
     #[inline(always)]
     fn next_instr_ptr(self) -> Option<usize> {
-        (self.0 != Self::COMPLETE.0).then_some(self.0)
+        match self {
+            Self::Next(instr_ptr) | Self::Switch(instr_ptr) => Some(instr_ptr),
+            Self::Complete => None,
+        }
     }
 }
 
@@ -560,7 +559,7 @@ impl<'store> Executor<'store> {
 
     fn throw_exception(&mut self, exception: ValueRef, protected_ip: usize) -> ExecResult<ExecFlow> {
         match self.dispatch_exception(exception, protected_ip)? {
-            Some(landing_pad) => Ok(ExecFlow::next(landing_pad)),
+            Some(flow) => Ok(flow),
             None => Err(Error::Exception(self.store.root_exception(exception)?).into()),
         }
     }
@@ -595,11 +594,12 @@ impl<'store> Executor<'store> {
         }
     }
 
-    fn dispatch_exception(&mut self, exception: ValueRef, mut protected_ip: usize) -> Result<Option<usize>, Trap> {
+    fn dispatch_exception(&mut self, exception: ValueRef, mut protected_ip: usize) -> Result<Option<ExecFlow>, Trap> {
         let object = self.store.state.gc.get(exception).ok_or(Trap::InvalidReference)?;
         let crate::store::GcObjectKind::Exception(tag_addr) = object.kind else {
             return Err(Trap::InvalidReference);
         };
+        let mut switched = false;
         loop {
             if let Some(catch) = self.matching_catch(protected_ip, tag_addr) {
                 let (landing_pad, base, with_ref, include_payload) = match catch {
@@ -623,13 +623,18 @@ impl<'store> Executor<'store> {
                 if with_ref {
                     ValueRef::stack_push(&mut self.store.value_stack, exception)?;
                 }
-                return Ok(Some(landing_pad as usize));
+                return Ok(Some(if switched {
+                    ExecFlow::Switch(landing_pad as usize)
+                } else {
+                    ExecFlow::Next(landing_pad as usize)
+                }));
             }
 
             self.store.value_stack.truncate_to_base(self.cf.locals_base);
             let Some(caller) = self.store.call_stack.pop_frame(self.call_stack_base) else {
                 return Ok(None);
             };
+            switched = true;
             self.switch_to_frame(caller);
             protected_ip = self.cf.instr_ptr.checked_sub(1).expect("invalid caller IP");
         }
@@ -647,7 +652,7 @@ impl<'store> Executor<'store> {
             if TAIL {
                 return Ok(self.exec_return());
             }
-            return Ok(ExecFlow::next(return_instr_ptr));
+            return Ok(ExecFlow::Next(return_instr_ptr));
         }
 
         let (param_count, result_count, base) = {
@@ -672,7 +677,7 @@ impl<'store> Executor<'store> {
                 Error::Trap(trap) => trap,
                 other => Trap::HostFunction(Box::new(other)),
             })?;
-        if TAIL { Ok(self.exec_return()) } else { Ok(ExecFlow::next(return_instr_ptr)) }
+        if TAIL { Ok(self.exec_return()) } else { Ok(ExecFlow::Next(return_instr_ptr)) }
     }
 
     fn exec_call_direct(&mut self, v: u32, return_instr_ptr: usize) -> ExecResult<ExecFlow> {
@@ -790,13 +795,14 @@ impl<'store> Executor<'store> {
             self.store.call_stack.push(self.cf, return_instr_ptr)?;
             self.cf = CallFrame::new(func_addr, locals_base, locals);
         }
+        let switched = next_func.is_some();
         if let Some(next_func) = next_func {
             self.func = next_func;
         }
         if owner != self.module.id() {
             self.set_module(owner);
         }
-        Ok(ExecFlow::next(0))
+        Ok(if switched { ExecFlow::Switch(0) } else { ExecFlow::Next(0) })
     }
 
     fn exec_call_ref<const TAIL: bool>(&mut self, type_addr: u32, return_instr_ptr: usize) -> ExecResult<ExecFlow> {
@@ -817,15 +823,16 @@ impl<'store> Executor<'store> {
     #[inline(always)]
     fn finish_return(&mut self) -> ExecFlow {
         let Some(caller) = self.store.call_stack.pop_frame(self.call_stack_base) else {
-            return ExecFlow::COMPLETE;
+            return ExecFlow::Complete;
         };
         let instr_ptr = caller.instr_ptr;
         if caller.func_addr == self.cf.func_addr {
             self.cf = caller;
+            ExecFlow::Next(instr_ptr)
         } else {
             self.switch_to_frame(caller);
+            ExecFlow::Switch(instr_ptr)
         }
-        ExecFlow::next(instr_ptr)
     }
 
     fn exec_return_void(&mut self) -> ExecFlow {

crates/tinywasm/src/store/mod.rs

@@ -824,7 +824,8 @@ impl Store {
                     };
                     let offset = usize::try_from(offset).unwrap_or(usize::MAX);
                     with_memory!(self.state, *mem_addr, |mem, kind| {
-                        mem.write_all(offset, &data.data).ok_or_else(|| memory::memory_oob(offset, data.data.len(), mem.len()))
+                        mem.write_all(offset, &data.data)
+                            .ok_or_else(|| memory::memory_oob(offset, data.data.len(), mem.len()))
                     })?;
                     self.state.data[data_addrs[i] as usize].drop();
                 }

crates/tinywasm/tests/borrow_instruction_stream.rs

@@ -0,0 +1,64 @@
+#![cfg(feature = "nightly-tail-calls")]
+
+use tinywasm::{HostFunction, Imports, ModuleInstance, Result, Store};
+
+const ADDER: &str = r#"
+    (module
+      (import "host" "bump" (func $bump (param i32) (result i32)))
+      (func (export "add") (param i32) (result i32)
+        local.get 0
+        call $bump
+        i32.const 2
+        i32.add))
+"#;
+
+const CALLER: &str = r#"
+    (module
+      (type $unary (func (param i32) (result i32)))
+      (import "adder" "add" (func $add (type $unary)))
+      (table 1 funcref)
+      (elem (i32.const 0) $add)
+      (func $local (param i32) (result i32)
+        local.get 0
+        i32.const 4
+        i32.add)
+      (func (export "main") (param i32) (result i32)
+        (local $acc i32)
+        (local $i i32)
+        local.get 0
+        local.set $acc
+        (loop $again
+          local.get $acc
+          i32.const 0
+          call_indirect (type $unary)
+          call $local
+          call $add
+          local.set $acc
+          local.get $i
+          i32.const 1
+          i32.add
+          local.tee $i
+          i32.const 64
+          i32.lt_u
+          br_if $again)
+        local.get $acc))
+"#;
+
+#[test]
+fn linked_calls_switch_instruction_stream() -> Result<()> {
+    let adder = tinywasm::parse_bytes(&wat::parse_str(ADDER).unwrap())?;
+    let caller = tinywasm::parse_bytes(&wat::parse_str(CALLER).unwrap())?;
+    let mut store = Store::default();
+
+    let mut adder_imports = Imports::new();
+    adder_imports.define("host", "bump", HostFunction::from(|_ctx, value: i32| -> Result<i32> { Ok(value + 1) }));
+    let adder_instance = ModuleInstance::instantiate(&mut store, &adder, Some(&adder_imports))?;
+
+    let mut caller_imports = Imports::new();
+    caller_imports.link_module("adder", adder_instance)?;
+    let caller_instance = ModuleInstance::instantiate(&mut store, &caller, Some(&caller_imports))?;
+    let main = caller_instance.func::<i32, i32>(&store, "main")?;
+
+    assert_eq!(main.call(&mut store, 32)?, 672);
+    Ok(())
+}