feat: new new instruction types for blocks/ends

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

7 files changed, 183 insertions, 71 deletions

diff --git a/crates/parser/src/conversion.rs b/crates/parser/src/conversion.rs
index 19b7586..7803d08 100644
--- a/crates/parser/src/conversion.rs
+++ b/crates/parser/src/conversion.rs
@@ -31,22 +31,17 @@ pub(crate) fn convert_module_code(
 ) -> Result<CodeSection> {
     let locals_reader = func.get_locals_reader()?;
     let count = locals_reader.get_count();
-
+    let pos = locals_reader.original_position();
     let mut locals = Vec::with_capacity(count as usize);
 
     for (i, local) in locals_reader.into_iter().enumerate() {
         let local = local?;
-        validator.define_locals(i, local.0, local.1)?;
-
+        validator.define_locals(pos + i, local.0, local.1)?;
         for _ in 0..local.0 {
             locals.push(convert_valtype(&local.1));
         }
     }
 
-    if locals.len() != count as usize {
-        return Err(crate::ParseError::Other("Invalid local index".to_string()));
-    }
-
     let body_reader = func.get_operators_reader()?;
     let body = process_operators(body_reader.original_position(), body_reader.into_iter(), validator)?;
 
@@ -118,8 +113,9 @@ pub fn process_operators<'a>(
     mut validator: FuncValidator<ValidatorResources>,
 ) -> Result<Box<[Instruction]>> {
     let mut instructions = Vec::new();
+    let mut last_label_pointer = Option::None;
 
-    for op in ops {
+    for (i, op) in ops.enumerate() {
         let op = op?;
         validator.op(offset, &op)?;
         offset += 1;
@@ -137,11 +133,63 @@ pub fn process_operators<'a>(
             }
             Unreachable => Instruction::Unreachable,
             Nop => Instruction::Nop,
-            Block { blockty } => Instruction::Block(convert_blocktype(blockty)),
-            Loop { blockty } => Instruction::Loop(convert_blocktype(blockty)),
-            If { blockty } => Instruction::If(convert_blocktype(blockty)),
-            Else => Instruction::Else,
-            End => Instruction::End,
+            Block { blockty } => {
+                last_label_pointer = Some(i);
+                Instruction::Block(convert_blocktype(blockty), 0)
+            }
+            Loop { blockty } => {
+                last_label_pointer = Some(i);
+                Instruction::Loop(convert_blocktype(blockty), 0)
+            }
+            If { blockty } => {
+                last_label_pointer = Some(i);
+                Instruction::If(convert_blocktype(blockty), 0)
+            }
+            End => {
+                if let Some(label_pointer) = last_label_pointer {
+                    // last_label_pointer is Some if we're ending a block
+                    match instructions[label_pointer] {
+                        Instruction::Block(_, ref mut end)
+                        | Instruction::Loop(_, ref mut end)
+                        | Instruction::Else(ref mut end)
+                        | Instruction::If(_, ref mut end) => {
+                            *end = i + 1; // Set the end position to be one after the End instruction
+                            last_label_pointer = None;
+                        }
+                        _ => {
+                            return Err(crate::ParseError::UnsupportedOperator(
+                                "Expected to end a block, but the last label was not a block".to_string(),
+                            ))
+                        }
+                    }
+
+                    Instruction::EndBlockFrame
+                } else {
+                    // last_label_pointer is None if we're ending the function
+                    Instruction::EndFunc
+                }
+            }
+            Else => {
+                let Some(label_pointer) = last_label_pointer else {
+                    return Err(crate::ParseError::UnsupportedOperator(
+                        "Expected to end an if block, but the last label was None".to_string(),
+                    ));
+                };
+
+                match instructions[label_pointer] {
+                    Instruction::If(_, ref mut end) => {
+                        *end = i + 1; // Set the end position to be one after the Else instruction
+                    }
+                    _ => {
+                        return Err(crate::ParseError::UnsupportedOperator(
+                            "Expected to end an if block, but the last label was not an if".to_string(),
+                        ));
+                    }
+                }
+
+                last_label_pointer = Some(i);
+                Instruction::Else(0)
+            }
             Br { relative_depth } => Instruction::Br(relative_depth),
             BrIf { relative_depth } => Instruction::BrIf(relative_depth),
             Return => Instruction::Return,
diff --git a/crates/tinywasm/src/runtime/executor/macros.rs b/crates/tinywasm/src/runtime/executor/macros.rs
index 7a999a2..01a7186 100644
--- a/crates/tinywasm/src/runtime/executor/macros.rs
+++ b/crates/tinywasm/src/runtime/executor/macros.rs
@@ -31,7 +31,7 @@ macro_rules! div_instr {
 /// Less than signed instruction
 macro_rules! lts_instr {
     ($ty:ty, $stack:ident) => {{
-        let [a, b] = $stack.values.pop_n_const::<2>()?;
+        let [b, a] = $stack.values.pop_n_const::<2>()?;
         let a: $ty = a.into();
         let b: $ty = b.into();
         $stack.values.push(((a < b) as i32).into());
diff --git a/crates/tinywasm/src/runtime/executor/mod.rs b/crates/tinywasm/src/runtime/executor/mod.rs
index 4119d20..662b080 100644
--- a/crates/tinywasm/src/runtime/executor/mod.rs
+++ b/crates/tinywasm/src/runtime/executor/mod.rs
@@ -1,7 +1,7 @@
 use super::{DefaultRuntime, Stack};
 use crate::{
     log::debug,
-    runtime::{BlockFrame, BlockFrameType, RawWasmValue},
+    runtime::{BlockFrame, BlockFrameInner, RawWasmValue},
     CallFrame, Error, ModuleInstance, Result, Store,
 };
 use alloc::vec::Vec;
@@ -116,12 +116,24 @@ fn exec_one(
             return Ok(ExecResult::Call);
         }
 
-        Loop(args) => {
-            cf.blocks.push(BlockFrame {
+        Loop(args, end_offset) => {
+            cf.block_frames.push(BlockFrame {
                 instr_ptr: cf.instr_ptr,
+                end_instr_ptr: cf.instr_ptr + *end_offset,
                 stack_ptr: stack.values.len(),
                 args: *args,
-                ty: BlockFrameType::Loop,
+                block: BlockFrameInner::Loop,
+            });
+            stack.values.block_args(*args)?;
+        }
+
+        Block(args, end_offset) => {
+            cf.block_frames.push(BlockFrame {
+                instr_ptr: cf.instr_ptr,
+                end_instr_ptr: cf.instr_ptr + *end_offset,
+                stack_ptr: stack.values.len(),
+                args: *args,
+                block: BlockFrameInner::Block,
             });
             stack.values.block_args(*args)?;
         }
@@ -144,42 +156,64 @@ fn exec_one(
         Br(v) => cf.break_to(*v, &mut stack.values)?,
         BrIf(v) => {
             let val: i32 = stack.values.pop().ok_or(Error::StackUnderflow)?.into();
+            debug!("br_if: {}", val);
             if val > 0 {
                 cf.break_to(*v, &mut stack.values)?
             };
         }
-        End => {
-            let blocks = &mut cf.blocks;
+
+        EndFunc => {
+            if cf.block_frames.len() > 0 {
+                panic!("endfunc: block frames not empty, this should have been validated by the parser");
+            }
+
+            if stack.call_stack.is_empty() {
+                debug!("end: no block to end and no parent call frame, returning");
+                return Ok(ExecResult::Return);
+            } else {
+                debug!("end: no block to end, returning to parent call frame");
+                *cf = stack.call_stack.pop()?;
+                return Ok(ExecResult::Call);
+            }
+        }
+
+        EndBlockFrame => {
+            let blocks = &mut cf.block_frames;
             let Some(block) = blocks.pop() else {
-                if stack.call_stack.is_empty() {
-                    debug!("end: no block to end and no parent call frame, returning");
-                    return Ok(ExecResult::Return);
-                } else {
-                    debug!("end: no block to end, returning to parent call frame");
-                    *cf = stack.call_stack.pop()?;
-                    return Ok(ExecResult::Call);
-                }
+                panic!("end: no block to end, this should have been validated by the parser");
             };
+
             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!(),
-                    };
+            let res: &[RawWasmValue] = match block.args {
+                BlockArgs::Empty => &[],
+                BlockArgs::Type(_t) => todo!(),
+                BlockArgs::FuncType(_t) => todo!(),
+            };
+
+            match block.block {
+                BlockFrameInner::Loop => {
+                    debug!("end(loop): continue loop");
 
                     // remove the loop values from the stack
                     stack.values.trim(block.stack_ptr);
 
-                    // push the loop result values to the stack
+                    // set the instruction pointer to the start of the loop
+                    cf.instr_ptr = block.instr_ptr;
+
+                    // push the loop back onto the stack
+                    blocks.push(block);
+                }
+                BlockFrameInner::Block => {
+                    // remove the block values from the stack
+                    stack.values.trim(block.stack_ptr);
+
+                    // push the block result values to the stack
                     stack.values.extend(res.iter().copied());
                 }
                 _ => {
-                    panic!("end: unimplemented block type end: {:?}", block.ty);
+                    panic!("end: unimplemented block type end: {:?}", block.block);
                 }
             }
         }
@@ -210,7 +244,13 @@ fn exec_one(
         F32Sub => sub_instr!(f32, stack),
         F64Sub => sub_instr!(f64, stack),
 
-        I32LtS => lts_instr!(i32, stack),
+        I32LtS => {
+            let [b, a] = stack.values.pop_n_const::<2>()?;
+            let a: i32 = a.into();
+            let b: i32 = b.into();
+            stack.values.push(((a < b) as i32).into());
+            debug!("i32.lt_s: {} < {} = {}", a, b, a < b);
+        }
         I64LtS => lts_instr!(i64, stack),
         F32Lt => lts_instr!(f32, stack),
         F64Lt => lts_instr!(f64, stack),
diff --git a/crates/tinywasm/src/runtime/stack.rs b/crates/tinywasm/src/runtime/stack.rs
index d7996b7..8e921b6 100644
--- a/crates/tinywasm/src/runtime/stack.rs
+++ b/crates/tinywasm/src/runtime/stack.rs
@@ -3,7 +3,7 @@ mod call_stack;
 mod value_stack;
 
 use self::{call_stack::CallStack, value_stack::ValueStack};
-pub(crate) use blocks::{BlockFrame, BlockFrameType};
+pub(crate) use blocks::{BlockFrame, BlockFrameInner};
 pub(crate) use call_stack::CallFrame;
 
 /// A WebAssembly Stack
diff --git a/crates/tinywasm/src/runtime/stack/blocks.rs b/crates/tinywasm/src/runtime/stack/blocks.rs
index d2bb6d6..4979fa2 100644
--- a/crates/tinywasm/src/runtime/stack/blocks.rs
+++ b/crates/tinywasm/src/runtime/stack/blocks.rs
@@ -6,6 +6,10 @@ use tinywasm_types::BlockArgs;
 pub(crate) struct Blocks(Vec<BlockFrame>);
 
 impl Blocks {
+    pub(crate) fn len(&self) -> usize {
+        self.0.len()
+    }
+
     #[inline]
     pub(crate) fn push(&mut self, block: BlockFrame) {
         self.0.push(block);
@@ -13,7 +17,7 @@ impl Blocks {
 
     #[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> {
+    pub(crate) fn get_relative_to_top(&self, index: usize) -> Option<&BlockFrame> {
         info!("get block: {}", index);
         info!("blocks: {:?}", self.0);
         self.0.get(self.0.len() - index - 1)
@@ -24,26 +28,29 @@ impl Blocks {
         self.0.pop()
     }
 
-    /// remove all blocks after the given index
+    /// keep the top `len` blocks and discard the rest
     #[inline]
-    pub(crate) fn trim(&mut self, index: usize) {
-        self.0.truncate(index + 1);
+    pub(crate) fn trim(&mut self, len: usize) {
+        self.0.truncate(len);
     }
 }
 
 #[derive(Debug, Clone)]
 pub(crate) struct BlockFrame {
-    // where to resume execution when the block is broken
+    // position of the instruction pointer when the block was entered
     pub(crate) instr_ptr: usize,
+    // position of the end instruction of the block
+    pub(crate) end_instr_ptr: usize,
+
     // position of the stack pointer when the block was entered
     pub(crate) stack_ptr: usize,
     pub(crate) args: BlockArgs,
-    pub(crate) ty: BlockFrameType,
+    pub(crate) block: BlockFrameInner,
 }
 
 #[derive(Debug, Copy, Clone)]
 #[allow(dead_code)]
-pub(crate) enum BlockFrameType {
+pub(crate) enum BlockFrameInner {
     Loop,
     If,
     Else,
diff --git a/crates/tinywasm/src/runtime/stack/call_stack.rs b/crates/tinywasm/src/runtime/stack/call_stack.rs
index 4493dd8..0fc1d12 100644
--- a/crates/tinywasm/src/runtime/stack/call_stack.rs
+++ b/crates/tinywasm/src/runtime/stack/call_stack.rs
@@ -1,8 +1,9 @@
 use crate::{runtime::RawWasmValue, Error, Result};
 use alloc::{boxed::Box, vec::Vec};
+use log::info;
 use tinywasm_types::{ValType, WasmValue};
 
-use super::{blocks::Blocks, BlockFrameType};
+use super::blocks::Blocks;
 
 // minimum call stack size
 const CALL_STACK_SIZE: usize = 1024;
@@ -68,10 +69,11 @@ impl CallStack {
 
 #[derive(Debug, Clone)]
 pub(crate) struct CallFrame {
+    // having real pointers here would be nice :( but we can't really do that in safe rust
     pub(crate) instr_ptr: usize,
     pub(crate) func_ptr: usize,
 
-    pub(crate) blocks: Blocks,
+    pub(crate) block_frames: Blocks,
     pub(crate) locals: Box<[RawWasmValue]>,
     pub(crate) local_count: usize,
 }
@@ -79,25 +81,33 @@ pub(crate) struct CallFrame {
 impl CallFrame {
     /// Break to a block at the given index (relative to the current frame)
     #[inline]
-    pub(crate) fn break_to(&mut self, block_index: u32, value_stack: &mut super::ValueStack) -> Result<()> {
-        let block = self
-            .blocks
-            .get(block_index as usize)
+    pub(crate) fn break_to(&mut self, break_to_relative: u32, value_stack: &mut super::ValueStack) -> Result<()> {
+        info!("we're in block_index: {}", self.block_frames.len());
+        info!("block_frames: {:?}", self.block_frames);
+        info!("break_to_relative: {}", break_to_relative);
+
+        let block_frame = self
+            .block_frames
+            .get_relative_to_top(break_to_relative as usize)
             .ok_or(Error::BlockStackUnderflow)?;
 
-        self.instr_ptr = block.instr_ptr;
-        value_stack.trim(block.stack_ptr);
+        info!("so we're breaking to: {:?} ?", block_frame);
+
+        todo!("break based on the type of the block we're breaking to");
+
+        // self.instr_ptr = block_frame.instr_ptr;
+        // value_stack.trim(block_frame.stack_ptr);
 
-        // Adjusting how to trim the blocks stack based on the block type
-        let trim_index = match block.ty {
-            // if we are breaking to a loop, we want to jump back to the start of the loop
-            BlockFrameType::Loop => block_index as usize - 2,
-            // if we are breaking to any other block, we want to jump to the end of the block
-            // TODO: check if this is correct
-            BlockFrameType::If | BlockFrameType::Else | BlockFrameType::Block => block_index as usize - 1,
-        };
+        // // // Adjusting how to trim the blocks stack based on the block type
+        // // let trim_index = match block_frame.block {
+        // //     // if we are breaking to a loop, we want to jump back to the start of the loop
+        // //     BlockFrameInner::Loop => block_index as usize - 1,
+        // //     // if we are breaking to any other block, we want to jump to the end of the block
+        // //     // TODO: check if this is correct
+        // //     BlockFrameInner::If | BlockFrameInner::Else | BlockFrameInner::Block => block_index as usize - 1,
+        // // };
 
-        self.blocks.trim(trim_index);
+        // self.block_frames.trim(block_index as usize);
         Ok(())
     }
 
@@ -111,7 +121,7 @@ impl CallFrame {
             func_ptr,
             local_count: locals.len(),
             locals: locals.into_boxed_slice(),
-            blocks: Blocks::default(),
+            block_frames: Blocks::default(),
         }
     }
 
diff --git a/crates/types/src/instructions.rs b/crates/types/src/instructions.rs
index 6db5e6b..6f11817 100644
--- a/crates/types/src/instructions.rs
+++ b/crates/types/src/instructions.rs
@@ -16,12 +16,18 @@ pub struct MemArg {
 
 type BrTableDefault = u32;
 type BrTableLen = usize;
+type EndOffset = usize;
 
 /// A WebAssembly Instruction
 ///
 /// These are our own internal bytecode instructions so they may not match the spec exactly.
 /// Wasm Bytecode can map to multiple of these instructions.
-/// For example, `br_table` stores the jump lables in the following `br_label` instructions to keep this enum small.
+///
+/// # Differences to the spec
+/// * `br_table` stores the jump lables in the following `br_label` instructions to keep this enum small.
+/// * Lables/Blocks: we store the label end offset in the instruction itself and
+///   have seperate EndBlockFrame and EndFunc instructions to mark the end of a block or function.
+///   This makes it easier to implement the label stack (we call it BlockFrameStack) iteratively.
 ///
 /// See <https://webassembly.github.io/spec/core/binary/instructions.html>
 #[derive(Debug, Clone, Copy, PartialEq)]
@@ -33,11 +39,12 @@ pub enum Instruction {
     // See <https://webassembly.github.io/spec/core/binary/instructions.html#control-instructions>
     Unreachable,
     Nop,
-    Block(BlockArgs),
-    Loop(BlockArgs),
-    If(BlockArgs),
-    Else,
-    End,
+    Block(BlockArgs, EndOffset),
+    Loop(BlockArgs, EndOffset),
+    If(BlockArgs, EndOffset),
+    Else(EndOffset),
+    EndBlockFrame,
+    EndFunc,
     Br(LabelAddr),
     BrIf(LabelAddr),
     BrTable(BrTableDefault, BrTableLen), // has to be followed by multiple BrLabel instructions