chore: refactor stack structs

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

crates/tinywasm/src/error.rs

@@ -10,6 +10,7 @@ pub enum Error {
 
     FuncDidNotReturn,
     StackUnderflow,
+    CallStackEmpty,
 
     InvalidStore,
 
@@ -24,6 +25,7 @@ impl Display for Error {
             Self::StackUnderflow => write!(f, "stack underflow"),
             Self::ParseError(err) => write!(f, "error parsing module: {:?}", err),
             Self::UnsupportedFeature(feature) => write!(f, "unsupported feature: {}", feature),
+            Self::CallStackEmpty => write!(f, "call stack empty"),
             Self::Other(message) => write!(f, "unknown error: {}", message),
             Self::InvalidStore => write!(f, "invalid store"),
             #[cfg(feature = "std")]

crates/tinywasm/src/func.rs

@@ -1,7 +1,10 @@
 use alloc::{format, string::String, string::ToString, vec, vec::Vec};
-use tinywasm_types::{FuncAddr, FuncType, ValType, WasmValue};
+use tinywasm_types::{FuncAddr, FuncType, WasmValue};
 
-use crate::{runtime::Stack, Error, ModuleInstance, Result, Store};
+use crate::{
+    runtime::{CallFrame, Stack},
+    Error, ModuleInstance, Result, Store,
+};
 
 #[derive(Debug)]
 pub struct FuncHandle {
@@ -12,45 +15,66 @@ pub struct FuncHandle {
 }
 impl FuncHandle {
     /// Call a function
+    /// See https://webassembly.github.io/spec/core/exec/modules.html#invocation
     pub fn call(&self, store: &mut Store, params: &[WasmValue]) -> Result<Vec<WasmValue>> {
-        let func = store
+        let mut stack = Stack::default();
+
+        // 1. Assert: funcs[func_addr] exists
+        // 2. let func_inst be the functiuon instance funcs[func_addr]
+        let func_inst = store
             .data
             .funcs
             .get(self.addr as usize)
             .ok_or(Error::Other(format!("function {} not found", self.addr)))?;
 
+        // 3. Let func_ty be the function type
         let func_ty = &self.ty;
 
-        // check that params match func_ty params
-        for (ty, param) in func_ty.params.iter().zip(params) {
+        // 4. If the length of the provided argument values is different from the number of expected arguments, then fail
+        if func_ty.params.len() != params.len() {
+            return Err(Error::Other(format!(
+                "param count mismatch: expected {}, got {}",
+                func_ty.params.len(),
+                params.len()
+            )));
+        }
+
+        // 5. For each value type and the corresponding value, check if types match
+        for (i, (ty, param)) in func_ty.params.iter().zip(params).enumerate() {
             if ty != &param.val_type() {
                 return Err(Error::Other(format!(
-                    "param type mismatch: expected {:?}, got {:?}",
-                    ty, param
+                    "param type mismatch at index {}: expected {:?}, got {:?}",
+                    i, ty, param
                 )));
             }
         }
 
-        let mut local_types: Vec<ValType> = Vec::new();
-        local_types.extend(func_ty.params.iter());
-        local_types.extend(func.locals().iter());
+        // 6. Let f be the dummy frame
+        let call_frame = CallFrame::new(self.addr as usize, params, func_inst.locals().iter());
 
-        // let runtime = &mut store.runtime;
+        // 7. Push the frame f to the call stack
+        stack.call_stack.push(call_frame);
 
-        let mut stack = Stack::default();
-        stack.locals.extend(params.iter().cloned());
+        // 8. Push the values to the stack (Not needed since the call frame owns the values)
 
-        let instrs = func.instructions().iter();
+        // 9. Invoke the function instance
+        let instrs = func_inst.instructions().iter();
         store.runtime.exec(&mut stack, instrs)?;
 
-        let res = func_ty
+        // Once the function returns:
+        let result_m = func_ty.results.len();
+        let res = stack.values.pop_n(result_m)?;
+        func_ty
             .results
             .iter()
-            .map(|_| stack.value_stack.pop())
-            .collect::<Option<Vec<_>>>()
-            .ok_or(Error::Other(
-                "function did not return the correct number of values".into(),
-            ))?;
+            .zip(res.iter())
+            .try_for_each(|(ty, val)| match ty == &val.val_type() {
+                true => Ok(()),
+                false => Err(Error::Other(format!(
+                    "result type mismatch: expected {:?}, got {:?}",
+                    ty, val
+                ))),
+            })?;
 
         Ok(res)
     }
@@ -103,6 +127,9 @@ impl_into_wasm_value_tuple!(T1, T2);
 impl_into_wasm_value_tuple!(T1, T2, T3);
 impl_into_wasm_value_tuple!(T1, T2, T3, T4);
 impl_into_wasm_value_tuple!(T1, T2, T3, T4, T5);
+impl_into_wasm_value_tuple!(T1, T2, T3, T4, T5, T6);
+impl_into_wasm_value_tuple!(T1, T2, T3, T4, T5, T6, T7);
+impl_into_wasm_value_tuple!(T1, T2, T3, T4, T5, T6, T7, T8);
 
 macro_rules! impl_from_wasm_value_tuple {
     ($($T:ident),*) => {
@@ -131,3 +158,6 @@ impl_from_wasm_value_tuple!(T1, T2);
 impl_from_wasm_value_tuple!(T1, T2, T3);
 impl_from_wasm_value_tuple!(T1, T2, T3, T4);
 impl_from_wasm_value_tuple!(T1, T2, T3, T4, T5);
+impl_from_wasm_value_tuple!(T1, T2, T3, T4, T5, T6);
+impl_from_wasm_value_tuple!(T1, T2, T3, T4, T5, T6, T7);
+impl_from_wasm_value_tuple!(T1, T2, T3, T4, T5, T6, T7, T8);

crates/tinywasm/src/runtime/mod.rs

@@ -9,45 +9,63 @@ use crate::{Error, Result};
 
 /// A WebAssembly Runtime.
 /// See https://webassembly.github.io/spec/core/exec/runtime.html
+///
+/// Generic over `CheckTypes` to enable type checking at runtime.
+/// This is useful for debugging, but should be disabled if you know
+/// that the module is valid.
 #[derive(Debug, Default)]
-pub struct Runtime {}
+pub struct Runtime<const CHECK_TYPES: bool> {}
 
-impl Runtime {
+impl<const CHECK_TYPES: bool> Runtime<CHECK_TYPES> {
     pub(crate) fn exec(
         &self,
         stack: &mut Stack,
         instrs: core::slice::Iter<Instruction>,
     ) -> Result<()> {
-        let locals = &mut stack.locals;
-        let value_stack = &mut stack.value_stack;
+        let call_frame = stack.call_stack.top_mut()?;
 
         for instr in instrs {
             use tinywasm_types::Instruction::*;
             match instr {
                 LocalGet(local_index) => {
-                    let val = &locals[*local_index as usize];
+                    let val = call_frame.get_local(*local_index as usize);
                     debug!("local: {:#?}", val);
-                    value_stack.push(val.clone());
+                    stack.values.push(val);
+                }
+                LocalSet(local_index) => {
+                    let val = stack.values.pop().ok_or(Error::StackUnderflow)?;
+                    call_frame.set_local(*local_index as usize, val);
                 }
                 I64Add => {
-                    let a = value_stack.pop().unwrap();
-                    let b = value_stack.pop().unwrap();
+                    let [a, b] = stack.values.pop_n_const::<2>()?;
+
                     let (WasmValue::I64(a), WasmValue::I64(b)) = (a, b) else {
                         panic!("Invalid type");
                     };
                     let c = WasmValue::I64(a + b);
-                    value_stack.push(c);
+                    stack.values.push(c);
                 }
                 I32Add => {
-                    let a = value_stack.pop().unwrap();
-                    let b = value_stack.pop().unwrap();
+                    let [a, b] = stack.values.pop_n_const::<2>()?;
+                    debug!("i64.add: {:?} + {:?}", a, b);
+
                     let (WasmValue::I32(a), WasmValue::I32(b)) = (a, b) else {
                         panic!("Invalid type");
                     };
                     let c = WasmValue::I32(a + b);
-                    value_stack.push(c);
+                    debug!("i64.add: {:?}", c);
+                    stack.values.push(c);
+                }
+                I32Sub => {
+                    let [a, b] = stack.values.pop_n_const::<2>()?;
+                    let (WasmValue::I32(a), WasmValue::I32(b)) = (a, b) else {
+                        panic!("Invalid type");
+                    };
+                    let c = WasmValue::I32(a - b);
+                    stack.values.push(c);
                 }
                 End => {
+                    debug!("stack: {:?}", stack);
                     return Ok(());
                 }
                 _ => todo!(),

crates/tinywasm/src/runtime/stack/call.rs

@@ -1,9 +1,48 @@
-use alloc::vec::Vec;
+use alloc::boxed::Box;
+use tinywasm_types::{ValType, WasmValue};
 
 #[derive(Debug)]
 pub struct CallFrame {
     pub instr_ptr: usize,
     pub func_ptr: usize,
 
-    pub local_addrs: Vec<usize>,
+    pub locals: Box<[WasmValue]>,
+    pub local_count: usize,
+}
+
+impl CallFrame {
+    pub fn new<'a>(
+        func_ptr: usize,
+        params: &[WasmValue],
+        local_types: impl Iterator<Item = &'a ValType>,
+    ) -> Self {
+        let mut locals = params.to_vec();
+        locals.extend(local_types.map(|ty| WasmValue::default_for(*ty)));
+        let locals = locals.into_boxed_slice();
+
+        Self {
+            instr_ptr: 0,
+            func_ptr,
+            local_count: locals.len(),
+            locals,
+        }
+    }
+
+    #[inline]
+    pub(crate) fn set_local(&mut self, local_index: usize, value: WasmValue) {
+        if local_index >= self.local_count {
+            panic!("Invalid local index");
+        }
+
+        self.locals[local_index] = value;
+    }
+
+    #[inline]
+    pub(crate) fn get_local(&self, local_index: usize) -> WasmValue {
+        if local_index >= self.local_count {
+            panic!("Invalid local index");
+        }
+
+        self.locals[local_index]
+    }
 }

crates/tinywasm/src/runtime/stack/mod.rs

@@ -4,6 +4,8 @@ mod call;
 pub use call::CallFrame;
 use tinywasm_types::WasmValue;
 
+use crate::{Error, Result};
+
 // minimum stack size
 pub const STACK_SIZE: usize = 1024;
 // minimum call stack size
@@ -12,28 +14,107 @@ pub const CALL_STACK_SIZE: usize = 1024;
 /// A WebAssembly Stack
 #[derive(Debug)]
 pub struct Stack {
-    /// Locals
-    // TODO: maybe store the locals on the stack instead?
-    pub locals: Vec<WasmValue>,
-
-    /// The value stack
-    // TODO: Split into Vec<u8> and Vec<ValType> for better memory usage?
-    pub value_stack: Vec<WasmValue>, // keeping this typed for now to make it easier to debug
-    pub value_stack_top: usize,
+    // keeping this typed for now to make it easier to debug
+    // TODO: Maybe split into Vec<u8> and Vec<ValType> for better memory usage?
+    pub(crate) values: ValueStack,
 
     /// The call stack
-    pub call_stack: Vec<CallFrame>,
-    pub call_stack_top: usize,
+    pub(crate) call_stack: CallStack,
+}
+
+#[derive(Debug)]
+pub struct CallStack {
+    stack: Vec<CallFrame>,
+    top: usize,
+}
+
+#[derive(Debug)]
+pub struct ValueStack {
+    stack: Vec<WasmValue>,
+    top: usize,
+}
+
+impl ValueStack {
+    #[inline]
+    pub(crate) fn _extend(&mut self, values: &[WasmValue]) {
+        self.top += values.len();
+        self.stack.extend(values.iter().cloned());
+    }
+
+    #[inline]
+    pub(crate) fn push(&mut self, value: WasmValue) {
+        self.top += 1;
+        self.stack.push(value);
+    }
+
+    #[inline]
+    pub(crate) fn pop(&mut self) -> Option<WasmValue> {
+        self.top -= 1;
+        self.stack.pop()
+    }
+
+    #[inline]
+    pub(crate) fn pop_n(&mut self, n: usize) -> Result<Vec<WasmValue>> {
+        if self.top < n {
+            return Err(Error::StackUnderflow);
+        }
+        self.top -= n;
+        let res = self.stack.drain(self.top..).rev().collect::<Vec<_>>();
+        Ok(res)
+    }
+
+    #[inline]
+    pub(crate) fn pop_n_const<const N: usize>(&mut self) -> Result<[WasmValue; N]> {
+        if self.top < N {
+            return Err(Error::StackUnderflow);
+        }
+        self.top -= N;
+        let mut res = [WasmValue::I32(0); N];
+        for i in res.iter_mut().rev() {
+            *i = self.stack.pop().ok_or(Error::InvalidStore)?;
+        }
+
+        Ok(res)
+    }
+}
+
+impl CallStack {
+    #[inline]
+    pub(crate) fn _top(&self) -> Result<&CallFrame> {
+        assert!(self.top <= self.stack.len());
+        if self.top == 0 {
+            return Err(Error::CallStackEmpty);
+        }
+        Ok(&self.stack[self.top - 1])
+    }
+
+    #[inline]
+    pub(crate) fn top_mut(&mut self) -> Result<&mut CallFrame> {
+        assert!(self.top <= self.stack.len());
+        if self.top == 0 {
+            return Err(Error::CallStackEmpty);
+        }
+        Ok(&mut self.stack[self.top - 1])
+    }
+
+    #[inline]
+    pub(crate) fn push(&mut self, call_frame: CallFrame) {
+        self.top += 1;
+        self.stack.push(call_frame);
+    }
 }
 
 impl Default for Stack {
     fn default() -> Self {
         Self {
-            locals: Vec::new(),
-            value_stack: Vec::with_capacity(STACK_SIZE),
-            value_stack_top: 0,
-            call_stack: Vec::with_capacity(CALL_STACK_SIZE),
-            call_stack_top: 0,
+            values: ValueStack {
+                stack: Vec::with_capacity(STACK_SIZE),
+                top: 0,
+            },
+            call_stack: CallStack {
+                stack: Vec::with_capacity(CALL_STACK_SIZE),
+                top: 0,
+            },
         }
     }
 }

crates/tinywasm/src/store.rs

@@ -24,7 +24,7 @@ pub struct Store {
     module_instance_count: usize,
 
     pub(crate) data: StoreData,
-    pub(crate) runtime: Runtime,
+    pub(crate) runtime: Runtime<true>,
 }
 
 impl PartialEq for Store {