chore: improve breaking

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

Author: explodingcamera

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

@@ -1,5 +1,6 @@
 use super::{DefaultRuntime, Stack};
 use crate::{
+    get_label_args,
     log::debug,
     runtime::{BlockType, LabelFrame, RawWasmValue},
     CallFrame, Error, ModuleInstance, Result, Store,
@@ -137,37 +138,45 @@ 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(), // - params,
-                    args: *args,
-                    ty: BlockType::If,
-                });
+                cf.enter_label(
+                    LabelFrame {
+                        instr_ptr: cf.instr_ptr,
+                        end_instr_ptr: cf.instr_ptr + *end_offset,
+                        stack_ptr: stack.values.len(), // - params,
+                        args: get_label_args(*args, &module)?,
+                        ty: BlockType::If,
+                    },
+                    &mut stack.values,
+                )
             }
         }
 
         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(), // - params,
-                args: *args,
-                ty: BlockType::Loop,
-            });
+            cf.enter_label(
+                LabelFrame {
+                    instr_ptr: cf.instr_ptr,
+                    end_instr_ptr: cf.instr_ptr + *end_offset,
+                    stack_ptr: stack.values.len(), // - params,
+                    args: get_label_args(*args, &module)?,
+                    ty: BlockType::Loop,
+                },
+                &mut stack.values,
+            );
         }
 
         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(), //- params,
