feat: support basic loops, br_if

Signed-off-by: Henry Gressmann <[email protected]>

crates/parser/src/conversion.rs

@@ -158,6 +158,10 @@ pub fn process_operators<'a>(
             GlobalSet { global_index } => Instruction::GlobalSet(global_index),
             MemorySize { .. } => Instruction::MemorySize,
             MemoryGrow { .. } => Instruction::MemoryGrow,
+            I32Const { value } => Instruction::I32Const(value),
+            I64Const { value } => Instruction::I64Const(value),
+            F32Const { value } => Instruction::F32Const(f32::from_bits(value.bits())), // TODO: check if this is correct
+            F64Const { value } => Instruction::F64Const(f64::from_bits(value.bits())), // TODO: check if this is correct
             I32Load { memarg } => Instruction::I32Load(convert_memarg(memarg)),
             I64Load { memarg } => Instruction::I64Load(convert_memarg(memarg)),
             F32Load { memarg } => Instruction::F32Load(convert_memarg(memarg)),

crates/tinywasm/src/func.rs

@@ -1,4 +1,5 @@
 use alloc::{format, string::String, string::ToString, vec, vec::Vec};
+use log::debug;
 use tinywasm_types::{FuncAddr, FuncType, WasmValue};
 
 use crate::{
@@ -50,6 +51,7 @@ impl FuncHandle {
         }
 
         // 6. Let f be the dummy frame
+        debug!("locals: {:?}", func_inst.locals());
         let call_frame = CallFrame::new(self.addr as usize, params, func_inst.locals().to_vec());
 
         // 7. Push the frame f to the call stack
@@ -58,7 +60,7 @@ impl FuncHandle {
         // 8. Push the values to the stack (Not needed since the call frame owns the values)
 
         // 9. Invoke the function instance
-        let instrs = func_inst.instructions().iter();
+        let instrs = func_inst.instructions();
         store.runtime.exec(&mut stack, instrs)?;
 
         // Once the function returns:

crates/tinywasm/src/runtime/executer.rs

@@ -1,35 +1,100 @@
 use super::{Runtime, Stack};
-use crate::{log::debug, Error, Result};
+use crate::{log::debug, runtime::RawWasmValue, Error, Result};
 use alloc::vec;
-use tinywasm_types::Instruction;
+use tinywasm_types::{BlockArgs, Instruction};
 
+#[derive(Debug)]
 #[allow(dead_code)]
 enum BlockMarker {
     Top,
-    Loop,
+    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: core::slice::Iter<Instruction>) -> Result<()> {
+    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);
 
-        for instr in 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)?;
-
-                    use BlockMarker::*;
                     match block {
-                        Top => return Ok(()),
-                        Block => todo!(),
-                        Loop => todo!(),
-                        If => todo!(),
-                        Else => todo!(),
+                        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) => {
@@ -41,6 +106,9 @@ impl<const CHECK_TYPES: bool> Runtime<CHECK_TYPES> {
                     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();
@@ -70,8 +138,16 @@ impl<const CHECK_TYPES: bool> Runtime<CHECK_TYPES> {
                     let b: i32 = b.into();
                     stack.values.push((a - b).into());
                 }
-                _ => todo!(),
+                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)

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

@@ -59,6 +59,7 @@ 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| RawWasmValue::from(*v)));
+        locals.extend(local_types.iter().map(|_| RawWasmValue::default()));
 
         Self {
             instr_ptr: 0,

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

@@ -1,12 +1,15 @@
 use crate::{runtime::RawWasmValue, Error, Result};
 use alloc::vec::Vec;
+use tinywasm_types::BlockArgs;
 
 // minimum stack size
 pub const STACK_SIZE: usize = 1024;
 
 #[derive(Debug)]
 pub struct ValueStack {
     stack: Vec<RawWasmValue>,
+
+    // TODO: don't pop the stack, just keep track of the top for better performance
     top: usize,
 }
 
@@ -21,7 +24,29 @@ impl Default for ValueStack {
 
 impl ValueStack {
     #[inline]
-    pub(crate) fn _extend(&mut self, values: impl IntoIterator<Item = RawWasmValue> + ExactSizeIterator) {
+    pub(crate) fn len(&self) -> usize {
+        assert!(self.top <= self.stack.len());
+        self.top
+    }
+
+    #[inline]
+    pub(crate) fn trim(&mut self, n: usize) {
+        assert!(self.top <= self.stack.len());
+        self.top -= n;
+        self.stack.truncate(self.top);
+    }
+
+    #[inline]
+    pub(crate) fn block_args(&self, args: BlockArgs) -> Result<()> {
+        match args {
+            BlockArgs::Empty => Ok(()),
+            BlockArgs::Type(_t) => todo!(),
+            BlockArgs::FuncType(_t) => todo!(),
+        }
+    }
+
+    #[inline]
+    pub(crate) fn extend(&mut self, values: impl IntoIterator<Item = RawWasmValue> + ExactSizeIterator) {
         self.top += values.len();
         self.stack.extend(values);
     }

crates/tinywasm/src/runtime/value.rs

@@ -10,7 +10,7 @@ pub struct RawWasmValue(u64);
 
 impl Debug for RawWasmValue {
     fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
-        write!(f, "RawWasmValue({})", self.0)
+        write!(f, "RawWasmValue({})", self.0 as i64) // cast to i64 so at least negative numbers for i32 and i64 are printed correctly
     }
 }
 

examples/wasm/loop.wat

@@ -1,5 +1,5 @@
 (module
-  (func $loopExample (export "loopExample")
+  (func $loop (export "loop") (result i32)
     (local i32)  ;; Declare a local i32 variable, let's call it 'i'
     (i32.const 0)  ;; Initialize 'i' to 0
     (local.set 0)
@@ -15,6 +15,8 @@
       (i32.lt_s)     ;; Check if 'i' is less than 10
       (br_if $loopStart)  ;; If 'i' < 10, continue the loop
     )
+
+    (local.get 0)  ;; After the loop, get the value of 'i' to be returned
+    ;; The function will return the value of 'i' here
   )
 )
-