feat: stack is now untyped, funcs is a boxed array again

Signed-off-by: Henry Gressmann <[email protected]>

crates/parser/src/lib.rs

@@ -108,7 +108,7 @@ impl TryFrom<ModuleReader> for TinyWasmModule {
             version: reader.version,
             start_func: reader.start_func,
             types: reader.type_section.into_boxed_slice(),
-            funcs,
+            funcs: funcs.into_boxed_slice(),
             exports: reader.export_section.into_boxed_slice(),
         })
     }

crates/tinywasm/src/func.rs

@@ -50,7 +50,7 @@ impl FuncHandle {
         }
 
         // 6. Let f be the dummy frame
-        let call_frame = CallFrame::new(self.addr as usize, params, func_inst.locals().iter());
+        let call_frame = CallFrame::new(self.addr as usize, params, func_inst.locals().to_vec());
 
         // 7. Push the frame f to the call stack
         stack.call_stack.push(call_frame);
@@ -64,19 +64,12 @@ impl FuncHandle {
         // Once the function returns:
         let result_m = func_ty.results.len();
         let res = stack.values.pop_n(result_m)?;
-        func_ty
-            .results
+
+        Ok(res
             .iter()
-            .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)
+            .zip(func_ty.results.iter())
+            .map(|(v, ty)| v.into_typed(*ty))
+            .collect())
     }
 }
 

crates/tinywasm/src/module.rs

@@ -45,7 +45,7 @@ impl Module {
     ) -> Result<ModuleInstance> {
         let idx = store.next_module_instance_idx();
 
-        let func_addrs = store.add_funcs(self.data.funcs, idx);
+        let func_addrs = store.add_funcs(self.data.funcs.into(), idx);
         let instance = ModuleInstance::new(
             self.data.types,
             self.data.start_func,

crates/tinywasm/src/runtime/executer.rs

@@ -1 +1,65 @@
+use super::{Runtime, Stack};
+use crate::{Error, Result};
+use log::debug;
+use tinywasm_types::Instruction;
 
+impl<const CHECK_TYPES: bool> Runtime<CHECK_TYPES> {
+    pub(crate) fn exec(
+        &self,
+        stack: &mut Stack,
+        instrs: core::slice::Iter<Instruction>,
+    ) -> Result<()> {
+        let call_frame = stack.call_stack.top_mut()?;
+
+        for instr in instrs {
+            use tinywasm_types::Instruction::*;
+            match instr {
+                LocalGet(local_index) => {
+                    let val = call_frame.get_local(*local_index as usize);
+                    debug!("local: {:#?}", val);
+                    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, b] = stack.values.pop_n_const::<2>()?;
+                    let a: i64 = a.into();
+                    let b: i64 = b.into();
+                    // let (WasmValue::I64(a), WasmValue::I64(b)) = (a, b) else {
+                    //     panic!("Invalid type");
+                    // };
+                    let c = a + b;
+                    stack.values.push(c.into());
+                }
+                I32Add => {
+                    let [a, b] = stack.values.pop_n_const::<2>()?;
+                    debug!("i64.add: {:?} + {:?}", a, b);
+                    let a: i32 = a.into();
+                    let b: i32 = b.into();
+                    // let (WasmValue::I32(a), WasmValue::I32(b)) = (a, b) else {
+                    //     panic!("Invalid type");
+                    // };
+                    stack.values.push((a + b).into());
+                }
+                I32Sub => {
+                    let [a, b] = stack.values.pop_n_const::<2>()?;
+                    // let (WasmValue::I32(a), WasmValue::I32(b)) = (a, b) else {
+                    //     panic!("Invalid type");
+                    // };
+                    let a: i32 = a.into();
+                    let b: i32 = b.into();
+                    stack.values.push((a - b).into());
+                }
+                End => {
+                    debug!("stack: {:?}", stack);
+                    return Ok(());
+                }
+                _ => todo!(),
+            }
+        }
+
+        Err(Error::FuncDidNotReturn)
+    }
+}

crates/tinywasm/src/runtime/mod.rs

@@ -1,77 +1,16 @@
 mod executer;
 mod stack;
+mod value;
 
-use log::debug;
 pub use stack::*;
-use tinywasm_types::{Instruction, WasmValue};
-
-use crate::{Error, Result};
+pub use value::UntypedWasmValue;
 
 /// 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.
+// Execution is implemented in the `executer` module
 #[derive(Debug, Default)]
 pub struct Runtime<const CHECK_TYPES: bool> {}
-
-impl<const CHECK_TYPES: bool> Runtime<CHECK_TYPES> {
-    pub(crate) fn exec(
-        &self,
-        stack: &mut Stack,
-        instrs: core::slice::Iter<Instruction>,
-    ) -> Result<()> {
-        let call_frame = stack.call_stack.top_mut()?;
-
-        for instr in instrs {
-            use tinywasm_types::Instruction::*;
-            match instr {
-                LocalGet(local_index) => {
-                    let val = call_frame.get_local(*local_index as usize);
-                    debug!("local: {:#?}", val);
-                    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, 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);
-                    stack.values.push(c);
-                }
-                I32Add => {
-                    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);
-                    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!(),
-            }
-        }
-
-        Err(Error::FuncDidNotReturn)
-    }
-}

crates/tinywasm/src/runtime/stack.rs

@@ -0,0 +1,15 @@
+mod call_stack;
+mod value_stack;
+use self::{call_stack::CallStack, value_stack::ValueStack};
+pub use call_stack::CallFrame;
+
+/// A WebAssembly Stack
+#[derive(Debug, Default)]
+pub struct Stack {
+    // 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(crate) call_stack: CallStack,
+}

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

@@ -0,0 +1,88 @@
+use crate::{runtime::UntypedWasmValue, Error, Result};
+use alloc::{boxed::Box, vec::Vec};
+use tinywasm_types::{ValType, WasmValue};
+
+// minimum call stack size
+pub const CALL_STACK_SIZE: usize = 1024;
+
+#[derive(Debug)]
+pub struct CallStack {
+    stack: Vec<CallFrame<true>>,
+    top: usize,
+}
+
+impl Default for CallStack {
+    fn default() -> Self {
+        Self {
+            stack: Vec::with_capacity(CALL_STACK_SIZE),
+            top: 0,
+        }
+    }
+}
+
+impl CallStack {
+    #[inline]
+    pub(crate) fn _top(&self) -> Result<&CallFrame<true>> {
+        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<true>> {
+        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<true>) {
+        self.top += 1;
+        self.stack.push(call_frame);
+    }
+}
+
+#[derive(Debug)]
+pub struct CallFrame<const CHECK: bool> {
+    pub instr_ptr: usize,
+    pub func_ptr: usize,
+
+    pub locals: Box<[UntypedWasmValue]>,
+    pub local_count: usize,
+}
+
+impl<const CHECK: bool> CallFrame<CHECK> {
+    pub fn new(func_ptr: usize, params: &[WasmValue], local_types: Vec<ValType>) -> Self {
+        let mut locals = Vec::with_capacity(local_types.len() + params.len());
+        locals.extend(params.iter().map(|v| UntypedWasmValue::from(*v)));
+
+        Self {
+            instr_ptr: 0,
+            func_ptr,
+            local_count: locals.len(),
+            locals: locals.into_boxed_slice(),
+        }
+    }
+
+    #[inline]
+    pub(crate) fn set_local(&mut self, local_index: usize, value: UntypedWasmValue) {
+        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) -> UntypedWasmValue {
+        if local_index >= self.local_count {
+            panic!("Invalid local index");
+        }
+
+        self.locals[local_index]
+    }
+}