-                args: *args,
-                ty: BlockType::Block,
-            });
+            cf.enter_label(
+                LabelFrame {
+                    instr_ptr: cf.instr_ptr,
+                    end_instr_ptr: cf.instr_ptr + *end_offset,
+                    stack_ptr: stack.values.len(), //- params,
+                    args: get_label_args(*args, &module)?,
+                    ty: BlockType::Block,
+                },
+                &mut stack.values,
+            );
         }
 
         BrTable(_default, len) => {
@@ -186,10 +195,10 @@ fn exec_one(
             todo!()
         }
 
-        Br(v) => cf.break_to(*v, &mut stack.values)?,
+        Br(v) => cf.break_to(*v, &mut stack.values, module)?,
         BrIf(v) => {
             if stack.values.pop_t::<i32>()? > 0 {
-                cf.break_to(*v, &mut stack.values)?
+                cf.break_to(*v, &mut stack.values, module)?
             };
         }
 
@@ -215,12 +224,7 @@ fn exec_one(
                 panic!("end: no label to end, this should have been validated by the parser");
             };
 
-            let res_count = match block.args {
-                BlockArgs::Empty => 0,
-                BlockArgs::Type(_) => 1,
-                BlockArgs::FuncType(t) => module.func_ty(t).results.len(),
-            };
-
+            let res_count = block.args.results;
             info!("we want to keep {} values on the stack", res_count);
             info!("current block stack ptr: {}", block.stack_ptr);
             info!("stack: {:?}", stack.values);

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::{BlockType, LabelFrame};
+pub(crate) use blocks::{get_label_args, BlockType, LabelFrame};
 pub(crate) use call_stack::CallFrame;
 
 /// A WebAssembly Stack

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

@@ -2,6 +2,8 @@ use alloc::vec::Vec;
 use log::info;
 use tinywasm_types::{BlockArgs, FuncType};
 
+use crate::{ModuleInstance, Result};
+
 #[derive(Debug, Default, Clone)]
 pub(crate) struct Labels(Vec<LabelFrame>);
 
@@ -35,7 +37,7 @@ impl Labels {
 
     /// keep the top `len` blocks and discard the rest
     #[inline]
-    pub(crate) fn trim(&mut self, len: usize) {
+    pub(crate) fn truncate(&mut self, len: usize) {
         self.0.truncate(len);
     }
 }
@@ -49,7 +51,7 @@ pub(crate) struct LabelFrame {
 
     // position of the stack pointer when the block was entered
     pub(crate) stack_ptr: usize,
-    pub(crate) args: BlockArgs,
+    pub(crate) args: LabelArgs,
     pub(crate) ty: BlockType,
 }
 
@@ -61,3 +63,20 @@ pub(crate) enum BlockType {
     Else,
     Block,
 }
+
+#[derive(Debug, Clone)]
+pub struct LabelArgs {
+    pub params: usize,
+    pub results: usize,
+}
+
+pub fn get_label_args(args: BlockArgs, module: &ModuleInstance) -> Result<LabelArgs> {
+    Ok(match args {
+        BlockArgs::Empty => LabelArgs { params: 0, results: 0 },
+        BlockArgs::Type(t) => LabelArgs { params: 0, results: 1 },
+        BlockArgs::FuncType(t) => LabelArgs {
+            params: module.func_ty(t).params.len(),
+            results: module.func_ty(t).results.len(),
+        },
+    })
+}

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

@@ -1,7 +1,7 @@
-use crate::{runtime::RawWasmValue, BlockType, Error, Result};
+use crate::{runtime::RawWasmValue, BlockType, Error, LabelFrame, ModuleInstance, Result};
 use alloc::{boxed::Box, vec::Vec};
 use log::{debug, info};
-use tinywasm_types::{ValType, WasmValue};
+use tinywasm_types::{BlockArgs, ValType, WasmValue};
 
 use super::blocks::Labels;
 
@@ -79,38 +79,50 @@ pub(crate) struct CallFrame {
 }
 
 impl CallFrame {
+    #[inline]
+    /// Push a new label to the label stack and ensure the stack has the correct values
+    pub(crate) fn enter_label(&mut self, label_frame: LabelFrame, stack: &mut super::ValueStack) {
+        if label_frame.args.params > 0 {
+            stack.extend_from_within((label_frame.stack_ptr - label_frame.args.params)..label_frame.stack_ptr);
+        }
+
+        self.labels.push(label_frame);
+    }
+
     /// Break to a block at the given index (relative to the current frame)
     #[inline]
-    pub(crate) fn break_to(&mut self, break_to_relative: u32, value_stack: &mut super::ValueStack) -> Result<()> {
+    pub(crate) fn break_to(
+        &mut self,
+        break_to_relative: u32,
+        value_stack: &mut super::ValueStack,
+        module: &ModuleInstance,
+    ) -> Result<()> {
         let current_label = self.labels.top().ok_or(Error::LabelStackUnderflow)?;
         let break_to = self
             .labels
             .get_relative_to_top(break_to_relative as usize)
             .ok_or(Error::LabelStackUnderflow)?;
 
-        info!("current label: {:?}", current_label);
-        info!("we're breaking to: {:?} ?", break_to);
+        // trim the lable's stack from the stack
+        value_stack.truncate_keep(break_to.stack_ptr, break_to.args.results);
 
         // instr_ptr points to the label instruction, but the next step
         // will increment it by 1 since we're changing the "current" instr_ptr
         match break_to.ty {
             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.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);
+                self.labels.truncate(self.labels.len() - break_to_relative as usize);
             }
             BlockType::Block => {
-                debug!("current instr_ptr: {}", self.instr_ptr);
-
-                // 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.truncate(break_to.stack_ptr);
+                // this is a block, so we want to jump to the next instruction after the block ends (the inst_ptr will be incremented by 1 before the next instruction is executed)
+                self.instr_ptr = break_to.end_instr_ptr;
 
                 // we also want to trim the label stack, including the block
-                self.labels.trim(self.labels.len() - break_to_relative as usize + 1);
+                self.labels
+                    .truncate(self.labels.len() - (break_to_relative as usize + 1));
             }
             _ => unimplemented!("break to block type: {:?}", current_label.ty),
         }

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

@@ -31,6 +31,12 @@ impl ValueStack {
         &self.stack
     }
 
+    #[inline]
+    pub(crate) fn extend_from_within(&mut self, range: Range<usize>) {
+        self.top += range.len();
+        self.stack.extend_from_within(range);
+    }
+
     #[inline]
     pub(crate) fn len(&self) -> usize {
         assert!(self.top <= self.stack.len());