pref: more callstack/callframe improvements

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

crates/tinywasm/src/func.rs

@@ -1,4 +1,4 @@
-use crate::log;
+use crate::{log, runtime::RawWasmValue};
 use alloc::{boxed::Box, format, string::String, string::ToString, vec, vec::Vec};
 use tinywasm_types::{FuncAddr, FuncType, ValType, WasmValue};
 
@@ -63,7 +63,7 @@ impl FuncHandle {
 
         // 6. Let f be the dummy frame
         log::debug!("locals: {:?}", locals);
-        let call_frame = CallFrame::new(func_inst, params, locals);
+        let call_frame = CallFrame::new(func_inst, params.iter().map(|v| RawWasmValue::from(*v)), locals);
 
         // 7. Push the frame f to the call stack
         // & 8. Push the values to the stack (Not needed since the call frame owns the values)

crates/tinywasm/src/imports.rs

@@ -154,7 +154,7 @@ impl Extern {
         let inner_func = move |ctx: FuncContext<'_>, args: &[WasmValue]| -> Result<Vec<WasmValue>> {
             let args = P::from_wasm_value_tuple(args)?;
             let result = func(ctx, args)?;
-            Ok(result.into_wasm_value_tuple())
+            Ok(result.into_wasm_value_tuple().to_vec())
         };
 
         let ty = tinywasm_types::FuncType { params: P::val_types(), results: R::val_types() };
@@ -383,7 +383,7 @@ impl Imports {
                             .ok_or_else(|| LinkingError::incompatible_import_type(import))?;
 
                         Self::compare_types(import, extern_func.ty(), import_func_type)?;
-                        imports.funcs.push(store.add_func(extern_func, *ty, idx)?);
+                        imports.funcs.push(store.add_func(extern_func, idx)?);
                     }
                     _ => return Err(LinkingError::incompatible_import_type(import).into()),
                 },

crates/tinywasm/src/reference.rs

