feat: improve lable/block handling, add support for more arithmetic opcodes

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

10 files changed, 294 insertions, 173 deletions

diff --git a/crates/parser/src/conversion.rs b/crates/parser/src/conversion.rs
index c80a7fc..432909c 100644
--- a/crates/parser/src/conversion.rs
+++ b/crates/parser/src/conversion.rs
@@ -129,7 +129,7 @@ pub fn process_operators<'a>(
                 let targets = targets
                     .targets()
                     .collect::<Result<Vec<u32>, wasmparser::BinaryReaderError>>()?;
-                instructions.push(Instruction::BrTable(def, targets.len() as u32));
+                instructions.push(Instruction::BrTable(def, targets.len()));
                 instructions.extend(targets.into_iter().map(Instruction::BrLabel));
                 continue;
             }
diff --git a/crates/tinywasm/src/error.rs b/crates/tinywasm/src/error.rs
index c186e46..c0e8d10 100644
--- a/crates/tinywasm/src/error.rs
+++ b/crates/tinywasm/src/error.rs
@@ -5,6 +5,11 @@ use core::fmt::Display;
 use tinywasm_parser::ParseError;
 
 #[derive(Debug)]
+pub enum Trap {
+    Unreachable,
+}
+
+#[derive(Debug)]
 pub enum Error {
     #[cfg(feature = "parser")]
     ParseError(ParseError),
@@ -15,6 +20,8 @@ pub enum Error {
     UnsupportedFeature(String),
     Other(String),
 
+    Trap(Trap),
+
     FuncDidNotReturn,
     StackUnderflow,
     BlockStackUnderflow,
@@ -31,6 +38,8 @@ impl Display for Error {
             #[cfg(feature = "std")]
             Self::Io(err) => write!(f, "I/O error: {}", err),
 
+            Self::Trap(trap) => write!(f, "trap: {:?}", trap),
+
             Self::Other(message) => write!(f, "unknown error: {}", message),
             Self::UnsupportedFeature(feature) => write!(f, "unsupported feature: {}", feature),
             Self::FuncDidNotReturn => write!(f, "function did not return"),
diff --git a/crates/tinywasm/src/runtime/executer.rs b/crates/tinywasm/src/runtime/executer.rs
deleted file mode 100644
index ab11fdd..0000000
--- a/crates/tinywasm/src/runtime/executer.rs
+++ /dev/null
@@ -1,155 +0,0 @@
-use super::{Runtime, Stack};
-use crate::{log::debug, runtime::RawWasmValue, Error, Result};
-use alloc::vec;
-use tinywasm_types::{BlockArgs, Instruction};
-
-#[derive(Debug)]
-#[allow(dead_code)]
-enum BlockMarker {
-    Top,
-    Loop {
-        instr_ptr: usize,
-        stack_ptr: usize,
-        args: BlockArgs,
-    },
-    If,
-    Else,
-    Block,
-}
-
-impl<const CHECK_TYPES: bool> Runtime<CHECK_TYPES> {
-    pub(crate) fn exec(&self, stack: &mut Stack, instrs: &[Instruction]) -> Result<()> {
-        let call_frame = stack.call_stack.top_mut()?;
-        let mut instr_ptr = call_frame.instr_ptr;
-
-        let mut blocks = vec![BlockMarker::Top];
-        debug!("locals: {:?}", call_frame.locals);
-
-        debug!("instrs: {:?}", instrs);
-
-        // TODO: maybe we don't need to check if the instr_ptr is valid since
-        // it should be validated by the parser
-        while let Some(instr) = instrs.get(instr_ptr) {
-            use tinywasm_types::Instruction::*;
-            match instr {
-                Loop(args) => {
-                    blocks.push(BlockMarker::Loop {
-                        instr_ptr,
-                        stack_ptr: stack.values.len(),
-                        args: *args,
-                    });
-                    debug!("loop: {:?}", args);
-                    stack.values.block_args(*args)?;
-                }
-                BrIf(v) => {
-                    // get block
-                    let block = blocks
-                        .get(blocks.len() - *v as usize - 1)
-                        .ok_or(Error::BlockStackUnderflow)?;
-
-                    match block {
-                        BlockMarker::Loop {
-                            instr_ptr: loop_instr_ptr,
-                            stack_ptr: stack_size,
-                            args: _,
-                        } => {
-                            let val = stack.values.pop().ok_or(Error::StackUnderflow)?;
-                            let val: i32 = val.into();
-
-                            // if val == 0 -> continue the loop
-                            if val != 0 {
-                                debug!("br_if: continue loop");
-                                instr_ptr = *loop_instr_ptr;
-                                stack.values.trim(*stack_size); // remove the loop values from the stack
-                            }
-
-                            // otherwise -> continue to loop end
-                        }
-                        _ => todo!(),
-                    }
-                }
-                End => {
-                    debug!("end, blocks: {:?}", blocks);
-                    debug!("     stack: {:?}", stack.values);
-                    let block = blocks.pop().ok_or(Error::BlockStackUnderflow)?;
-                    match block {
-                        BlockMarker::Top => {
-                            debug!("end: return");
-                            return Ok(());
-                        }
-                        BlockMarker::Loop {
-                            instr_ptr: _loop_instr_ptr,
-                            stack_ptr: stack_size,
-                            args,
-                        } => {
-                            debug!("end(loop): break loop");
-                            let res: &[RawWasmValue] = match args {
-                                BlockArgs::Empty => &[],
-                                BlockArgs::Type(_t) => todo!(),
-                                BlockArgs::FuncType(_t) => todo!(),
-                            };
-
-                            stack.values.trim(stack_size); // remove the loop values from the stack
-                            stack.values.extend(res.iter().copied()); // push the loop result values to the stack
-                        }
-                        _ => {
-                            panic!("Attempted to end a block that is not the top block");
-                        }
-                    }
-                }
-                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);
-                }
-                I32Const(val) => {
-                    stack.values.push((*val).into());
-                }
-                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());
-                }
-                I32LtS => {
-                    let [a, b] = stack.values.pop_n_const::<2>()?;
-                    let a: i32 = a.into();
-                    let b: i32 = b.into();
-                    stack.values.push(((a < b) as i32).into());
-                }
-                i => todo!("{:?}", i),
-            }
-
-            instr_ptr += 1;
-        }
-
-        Err(Error::FuncDidNotReturn)
-    }
-}
diff --git a/crates/tinywasm/src/runtime/executer/macros.rs b/crates/tinywasm/src/runtime/executer/macros.rs
new file mode 100644
index 0000000..154d7f7
--- /dev/null
+++ b/crates/tinywasm/src/runtime/executer/macros.rs
@@ -0,0 +1,44 @@
+/// Add two values from the stack
+macro_rules! add_instr {
+    ($ty:ty, $stack:ident) => {{
+        let [a, b] = $stack.values.pop_n_const::<2>()?;
+        let a: $ty = a.into();
+        let b: $ty = b.into();
+        $stack.values.push((a + b).into());
+    }};
+}
+
+/// Subtract the top two values on the stack
+macro_rules! sub_instr {
+    ($ty:ty, $stack:ident) => {{
+        let [a, b] = $stack.values.pop_n_const::<2>()?;
+        let a: $ty = a.into();
+        let b: $ty = b.into();
+        $stack.values.push((a - b).into());
+    }};
+}
+
+/// Divide the top two values on the stack
+macro_rules! div_instr {
+    ($ty:ty, $stack:ident) => {{
+        let [a, b] = $stack.values.pop_n_const::<2>()?;
+        let a: $ty = a.into();
+        let b: $ty = b.into();
+        $stack.values.push((a / b).into());
+    }};
+}
+
+/// Less than signed instruction
+macro_rules! lts_instr {
+    ($ty:ty, $stack:ident) => {{
+        let [a, b] = $stack.values.pop_n_const::<2>()?;
+        let a: $ty = a.into();
+        let b: $ty = b.into();
+        $stack.values.push(((a < b) as i32).into());
+    }};
+}
+
+pub(super) use add_instr;
+pub(super) use div_instr;
+pub(super) use lts_instr;
+pub(super) use sub_instr;
diff --git a/crates/tinywasm/src/runtime/executer/mod.rs b/crates/tinywasm/src/runtime/executer/mod.rs
new file mode 100644
index 0000000..28fc468
--- /dev/null
+++ b/crates/tinywasm/src/runtime/executer/mod.rs
@@ -0,0 +1,130 @@
+use super::{Runtime, Stack};
+use crate::{
+    log::debug,
+    runtime::{BlockFrame, BlockFrameType, RawWasmValue},
+    Error, Result,
+};
+use alloc::vec::Vec;
+use tinywasm_types::{BlockArgs, Instruction};
+
+mod macros;
+use macros::*;
+
+impl<const CHECK_TYPES: bool> Runtime<CHECK_TYPES> {
+    pub(crate) fn exec(&self, stack: &mut Stack, instrs: &[Instruction]) -> Result<()> {
+        let cf = stack.call_stack.top_mut()?;
+
+        // TODO: maybe we don't need to check if the instr_ptr is valid since
+        // it should be validated by the parser
+        while let Some(instr) = instrs.get(cf.instr_ptr) {
+            use tinywasm_types::Instruction::*;
+            match instr {
+                Nop => {} // do nothing
+                Unreachable => return Err(Error::Trap(crate::Trap::Unreachable)),
+                Loop(args) => {
+                    cf.blocks.push(BlockFrame {
+                        instr_ptr: cf.instr_ptr,
+                        stack_ptr: stack.values.len(),
+                        args: *args,
+                        ty: BlockFrameType::Loop,
+                    });
+                    stack.values.block_args(*args)?;
+                }
+                BrTable(_default, len) => {
+                    let instr = instrs[cf.instr_ptr + 1..cf.instr_ptr + 1 + *len]
+                        .iter()
+                        .map(|i| match i {
+                            BrLabel(l) => Ok(*l),
+                            _ => panic!("Expected BrLabel, this should have been validated by the parser"),
+                        })
+                        .collect::<Result<Vec<_>>>()?;
+
+                    if instr.len() != *len {
+                        panic!("Expected {} BrLabel instructions, got {}", len, instr.len());
+                    }
+
+                    todo!()
+                }
+                Br(v) => cf.break_to(*v, &mut stack.values)?,
+                BrIf(v) => {
+                    let val: i32 = stack.values.pop().ok_or(Error::StackUnderflow)?.into();
+                    if val > 0 {
+                        cf.break_to(*v, &mut stack.values)?
+                    };
+                }
+                End => {
+                    let blocks = &mut cf.blocks;
+                    let Some(block) = blocks.pop() else {
+                        debug!("end: no block to end, returning");
+                        return Ok(());
+                    };
+                    debug!("end, blocks: {:?}", blocks);
+                    debug!("     instr_ptr: {}", cf.instr_ptr);
+
+                    match block.ty {
+                        BlockFrameType::Loop => {
+                            debug!("end(loop): break loop");
+                            let res: &[RawWasmValue] = match block.args {
+                                BlockArgs::Empty => &[],
+                                BlockArgs::Type(_t) => todo!(),
+                                BlockArgs::FuncType(_t) => todo!(),
+                            };
+
+                            // remove the loop values from the stack
+                            stack.values.trim(block.stack_ptr);
+
+                            // push the loop result values to the stack
+                            stack.values.extend(res.iter().copied());
+                        }
+                        _ => {
+                            panic!("Attempted to end a block that is not the top block");
+                        }
+                    }
+                }
+                LocalGet(local_index) => {
+                    debug!("local.get: {:?}", local_index);
+                    let val = cf.get_local(*local_index as usize);
+                    debug!("local: {:#?}", val);
+                    stack.values.push(val);
+                }
+                LocalSet(local_index) => {
+                    debug!("local.set: {:?}", local_index);
+                    let val = stack.values.pop().ok_or(Error::StackUnderflow)?;
+                    cf.set_local(*local_index as usize, val);
+                }
+                I32Const(val) => {
+                    stack.values.push((*val).into());
+                }
+                I64Add => add_instr!(i64, stack),
+                I32Add => add_instr!(i32, stack),
+                F32Add => add_instr!(f32, stack),
+                F64Add => add_instr!(f64, stack),
+
+                I32Sub => sub_instr!(i32, stack),
+                I64Sub => sub_instr!(i64, stack),
+                F32Sub => sub_instr!(f32, stack),
+                F64Sub => sub_instr!(f64, stack),
+
+                I32LtS => lts_instr!(i32, stack),
+                I64LtS => lts_instr!(i64, stack),
+                F32Lt => lts_instr!(f32, stack),
+                F64Lt => lts_instr!(f64, stack),
+
+                I32DivS => div_instr!(i32, stack),
+                I64DivS => div_instr!(i64, stack),
+                F32Div => div_instr!(f32, stack),
+                F64Div => div_instr!(f64, stack),
+
+                i => todo!("{:?}", i),
+            }
+
+            cf.instr_ptr += 1;
+        }
+
+        debug!("end of exec");
+        debug!("stack: {:?}", stack.values);
+        debug!("insts: {:?}", instrs);
+        debug!("instr_ptr: {}", cf.instr_ptr);
+        Err(Error::FuncDidNotReturn)
+    }
+}
diff --git a/crates/tinywasm/src/runtime/stack.rs b/crates/tinywasm/src/runtime/stack.rs
index 113cfc6..d7996b7 100644
--- a/crates/tinywasm/src/runtime/stack.rs
+++ b/crates/tinywasm/src/runtime/stack.rs
@@ -1,7 +1,10 @@
+mod blocks;
 mod call_stack;
 mod value_stack;
