chore: improve value stack

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

10 files changed, 191 insertions, 45 deletions

diff --git a/benchmarks/benches/fibonacci.rs b/benchmarks/benches/fibonacci.rs
index 15c09ac..61b9a68 100644
--- a/benchmarks/benches/fibonacci.rs
+++ b/benchmarks/benches/fibonacci.rs
@@ -46,21 +46,21 @@ fn run_native_recursive(n: i32) -> i32 {
 
 const FIBONACCI: &[u8] = include_bytes!("../../examples/rust/out/fibonacci.wasm");
 fn criterion_benchmark(c: &mut Criterion) {
-    {
-        let mut group = c.benchmark_group("fibonacci");
-        group.bench_function("native", |b| b.iter(|| run_native(black_box(60))));
-        group.bench_function("tinywasm", |b| b.iter(|| run_tinywasm(FIBONACCI, black_box(60), "fibonacci")));
-        group.bench_function("wasmi", |b| b.iter(|| run_wasmi(FIBONACCI, black_box(60), "fibonacci")));
-        group.bench_function("wasmer", |b| b.iter(|| run_wasmer(FIBONACCI, black_box(60), "fibonacci")));
-    }
+    // {
+    //     let mut group = c.benchmark_group("fibonacci");
+    //     group.bench_function("native", |b| b.iter(|| run_native(black_box(60))));
+    //     group.bench_function("tinywasm", |b| b.iter(|| run_tinywasm(FIBONACCI, black_box(60), "fibonacci")));
+    //     group.bench_function("wasmi", |b| b.iter(|| run_wasmi(FIBONACCI, black_box(60), "fibonacci")));
+    //     group.bench_function("wasmer", |b| b.iter(|| run_wasmer(FIBONACCI, black_box(60), "fibonacci")));
+    // }
 
     {
         let mut group = c.benchmark_group("fibonacci-recursive");
         group.measurement_time(std::time::Duration::from_secs(5));
-        group.bench_function("native", |b| b.iter(|| run_native_recursive(black_box(26))));
+        // group.bench_function("native", |b| b.iter(|| run_native_recursive(black_box(26))));
         group.bench_function("tinywasm", |b| b.iter(|| run_tinywasm(FIBONACCI, black_box(26), "fibonacci_recursive")));
-        group.bench_function("wasmi", |b| b.iter(|| run_wasmi(FIBONACCI, black_box(26), "fibonacci_recursive")));
-        group.bench_function("wasmer", |b| b.iter(|| run_wasmer(FIBONACCI, black_box(26), "fibonacci_recursive")));
+        // group.bench_function("wasmi", |b| b.iter(|| run_wasmi(FIBONACCI, black_box(26), "fibonacci_recursive")));
+        // group.bench_function("wasmer", |b| b.iter(|| run_wasmer(FIBONACCI, black_box(26), "fibonacci_recursive")));
     }
 }
 
diff --git a/crates/parser/Cargo.toml b/crates/parser/Cargo.toml
index 03d7c9b..6d48a24 100644
--- a/crates/parser/Cargo.toml
+++ b/crates/parser/Cargo.toml
@@ -16,3 +16,4 @@ tinywasm-types={version="0.7.0", path="../types", default-features=false}
 default=["std", "logging"]
 logging=["log"]
 std=["tinywasm-types/std", "wasmparser/std"]
+nightly=[]
diff --git a/crates/tinywasm/Cargo.toml b/crates/tinywasm/Cargo.toml
index f5aaf50..78c1b42 100644
--- a/crates/tinywasm/Cargo.toml
+++ b/crates/tinywasm/Cargo.toml
@@ -34,7 +34,7 @@ std=["tinywasm-parser?/std", "tinywasm-types/std"]
 parser=["tinywasm-parser"]
 archive=["tinywasm-types/archive"]
 simd=[]
-nightly=[]
+nightly=["tinywasm-parser?/nightly"]
 
 [[test]]
 name="test-mvp"
diff --git a/crates/tinywasm/src/boxvec.rs b/crates/tinywasm/src/boxvec.rs
new file mode 100644
index 0000000..bf7cd8b
--- /dev/null
+++ b/crates/tinywasm/src/boxvec.rs
@@ -0,0 +1,120 @@
+use crate::unlikely;
+use alloc::{borrow::Cow, boxed::Box, vec};
+use core::ops::RangeBounds;
+
+// A Vec-like type that doesn't deallocate memory when popping elements.
+#[derive(Debug)]
+pub(crate) struct BoxVec<T> {
+    pub(crate) data: Box<[T]>,
+    pub(crate) end: usize,
+}
+
+impl<T: Copy + Default> BoxVec<T> {
+    #[inline(always)]
+    pub(crate) fn with_capacity(capacity: usize) -> Self {
+        Self { data: vec![T::default(); capacity].into_boxed_slice(), end: 0 }
+    }
+
+    #[inline(always)]
+    pub(crate) fn push(&mut self, value: T) {
+        assert!(self.end <= self.data.len(), "stack overflow");
+        self.data[self.end] = value;
+        self.end += 1;
+    }
+
+    #[inline(always)]
+    pub(crate) fn pop(&mut self) -> Option<T> {
+        assert!(self.end <= self.data.len(), "invalid stack state (should be impossible)");
+        if unlikely(self.end == 0) {
+            None
+        } else {
+            self.end -= 1;
+            Some(self.data[self.end])
+        }
+    }
+
+    #[inline(always)]
+    pub(crate) fn len(&self) -> usize {
+        self.end
+    }
+
+    #[inline(always)]
+    pub(crate) fn extend_from_slice(&mut self, values: &[T]) {
+        let new_end = self.end + values.len();
+        assert!(new_end <= self.data.len(), "stack overflow");
+        self.data[self.end..new_end].copy_from_slice(values);
+        self.end = new_end;
+    }
+
+    #[inline(always)]
+    pub(crate) fn last_mut(&mut self) -> Option<&mut T> {
+        assert!(self.end <= self.data.len(), "invalid stack state (should be impossible)");
+        if unlikely(self.end == 0) {
+            None
+        } else {
+            Some(&mut self.data[self.end - 1])
+        }
+    }
+
+    #[inline(always)]
+    pub(crate) fn last(&self) -> Option<&T> {
+        assert!(self.end <= self.data.len(), "invalid stack state (should be impossible)");
+        if unlikely(self.end == 0) {
+            None
+        } else {
+            Some(&self.data[self.end - 1])
+        }
+    }
+
+    #[inline(always)]
+    pub(crate) fn drain(&mut self, range: impl RangeBounds<usize>) -> Cow<'_, [T]> {
+        let start = match range.start_bound() {
+            core::ops::Bound::Included(&start) => start,
+            core::ops::Bound::Excluded(&start) => start + 1,
+            core::ops::Bound::Unbounded => 0,
+        };
+        let end = match range.end_bound() {
+            core::ops::Bound::Included(&end) => end + 1,
+            core::ops::Bound::Excluded(&end) => end,
+            core::ops::Bound::Unbounded => self.end,
+        };
+
+        assert!(start <= end);
+        assert!(end <= self.end);
+
+        if end == self.end {
+            self.end = start;
+            return Cow::Borrowed(&self.data[start..end]);
+        }
+
+        let drain = self.data[start..end].to_vec();
+        self.data.copy_within(end..self.end, start);
+        self.end -= end - start;
+        Cow::Owned(drain)
+    }
+}
+
+impl<T> core::ops::Index<usize> for BoxVec<T> {
+    type Output = T;
+
+    #[inline(always)]
+    fn index(&self, index: usize) -> &T {
+        &self.data[index]
+    }
+}
+
+impl<T> core::ops::Index<core::ops::Range<usize>> for BoxVec<T> {
+    type Output = [T];
+
+    #[inline(always)]
+    fn index(&self, index: core::ops::Range<usize>) -> &[T] {
+        &self.data[index]
+    }
+}
+
+impl<T> core::ops::IndexMut<usize> for BoxVec<T> {
+    #[inline(always)]
+    fn index_mut(&mut self, index: usize) -> &mut T {
+        &mut self.data[index]
+    }
+}
diff --git a/crates/tinywasm/src/func.rs b/crates/tinywasm/src/func.rs
index e43b680..a6e16ad 100644
--- a/crates/tinywasm/src/func.rs
+++ b/crates/tinywasm/src/func.rs
@@ -59,8 +59,8 @@ impl FuncHandle {
         };
 
         // 6. Let f be the dummy frame
