feat: new new instruction types for blocks/ends

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

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,

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());

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),

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

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

@@ -6,14 +6,18 @@ 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);
     }
 
     #[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,