@@ -83,7 +83,7 @@ impl MemoryRef {
     /// Load a UTF-8 string from memory
     pub fn load_string(&self, offset: usize, len: usize) -> Result<String> {
         let bytes = self.load_vec(offset, len)?;
-        Ok(String::from_utf8(bytes).map_err(|_| crate::Error::Other("Invalid UTF-8 string".to_string()))?)
+        String::from_utf8(bytes).map_err(|_| crate::Error::Other("Invalid UTF-8 string".to_string()))
     }
 
     /// Load a C-style string from memory

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

@@ -1,7 +1,7 @@
 use super::{InterpreterRuntime, Stack};
 use crate::{cold, log, unlikely};
 use crate::{
-    runtime::{BlockType, CallFrame, LabelArgs, LabelFrame},
+    runtime::{BlockType, CallFrame, LabelFrame},
     Error, FuncContext, ModuleInstance, Result, Store, Trap,
 };
 use alloc::format;
@@ -29,13 +29,12 @@ impl InterpreterRuntime {
         let mut cf = stack.call_stack.pop()?;
 
         let mut func_inst = cf.func_instance.clone();
-        let mut wasm_func = func_inst.assert_wasm().expect("exec expected wasm function");
+        let mut wasm_func = func_inst.assert_wasm()?;
 
         // The function to execute, gets updated from ExecResult::Call
         let mut instrs = &wasm_func.instructions;
         let mut instr_count = instrs.len();
-
-        let mut current_module = store.get_module_instance(func_inst.owner).unwrap().clone();
+        let mut current_module = store.get_module_instance_raw(func_inst.owner);
 
         loop {
             if unlikely(cf.instr_ptr >= instr_count) {
@@ -164,8 +163,8 @@ fn exec_one(
                 }
             };
 
-            let params = stack.values.pop_n_rev(ty.params.len())?.collect::<Vec<_>>();
-            let call_frame = CallFrame::new_raw(func_inst, &params, locals);
+            let params = stack.values.pop_n_rev(ty.params.len())?;
+            let call_frame = CallFrame::new(func_inst, params, locals);
 
             // push the call frame
             cf.instr_ptr += 1; // skip the call instruction
@@ -213,8 +212,8 @@ fn exec_one(
                 }
             };
 
-            let params = stack.values.pop_n_rev(func_ty.params.len())?.collect::<Vec<_>>();
-            let call_frame = CallFrame::new_raw(func_inst, &params, locals);
+            let params = stack.values.pop_n_rev(func_ty.params.len())?;
+            let call_frame = CallFrame::new(func_inst, params, locals);
 
             // push the call frame
             cf.instr_ptr += 1; // skip the call instruction
@@ -229,13 +228,14 @@ fn exec_one(
             // truthy value is on the top of the stack, so enter the then block
             if stack.values.pop_t::<i32>()? != 0 {
                 cf.enter_label(
-                    LabelFrame {
-                        instr_ptr: cf.instr_ptr,
-                        end_instr_ptr: cf.instr_ptr + *end_offset,
-                        stack_ptr: stack.values.len(), // - params,
-                        args: LabelArgs::new(*args, module)?,
-                        ty: BlockType::If,
-                    },
+                    LabelFrame::new(
+                        cf.instr_ptr,
+                        cf.instr_ptr + *end_offset,
+                        stack.values.len(), // - params,
+                        BlockType::If,
+                        args,
+                        module,
+                    ),
                     &mut stack.values,
                 );
                 return Ok(ExecResult::Ok);
@@ -244,13 +244,14 @@ fn exec_one(
             // falsy value is on the top of the stack
             if let Some(else_offset) = else_offset {
                 cf.enter_label(
-                    LabelFrame {
-                        instr_ptr: cf.instr_ptr + *else_offset,
-                        end_instr_ptr: cf.instr_ptr + *end_offset,
-                        stack_ptr: stack.values.len(), // - params,
-                        args: LabelArgs::new(*args, module)?,
-                        ty: BlockType::Else,
-                    },
+                    LabelFrame::new(
+                        cf.instr_ptr + *else_offset,
+                        cf.instr_ptr + *end_offset,
+                        stack.values.len(), // - params,
+                        BlockType::Else,
+                        args,
+                        module,
+                    ),
                     &mut stack.values,
                 );
                 cf.instr_ptr += *else_offset;
@@ -261,26 +262,28 @@ fn exec_one(
 
         Loop(args, end_offset) => {
             cf.enter_label(
-                LabelFrame {
-                    instr_ptr: cf.instr_ptr,
-                    end_instr_ptr: cf.instr_ptr + *end_offset,
-                    stack_ptr: stack.values.len(), // - params,
-                    args: LabelArgs::new(*args, module)?,
-                    ty: BlockType::Loop,
-                },
+                LabelFrame::new(
+                    cf.instr_ptr,
+                    cf.instr_ptr + *end_offset,
+                    stack.values.len(), // - params,
+                    BlockType::Loop,
+                    args,
+                    module,
+                ),
                 &mut stack.values,
             );
         }
 
         Block(args, end_offset) => {
             cf.enter_label(
-                LabelFrame {
-                    instr_ptr: cf.instr_ptr,
-                    end_instr_ptr: cf.instr_ptr + *end_offset,
-                    stack_ptr: stack.values.len(), //- params,
-                    args: LabelArgs::new(*args, module)?,
-                    ty: BlockType::Block,
-                },
+                LabelFrame::new(
+                    cf.instr_ptr,
+                    cf.instr_ptr + *end_offset,
+                    stack.values.len(), // - params,
+                    BlockType::Block,
+                    args,
+                    module,
+                ),
                 &mut stack.values,
             );
         }
@@ -341,7 +344,7 @@ fn exec_one(
                 panic!("else: no label to end, this should have been validated by the parser");
             };
 
-            let res_count = block.args.results;
+            let res_count = block.results;
             stack.values.truncate_keep(block.stack_ptr, res_count);
             cf.instr_ptr += *end_offset;
         }
@@ -352,7 +355,7 @@ fn exec_one(
                 cold();
                 panic!("end: no label to end, this should have been validated by the parser");
             };
-            stack.values.truncate_keep(block.stack_ptr, block.args.results)
+            stack.values.truncate_keep(block.stack_ptr, block.results)
         }
 
         LocalGet(local_index) => stack.values.push(cf.get_local(*local_index as usize)),

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

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

@@ -1,12 +1,17 @@
 use alloc::vec::Vec;
 use tinywasm_types::BlockArgs;
 
-use crate::{ModuleInstance, Result};
+use crate::{unlikely, ModuleInstance};
 
-#[derive(Debug, Clone, Default)]
+#[derive(Debug, Clone)]
 pub(crate) struct Labels(Vec<LabelFrame>);
 
 impl Labels {
+    pub(crate) fn new() -> Self {
+        // this is somehow a lot faster than Vec::with_capacity(128) or even using Default::default() in the benchmarks
+        Self(Vec::new())
+    }
+
     pub(crate) fn len(&self) -> usize {
         self.0.len()
     }
@@ -19,7 +24,12 @@ impl Labels {
     #[inline]
     /// get the label at the given index, where 0 is the top of the stack
     pub(crate) fn get_relative_to_top(&self, index: usize) -> Option<&LabelFrame> {
-        self.0.get(self.0.len() - index - 1)
+        // the vast majority of wasm functions don't use break to return
+        if unlikely(index >= self.0.len()) {
+            return None;
+        }
+
+        Some(&self.0[self.0.len() - index - 1])
     }
 
     #[inline]
@@ -43,10 +53,34 @@ pub(crate) struct LabelFrame {
 
     // position of the stack pointer when the block was entered
     pub(crate) stack_ptr: usize,
-    pub(crate) args: LabelArgs,
+    pub(crate) results: usize,
+    pub(crate) params: usize,
     pub(crate) ty: BlockType,
 }
 
+impl LabelFrame {
+    #[inline]
+    pub(crate) fn new(
+        instr_ptr: usize,
+        end_instr_ptr: usize,
+        stack_ptr: usize,
+        ty: BlockType,
+        args: &BlockArgs,
+        module: &ModuleInstance,
+    ) -> Self {
+        let (params, results) = match args {
+            BlockArgs::Empty => (0, 0),
+            BlockArgs::Type(_) => (0, 1),
+            BlockArgs::FuncType(t) => {
+                let ty = module.func_ty(*t);
+                (ty.params.len(), ty.results.len())
+            }
+        };
+
+        Self { instr_ptr, end_instr_ptr, stack_ptr, results, params, ty }
+    }
+}
+
 #[derive(Debug, Copy, Clone)]
 #[allow(dead_code)]
 pub(crate) enum BlockType {
@@ -55,21 +89,3 @@ pub(crate) enum BlockType {
     Else,
     Block,
 }
-
-#[derive(Debug, Clone, Default)]
-pub(crate) struct LabelArgs {
-    pub(crate) params: usize,
-    pub(crate) results: usize,
-}
-
-impl LabelArgs {
-    pub(crate) fn new(args: BlockArgs, module: &ModuleInstance) -> Result<Self> {
-        Ok(match args {
-            BlockArgs::Empty => LabelArgs { params: 0, results: 0 },
-            BlockArgs::Type(_) => LabelArgs { params: 0, results: 1 },
-            BlockArgs::FuncType(t) => {
-                LabelArgs { params: module.func_ty(t).params.len(), results: module.func_ty(t).results.len() }
-            }
-        })
-    }
-}