chore: benchmarking + major performance improvements (50%+ faster)

Signed-off-by: Henry Gressmann <mail@henrygressmann.de>

6 files changed, 94 insertions, 99 deletions

diff --git a/crates/tinywasm/src/imports.rs b/crates/tinywasm/src/imports.rs
index 9c5086a..294e547 100644
--- a/crates/tinywasm/src/imports.rs
+++ b/crates/tinywasm/src/imports.rs
@@ -255,6 +255,13 @@ impl Imports {
         Imports { values: BTreeMap::new(), modules: BTreeMap::new() }
     }
 
+    /// Merge two import sets
+    pub fn merge(mut self, other: Self) -> Self {
+        self.values.extend(other.values);
+        self.modules.extend(other.modules);
+        self
+    }
+
     /// Link a module
     ///
     /// This will automatically link all imported values on instantiation
diff --git a/crates/tinywasm/src/lib.rs b/crates/tinywasm/src/lib.rs
index f2951b6..bc64ac4 100644
--- a/crates/tinywasm/src/lib.rs
+++ b/crates/tinywasm/src/lib.rs
@@ -122,3 +122,13 @@ pub mod parser {
 pub mod types {
     pub use tinywasm_types::*;
 }
+
+#[cold]
+pub(crate) fn cold() {}
+
+pub(crate) fn unlikely(b: bool) -> bool {
+    if b {
+        cold();
+    }
+    b
+}
diff --git a/crates/tinywasm/src/runtime/interpreter/macros.rs b/crates/tinywasm/src/runtime/interpreter/macros.rs
index cc65812..eba866c 100644
--- a/crates/tinywasm/src/runtime/interpreter/macros.rs
+++ b/crates/tinywasm/src/runtime/interpreter/macros.rs
@@ -146,13 +146,8 @@ macro_rules! comp {
     }};
 
     ($op:tt, $intermediate:ty, $to:ty, $stack:ident) => {{
-        let [a, b] = $stack.values.pop_n_const::<2>()?;
-        let a: $intermediate = a.into();
-        let b: $intermediate = b.into();
-
-        // Cast to unsigned type before comparison
-        let a = a as $to;
-        let b = b as $to;
+        let b = $stack.values.pop_t::<$intermediate>()? as $to;
+        let a = $stack.values.pop_t::<$intermediate>()? as $to;
         $stack.values.push(((a $op b) as i32).into());
     }};
 }