+
 use self::{call_stack::CallStack, value_stack::ValueStack};
-pub use call_stack::CallFrame;
+pub(crate) use blocks::{BlockFrame, BlockFrameType};
+pub(crate) use call_stack::CallFrame;
 
 /// A WebAssembly Stack
 #[derive(Debug, Default)]
diff --git a/crates/tinywasm/src/runtime/stack/blocks.rs b/crates/tinywasm/src/runtime/stack/blocks.rs
new file mode 100644
index 0000000..10003d9
--- /dev/null
+++ b/crates/tinywasm/src/runtime/stack/blocks.rs
@@ -0,0 +1,51 @@
+use alloc::vec::Vec;
+use log::info;
+use tinywasm_types::BlockArgs;
+
+#[derive(Debug, Default)]
+pub(crate) struct Blocks(Vec<BlockFrame>);
+
+impl Blocks {
+    #[inline]
+    pub(crate) fn push(&mut self, block: BlockFrame) {
+        self.0.push(block);
+    }
+
+    #[inline]
+    /// get the block at the given index, where 0 is the top of the stack
+    pub(crate) fn get(&self, index: usize) -> Option<&BlockFrame> {
+        info!("get block: {}", index);
+        info!("blocks: {:?}", self.0);
+        self.0.get(self.0.len() - index - 1)
+    }
+
+    #[inline]
+    pub(crate) fn pop(&mut self) -> Option<BlockFrame> {
+        self.0.pop()
+    }
+
+    /// remove all blocks after the given index
+    #[inline]
+    pub(crate) fn trim(&mut self, index: usize) {
+        self.0.truncate(index + 1);
+    }
+}
+
+#[derive(Debug)]
+pub(crate) struct BlockFrame {
+    // position of the instruction pointer when the block was entered
+    pub instr_ptr: usize,
+    // position of the stack pointer when the block was entered
+    pub stack_ptr: usize,
+    pub args: BlockArgs,
+    pub ty: BlockFrameType,
+}
+
+#[derive(Debug, Copy, Clone)]
+#[allow(dead_code)]
+pub(crate) enum BlockFrameType {
+    Loop,
+    If,
+    Else,
+    Block,
+}
diff --git a/crates/tinywasm/src/runtime/stack/call_stack.rs b/crates/tinywasm/src/runtime/stack/call_stack.rs
index d1ae18d..9dbca46 100644
--- a/crates/tinywasm/src/runtime/stack/call_stack.rs
+++ b/crates/tinywasm/src/runtime/stack/call_stack.rs
@@ -2,12 +2,14 @@ use crate::{runtime::RawWasmValue, Error, Result};
 use alloc::{boxed::Box, vec::Vec};
 use tinywasm_types::{ValType, WasmValue};
 