-        let call_frame_params = params.iter().map(|v| RawWasmValue::from(*v));
-        let call_frame = CallFrame::new(wasm_func.clone(), func_inst.owner, call_frame_params, 0);
+        let call_frame_params = params.iter().map(|v| RawWasmValue::from(*v)).collect::<Vec<_>>();
+        let call_frame = CallFrame::new(wasm_func.clone(), func_inst.owner, &call_frame_params, 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/lib.rs b/crates/tinywasm/src/lib.rs
index 9c48ee0..0ab61f8 100644
--- a/crates/tinywasm/src/lib.rs
+++ b/crates/tinywasm/src/lib.rs
@@ -100,6 +100,7 @@ pub use module::Module;
 pub use reference::*;
 pub use store::*;
 
+mod boxvec;
 mod func;
 mod imports;
 mod instance;
diff --git a/crates/tinywasm/src/runtime/interpreter/mod.rs b/crates/tinywasm/src/runtime/interpreter/mod.rs
index 5667089..b98d895 100644
--- a/crates/tinywasm/src/runtime/interpreter/mod.rs
+++ b/crates/tinywasm/src/runtime/interpreter/mod.rs
@@ -64,13 +64,13 @@ impl<'store, 'stack> Executor<'store, 'stack> {
     }
 
     #[inline(always)]
-    pub(crate) fn exec_next(&mut self) -> Result<ControlFlow<()>> {
+    fn exec_next(&mut self) -> Result<ControlFlow<()>> {
         use tinywasm_types::Instruction::*;
         match self.cf.fetch_instr() {
-            Nop => cold(),
+            Nop => self.exec_noop(),
             Unreachable => self.exec_unreachable()?,
 
-            Drop => self.stack.values.pop().map(|_| ())?,
+            Drop => self.exec_drop()?,
             Select(_valtype) => self.exec_select()?,
 
             Call(v) => return self.exec_call_direct(*v),
@@ -80,7 +80,7 @@ impl<'store, 'stack> Executor<'store, 'stack> {
             Else(end_offset) => self.exec_else(*end_offset)?,
             Loop(args, end) => self.enter_block(self.cf.instr_ptr, *end, BlockType::Loop, *args),
             Block(args, end) => self.enter_block(self.cf.instr_ptr, *end, BlockType::Block, *args),
-            Br(v) => break_to!(*v, self),
+            Br(v) => return self.exec_br(*v),
             BrIf(v) => return self.exec_br_if(*v),
             BrTable(default, len) => return self.exec_brtable(*default, *len),
             Return => return self.exec_return(),
@@ -310,7 +310,6 @@ impl<'store, 'stack> Executor<'store, 'stack> {
             I32StoreLocal { local, const_i32, offset, mem_addr } => {
                 self.exec_i32_store_local(*local, *const_i32, *offset, *mem_addr)?
             }
-
             i => {
                 cold();
                 return Err(Error::UnsupportedFeature(format!("unimplemented instruction: {:?}", i)));
@@ -345,6 +344,13 @@ impl<'store, 'stack> Executor<'store, 'stack> {
     }
 
     #[inline(always)]
+    fn exec_br(&mut self, to: u32) -> Result<ControlFlow<()>> {
+        break_to!(to, self);
+        self.cf.instr_ptr += 1;
+        Ok(ControlFlow::Continue(()))
+    }
+
+    #[inline(always)]
     fn exec_br_if(&mut self, to: u32) -> Result<ControlFlow<()>> {
         let val: i32 = self.stack.values.pop()?.into();
         if val != 0 {
@@ -397,6 +403,9 @@ impl<'store, 'stack> Executor<'store, 'stack> {
     }
 
     #[inline(always)]
+    fn exec_noop(&self) {}
+
+    #[inline(always)]
     fn exec_ref_is_null(&mut self) -> Result<()> {
         self.stack.values.replace_top(|val| ((i32::from(val) == -1) as i32).into())
     }
@@ -551,19 +560,25 @@ impl<'store, 'stack> Executor<'store, 'stack> {
         let table_idx = self.module.resolve_table_addr(table_index);
         let table = self.store.get_table(table_idx)?;
         let delta: i32 = self.stack.values.pop()?.into();
-        let prev_size = table.borrow().size() as i32;
+        let prev_size = table.borrow().size();
         table.borrow_mut().grow_to_fit((prev_size + delta) as usize)?;
         self.stack.values.push(prev_size.into());
         Ok(())
     }
 
     #[inline(always)]
-    fn exec_table_fill(&mut self, table_index: u32) -> Result<()> {
+    fn exec_table_fill(&mut self, _table_index: u32) -> Result<()> {
         // TODO: implement
         Ok(())
     }
 
     #[inline(always)]
+    fn exec_drop(&mut self) -> Result<()> {
+        self.stack.values.pop()?;
+        Ok(())
+    }
+
+    #[inline(always)]
     fn exec_select(&mut self) -> Result<()> {
         let cond: i32 = self.stack.values.pop()?.into();
         let val2 = self.stack.values.pop()?;
@@ -695,7 +710,7 @@ impl<'store, 'stack> Executor<'store, 'stack> {
     #[inline(always)]
     fn exec_call(&mut self, wasm_func: Rc<WasmFunction>, owner: ModuleInstanceAddr) -> Result<ControlFlow<()>> {
         let params = self.stack.values.pop_n_rev(wasm_func.ty.params.len())?;
-        let new_call_frame = CallFrame::new(wasm_func, owner, params, self.stack.blocks.len() as u32);
+        let new_call_frame = CallFrame::new(wasm_func, owner, &params, self.stack.blocks.len() as u32);
         self.cf.instr_ptr += 1; // skip the call instruction
         self.stack.call_stack.push(core::mem::replace(&mut self.cf, new_call_frame))?;
         self.module.swap_with(self.cf.module_addr, self.store);
diff --git a/crates/tinywasm/src/runtime/stack/call_stack.rs b/crates/tinywasm/src/runtime/stack/call_stack.rs
index 93f2c03..4bf678e 100644
--- a/crates/tinywasm/src/runtime/stack/call_stack.rs
+++ b/crates/tinywasm/src/runtime/stack/call_stack.rs
@@ -103,7 +103,7 @@ impl CallFrame {
     pub(crate) fn new(
         wasm_func_inst: Rc<WasmFunction>,
         owner: ModuleInstanceAddr,
-        params: impl ExactSizeIterator<Item = RawWasmValue>,
+        params: &[RawWasmValue],
         block_ptr: u32,
     ) -> Self {
         let locals = {
diff --git a/crates/tinywasm/src/runtime/stack/value_stack.rs b/crates/tinywasm/src/runtime/stack/value_stack.rs
index 6d0e21d..49710af 100644
--- a/crates/tinywasm/src/runtime/stack/value_stack.rs
+++ b/crates/tinywasm/src/runtime/stack/value_stack.rs
@@ -1,5 +1,5 @@
-use crate::{cold, runtime::RawWasmValue, unlikely, Error, Result};
-use alloc::vec::Vec;
+use crate::{boxvec::BoxVec, cold, runtime::RawWasmValue, unlikely, Error, Result};
+use alloc::{borrow::Cow, vec::Vec};
 use tinywasm_types::{ValType, WasmValue};
 
 use super::BlockFrame;
@@ -18,19 +18,19 @@ use crate::runtime::raw_simd::RawSimdWasmValue;
 
 #[derive(Debug)]
 pub(crate) struct ValueStack {
-    stack: Vec<RawWasmValue>,
+    pub(crate) stack: BoxVec<RawWasmValue>,
 
     #[cfg(feature = "simd")]
-    simd_stack: Vec<RawSimdWasmValue>,
+    simd_stack: BoxVec<RawSimdWasmValue>,
 }
 
 impl Default for ValueStack {
     fn default() -> Self {
         Self {
-            stack: Vec::with_capacity(MIN_VALUE_STACK_SIZE),
+            stack: BoxVec::with_capacity(MIN_VALUE_STACK_SIZE),
 
             #[cfg(feature = "simd")]
-            simd_stack: Vec::with_capacity(MIN_SIMD_VALUE_STACK_SIZE),
+            simd_stack: BoxVec::with_capacity(MIN_SIMD_VALUE_STACK_SIZE),
         }
     }
 }
@@ -59,9 +59,20 @@ impl ValueStack {
 
     #[inline(always)]
     pub(crate) fn calculate(&mut self, func: fn(RawWasmValue, RawWasmValue) -> RawWasmValue) -> Result<()> {
-        let v2 = self.pop()?;
-        let v1 = self.last_mut()?;
-        *v1 = func(*v1, v2);
+        if self.stack.end < 2 {
+            cold(); // cold in here instead of the stack makes a huge performance difference
+            return Err(Error::ValueStackUnderflow);
+        }
+
+        assert!(
+            self.stack.end >= 2 && self.stack.end <= self.stack.data.len(),
+            "invalid stack state (should be impossible)"
+        );
+
+        self.stack.data[self.stack.end - 2] =
+            func(self.stack.data[self.stack.end - 2], self.stack.data[self.stack.end - 1]);
+
+        self.stack.end -= 1;
         Ok(())
     }
 
@@ -113,7 +124,7 @@ impl ValueStack {
         match self.stack.last_mut() {
             Some(v) => Ok(v),
             None => {
-                cold();
+                cold(); // cold in here instead of the stack makes a huge performance difference
                 Err(Error::ValueStackUnderflow)
             }
         }
@@ -124,7 +135,7 @@ impl ValueStack {
         match self.stack.last() {
             Some(v) => Ok(v),
             None => {
-                cold();
+                cold(); // cold in here instead of the stack makes a huge performance difference
                 Err(Error::ValueStackUnderflow)
             }
         }
@@ -135,7 +146,7 @@ impl ValueStack {
         match self.stack.pop() {
             Some(v) => Ok(v),
             None => {
-                cold();
+                cold(); // cold in here instead of the stack makes a huge performance difference
                 Err(Error::ValueStackUnderflow)
             }
         }
@@ -165,10 +176,9 @@ impl ValueStack {
         self.stack.drain(bf.stack_ptr as usize..end);
 
         #[cfg(feature = "simd")]
-        {
-            let end = self.simd_stack.len() - bf.simd_results as usize;
-            self.simd_stack.drain(bf.simd_stack_ptr as usize..end);
-        }
+        let end = self.simd_stack.len() - bf.simd_results as usize;
+        #[cfg(feature = "simd")]
+        self.simd_stack.drain(bf.simd_stack_ptr as usize..end);
     }
 
     #[inline]
@@ -177,10 +187,9 @@ impl ValueStack {
         self.stack.drain(bf.stack_ptr as usize..end);
 
         #[cfg(feature = "simd")]
-        {
-            let end = self.simd_stack.len() - bf.simd_params as usize;
-            self.simd_stack.drain(bf.simd_stack_ptr as usize..end);
-        }
+        let end = self.simd_stack.len() - bf.simd_params as usize;
+        #[cfg(feature = "simd")]
+        self.simd_stack.drain(bf.simd_stack_ptr as usize..end);
     }
 
     #[inline]
@@ -193,7 +202,7 @@ impl ValueStack {
     }
 
     #[inline]
-    pub(crate) fn pop_n_rev(&mut self, n: usize) -> Result<alloc::vec::Drain<'_, RawWasmValue>> {
+    pub(crate) fn pop_n_rev(&mut self, n: usize) -> Result<Cow<'_, [RawWasmValue]>> {
         if unlikely(self.stack.len() < n) {
             return Err(Error::ValueStackUnderflow);
         }
@@ -202,7 +211,7 @@ impl ValueStack {
 }
 
 #[inline(always)]
-fn truncate_keep<T>(data: &mut Vec<T>, n: u32, end_keep: u32) {
+fn truncate_keep<T: Copy + Default>(data: &mut BoxVec<T>, n: u32, end_keep: u32) {
     let total_to_keep = n + end_keep;
     let len = data.len() as u32;
     assert!(len >= total_to_keep, "RawWasmValueotal to keep should be less than or equal to self.top");
diff --git a/rust-toolchain.toml b/rust-toolchain.toml
index a84b93a..a6f0132 100644
--- a/rust-toolchain.toml
+++ b/rust-toolchain.toml
@@ -1,2 +1,2 @@
 [toolchain]
-channel="nightly-2024-04-15"
+channel="nightly-2024-05-25"