From b8785f3951166266008876c01ec29475ad688cac Mon Sep 17 00:00:00 2001 From: Henry Date: Sun, 27 Sep 2026 18:00:30 +0200 Subject: [PATCH] pref: borrow instructions to improve become codegen Signed-off-by: Henry --- Cargo.lock | 8 +-- .../src/interpreter/executor/dispatch.rs | 4 +- .../interpreter/executor/dispatch_become.rs | 32 +++++++--- .../tinywasm/src/interpreter/executor/mod.rs | 43 +++++++------ crates/tinywasm/src/store/mod.rs | 3 +- .../tests/borrow_instruction_stream.rs | 64 +++++++++++++++++++ 6 files changed, 119 insertions(+), 35 deletions(-) create mode 100644 crates/tinywasm/tests/borrow_instruction_stream.rs diff --git a/Cargo.lock b/Cargo.lock index fad4ef6f..bd732b24 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1365,18 +1365,18 @@ dependencies = [ [[package]] name = "zerocopy" -version = "0.8.58" +version = "0.8.59" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c17e8fafad82b542ff3717217ecdc736231b59e387768c9630123b4ce4d2db44" +checksum = "6df92bf3d9227be3d53173901ddbffac2babc27ae50f397776ffd6dc33f800cb" dependencies = [ "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.8.58" +version = "0.8.59" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "595f56e044df4f46a0c9a626f65c3d99eb8488f7e8a8baa12dd76326d9710bf2" +checksum = "ac4f328cf2f05d084e496c3e9c3f33ed0a183656a16e1fcec4d464d8373aec82" dependencies = [ "proc-macro2", "quote", diff --git a/crates/tinywasm/src/interpreter/executor/dispatch.rs b/crates/tinywasm/src/interpreter/executor/dispatch.rs index 57efc54c..ac520475 100644 --- a/crates/tinywasm/src/interpreter/executor/dispatch.rs +++ b/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 { 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)) } }; } diff --git a/crates/tinywasm/src/interpreter/executor/dispatch_become.rs b/crates/tinywasm/src/interpreter/executor/dispatch_become.rs index c24f6bab..d5c86442 100644 --- a/crates/tinywasm/src/interpreter/executor/dispatch_become.rs +++ b/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")] diff --git a/crates/tinywasm/src/interpreter/executor/mod.rs b/crates/tinywasm/src/interpreter/executor/mod.rs index c68e0aa8..149ebc3a 100644 --- a/crates/tinywasm/src/interpreter/executor/mod.rs +++ b/crates/tinywasm/src/interpreter/executor/mod.rs @@ -55,20 +55,19 @@ impl From 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 { - (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 { 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, Trap> { + fn dispatch_exception(&mut self, exception: ValueRef, mut protected_ip: usize) -> Result, 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 { @@ -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(&mut self, type_addr: u32, return_instr_ptr: usize) -> ExecResult { @@ -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 { diff --git a/crates/tinywasm/src/store/mod.rs b/crates/tinywasm/src/store/mod.rs index af63a334..7f5ae61d 100644 --- a/crates/tinywasm/src/store/mod.rs +++ b/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(); } diff --git a/crates/tinywasm/tests/borrow_instruction_stream.rs b/crates/tinywasm/tests/borrow_instruction_stream.rs new file mode 100644 index 00000000..2bd8d68d --- /dev/null +++ b/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 { 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::(&store, "main")?; + + assert_eq!(main.call(&mut store, 32)?, 672); + Ok(()) +}