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

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

13 files changed, 324 insertions, 249 deletions

diff --git a/crates/parser/src/lib.rs b/crates/parser/src/lib.rs
index b1729f4..fdfea49 100644
--- a/crates/parser/src/lib.rs
+++ b/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(),
         })
     }
diff --git a/crates/tinywasm/src/func.rs b/crates/tinywasm/src/func.rs
index 1fda75e..d967a45 100644
--- a/crates/tinywasm/src/func.rs
+++ b/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
-            .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)
+        Ok(res
+            .iter()
+            .zip(func_ty.results.iter())
+            .map(|(v, ty)| v.into_typed(*ty))
+            .collect())
     }
 }
 
diff --git a/crates/tinywasm/src/module.rs b/crates/tinywasm/src/module.rs
index 04320a2..41ad860 100644
--- a/crates/tinywasm/src/module.rs
+++ b/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,
diff --git a/crates/tinywasm/src/runtime/executer.rs b/crates/tinywasm/src/runtime/executer.rs
index 8b13789..19e0a29 100644
--- a/crates/tinywasm/src/runtime/executer.rs
+++ b/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)
+    }
+}
diff --git a/crates/tinywasm/src/runtime/mod.rs b/crates/tinywasm/src/runtime/mod.rs
index f6c0596..ed5701c 100644
--- a/crates/tinywasm/src/runtime/mod.rs
+++ b/crates/tinywasm/src/runtime/mod.rs
@@ -1,11 +1,9 @@
 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
@@ -13,65 +11,6 @@ use crate::{Error, Result};
 /// 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)
