perf: reduce block overhead (#3)

10 files changed, 99 insertions, 74 deletions

diff --git a/crates/parser/src/conversion.rs b/crates/parser/src/conversion.rs
index cd8840a..68567f3 100644
--- a/crates/parser/src/conversion.rs
+++ b/crates/parser/src/conversion.rs
@@ -291,8 +291,6 @@ pub(crate) fn process_operators(
             }
             End => {
                 if let Some(label_pointer) = labels_ptrs.pop() {
-                    log::debug!("ending block: {:?}", instructions[label_pointer]);
-
                     let current_instr_ptr = instructions.len();
 
                     // last_label_pointer is Some if we're ending a block
diff --git a/crates/tinywasm/src/error.rs b/crates/tinywasm/src/error.rs
index 6b8119f..85ffdf6 100644
--- a/crates/tinywasm/src/error.rs
+++ b/crates/tinywasm/src/error.rs
@@ -24,17 +24,17 @@ pub enum Error {
     FuncDidNotReturn,
 
     /// The stack is empty
-    StackUnderflow,
+    ValueStackUnderflow,
 
     /// The label stack is empty
-    LabelStackUnderflow,
+    BlockStackUnderflow,
+
+    /// The call stack is empty
+    CallStackUnderflow,
 
     /// An invalid label type was encountered
     InvalidLabelType,
 
-    /// The call stack is empty
-    CallStackEmpty,
-
     /// The store is not the one that the module instance was instantiated in
     InvalidStore,
 
@@ -189,13 +189,13 @@ impl Display for Error {
 
             Self::Trap(trap) => write!(f, "trap: {}", trap),
             Self::Linker(err) => write!(f, "linking error: {}", err),
-            Self::CallStackEmpty => write!(f, "call stack empty"),
+            Self::CallStackUnderflow => write!(f, "call stack empty"),
             Self::InvalidLabelType => write!(f, "invalid label type"),
             Self::Other(message) => write!(f, "unknown error: {}", message),
             Self::UnsupportedFeature(feature) => write!(f, "unsupported feature: {}", feature),
             Self::FuncDidNotReturn => write!(f, "function did not return"),
-            Self::LabelStackUnderflow => write!(f, "label stack underflow"),
-            Self::StackUnderflow => write!(f, "stack underflow"),
+            Self::BlockStackUnderflow => write!(f, "label stack underflow"),
+            Self::ValueStackUnderflow => write!(f, "value stack underflow"),
             Self::InvalidStore => write!(f, "invalid store"),
         }
     }
diff --git a/crates/tinywasm/src/func.rs b/crates/tinywasm/src/func.rs
index faf7931..757cb85 100644
--- a/crates/tinywasm/src/func.rs
+++ b/crates/tinywasm/src/func.rs
@@ -58,7 +58,7 @@ impl FuncHandle {
 
         // 6. Let f be the dummy frame
         let call_frame =
-            CallFrame::new(wasm_func.clone(), func_inst.owner, params.iter().map(|v| RawWasmValue::from(*v)));
+            CallFrame::new(wasm_func.clone(), func_inst.owner, params.iter().map(|v| RawWasmValue::from(*v)), 0);
 
         // 7. Push the frame f to the call stack
         // & 8. Push the values to the stack (Not needed since the call frame owns the values)
diff --git a/crates/tinywasm/src/instance.rs b/crates/tinywasm/src/instance.rs
index 6bfe99c..3fc4fe0 100644
--- a/crates/tinywasm/src/instance.rs
+++ b/crates/tinywasm/src/instance.rs
@@ -61,7 +61,7 @@ impl ModuleInstance {
         // don't need to create a auxiliary frame etc.
 
         let idx = store.next_module_instance_idx();
-        log::error!("Instantiating module at index {}", idx);
+        log::info!("Instantiating module at index {}", idx);
         let imports = imports.unwrap_or_default();
 
         let mut addrs = imports.link(store, &module, idx)?;
diff --git a/crates/tinywasm/src/runtime/interpreter/macros.rs b/crates/tinywasm/src/runtime/interpreter/macros.rs
index d6c0553..824666d 100644
--- a/crates/tinywasm/src/runtime/interpreter/macros.rs
+++ b/crates/tinywasm/src/runtime/interpreter/macros.rs
@@ -11,7 +11,7 @@
 // from a function, so we need to check if the label stack is empty
 macro_rules! break_to {
     ($cf:ident, $stack:ident, $break_to_relative:ident) => {{
-        if $cf.break_to(*$break_to_relative, &mut $stack.values).is_none() {
+        if $cf.break_to(*$break_to_relative, &mut $stack.values, &mut $stack.blocks).is_none() {
             if $stack.call_stack.is_empty() {
                 return Ok(ExecResult::Return);
             } else {
@@ -53,7 +53,6 @@ macro_rules! mem_load {
 
         const LEN: usize = core::mem::size_of::<$load_type>();
         let val = mem_ref.load_as::<LEN, $load_type>(addr, $arg.align as usize)?;
-        // let loaded_value = mem_ref.load_as::<$load_type>(addr, $arg.align as usize)?;
         $stack.values.push((val as $target_type).into());
     }};
 }
diff --git a/crates/tinywasm/src/runtime/interpreter/mod.rs b/crates/tinywasm/src/runtime/interpreter/mod.rs
index 7e34ac1..78f6ed2 100644
--- a/crates/tinywasm/src/runtime/interpreter/mod.rs
+++ b/crates/tinywasm/src/runtime/interpreter/mod.rs
@@ -4,7 +4,7 @@ use core::ops::{BitAnd, BitOr, BitXor, Neg};
 use tinywasm_types::{ElementKind, ValType};
 
 use super::{InterpreterRuntime, Stack};
-use crate::runtime::{BlockType, CallFrame, LabelFrame};
+use crate::runtime::{BlockFrame, BlockType, CallFrame};
 use crate::{cold, log, unlikely};
 use crate::{Error, FuncContext, ModuleInstance, Result, Store, Trap};
 
@@ -32,8 +32,13 @@ impl InterpreterRuntime {
             match exec_one(&mut cf, stack, store, &current_module) {
                 // Continue execution at the new top of the call stack
                 Ok(ExecResult::Call) => {
+                    let old = cf.block_ptr;
                     cf = stack.call_stack.pop()?;
 
+                    if old > cf.block_ptr {
+                        stack.blocks.truncate(old);
+                    }
+
                     // keeping the pointer seperate from the call frame is about 2% faster
                     // than storing it in the call frame
                     if cf.func_instance.1 != current_module.id() {
@@ -123,7 +128,7 @@ fn exec_one(cf: &mut CallFrame, stack: &mut Stack, store: &mut Store, module: &M
             };
 
             let params = stack.values.pop_n_rev(wasm_func.ty.params.len())?;
-            let call_frame = CallFrame::new(wasm_func, func_inst.owner, params);
+            let call_frame = CallFrame::new(wasm_func, func_inst.owner, params, stack.blocks.len());
 
             // push the call frame
             cf.instr_ptr += 1; // skip the call instruction
@@ -180,7 +185,7 @@ fn exec_one(cf: &mut CallFrame, stack: &mut Stack, store: &mut Store, module: &M
             }
 
             let params = stack.values.pop_n_rev(wasm_func.ty.params.len())?;
-            let call_frame = CallFrame::new(wasm_func, func_inst.owner, params);
+            let call_frame = CallFrame::new(wasm_func, func_inst.owner, params, stack.blocks.len());
 
             // push the call frame
             cf.instr_ptr += 1; // skip the call instruction
@@ -194,8 +199,8 @@ fn exec_one(cf: &mut CallFrame, stack: &mut Stack, store: &mut Store, module: &M
         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 {
-                cf.enter_label(
-                    LabelFrame::new(
+                cf.enter_block(
+                    BlockFrame::new(
                         cf.instr_ptr,
                         cf.instr_ptr + *end_offset,
                         stack.values.len(), // - params,
@@ -204,13 +209,14 @@ fn exec_one(cf: &mut CallFrame, stack: &mut Stack, store: &mut Store, module: &M
                         module,
                     ),
                     &mut stack.values,
+                    &mut stack.blocks,
                 );
                 return Ok(ExecResult::Ok);
             }
 
             // falsy value is on the top of the stack
             if let Some(else_offset) = else_offset {
-                let label = LabelFrame::new(
+                let label = BlockFrame::new(
                     cf.instr_ptr + *else_offset,
                     cf.instr_ptr + *end_offset,
                     stack.values.len(), // - params,
@@ -219,15 +225,15 @@ fn exec_one(cf: &mut CallFrame, stack: &mut Stack, store: &mut Store, module: &M
                     module,
                 );
                 cf.instr_ptr += *else_offset;
-                cf.enter_label(label, &mut stack.values);
+                cf.enter_block(label, &mut stack.values, &mut stack.blocks);
             } else {
                 cf.instr_ptr += *end_offset;
             }
         }
 
         Loop(args, end_offset) => {
-            cf.enter_label(
-                LabelFrame::new(
+            cf.enter_block(
+                BlockFrame::new(
                     cf.instr_ptr,
                     cf.instr_ptr + *end_offset,
                     stack.values.len(), // - params,
@@ -236,12 +242,13 @@ fn exec_one(cf: &mut CallFrame, stack: &mut Stack, store: &mut Store, module: &M
                     module,
                 ),
                 &mut stack.values,
+                &mut stack.blocks,
             );
         }
 
         Block(args, end_offset) => {
-            cf.enter_label(
-                LabelFrame::new(
+            cf.enter_block(
+                BlockFrame::new(
                     cf.instr_ptr,
                     cf.instr_ptr + *end_offset,
                     stack.values.len(), // - params,
@@ -250,6 +257,7 @@ fn exec_one(cf: &mut CallFrame, stack: &mut Stack, store: &mut Store, module: &M
                     module,
                 ),
                 &mut stack.values,
+                &mut stack.blocks,
             );
         }
 
@@ -291,10 +299,10 @@ fn exec_one(cf: &mut CallFrame, stack: &mut Stack, store: &mut Store, module: &M
         },
 
         EndFunc => {
-            assert!(
-                cf.labels.len() == 0,
-                "endfunc: block frames not empty, this should have been validated by the parser"
-            );
+            if stack.blocks.len() != cf.block_ptr {
+                cold();
+                panic!("endfunc: block frames not empty, this should have been validated by the parser");
+            }
 
             match stack.call_stack.is_empty() {
                 true => return Ok(ExecResult::Return),
@@ -304,7 +312,7 @@ fn exec_one(cf: &mut CallFrame, stack: &mut Stack, store: &mut Store, module: &M
 
         // We're essentially using else as a EndBlockFrame instruction for if blocks
         Else(end_offset) => {
-            let Some(block) = cf.labels.pop() else {
+            let Some(block) = stack.blocks.pop() else {
                 cold();
                 panic!("else: no label to end, this should have been validated by the parser");
             };
@@ -316,16 +324,28 @@ fn exec_one(cf: &mut CallFrame, stack: &mut Stack, store: &mut Store, module: &M
 
         EndBlockFrame => {
             // remove the label from the label stack
-            let Some(block) = cf.labels.pop() else {
+            let Some(block) = stack.blocks.pop() else {
                 cold();
-                panic!("end: no label to end, this should have been validated by the parser");
+                panic!("end blockframe: no label to end, this should have been validated by the parser");
             };
-            stack.values.truncate_keep(block.stack_ptr, block.results)
+
+            stack.values.truncate_keep(block.stack_ptr, block.results);
         }
 
         LocalGet(local_index) => stack.values.push(cf.get_local(*local_index as usize)),
         LocalSet(local_index) => cf.set_local(*local_index as usize, stack.values.pop()?),
-        LocalTee(local_index) => cf.set_local(*local_index as usize, *stack.values.last()?),
+        LocalTee(local_index) => {
+            let last_val = match stack.values.last() {
+                Ok(val) => val,
+                Err(_) => {
+                    log::error!("index: {}", local_index);
+                    log::error!("stack: {:?}", stack.values);
+
+                    panic!();
+                }
+            };
+            cf.set_local(*local_index as usize, *last_val)
+        }
 
         GlobalGet(global_index) => {
             let idx = module.resolve_global_addr(*global_index);
diff --git a/crates/tinywasm/src/runtime/stack.rs b/crates/tinywasm/src/runtime/stack.rs
index 31c41f0..3db3c4b 100644
--- a/crates/tinywasm/src/runtime/stack.rs
+++ b/crates/tinywasm/src/runtime/stack.rs
@@ -1,20 +1,21 @@
-mod blocks;
+mod block_stack;
 mod call_stack;
 mod value_stack;
 
 use self::{call_stack::CallStack, value_stack::ValueStack};
-pub(crate) use blocks::{BlockType, LabelFrame};
+pub(crate) use block_stack::{BlockFrame, BlockStack, BlockType};
 pub(crate) use call_stack::CallFrame;
 
 /// A WebAssembly Stack
 #[derive(Debug)]
 pub struct Stack {
     pub(crate) values: ValueStack,
+    pub(crate) blocks: BlockStack,
     pub(crate) call_stack: CallStack,
 }
 
 impl Stack {
     pub(crate) fn new(call_frame: CallFrame) -> Self {
-        Self { values: ValueStack::default(), call_stack: CallStack::new(call_frame) }
+        Self { values: ValueStack::default(), blocks: BlockStack::default(), call_stack: CallStack::new(call_frame) }
     }
 }
diff --git a/crates/tinywasm/src/runtime/stack/blocks.rs b/crates/tinywasm/src/runtime/stack/block_stack.rs
index c04632c..cf7a0d4 100644
--- a/crates/tinywasm/src/runtime/stack/blocks.rs
+++ b/crates/tinywasm/src/runtime/stack/block_stack.rs
@@ -1,33 +1,28 @@
+use crate::{unlikely, ModuleInstance};
 use alloc::vec::Vec;
 use tinywasm_types::BlockArgs;
 
-use crate::{unlikely, ModuleInstance};
-
-#[derive(Debug, Clone)]
-pub(crate) struct Labels(Vec<LabelFrame>); // TODO: maybe Box<[LabelFrame]> by analyzing the lable count when parsing the module?
-
-impl Labels {
-    #[inline]
-    pub(crate) fn new() -> Self {
-        // this is somehow a lot faster than Vec::with_capacity(128) or even using Default::default() in the benchmarks
-        Self(Vec::new())
-    }
+#[derive(Debug, Clone, Default)]
+pub(crate) struct BlockStack(Vec<BlockFrame>); // TODO: maybe Box<[LabelFrame]> by analyzing the lable count when parsing the module?
 
+impl BlockStack {
     #[inline]
     pub(crate) fn len(&self) -> usize {
         self.0.len()
     }
 
     #[inline]
-    pub(crate) fn push(&mut self, label: LabelFrame) {
-        self.0.push(label);
+    pub(crate) fn push(&mut self, block: BlockFrame) {
+        self.0.push(block);
     }
 
     #[inline]
     /// get the label at the given index, where 0 is the top of the stack
-    pub(crate) fn get_relative_to_top(&self, index: usize) -> Option<&LabelFrame> {
+    pub(crate) fn get_relative_to(&self, index: usize, offset: usize) -> Option<&BlockFrame> {
+        let len = self.0.len() - offset;
+
         // the vast majority of wasm functions don't use break to return
-        if unlikely(index >= self.0.len()) {
+        if unlikely(index >= len) {
             return None;
         }
 
@@ -35,7 +30,7 @@ impl Labels {
     }
 
     #[inline]
-    pub(crate) fn pop(&mut self) -> Option<LabelFrame> {
+    pub(crate) fn pop(&mut self) -> Option<BlockFrame> {
         self.0.pop()
     }
 
@@ -47,7 +42,7 @@ impl Labels {
 }
 
 #[derive(Debug, Clone)]
-pub(crate) struct LabelFrame {
+pub(crate) struct BlockFrame {
     // position of the instruction pointer when the block was entered
     pub(crate) instr_ptr: usize,
     // position of the end instruction of the block
@@ -60,7 +55,7 @@ pub(crate) struct LabelFrame {
     pub(crate) ty: BlockType,
 }
 
-impl LabelFrame {
+impl BlockFrame {
     #[inline]
     pub(crate) fn new(
         instr_ptr: usize,
diff --git a/crates/tinywasm/src/runtime/stack/call_stack.rs b/crates/tinywasm/src/runtime/stack/call_stack.rs
index a902833..1c441a8 100644
--- a/crates/tinywasm/src/runtime/stack/call_stack.rs
+++ b/crates/tinywasm/src/runtime/stack/call_stack.rs
@@ -1,11 +1,12 @@
 use alloc::{boxed::Box, rc::Rc, vec::Vec};
 use tinywasm_types::{ModuleInstanceAddr, WasmFunction};
 
-use super::{blocks::Labels, LabelFrame};
 use crate::runtime::{BlockType, RawWasmValue};
 use crate::unlikely;
 use crate::{Error, Result, Trap};
 
+use super::BlockFrame;
+
 const CALL_STACK_SIZE: usize = 128;
 const CALL_STACK_MAX_SIZE: usize = 1024;
 
@@ -31,7 +32,7 @@ impl CallStack {
     pub(crate) fn pop(&mut self) -> Result<CallFrame> {
         match self.stack.pop() {
             Some(frame) => Ok(frame),
-            None => Err(Error::CallStackEmpty),
+            None => Err(Error::CallStackUnderflow),
         }
     }
 
@@ -48,25 +49,35 @@ impl CallStack {
 #[derive(Debug, Clone)]
 pub(crate) struct CallFrame {
     pub(crate) instr_ptr: usize,
+    pub(crate) block_ptr: usize,
     pub(crate) func_instance: (Rc<WasmFunction>, ModuleInstanceAddr),
-    pub(crate) labels: Labels,
     pub(crate) locals: Box<[RawWasmValue]>,
 }
 
 impl CallFrame {
     /// Push a new label to the label stack and ensure the stack has the correct values
-    pub(crate) fn enter_label(&mut self, label_frame: LabelFrame, stack: &mut super::ValueStack) {
-        if label_frame.params > 0 {
-            stack.extend_from_within((label_frame.stack_ptr - label_frame.params)..label_frame.stack_ptr);
+    pub(crate) fn enter_block(
+        &mut self,
+        block_frame: BlockFrame,
+        values: &mut super::ValueStack,
+        blocks: &mut super::BlockStack,
+    ) {
+        if block_frame.params > 0 {
+            values.extend_from_within((block_frame.stack_ptr - block_frame.params)..block_frame.stack_ptr);
         }
 
-        self.labels.push(label_frame);
+        blocks.push(block_frame);
     }
 
     /// 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)
-    pub(crate) fn break_to(&mut self, break_to_relative: u32, value_stack: &mut super::ValueStack) -> Option<()> {
-        let break_to = self.labels.get_relative_to_top(break_to_relative as usize)?;
+    pub(crate) fn break_to(
+        &mut self,
+        break_to_relative: u32,
+        values: &mut super::ValueStack,
+        blocks: &mut super::BlockStack,
+    ) -> Option<()> {
+        let break_to = blocks.get_relative_to(break_to_relative as usize, self.block_ptr)?;
 
         // instr_ptr points to the label instruction, but the next step
         // will increment it by 1 since we're changing the "current" instr_ptr
@@ -76,12 +87,12 @@ impl CallFrame {
                 self.instr_ptr = break_to.instr_ptr;
 
                 // We also want to push the params to the stack
-                value_stack.break_to(break_to.stack_ptr, break_to.params);
+                values.break_to(break_to.stack_ptr, break_to.params);
 
                 // check if we're breaking to the loop
                 if break_to_relative != 0 {
                     // we also want to trim the label stack to the loop (but not including the loop)
-                    self.labels.truncate(self.labels.len() - break_to_relative as usize);
+                    blocks.truncate(blocks.len() - break_to_relative as usize);
                     return Some(());
                 }
             }
@@ -89,13 +100,13 @@ impl CallFrame {
             BlockType::Block | BlockType::If | BlockType::Else => {
                 // this is a block, so we want to jump to the next instruction after the block ends
                 // We also want to push the block's results to the stack
-                value_stack.break_to(break_to.stack_ptr, break_to.results);
+                values.break_to(break_to.stack_ptr, break_to.results);
 
                 // (the inst_ptr will be incremented by 1 before the next instruction is executed)
                 self.instr_ptr = break_to.end_instr_ptr;
 
                 // we also want to trim the label stack, including the block
-                self.labels.truncate(self.labels.len() - (break_to_relative as usize + 1));
+                blocks.truncate(blocks.len() - (break_to_relative as usize + 1));
             }
         }
 
@@ -108,6 +119,7 @@ impl CallFrame {
         wasm_func_inst: Rc<WasmFunction>,
         owner: ModuleInstanceAddr,
         params: impl Iterator<Item = RawWasmValue> + ExactSizeIterator,
+        block_ptr: usize,
     ) -> Self {
         let locals = {
             let local_types = &wasm_func_inst.locals;
@@ -118,7 +130,7 @@ impl CallFrame {
             locals.into_boxed_slice()
         };
 
-        Self { instr_ptr: 0, func_instance: (wasm_func_inst, owner), locals, labels: Labels::new() }
+        Self { instr_ptr: 0, func_instance: (wasm_func_inst, owner), locals, block_ptr }
     }
 
     #[inline]
diff --git a/crates/tinywasm/src/runtime/stack/value_stack.rs b/crates/tinywasm/src/runtime/stack/value_stack.rs
index cc35bc2..9649177 100644
--- a/crates/tinywasm/src/runtime/stack/value_stack.rs
+++ b/crates/tinywasm/src/runtime/stack/value_stack.rs
@@ -63,7 +63,7 @@ impl ValueStack {
             Some(v) => Ok(v),
             None => {
                 cold();
-                Err(Error::StackUnderflow)
+                Err(Error::ValueStackUnderflow)
             }
         }
     }
@@ -74,7 +74,7 @@ impl ValueStack {
             Some(v) => Ok(v.into()),
             None => {
                 cold();
-                Err(Error::StackUnderflow)
+                Err(Error::ValueStackUnderflow)
             }
         }
     }
@@ -85,7 +85,7 @@ impl ValueStack {
             Some(v) => Ok(v),
             None => {
                 cold();
-                Err(Error::StackUnderflow)
+                Err(Error::ValueStackUnderflow)
             }
         }
     }
@@ -105,7 +105,7 @@ impl ValueStack {
     pub(crate) fn last_n(&self, n: usize) -> Result<&[RawWasmValue]> {
         let len = self.stack.len();
         if unlikely(len < n) {
-            return Err(Error::StackUnderflow);
+            return Err(Error::ValueStackUnderflow);
         }
         Ok(&self.stack[len - n..len])
     }
@@ -114,7 +114,7 @@ impl ValueStack {
     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);
+            return Err(Error::ValueStackUnderflow);
         }
         let res = self.stack.drain((len - n)..);
         Ok(res)