chore: cleanup

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

5 files changed, 70 insertions, 78 deletions

diff --git a/crates/parser/src/module.rs b/crates/parser/src/module.rs
index a4afad0..27fcb1b 100644
--- a/crates/parser/src/module.rs
+++ b/crates/parser/src/module.rs
@@ -210,6 +210,12 @@ impl ModuleReader {
             .map(|(func_idx, ((instructions, data, locals), ty_idx))| {
                 let ty = self.func_types.get(ty_idx as usize).expect("No func type for func, this is a bug").clone();
                 let params = ValueCountsSmall::from(&ty.params);
+                let locals = ValueCountsSmall {
+                    c32: u16::try_from(locals.c32).unwrap_or_else(|_| unreachable!("local count exceeds u16")),
+                    c64: u16::try_from(locals.c64).unwrap_or_else(|_| unreachable!("local count exceeds u16")),
+                    c128: u16::try_from(locals.c128).unwrap_or_else(|_| unreachable!("local count exceeds u16")),
+                    cref: u16::try_from(locals.cref).unwrap_or_else(|_| unreachable!("local count exceeds u16")),
+                };
                 let self_func_addr = imported_func_count + func_idx as u32;
                 let mut instructions = instructions.to_vec();
                 Self::apply_instruction_rewrites(&mut instructions, self_func_addr);
diff --git a/crates/tinywasm/src/func.rs b/crates/tinywasm/src/func.rs
index 98de9a2..fd1279f 100644
--- a/crates/tinywasm/src/func.rs
+++ b/crates/tinywasm/src/func.rs
@@ -66,8 +66,8 @@ impl FuncHandle {
         // Reset stack, push args, allocate locals, create entry frame.
         store.stack.clear();
         store.stack.values.extend_from_wasmvalues(params)?;
-        let (locals_base, _stack_base, stack_offset) =
-            store.stack.values.enter_locals(wasm_func.params, wasm_func.locals)?;
+        let locals_base = store.stack.values.enter_locals(&wasm_func.params, &wasm_func.locals)?;
+        let stack_offset = wasm_func.locals;
         let callframe = CallFrame::new(self.addr, func_inst.owner, locals_base, stack_offset);
 
         // Execute until completion and then collect result values from the stack.
@@ -100,8 +100,8 @@ impl FuncHandle {
             Function::Wasm(wasm_func) => {
                 store.stack.clear();
                 store.stack.values.extend_from_wasmvalues(params)?;
-                let (locals_base, _stack_base, stack_offset) =
-                    store.stack.values.enter_locals(wasm_func.params, wasm_func.locals)?;
+                let locals_base = store.stack.values.enter_locals(&wasm_func.params, &wasm_func.locals)?;
+                let stack_offset = wasm_func.locals;
                 let callframe = CallFrame::new(self.addr, func_inst_owner, locals_base, stack_offset);
 
                 Ok(FuncExecution {
diff --git a/crates/tinywasm/src/interpreter/executor.rs b/crates/tinywasm/src/interpreter/executor.rs
index 76464cb..85208bf 100644
--- a/crates/tinywasm/src/interpreter/executor.rs
+++ b/crates/tinywasm/src/interpreter/executor.rs
@@ -49,10 +49,7 @@ impl<'store, const BUDGETED: bool> Executor<'store, BUDGETED> {
 
     #[inline(always)]
     fn exec<const ITERATIONS: usize>(&mut self) -> Result<Option<()>> {
-        for _ in 0..ITERATIONS {
-            use tinywasm_types::Instruction::*;
-
-            macro_rules! stack_op {
+        macro_rules! stack_op {
             (simd_unary $method:ident) => { stack_op!(unary Value128, |v| v.$method()) };
             (simd_binary $method:ident) => { stack_op!(binary Value128, |a, b| a.$method(b)) };
             (unary $ty:ty, |$v:ident| $expr:expr) => { self.store.stack.values.unary::<$ty>(|$v| Ok($expr))? };
@@ -73,6 +70,9 @@ impl<'store, const BUDGETED: bool> Executor<'store, BUDGETED> {
             }};
         }
 
+        for _ in 0..ITERATIONS {
+            use tinywasm_types::Instruction::*;
+
             let next = match self.func.instructions.0.get(self.cf.instr_ptr as usize) {
                 Some(instr) => instr,
                 None => unreachable!(
@@ -164,9 +164,7 @@ impl<'store, const BUDGETED: bool> Executor<'store, BUDGETED> {
                     let mem = self.store.state.get_mem_mut(self.module.resolve_mem_addr(m.mem_addr()));
                     let addr = u64::from(self.store.stack.values.local_get::<u32>(&self.cf, *addr_local));
                     let value = self.store.stack.values.local_get::<u32>(&self.cf, *value_local).to_mem_bytes();
-                    if let Err(e) = mem.store((m.offset() + addr) as usize, value.len(), &value) {
-                        return Err(e);
-                    }
+                    mem.store((m.offset() + addr) as usize, value.len(), &value)?;
                 }
                 I32LoadLocalTee(m, addr_local, dst_local) => {
                     let mem = self.store.state.get_mem(self.module.resolve_mem_addr(m.mem_addr()));
@@ -677,24 +675,15 @@ impl<'store, const BUDGETED: bool> Executor<'store, BUDGETED> {
             self.store.stack.values.truncate_keep_counts(self.cf.locals_base, wasm_func.params);
         }
 
-        let (locals_base, _stack_base, stack_offset) =
-            match self.store.stack.values.enter_locals(wasm_func.params, wasm_func.locals) {
-                Ok(v) => v,
-                Err(Error::Trap(Trap::ValueStackOverflow)) if !IS_RETURN_CALL => {
-                    return Err(Trap::CallStackOverflow.into());
-                }
-                Err(err) => return Err(err),
-            };
-
-        let new_call_frame = CallFrame::new(func_addr, owner, locals_base, stack_offset);
+        let res = self.store.stack.values.enter_locals(&wasm_func.params, &wasm_func.locals);
+        let locals_base = res.map_err(|err| if IS_RETURN_CALL { err } else { Error::Trap(Trap::CallStackOverflow) })?;
+        let new_call_frame = CallFrame::new(func_addr, owner, locals_base, wasm_func.locals);
 
-        if IS_RETURN_CALL {
-            self.cf = new_call_frame;
-        } else {
+        if !IS_RETURN_CALL {
             self.cf.incr_instr_ptr(); // skip the call instruction
             self.store.stack.call_stack.push(self.cf)?;
-            self.cf = new_call_frame;
         }
+        self.cf = new_call_frame;
 
         if self.cf.module_addr != self.module.idx {
             self.module = self.store.get_module_instance_raw(self.cf.module_addr).clone();
@@ -730,24 +719,15 @@ impl<'store, const BUDGETED: bool> Executor<'store, BUDGETED> {
             self.store.stack.values.truncate_keep_counts(self.cf.locals_base, params);
         }
 
-        let (locals_base, _stack_base, stack_offset) = match self.store.stack.values.enter_locals(params, locals) {
-            Ok(v) => v,
-            Err(Error::Trap(Trap::ValueStackOverflow)) if !IS_RETURN_CALL => {
-                return Err(Trap::CallStackOverflow.into());
-            }
-            Err(err) => return Err(err),
-        };
-
-        let new_call_frame = CallFrame::new(self.cf.func_addr, self.cf.module_addr, locals_base, stack_offset);
+        let res = self.store.stack.values.enter_locals(&params, &locals);
+        let locals_base = res.map_err(|err| if IS_RETURN_CALL { err } else { Error::Trap(Trap::CallStackOverflow) })?;
+        let new_call_frame = CallFrame::new(self.cf.func_addr, self.cf.module_addr, locals_base, locals);
 
-        if IS_RETURN_CALL {
-            self.cf = new_call_frame;
-        } else {
+        if !IS_RETURN_CALL {
             self.cf.incr_instr_ptr();
             self.store.stack.call_stack.push(self.cf)?;
-            self.cf = new_call_frame;
         }
-
+        self.cf = new_call_frame;
         Ok(())
     }
 
@@ -987,9 +967,7 @@ impl<'store, const BUDGETED: bool> Executor<'store, BUDGETED> {
             false => self.store.stack.values.pop::<i32>() as u32 as u64,
         };
 
-        if let Err(e) = mem.store((offset + addr) as usize, val.len(), &val) {
-            return Err(e);
-        }
+        mem.store((offset + addr) as usize, val.len(), &val)?;
 
         Ok(())
     }
@@ -1009,9 +987,7 @@ impl<'store, const BUDGETED: bool> Executor<'store, BUDGETED> {
             false => u64::from(self.store.stack.values.pop::<i32>() as u32),
         };
 
-        if let Err(e) = mem.store((offset + addr) as usize, val.len(), &val) {
-            return Err(e);
-        }
+        mem.store((offset + addr) as usize, val.len(), &val)?;
 
         Ok(())
     }
diff --git a/crates/tinywasm/src/interpreter/stack/value_stack.rs b/crates/tinywasm/src/interpreter/stack/value_stack.rs
index 395f0a4..1ec20fe 100644
--- a/crates/tinywasm/src/interpreter/stack/value_stack.rs
+++ b/crates/tinywasm/src/interpreter/stack/value_stack.rs
@@ -1,6 +1,6 @@
 use alloc::boxed::Box;
 use alloc::vec::Vec;
-use tinywasm_types::{ExternRef, FuncRef, LocalAddr, ValType, ValueCounts, ValueCountsSmall, WasmValue};
+use tinywasm_types::{ExternRef, FuncRef, LocalAddr, ValType, ValueCountsSmall, WasmValue};
 
 use crate::{Result, Trap, engine::Config, interpreter::*, unlikely};
 
@@ -84,20 +84,22 @@ impl<T: Copy + Default> Stack<T> {
 
     pub(crate) fn truncate_keep(&mut self, n: usize, end_keep: usize) {
         debug_assert!(n <= self.len);
-        if n >= self.len {
+        let len = self.len;
+        if n >= len {
             return;
         }
 
-        let keep = (self.len - n).min(end_keep);
-        if keep != 0 {
-            let src = self.len - keep;
-            self.data.copy_within(src..self.len, n);
+        if end_keep == 0 {
+            self.len = n;
+            return;
         }
 
+        let keep = (len - n).min(end_keep);
+        self.data.copy_within((len - keep)..len, n);
         self.len = n + keep;
     }
 
-    pub(crate) fn enter_locals(&mut self, param_count: usize, local_count: usize) -> Result<(u32, u32)> {
+    pub(crate) fn enter_locals(&mut self, param_count: usize, local_count: usize) -> Result<u32> {
         debug_assert!(param_count <= local_count);
         let start = self.len - param_count;
         let end = start + local_count;
@@ -107,10 +109,11 @@ impl<T: Copy + Default> Stack<T> {
         }
 
         let init_start = start + param_count;
-        self.data[init_start..end].fill(T::default());
+        if init_start != end {
+            self.data[init_start..end].fill(T::default());
+        }
         self.len = end;
-
-        Ok((start as u32, end as u32))
+        Ok(start as u32)
     }
 
     pub(crate) fn select_many(&mut self, count: usize, condition: bool) {
@@ -246,33 +249,40 @@ impl ValueStack {
         val_types.into_iter().map(|val_type| self.pop_wasmvalue(*val_type))
     }
 
-    pub(crate) fn enter_locals(
-        &mut self,
-        params: ValueCountsSmall,
-        locals: ValueCounts,
-    ) -> Result<(StackBase, StackBase, ValueCountsSmall)> {
-        let stack_offset = ValueCountsSmall {
-            c32: u16::try_from(locals.c32).unwrap_or_else(|_| unreachable!("local count exceeds u16")),
-            c64: u16::try_from(locals.c64).unwrap_or_else(|_| unreachable!("local count exceeds u16")),
-            c128: u16::try_from(locals.c128).unwrap_or_else(|_| unreachable!("local count exceeds u16")),
-            cref: u16::try_from(locals.cref).unwrap_or_else(|_| unreachable!("local count exceeds u16")),
+    pub(crate) fn enter_locals(&mut self, params: &ValueCountsSmall, locals: &ValueCountsSmall) -> Result<StackBase> {
+        let locals_base32 = if params.c32 == 0 && locals.c32 == 0 {
+            self.stack_32.len as u32
+        } else {
+            self.stack_32.enter_locals(params.c32 as usize, locals.c32 as usize)?
+        };
+        let locals_base64 = if params.c64 == 0 && locals.c64 == 0 {
+            self.stack_64.len as u32
+        } else {
+            self.stack_64.enter_locals(params.c64 as usize, locals.c64 as usize)?
+        };
+        let locals_base128 = if params.c128 == 0 && locals.c128 == 0 {
+            self.stack_128.len as u32
+        } else {
+            self.stack_128.enter_locals(params.c128 as usize, locals.c128 as usize)?
+        };
+        let locals_baseref = if params.cref == 0 && locals.cref == 0 {
+            self.stack_ref.len as u32
+        } else {
+            self.stack_ref.enter_locals(params.cref as usize, locals.cref as usize)?
         };
 
-        let (locals_base32, stack_base32) = self.stack_32.enter_locals(params.c32 as usize, locals.c32 as usize)?;
-        let (locals_base64, stack_base64) = self.stack_64.enter_locals(params.c64 as usize, locals.c64 as usize)?;
-        let (locals_base128, stack_base128) =
-            self.stack_128.enter_locals(params.c128 as usize, locals.c128 as usize)?;
-        let (locals_baseref, stack_baseref) =
-            self.stack_ref.enter_locals(params.cref as usize, locals.cref as usize)?;
-
-        Ok((
-            StackBase { s32: locals_base32, s64: locals_base64, s128: locals_base128, sref: locals_baseref },
-            StackBase { s32: stack_base32, s64: stack_base64, s128: stack_base128, sref: stack_baseref },
-            stack_offset,
-        ))
+        Ok(StackBase { s32: locals_base32, s64: locals_base64, s128: locals_base128, sref: locals_baseref })
     }
 
     pub(crate) fn truncate_keep_counts(&mut self, base: StackBase, keep: ValueCountsSmall) {
+        if keep.c32 == 0 && keep.c64 == 0 && keep.c128 == 0 && keep.cref == 0 {
+            self.stack_32.len = base.s32 as usize;
+            self.stack_64.len = base.s64 as usize;
+            self.stack_128.len = base.s128 as usize;
+            self.stack_ref.len = base.sref as usize;
+            return;
+        }
+
         self.stack_32.truncate_keep(base.s32 as usize, keep.c32 as usize);
         self.stack_64.truncate_keep(base.s64 as usize, keep.c64 as usize);
         self.stack_128.truncate_keep(base.s128 as usize, keep.c128 as usize);
diff --git a/crates/types/src/lib.rs b/crates/types/src/lib.rs
index 8a815cb..2a167a4 100644
--- a/crates/types/src/lib.rs
+++ b/crates/types/src/lib.rs
@@ -262,7 +262,7 @@ impl<'a, T: IntoIterator<Item = &'a ValType>> From<T> for ValueCountsSmall {
 pub struct WasmFunction {
     pub instructions: ArcSlice<Instruction>,
     pub data: WasmFunctionData,
-    pub locals: ValueCounts,
+    pub locals: ValueCountsSmall,
     pub params: ValueCountsSmall,
     pub ty: FuncType,
 }