perf: borrow instruction stream across tail dispatch

Author: matthargett

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

@@ -3,7 +3,8 @@ 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], FuncAddr, usize, Instruction) -> ExecResult<()>;
 type BoundedHandler = for<'store> fn(&mut Executor<'store>, usize, Instruction, u32) -> ExecResult<()>;
 
 #[cold]
@@ -29,22 +30,30 @@ 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],
+                func_addr: FuncAddr,
                 $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, func_addr, 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(())),
+                            Some(next_instr_ptr) => {
+                                if $executor.cf.func_addr != func_addr {
+                                    $executor.cf.instr_ptr = next_instr_ptr;
+                                    return Ok(());
+                                }
+                                $dispatch_next!(next_instr_ptr)
+                            },
+                            None => return cold!({ $executor.completed = true; Ok(()) }),
                         }
                     }};
                 }
@@ -142,10 +151,18 @@ 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 func_addr = self.cf.func_addr;
+            let instr_ptr = self.cf.instr_ptr;
+            let instruction = instructions[instr_ptr];
+            let handler = Unbudgeted::handler_for(instruction.opcode());
+            handler(&mut self, instructions, func_addr, instr_ptr, instruction)?;
+            if self.completed {
+                return Ok(());
+            }
+        }
     }
 
     #[cfg(feature = "std")]

crates/tinywasm/tests/borrow_instruction_stream.rs

@@ -0,0 +1,78 @@
+#![cfg(feature = "nightly-tail-calls")]
+
+use tinywasm::{ExecProgress, 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_preserve_instruction_stream_across_function_switches() -> 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);
+
+    let mut resumable = main.call_resumable(&mut store, 32)?;
+    let mut suspended = false;
+    loop {
+        match resumable.resume_with_fuel(8)? {
+            ExecProgress::Completed(value) => {
+                assert_eq!(value, 672);
+                break;
+            }
+            ExecProgress::Suspended => suspended = true,
+        }
+    }
+    assert!(suspended);
+
+    Ok(())
+}