feat: control flow improvements

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

3 files changed, 19 insertions, 7 deletions

diff --git a/crates/tinywasm/src/runtime/executer/mod.rs b/crates/tinywasm/src/runtime/executer/mod.rs
index 28fc468..c09a22e 100644
--- a/crates/tinywasm/src/runtime/executer/mod.rs
+++ b/crates/tinywasm/src/runtime/executer/mod.rs
@@ -92,9 +92,14 @@ impl<const CHECK_TYPES: bool> Runtime<CHECK_TYPES> {
                     let val = stack.values.pop().ok_or(Error::StackUnderflow)?;
                     cf.set_local(*local_index as usize, val);
                 }
-                I32Const(val) => {
-                    stack.values.push((*val).into());
+                LocalTee(local_index) => {
+                    debug!("local.tee: {:?}", local_index);
+                    let val = stack.values.pop().ok_or(Error::StackUnderflow)?;
+                    cf.set_local(*local_index as usize, val);
+                    stack.values.push(val);
                 }
+                I32Const(val) => stack.values.push((*val).into()),
+                I64Const(val) => stack.values.push((*val).into()),
                 I64Add => add_instr!(i64, stack),
                 I32Add => add_instr!(i32, stack),
                 F32Add => add_instr!(f32, stack),
diff --git a/crates/tinywasm/src/runtime/stack/blocks.rs b/crates/tinywasm/src/runtime/stack/blocks.rs
index 10003d9..6bf2237 100644
--- a/crates/tinywasm/src/runtime/stack/blocks.rs
+++ b/crates/tinywasm/src/runtime/stack/blocks.rs
@@ -33,7 +33,7 @@ impl Blocks {
 
 #[derive(Debug)]
 pub(crate) struct BlockFrame {
-    // position of the instruction pointer when the block was entered
+    // where to resume execution when the block is broken
     pub instr_ptr: usize,
     // position of the stack pointer when the block was entered
     pub stack_ptr: usize,
diff --git a/crates/tinywasm/src/runtime/stack/call_stack.rs b/crates/tinywasm/src/runtime/stack/call_stack.rs
index 9dbca46..3eaf886 100644
--- a/crates/tinywasm/src/runtime/stack/call_stack.rs
+++ b/crates/tinywasm/src/runtime/stack/call_stack.rs
@@ -2,7 +2,7 @@ use crate::{runtime::RawWasmValue, Error, Result};
 use alloc::{boxed::Box, vec::Vec};
 use tinywasm_types::{ValType, WasmValue};
 
-use super::blocks::Blocks;
+use super::{blocks::Blocks, BlockFrameType};
 
 // minimum call stack size
 pub const CALL_STACK_SIZE: usize = 1024;
@@ -70,9 +70,16 @@ impl CallFrame {
         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);
+        // 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,
+        };
+
+        self.blocks.trim(trim_index);
         Ok(())
     }