+use super::blocks::Blocks;
+
 // minimum call stack size
 pub const CALL_STACK_SIZE: usize = 1024;
 
 #[derive(Debug)]
 pub struct CallStack {
-    stack: Vec<CallFrame<true>>,
+    stack: Vec<CallFrame>,
     top: usize,
 }
 
@@ -22,7 +24,7 @@ impl Default for CallStack {
 
 impl CallStack {
     #[inline]
-    pub(crate) fn _top(&self) -> Result<&CallFrame<true>> {
+    pub(crate) fn _top(&self) -> Result<&CallFrame> {
         assert!(self.top <= self.stack.len());
         if self.top == 0 {
             return Err(Error::CallStackEmpty);
@@ -31,7 +33,7 @@ impl CallStack {
     }
 
     #[inline]
-    pub(crate) fn top_mut(&mut self) -> Result<&mut CallFrame<true>> {
+    pub(crate) fn top_mut(&mut self) -> Result<&mut CallFrame> {
         assert!(self.top <= self.stack.len());
         if self.top == 0 {
             return Err(Error::CallStackEmpty);
@@ -40,22 +42,40 @@ impl CallStack {
     }
 
     #[inline]
-    pub(crate) fn push(&mut self, call_frame: CallFrame<true>) {
+    pub(crate) fn push(&mut self, call_frame: CallFrame) {
         self.top += 1;
         self.stack.push(call_frame);
     }
 }
 
 #[derive(Debug)]
-pub struct CallFrame<const CHECK: bool> {
+pub struct CallFrame {
     pub instr_ptr: usize,
     pub func_ptr: usize,
 
+    pub blocks: Blocks,
     pub locals: Box<[RawWasmValue]>,
     pub local_count: usize,
 }
 
-impl<const CHECK: bool> CallFrame<CHECK> {
+impl CallFrame {
+    /// Break to a block at the given index (relative to the current frame)
+    #[inline]
+    pub fn break_to(&mut self, block_index: u32, value_stack: &mut super::ValueStack) -> Result<()> {
+        let block = self
+            .blocks
+            .get(block_index as usize)
+            .ok_or(Error::BlockStackUnderflow)?;
+
+        self.instr_ptr = block.instr_ptr;
+        value_stack.trim(block.stack_ptr);
+
+        // -2 because the block we're breaking to is still on the stack
+        // TODO: this might be wrong
+        self.blocks.trim(block_index as usize - 2);
+        Ok(())
+    }
+
     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| RawWasmValue::from(*v)));
@@ -66,6 +86,7 @@ impl<const CHECK: bool> CallFrame<CHECK> {
             func_ptr,
             local_count: locals.len(),
             locals: locals.into_boxed_slice(),
+            blocks: Blocks::default(),
         }
     }
 
diff --git a/crates/types/src/instructions.rs b/crates/types/src/instructions.rs
index 40bb9f4..005da36 100644
--- a/crates/types/src/instructions.rs
+++ b/crates/types/src/instructions.rs
@@ -14,6 +14,9 @@ pub struct MemArg {
     pub offset: u64,
 }
 
+type BrTableDefault = u32;
+type BrTableLen = usize;
+
 /// A WebAssembly Instruction
 /// See https://webassembly.github.io/spec/core/binary/instructions.html
 /// These are our own internal bytecode instructions so they may not match the spec exactly.
@@ -35,7 +38,7 @@ pub enum Instruction {
     End,
     Br(LabelAddr),
     BrIf(LabelAddr),
-    BrTable(u32, u32), // has to be followed by multiple BrLabel instructions
+    BrTable(BrTableDefault, BrTableLen), // has to be followed by multiple BrLabel instructions
     Return,
     Call(FuncAddr),
     CallIndirect(TypeAddr, TableAddr),
diff --git a/examples/wasm/loop.wat b/examples/wasm/loop.wat
index 02683e7..ecdf136 100644
--- a/examples/wasm/loop.wat
+++ b/examples/wasm/loop.wat
@@ -4,16 +4,31 @@
     (i32.const 0)  ;; Initialize 'i' to 0
     (local.set 0)
 
-    (loop $loopStart  ;; Start of the loop
-      (local.get 0)  ;; Get the current value of 'i'
-      (i32.const 1)  ;; Push 1 onto the stack
-      (i32.add)      ;; Add 'i' and 1
-      (local.set 0)  ;; Update 'i' with the new value
+    ;; (loop $test
+    ;;   (loop $test2
+    ;;     (local.get 0)
+    ;;     (i32.const 1)
+    ;;     (i32.add)
+    ;;     (local.set 0)
+    ;;     (br $test2)
+    ;;   )
+    ;; )
 
-      (local.get 0)  ;; Push the current value of 'i' to check the condition
-      (i32.const 10) ;; Push 10 onto the stack
-      (i32.lt_s)     ;; Check if 'i' is less than 10
-      (br_if $loopStart)  ;; If 'i' < 10, continue the loop
+
+    (loop $blockStart
+      (loop $blockStart2
+        (loop $loopStart  ;; Start of the loop
+          (local.get 0)  ;; Get the current value of 'i'
+          (i32.const 1)  ;; Push 1 onto the stack
+          (i32.add)      ;; Add 'i' and 1
+          (local.set 0)  ;; Update 'i' with the new value
+    
+          (local.get 0)  ;; Push the current value of 'i' to check the condition
+          (i32.const 10) ;; Push 10 onto the stack
+          (i32.lt_s)     ;; Check if 'i' is less than 10
+          (br_if $blockStart)  ;; If 'i' < 10, continue the loop
+        )
+      )
     )
 
     (local.get 0)  ;; After the loop, get the value of 'i' to be returned