@@ -160,7 +155,7 @@ macro_rules! comp {
 /// Compare a value on the stack to zero
 macro_rules! comp_zero {
     ($op:tt, $ty:ty, $stack:ident) => {{
-        let a: $ty = $stack.values.pop()?.into();
+        let a = $stack.values.pop_t::<$ty>()?;
         $stack.values.push(((a $op 0) as i32).into());
     }};
 }
@@ -173,20 +168,14 @@ macro_rules! arithmetic {
 
     // also allow operators such as +, -
     ($op:tt, $ty:ty, $stack:ident) => {{
-        let [a, b] = $stack.values.pop_n_const::<2>()?;
-        let a: $ty = a.into();
-        let b: $ty = b.into();
+        let b: $ty = $stack.values.pop_t()?;
+        let a: $ty = $stack.values.pop_t()?;
         $stack.values.push((a $op b).into());
     }};
 
     ($op:ident, $intermediate:ty, $to:ty, $stack:ident) => {{
-        let [a, b] = $stack.values.pop_n_const::<2>()?;
-        let a: $to = a.into();
-        let b: $to = b.into();
-
-        let a = a as $intermediate;
-        let b = b as $intermediate;
-
+        let b = $stack.values.pop_t::<$to>()? as $intermediate;
+        let a = $stack.values.pop_t::<$to>()? as $intermediate;
         let result = a.$op(b);
         $stack.values.push((result as $to).into());
     }};
@@ -215,19 +204,14 @@ macro_rules! checked_int_arithmetic {
     }};
 
     ($op:ident, $from:ty, $to:ty, $stack:ident) => {{
-        let [a, b] = $stack.values.pop_n_const::<2>()?;
-        let a: $from = a.into();
-        let b: $from = b.into();
-
-        let a_casted: $to = a as $to;
-        let b_casted: $to = b as $to;
+        let b = $stack.values.pop_t::<$from>()? as $to;
+        let a = $stack.values.pop_t::<$from>()? as $to;
 
-        if b_casted == 0 {
+        if b == 0 {
             return Err(Error::Trap(crate::Trap::DivisionByZero));
         }
 
-        let result = a_casted.$op(b_casted).ok_or_else(|| Error::Trap(crate::Trap::IntegerOverflow))?;
-
+        let result = a.$op(b).ok_or_else(|| Error::Trap(crate::Trap::IntegerOverflow))?;
         // Cast back to original type if different
         $stack.values.push((result as $from).into());
     }};
diff --git a/crates/tinywasm/src/runtime/interpreter/mod.rs b/crates/tinywasm/src/runtime/interpreter/mod.rs
index db013a0..ad261db 100644
--- a/crates/tinywasm/src/runtime/interpreter/mod.rs
+++ b/crates/tinywasm/src/runtime/interpreter/mod.rs
@@ -1,7 +1,6 @@
 use super::{InterpreterRuntime, Stack};
-use crate::log;
+use crate::{cold, log, unlikely};
 use crate::{
-    log::debug,
     runtime::{BlockType, CallFrame, LabelArgs, LabelFrame},
     Error, FuncContext, ModuleInstance, Result, Store, Trap,
 };
@@ -23,6 +22,7 @@ use macros::*;
 use traits::*;
 
 impl InterpreterRuntime {
+    #[inline(always)] // a small 2-3% performance improvement in some cases
     pub(crate) fn exec(&self, store: &mut Store, stack: &mut Stack) -> Result<()> {
         // The current call frame, gets updated inside of exec_one
         let mut cf = stack.call_stack.pop()?;
@@ -80,10 +80,10 @@ impl InterpreterRuntime {
             }
         }
 
-        debug!("end of exec");
-        debug!("stack: {:?}", stack.values);
-        debug!("insts: {:?}", instrs);
-        debug!("instr_ptr: {}", cf.instr_ptr);
+        log::debug!("end of exec");
+        log::debug!("stack: {:?}", stack.values);
+        log::debug!("insts: {:?}", instrs);
+        log::debug!("instr_ptr: {}", cf.instr_ptr);
         Err(Error::FuncDidNotReturn)
     }
 }
@@ -115,7 +115,7 @@ macro_rules! break_to {
 /// Run a single step of the interpreter
 /// A seperate function is used so later, we can more easily implement
 /// a step-by-step debugger (using generators once they're stable?)
-#[inline]
+#[inline(always)] // this improves performance by more than 20% in some cases
 fn exec_one(
     cf: &mut CallFrame,
     instr: &Instruction,
@@ -124,12 +124,13 @@ fn exec_one(
     store: &mut Store,
     module: &ModuleInstance,
 ) -> Result<ExecResult> {
-    debug!("ptr: {} instr: {:?}", cf.instr_ptr, instr);
-
     use tinywasm_types::Instruction::*;
     match instr {
         Nop => { /* do nothing */ }
-        Unreachable => return Ok(ExecResult::Trap(crate::Trap::Unreachable)), // we don't need to include the call frame here because it's already on the stack
+        Unreachable => {
+            cold();
+            return Ok(ExecResult::Trap(crate::Trap::Unreachable));
+        } // we don't need to include the call frame here because it's already on the stack
         Drop => stack.values.pop().map(|_| ())?,
 
         Select(
@@ -162,7 +163,7 @@ fn exec_one(
                 }
             };
 
-            let params = stack.values.pop_n_rev(ty.params.len())?;
+            let params = stack.values.pop_n_rev(ty.params.len())?.collect::<Vec<_>>();
             let call_frame = CallFrame::new_raw(func_inst, &params, locals);
 
             // push the call frame
@@ -191,15 +192,9 @@ fn exec_one(
 
             let func_inst = store.get_func(func_ref as usize)?.clone();
             let func_ty = func_inst.func.ty();
-
-            log::info!("type_addr: {}", type_addr);
-            log::info!("types: {:?}", module.func_tys());
             let call_ty = module.func_ty(*type_addr);
 
-            log::info!("call_indirect: current fn owner: {:?}", module.id());
-            log::info!("call_indirect: func owner: {:?}", func_inst.owner);
-
-            if func_ty != call_ty {
+            if unlikely(func_ty != call_ty) {
                 log::error!("indirect call type mismatch: {:?} != {:?}", func_ty, call_ty);
                 return Err(
                     Trap::IndirectCallTypeMismatch { actual: func_ty.clone(), expected: call_ty.clone() }.into()
@@ -217,7 +212,7 @@ fn exec_one(
                 }
             };
 
-            let params = stack.values.pop_n_rev(func_ty.params.len())?;
+            let params = stack.values.pop_n_rev(func_ty.params.len())?.collect::<Vec<_>>();
             let call_frame = CallFrame::new_raw(func_inst, &params, locals);
 
             // push the call frame
@@ -232,7 +227,6 @@ fn exec_one(
         If(args, else_offset, end_offset) => {
             // truthy value is on the top of the stack, so enter the then block
             if stack.values.pop_t::<i32>()? != 0 {
-                log::trace!("entering then");
                 cf.enter_label(
                     LabelFrame {
                         instr_ptr: cf.instr_ptr,
@@ -248,7 +242,6 @@ fn exec_one(
 
             // falsy value is on the top of the stack
             if let Some(else_offset) = else_offset {
-                log::debug!("entering else at {}", cf.instr_ptr + *else_offset);
                 cf.enter_label(
                     LabelFrame {
                         instr_ptr: cf.instr_ptr + *else_offset,
@@ -266,7 +259,6 @@ fn exec_one(
         }
 
         Loop(args, end_offset) => {
-            // let params = stack.values.pop_block_params(*args, &module)?;
             cf.enter_label(
                 LabelFrame {
                     instr_ptr: cf.instr_ptr,
@@ -297,11 +289,14 @@ fn exec_one(
                 .iter()
                 .map(|i| match i {
                     BrLabel(l) => Ok(*l),
-                    _ => panic!("Expected BrLabel, this should have been validated by the parser"),
+                    _ => {
+                        cold();
+                        panic!("Expected BrLabel, this should have been validated by the parser")
+                    }
                 })
                 .collect::<Result<Vec<_>>>()?;
 
-            if instr.len() != *len {
+            if unlikely(instr.len() != *len) {
                 panic!(
                     "Expected {} BrLabel instructions, got {}, this should have been validated by the parser",
                     len,
@@ -341,6 +336,7 @@ fn exec_one(
         // We're essentially using else as a EndBlockFrame instruction for if blocks
         Else(end_offset) => {
             let Some(block) = cf.labels.pop() else {
+                cold();
                 panic!("else: no label to end, this should have been validated by the parser");
             };
 
@@ -352,6 +348,7 @@ fn exec_one(
         EndBlockFrame => {
             // remove the label from the label stack
             let Some(block) = cf.labels.pop() else {
+                cold();
                 panic!("end: no label to end, this should have been validated by the parser");
             };
             stack.values.truncate_keep(block.stack_ptr, block.args.results)
@@ -379,7 +376,8 @@ fn exec_one(
 
         MemorySize(addr, byte) => {
             if *byte != 0 {
-                unimplemented!("memory.size with byte != 0");
+                cold();
+                return Err(Error::UnsupportedFeature("memory.size with byte != 0".to_string()));
             }
 
             let mem_idx = module.resolve_mem_addr(*addr);
@@ -389,6 +387,7 @@ fn exec_one(
 
         MemoryGrow(addr, byte) => {
             if *byte != 0 {
+                cold();
                 return Err(Error::UnsupportedFeature("memory.grow with byte != 0".to_string()));
             }
 
@@ -633,6 +632,7 @@ fn exec_one(
         I64TruncSatF64U => arithmetic_single!(trunc, f64, u64, stack),
 
         i => {
+            cold();
             log::error!("unimplemented instruction: {:?}", i);
             return Err(Error::UnsupportedFeature(alloc::format!("unimplemented instruction: {:?}", i)));
         }
diff --git a/crates/tinywasm/src/runtime/stack/call_stack.rs b/crates/tinywasm/src/runtime/stack/call_stack.rs
index 20de728..b19e332 100644
--- a/crates/tinywasm/src/runtime/stack/call_stack.rs
+++ b/crates/tinywasm/src/runtime/stack/call_stack.rs
@@ -78,7 +78,6 @@ impl CallFrame {
 
     /// Break to a block at the given index (relative to the current frame)
     /// Returns `None` if there is no block at the given index (e.g. if we need to return, this is handled by the caller)
-    #[inline]
     pub(crate) fn break_to(&mut self, break_to_relative: u32, value_stack: &mut super::ValueStack) -> Option<()> {
         log::debug!("break_to_relative: {}", break_to_relative);
         let break_to = self.labels.get_relative_to_top(break_to_relative as usize)?;
diff --git a/crates/tinywasm/src/runtime/stack/value_stack.rs b/crates/tinywasm/src/runtime/stack/value_stack.rs
index f8f0951..d36373b 100644
--- a/crates/tinywasm/src/runtime/stack/value_stack.rs
+++ b/crates/tinywasm/src/runtime/stack/value_stack.rs
@@ -1,6 +1,6 @@
 use core::ops::Range;
 
-use crate::{runtime::RawWasmValue, Error, Result};
+use crate::{cold, runtime::RawWasmValue, unlikely, Error, Result};
 use alloc::vec::Vec;
 use tinywasm_types::{ValType, WasmValue};
 
@@ -10,19 +10,17 @@ pub(crate) const STACK_SIZE: usize = 1024;
 #[derive(Debug)]
 pub(crate) struct ValueStack {
     stack: Vec<RawWasmValue>,
-    top: usize,
 }
 
 impl Default for ValueStack {
     fn default() -> Self {
-        Self { stack: Vec::with_capacity(STACK_SIZE), top: 0 }
+        Self { stack: Vec::with_capacity(STACK_SIZE) }
     }
 }
 
 impl ValueStack {
     #[inline]
     pub(crate) fn extend_from_within(&mut self, range: Range<usize>) {
-        self.top += range.len();
         self.stack.extend_from_within(range);
     }
 
@@ -32,98 +30,95 @@ impl ValueStack {
             return;
         }
 
-        self.top += values.len();
         self.stack.extend(values.iter().map(|v| RawWasmValue::from(*v)));
     }
 
     #[inline]
     pub(crate) fn len(&self) -> usize {
-        assert!(self.top <= self.stack.len());
-        self.top
+        self.stack.len()
     }
 
     pub(crate) fn truncate_keep(&mut self, n: usize, end_keep: usize) {
         let total_to_keep = n + end_keep;
-        assert!(self.top >= total_to_keep, "Total to keep should be less than or equal to self.top");
+        let len = self.stack.len();
+        assert!(len >= total_to_keep, "Total to keep should be less than or equal to self.top");
 
-        let current_size = self.stack.len();
-        if current_size <= total_to_keep {
+        if len <= total_to_keep {
             return; // No need to truncate if the current size is already less than or equal to total_to_keep
         }
 
-        let items_to_remove = current_size - total_to_keep;
-        let remove_start_index = self.top - items_to_remove - end_keep;
-        let remove_end_index = self.top - end_keep;
-
+        let items_to_remove = len - total_to_keep;
+        let remove_start_index = len - items_to_remove - end_keep;
+        let remove_end_index = len - end_keep;
         self.stack.drain(remove_start_index..remove_end_index);
-        self.top = total_to_keep; // Update top to reflect the new size
     }
 
     #[inline]
     pub(crate) fn push(&mut self, value: RawWasmValue) {
-        self.top += 1;
         self.stack.push(value);
     }
 
     #[inline]
     pub(crate) fn last(&self) -> Result<&RawWasmValue> {
-        self.stack.last().ok_or(Error::StackUnderflow)
+        match self.stack.last() {
+            Some(v) => Ok(v),
+            None => {
+                cold();
+                Err(Error::StackUnderflow)
+            }
+        }
     }
 
     #[inline]
     pub(crate) fn pop_t<T: From<RawWasmValue>>(&mut self) -> Result<T> {
-        self.top -= 1;
-        Ok(self.stack.pop().ok_or(Error::StackUnderflow)?.into())
+        match self.stack.pop() {
+            Some(v) => Ok(v.into()),
+            None => {
+                cold();
+                Err(Error::StackUnderflow)
+            }
+        }
     }
 
     #[inline]
     pub(crate) fn pop(&mut self) -> Result<RawWasmValue> {
-        self.top -= 1;
-        self.stack.pop().ok_or(Error::StackUnderflow)
+        match self.stack.pop() {
+            Some(v) => Ok(v),
+            None => {
+                cold();
+                Err(Error::StackUnderflow)
+            }
+        }
     }
 
     #[inline]
     pub(crate) fn pop_params(&mut self, types: &[ValType]) -> Result<Vec<WasmValue>> {
-        let res = self.pop_n_rev(types.len())?.iter().zip(types.iter()).map(|(v, ty)| v.attach_type(*ty)).collect();
+        let res = self.pop_n_rev(types.len())?.zip(types.iter()).map(|(v, ty)| v.attach_type(*ty)).collect();
         Ok(res)
     }
 
     pub(crate) fn break_to(&mut self, new_stack_size: usize, result_count: usize) {
-        assert!(self.top >= result_count);
-        self.stack.copy_within((self.top - result_count)..self.top, new_stack_size);
-        self.top = new_stack_size + result_count;
-        self.stack.truncate(self.top);
+        let len = self.stack.len();
+        self.stack.copy_within((len - result_count)..len, new_stack_size);
+        self.stack.truncate(new_stack_size + result_count);
     }
 
     #[inline]
     pub(crate) fn last_n(&self, n: usize) -> Result<&[RawWasmValue]> {
-        if self.top < n {
-            return Err(Error::StackUnderflow);
-        }
-        Ok(&self.stack[self.top - n..self.top])
-    }
-
-    #[inline]
-    pub(crate) fn pop_n_rev(&mut self, n: usize) -> Result<Vec<RawWasmValue>> {
-        if self.top < n {
+        let len = self.stack.len();
+        if unlikely(len < n) {
             return Err(Error::StackUnderflow);
         }
-        self.top -= n;
-        let res = self.stack.drain(self.top..).collect::<Vec<_>>();
-        Ok(res)
+        Ok(&self.stack[len - n..len])
     }
 
     #[inline]
-    pub(crate) fn pop_n_const<const N: usize>(&mut self) -> Result<[RawWasmValue; N]> {
-        if self.top < N {
+    pub(crate) fn pop_n_rev(&mut self, n: usize) -> Result<alloc::vec::Drain<'_, RawWasmValue>> {
+        let len = self.stack.len();
+        if unlikely(len < n) {
             return Err(Error::StackUnderflow);
         }
-        self.top -= N;
-        let mut res = [RawWasmValue::default(); N];
-        for i in res.iter_mut().rev() {
-            *i = self.stack.pop().ok_or(Error::InvalidStore)?;
-        }
-
+        let res = self.stack.drain((len - n)..);
         Ok(res)
     }
 }