chore: improve label handling

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

Author: explodingcamera

crates/parser/src/conversion.rs

@@ -81,10 +81,7 @@ pub(crate) fn convert_blocktype(blocktype: wasmparser::BlockType) -> BlockArgs {
         // TODO: maybe solve this differently so we can support 128-bit values
         // without having to increase the size of the WasmValue enum
         Type(ty) => BlockArgs::Type(convert_valtype(&ty)),
-
-        // Wasm 2.0
-        FuncType(_ty) => unimplemented!(),
-        // FuncType(ty) => BlockArgs::FuncType(*ty),
+        FuncType(ty) => BlockArgs::FuncType(ty),
     }
 }
 

crates/tinywasm/src/func.rs

@@ -1,5 +1,5 @@
 use alloc::{format, string::String, string::ToString, vec, vec::Vec};
-use log::debug;
+use log::{debug, info};
 use tinywasm_types::{FuncAddr, FuncType, WasmValue};
 
 use crate::{
@@ -34,6 +34,7 @@ impl FuncHandle {
 
         // 4. If the length of the provided argument values is different from the number of expected arguments, then fail
         if func_ty.params.len() != params.len() {
+            info!("func_ty.params: {:?}", func_ty.params);
             return Err(Error::Other(format!(
                 "param count mismatch: expected {}, got {}",
                 func_ty.params.len(),

crates/tinywasm/src/runtime/executor/mod.rs

@@ -6,7 +6,7 @@ use crate::{
 };
 use alloc::vec::Vec;
 use log::info;
-use tinywasm_types::{BlockArgs, Instruction};
+use tinywasm_types::{BlockArgs, FuncType, Instruction, ValType};
 
 mod macros;
 use macros::*;
@@ -137,37 +137,37 @@ fn exec_one(
             info!("end: {:?} (@{})", instrs[end_instr_ptr], end_instr_ptr);
 
             if stack.values.pop_t::<i32>()? != 0 {
+                // let params = stack.values.pop_block_params(*args, &module)?;
                 cf.labels.push(LabelFrame {
                     instr_ptr: cf.instr_ptr,
                     end_instr_ptr: cf.instr_ptr + *end_offset,
-                    stack_ptr: stack.values.len(),
+                    stack_ptr: stack.values.len(), // - params,
                     args: *args,
                     ty: BlockType::If,
                 });
-                stack.values.push_block_args(*args)?;
             }
         }
 
         Loop(args, end_offset) => {
+            // let params = stack.values.pop_block_params(*args, &module)?;
             cf.labels.push(LabelFrame {
                 instr_ptr: cf.instr_ptr,
                 end_instr_ptr: cf.instr_ptr + *end_offset,
-                stack_ptr: stack.values.len(),
+                stack_ptr: stack.values.len(), // - params,
                 args: *args,
                 ty: BlockType::Loop,
             });
-            stack.values.push_block_args(*args)?;
         }
 
         Block(args, end_offset) => {
+            // let params = stack.values.pop_block_params(*args, &module)?;
             cf.labels.push(LabelFrame {
                 instr_ptr: cf.instr_ptr,
                 end_instr_ptr: cf.instr_ptr + *end_offset,
-                stack_ptr: stack.values.len(),
+                stack_ptr: stack.values.len(), //- params,
                 args: *args,
                 ty: BlockType::Block,
             });
-            stack.values.push_block_args(*args)?;
         }
 
         BrTable(_default, len) => {
@@ -215,17 +215,18 @@ fn exec_one(
                 panic!("end: no label to end, this should have been validated by the parser");
             };
 
-            let res: &[RawWasmValue] = match block.args {
-                BlockArgs::Empty => &[],
-                BlockArgs::Type(_t) => todo!(),
-                BlockArgs::FuncType(_t) => todo!(),
+            let res_count = match block.args {
+                BlockArgs::Empty => 0,
+                BlockArgs::Type(_) => 1,
+                BlockArgs::FuncType(t) => module.func_ty(t).results.len(),
             };
 
-            // trim the lable's stack from the stack
-            stack.values.trim(block.stack_ptr);
+            info!("we want to keep {} values on the stack", res_count);
+            info!("current block stack ptr: {}", block.stack_ptr);
+            info!("stack: {:?}", stack.values);
 
-            // push the block result values to the stack
-            stack.values.extend(res.iter().copied());
+            // trim the lable's stack from the stack
+            stack.values.truncate_keep(block.stack_ptr, res_count)
         }
 
         LocalGet(local_index) => stack.values.push(cf.get_local(*local_index as usize)),

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

@@ -1,6 +1,6 @@
 use alloc::vec::Vec;
 use log::info;
-use tinywasm_types::BlockArgs;
+use tinywasm_types::{BlockArgs, FuncType};
 
 #[derive(Debug, Default, Clone)]
 pub(crate) struct Labels(Vec<LabelFrame>);

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

@@ -97,7 +97,7 @@ impl CallFrame {
             BlockType::Loop => {
                 // this is a loop, so we want to jump back to the start of the loop
                 self.instr_ptr = break_to.instr_ptr;
-                value_stack.trim(break_to.stack_ptr);
+                value_stack.truncate(break_to.stack_ptr);
 
                 // we also want to trim the label stack to the loop (but not including the loop)
                 self.labels.trim(self.labels.len() - break_to_relative as usize);
@@ -107,11 +107,10 @@ impl CallFrame {
 
                 // this is a block, so we want to jump to the next instruction after the block ends
                 self.instr_ptr = break_to.end_instr_ptr + 1;
-                value_stack.trim(break_to.stack_ptr);
+                value_stack.truncate(break_to.stack_ptr);
 
                 // we also want to trim the label stack, including the block
                 self.labels.trim(self.labels.len() - break_to_relative as usize + 1);
-                panic!()
             }
             _ => unimplemented!("break to block type: {:?}", current_label.ty),
         }

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

@@ -1,5 +1,8 @@
-use crate::{runtime::RawWasmValue, Error, Result};
+use core::ops::Range;
+
+use crate::{runtime::RawWasmValue, Error, ModuleInstance, Result};
 use alloc::vec::Vec;
+use log::info;
 use tinywasm_types::BlockArgs;
 
 // minimum stack size
@@ -23,26 +26,36 @@ impl Default for ValueStack {
 }
 
 impl ValueStack {
+    #[cfg(test)]
+    pub(crate) fn data(&self) -> &[RawWasmValue] {
+        &self.stack
+    }
+
     #[inline]
     pub(crate) fn len(&self) -> usize {
         assert!(self.top <= self.stack.len());
         self.top
     }
 
     #[inline]
-    pub(crate) fn trim(&mut self, n: usize) {
+    pub(crate) fn truncate(&mut self, n: usize) {
         assert!(self.top <= self.stack.len());
         self.top -= n;
         self.stack.truncate(self.top);
     }
 
     #[inline]
-    pub(crate) fn push_block_args(&self, args: BlockArgs) -> Result<()> {
-        match args {
-            BlockArgs::Empty => Ok(()),
-            BlockArgs::Type(_t) => todo!("support block args (type)"),
-            BlockArgs::FuncType(_t) => todo!("support block args (func type)"),
+    // example: [1, 2, 3] n=1, end_keep=1 => [1, 3]
+    // example: [1] n=1, end_keep=1 => [1]
+    pub(crate) fn truncate_keep(&mut self, n: usize, end_keep: usize) {
+        if n == end_keep || n == 0 {
+            return;
         }
+
+        assert!(self.top <= self.stack.len());
+        info!("removing from {} to {}", self.top - n, self.top - end_keep);
+        self.stack.drain(self.top - n..self.top - end_keep);
+        self.top -= n - end_keep;
     }
 
     #[inline]
@@ -98,3 +111,56 @@ impl ValueStack {
         Ok(res)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::std::panic;
+
+    fn crate_stack<T: Into<RawWasmValue> + Copy>(data: &[T]) -> ValueStack {
+        let mut stack = ValueStack::default();
+        stack.extend(data.iter().map(|v| (*v).into()));
+        stack
+    }
+
+    fn assert_truncate_keep<T: Into<RawWasmValue> + Copy>(data: &[T], n: usize, end_keep: usize, expected: &[T]) {
+        let mut stack = crate_stack(data);
+        stack.truncate_keep(n, end_keep);
+        assert_eq!(
+            stack.data(),
+            expected.iter().map(|v| (*v).into()).collect::<Vec<_>>().as_slice()
+        );
+    }
+
+    fn catch_unwind_silent<F: FnOnce() -> R + panic::UnwindSafe, R>(f: F) -> crate::std::thread::Result<R> {
+        let prev_hook = panic::take_hook();
+        panic::set_hook(alloc::boxed::Box::new(|_| {}));
+        let result = panic::catch_unwind(f);
+        panic::set_hook(prev_hook);
+        result
+    }
+
+    #[test]
+    fn test_truncate_keep() {
+        assert_truncate_keep(&[1, 2, 3], 1, 1, &[1, 2, 3]);
+        assert_truncate_keep(&[1], 1, 1, &[1]);
+        assert_truncate_keep(&[1, 2, 3], 2, 1, &[1, 3]);
+        assert_truncate_keep::<i32>(&[], 0, 0, &[]);
+        catch_unwind_silent(|| assert_truncate_keep(&[1, 2, 3], 4, 1, &[1, 3])).expect_err("should panic");
+    }
+
+    #[test]
+    fn test_value_stack() {
+        let mut stack = ValueStack::default();
+        stack.push(1.into());
+        stack.push(2.into());
+        stack.push(3.into());
+        assert_eq!(stack.len(), 3);
+        assert_eq!(stack.pop_t::<i32>().unwrap(), 3);
+        assert_eq!(stack.len(), 2);
+        assert_eq!(stack.pop_t::<i32>().unwrap(), 2);
+        assert_eq!(stack.len(), 1);
+        assert_eq!(stack.pop_t::<i32>().unwrap(), 1);
+        assert_eq!(stack.len(), 0);
+    }
+}