-    }
-}
diff --git a/crates/tinywasm/src/runtime/stack.rs b/crates/tinywasm/src/runtime/stack.rs
new file mode 100644
index 0000000..113cfc6
--- /dev/null
+++ b/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,
+}
diff --git a/crates/tinywasm/src/runtime/stack/call.rs b/crates/tinywasm/src/runtime/stack/call.rs
deleted file mode 100644
index 4682b16..0000000
--- a/crates/tinywasm/src/runtime/stack/call.rs
+++ /dev/null
@@ -1,48 +0,0 @@
-use alloc::boxed::Box;
-use tinywasm_types::{ValType, WasmValue};
-
-#[derive(Debug)]
-pub struct CallFrame {
-    pub instr_ptr: usize,
-    pub func_ptr: 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]
-    }
-}
diff --git a/crates/tinywasm/src/runtime/stack/call_stack.rs b/crates/tinywasm/src/runtime/stack/call_stack.rs
new file mode 100644
index 0000000..eb5256e
--- /dev/null
+++ b/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]
+    }
+}
diff --git a/crates/tinywasm/src/runtime/stack/mod.rs b/crates/tinywasm/src/runtime/stack/mod.rs
deleted file mode 100644
index 271cdd5..0000000
--- a/crates/tinywasm/src/runtime/stack/mod.rs
+++ /dev/null
@@ -1,120 +0,0 @@
-use alloc::vec::Vec;
-
-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
-pub const CALL_STACK_SIZE: usize = 1024;
-
-/// A WebAssembly Stack
-#[derive(Debug)]
-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,
-}
-
-#[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 {
-            values: ValueStack {
-                stack: Vec::with_capacity(STACK_SIZE),
-                top: 0,
-            },
-            call_stack: CallStack {
-                stack: Vec::with_capacity(CALL_STACK_SIZE),
-                top: 0,
-            },
-        }
-    }
-}
diff --git a/crates/tinywasm/src/runtime/stack/value_stack.rs b/crates/tinywasm/src/runtime/stack/value_stack.rs
new file mode 100644
index 0000000..83f54c9
--- /dev/null
+++ b/crates/tinywasm/src/runtime/stack/value_stack.rs
@@ -0,0 +1,67 @@
+use crate::{runtime::UntypedWasmValue, Error, Result};
+use alloc::vec::Vec;
+
+// minimum stack size
+pub const STACK_SIZE: usize = 1024;
+
+#[derive(Debug)]
+pub struct ValueStack {
+    stack: Vec<UntypedWasmValue>,
+    top: usize,
+}
+
+impl Default for ValueStack {
+    fn default() -> Self {
+        Self {
+            stack: Vec::with_capacity(STACK_SIZE),
+            top: 0,
+        }
+    }
+}
+
+impl ValueStack {
+    #[inline]
+    pub(crate) fn _extend(
+        &mut self,
+        values: impl IntoIterator<Item = UntypedWasmValue> + ExactSizeIterator,
+    ) {
+        self.top += values.len();
+        self.stack.extend(values);
+    }
+
+    #[inline]
+    pub(crate) fn push(&mut self, value: UntypedWasmValue) {
+        self.top += 1;
+        self.stack.push(value);
+    }
+
+    #[inline]
+    pub(crate) fn pop(&mut self) -> Option<UntypedWasmValue> {
+        self.top -= 1;
+        self.stack.pop()
+    }
+
+    #[inline]
+    pub(crate) fn pop_n(&mut self, n: usize) -> Result<Vec<UntypedWasmValue>> {
+        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<[UntypedWasmValue; N]> {
+        if self.top < N {
+            return Err(Error::StackUnderflow);
+        }
+        self.top -= N;
+        let mut res = [UntypedWasmValue::default(); N];
+        for i in res.iter_mut().rev() {
+            *i = self.stack.pop().ok_or(Error::InvalidStore)?;
+        }
+
+        Ok(res)
+    }
+}
diff --git a/crates/tinywasm/src/runtime/value.rs b/crates/tinywasm/src/runtime/value.rs
new file mode 100644
index 0000000..dcab276
--- /dev/null
+++ b/crates/tinywasm/src/runtime/value.rs
@@ -0,0 +1,77 @@
+use tinywasm_types::{ValType, WasmValue};
+
+#[derive(Debug, Clone, Copy, Default)]
+pub struct UntypedWasmValue(u64);
+
+impl UntypedWasmValue {
+    pub fn into_typed(self, ty: ValType) -> WasmValue {
+        match ty {
+            ValType::I32 => WasmValue::I32(self.0 as i32),
+            ValType::I64 => WasmValue::I64(self.0 as i64),
+            ValType::F32 => WasmValue::F32(f32::from_bits(self.0 as u32)),
+            ValType::F64 => WasmValue::F64(f64::from_bits(self.0)),
+            ValType::ExternRef => todo!(),
+            ValType::FuncRef => todo!(),
+            ValType::V128 => todo!(),
+        }
+    }
+}
+
+impl From<i32> for UntypedWasmValue {
+    fn from(i: i32) -> Self {
+        Self(i as u64)
+    }
+}
+
+impl From<UntypedWasmValue> for i32 {
+    fn from(v: UntypedWasmValue) -> Self {
+        v.0 as i32
+    }
+}
+
+impl From<i64> for UntypedWasmValue {
+    fn from(i: i64) -> Self {
+        Self(i as u64)
+    }
+}
+
+impl From<UntypedWasmValue> for i64 {
+    fn from(v: UntypedWasmValue) -> Self {
+        v.0 as i64
+    }
+}
+
+impl From<f32> for UntypedWasmValue {
+    fn from(i: f32) -> Self {
+        Self(i.to_bits() as u64)
+    }
+}
+
+impl From<UntypedWasmValue> for f32 {
+    fn from(v: UntypedWasmValue) -> Self {
+        f32::from_bits(v.0 as u32)
+    }
+}
+
+impl From<f64> for UntypedWasmValue {
+    fn from(i: f64) -> Self {
+        Self(i.to_bits())
+    }
+}
+
+impl From<UntypedWasmValue> for f64 {
+    fn from(v: UntypedWasmValue) -> Self {
+        f64::from_bits(v.0)
+    }
+}
+
+impl From<WasmValue> for UntypedWasmValue {
+    fn from(v: WasmValue) -> Self {
+        match v {
+            WasmValue::I32(i) => Self(i as u64),
+            WasmValue::I64(i) => Self(i as u64),
+            WasmValue::F32(i) => Self(i.to_bits() as u64),
+            WasmValue::F64(i) => Self(i.to_bits()),
+        }
+    }
+}
diff --git a/crates/tinywasm/src/store.rs b/crates/tinywasm/src/store.rs
index 85ffdbc..ec0d3b0 100644
--- a/crates/tinywasm/src/store.rs
+++ b/crates/tinywasm/src/store.rs
@@ -103,7 +103,7 @@ impl Store {
         idx: ModuleInstanceAddr,
     ) -> Vec<FuncAddr> {
         let mut func_addrs = Vec::with_capacity(funcs.len());
-        for func in funcs {
+        for func in funcs.into_iter() {
             self.data.funcs.push(FunctionInstance {
                 func,
                 _module_instance: idx,
diff --git a/crates/types/src/lib.rs b/crates/types/src/lib.rs
index 6430b03..875deaf 100644
--- a/crates/types/src/lib.rs
+++ b/crates/types/src/lib.rs
@@ -10,7 +10,7 @@ pub struct TinyWasmModule {
     pub version: Option<u16>,
     pub start_func: Option<FuncAddr>,
 
-    pub funcs: Vec<Function>,
+    pub funcs: Box<[Function]>,
     pub types: Box<[FuncType]>,
     pub exports: Box<[Export]>,
     // pub tables: Option<TableType>,