chore: refactor executor, add LocalCopy instructions

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

Author: explodingcamera

.vscode/settings.json

@@ -4,5 +4,5 @@
     },
     "rust-analyzer.linkedProjects": [
         "./Cargo.toml",
-    ]
+    ],
 }

Cargo.lock

@@ -176,7 +176,7 @@ pub(crate) fn convert_module_code(
 
     // maps a local's address to the index in the type's locals array
     let mut local_addr_map = Vec::with_capacity(count as usize);
-    let mut local_counts = LocalCounts::default();
+    let mut local_counts = ValueCounts::default();
 
     for (i, local) in locals_reader.into_iter().enumerate() {
         let local = local?;
@@ -186,20 +186,20 @@ pub(crate) fn convert_module_code(
     for i in 0..validator.len_locals() {
         match validator.get_local_type(i) {
             Some(wasmparser::ValType::I32) | Some(wasmparser::ValType::F32) => {
-                local_addr_map.push(local_counts.local_32);
-                local_counts.local_32 += 1;
+                local_addr_map.push(local_counts.c32);
+                local_counts.c32 += 1;
             }
             Some(wasmparser::ValType::I64) | Some(wasmparser::ValType::F64) => {
-                local_addr_map.push(local_counts.local_64);
-                local_counts.local_64 += 1;
+                local_addr_map.push(local_counts.c64);
+                local_counts.c64 += 1;
             }
             Some(wasmparser::ValType::V128) => {
-                local_addr_map.push(local_counts.local_128);
-                local_counts.local_128 += 1;
+                local_addr_map.push(local_counts.c128);
+                local_counts.c128 += 1;
             }
             Some(wasmparser::ValType::Ref(_)) => {
-                local_addr_map.push(local_counts.local_ref);
-                local_counts.local_ref += 1;
+                local_addr_map.push(local_counts.cref);
+                local_counts.cref += 1;
             }
             None => return Err(crate::ParseError::UnsupportedOperator("Unknown local type".to_string())),
         }

crates/parser/src/conversion.rs

@@ -59,6 +59,7 @@ impl Parser {
             sign_extension: true,
             saturating_float_to_int: true,
             function_references: true,
+            tail_call: true,
 
             component_model: false,
             component_model_nested_names: false,
@@ -71,7 +72,6 @@ impl Parser {
             memory_control: false,
             relaxed_simd: false,
             simd: false,
-            tail_call: false,
             threads: false,
             multi_memory: false, // should be working mostly
             custom_page_sizes: false,

crates/parser/src/lib.rs

@@ -3,12 +3,12 @@ use crate::{conversion, ParseError, Result};
 use alloc::string::ToString;
 use alloc::{boxed::Box, format, vec::Vec};
 use tinywasm_types::{
-    Data, Element, Export, FuncType, Global, Import, Instruction, LocalCounts, MemoryType, TableType, TinyWasmModule,
-    WasmFunction,
+    Data, Element, Export, FuncType, Global, Import, Instruction, MemoryType, TableType, TinyWasmModule, ValType,
+    ValueCounts, ValueCountsSmall, WasmFunction,
 };
 use wasmparser::{FuncValidatorAllocations, Payload, Validator};
 
-pub(crate) type Code = (Box<[Instruction]>, LocalCounts);
+pub(crate) type Code = (Box<[Instruction]>, ValueCounts);
 
 #[derive(Default)]
 pub(crate) struct ModuleReader {
@@ -193,10 +193,18 @@ impl ModuleReader {
             .code
             .into_iter()
             .zip(self.code_type_addrs)
-            .map(|((instructions, locals), ty_idx)| WasmFunction {
-                instructions,
-                locals,
-                ty: self.func_types.get(ty_idx as usize).expect("No func type for func, this is a bug").clone(),
+            .map(|((instructions, locals), ty_idx)| {
+                let mut params = ValueCountsSmall::default();
+                let ty = self.func_types.get(ty_idx as usize).expect("No func type for func, this is a bug").clone();
+                for param in ty.params.iter() {
+                    match param {
+                        ValType::I32 | ValType::F32 => params.c32 += 1,
+                        ValType::I64 | ValType::F64 => params.c64 += 1,
+                        ValType::V128 => params.c128 += 1,
+                        ValType::RefExtern | ValType::RefFunc => params.cref += 1,
+                    }
+                }
+                WasmFunction { instructions, params, locals, ty }
             })
             .collect::<Vec<_>>()
             .into_boxed_slice();

crates/parser/src/module.rs

@@ -122,6 +122,7 @@ macro_rules! impl_visit_operator {
     (@@sign_extension $($rest:tt)* ) => {};
     (@@saturating_float_to_int $($rest:tt)* ) => {};
     (@@bulk_memory $($rest:tt)* ) => {};
+    (@@tail_call $($rest:tt)* ) => {};
     (@@$proposal:ident $op:ident $({ $($arg:ident: $argty:ty),* })? => $visit:ident) => {
         #[cold]
         fn $visit(&mut self $($(,$arg: $argty)*)?) {
@@ -317,6 +318,14 @@ impl<'a, R: WasmModuleResources> wasmparser::VisitOperator<'a> for FunctionBuild
         visit_i64_trunc_sat_f64_u, Instruction::I64TruncSatF64U
     }
 
+    fn visit_return_call(&mut self, function_index: u32) -> Self::Output {
+        self.instructions.push(Instruction::ReturnCall(function_index));
+    }
+
+    fn visit_return_call_indirect(&mut self, type_index: u32, table_index: u32) -> Self::Output {
+        self.instructions.push(Instruction::ReturnCallIndirect(type_index, table_index));
+    }
+
     fn visit_global_set(&mut self, global_index: u32) -> Self::Output {
         match self.validator.get_operand_type(0) {
             Some(Some(t)) => self.instructions.push(match t {
@@ -385,6 +394,30 @@ impl<'a, R: WasmModuleResources> wasmparser::VisitOperator<'a> for FunctionBuild
             return;
         };
 
+        match self.instructions.last() {
+            Some(Instruction::LocalGet32(from))
+            | Some(Instruction::LocalGet64(from))
+            | Some(Instruction::LocalGet128(from))
+            | Some(Instruction::LocalGetRef(from)) => {
+                let from = *from;
+                self.instructions.pop();
+                // validation will ensure that the last instruction is the correct local.get
+                match self.validator.get_operand_type(0) {
+                    Some(Some(t)) => self.instructions.push(match t {
+                        wasmparser::ValType::I32 => Instruction::LocalCopy32(from, resolved_idx),
+                        wasmparser::ValType::F32 => Instruction::LocalCopy32(from, resolved_idx),
+                        wasmparser::ValType::I64 => Instruction::LocalCopy64(from, resolved_idx),
+                        wasmparser::ValType::F64 => Instruction::LocalCopy64(from, resolved_idx),
+                        wasmparser::ValType::V128 => Instruction::LocalCopy128(from, resolved_idx),
+                        wasmparser::ValType::Ref(_) => Instruction::LocalCopyRef(from, resolved_idx),
+                    }),
+                    _ => self.visit_unreachable(),
+                }
+                return;
+            }
+            _ => {}
+        }
+
         match self.validator.get_operand_type(0) {
             Some(Some(t)) => self.instructions.push(match t {
                 wasmparser::ValType::I32 => Instruction::LocalSet32(resolved_idx),

crates/parser/src/visit.rs

@@ -25,7 +25,7 @@ fn argon2id_run(module: TinyWasmModule) -> Result<()> {
     let mut store = Store::default();
     let instance = ModuleInstance::instantiate(&mut store, module.into(), None)?;
     let argon2 = instance.exported_func::<(i32, i32, i32), i32>(&store, "argon2id")?;
-    argon2.call(&mut store, (1000, 1, 1))?;
+    argon2.call(&mut store, (1000, 2, 1))?;
     Ok(())
 }
 

crates/tinywasm/benches/argon2id.rs

@@ -1,5 +1,5 @@
 use alloc::string::{String, ToString};
-use core::fmt::Display;
+use core::{fmt::Display, ops::ControlFlow};
 use tinywasm_types::FuncType;
 
 #[cfg(feature = "parser")]
@@ -23,15 +23,6 @@ pub enum Error {
     /// A function did not return a value
     FuncDidNotReturn,
 
-    /// The stack is empty
-    ValueStackUnderflow,
-
-    /// The label stack is empty
-    BlockStackUnderflow,
-
-    /// The call stack is empty
-    CallStackUnderflow,
-
     /// An invalid label type was encountered
     InvalidLabelType,
 
@@ -189,13 +180,10 @@ impl Display for Error {
 
             Self::Trap(trap) => write!(f, "trap: {}", trap),
             Self::Linker(err) => write!(f, "linking error: {}", err),
-            Self::CallStackUnderflow => write!(f, "call stack empty"),
             Self::InvalidLabelType => write!(f, "invalid label type"),
             Self::Other(message) => write!(f, "unknown error: {}", message),
             Self::UnsupportedFeature(feature) => write!(f, "unsupported feature: {}", feature),
             Self::FuncDidNotReturn => write!(f, "function did not return"),
-            Self::BlockStackUnderflow => write!(f, "label stack underflow"),
-            Self::ValueStackUnderflow => write!(f, "value stack underflow"),
             Self::InvalidStore => write!(f, "invalid store"),
         }
     }
@@ -249,3 +237,16 @@ impl From<tinywasm_parser::ParseError> for Error {
 
 /// A wrapper around [`core::result::Result`] for tinywasm operations
 pub type Result<T, E = Error> = crate::std::result::Result<T, E>;
+
+pub(crate) trait Controlify<T> {
+    fn to_cf(self) -> ControlFlow<Option<Error>, T>;
+}
+
+impl<T> Controlify<T> for Result<T, Error> {
+    fn to_cf(self) -> ControlFlow<Option<Error>, T> {
+        match self {
+            Ok(value) => ControlFlow::Continue(value),
+            Err(err) => ControlFlow::Break(Some(err)),
+        }
+    }
+}

crates/tinywasm/src/error.rs

@@ -48,7 +48,7 @@ impl FuncHandle {
             return Err(Error::Other("Type mismatch".into()));
         }
 
-        let func_inst = store.get_func(self.addr)?;
+        let func_inst = store.get_func(&self.addr);
         let wasm_func = match &func_inst.func {
             Function::Host(host_func) => {
                 let func = &host_func.clone().func;
@@ -76,7 +76,7 @@ impl FuncHandle {
         // assert!(stack.values.len() >= result_m);
 
         // 2. Pop m values from the stack
-        let res = stack.values.pop_results(&func_ty.results)?;
+        let res = stack.values.pop_results(&func_ty.results);
 
         // The values are returned as the results of the invocation.
         Ok(res)

crates/tinywasm/src/func.rs

@@ -143,8 +143,6 @@ impl Extern {
     }
 
     /// Create a new typed function import
-    // TODO: currently, this is slower than `Extern::func` because of the type conversions.
-    //       we should be able to optimize this and make it even faster than `Extern::func`.
     pub fn typed_func<P, R>(func: impl Fn(FuncContext<'_>, P) -> Result<R> + 'static) -> Self
     where
         P: FromWasmValueTuple + ValTypesFromTuple,
@@ -388,23 +386,23 @@ impl Imports {
 
                     match (val, &import.kind) {
                         (ExternVal::Global(global_addr), ImportKind::Global(ty)) => {
-                            let global = store.get_global(global_addr)?;
+                            let global = store.get_global(&global_addr);
                             Self::compare_types(import, &global.ty, ty)?;
                             imports.globals.push(global_addr);
                         }
                         (ExternVal::Table(table_addr), ImportKind::Table(ty)) => {
-                            let table = store.get_table(table_addr)?;
+                            let table = store.get_table(&table_addr);
                             Self::compare_table_types(import, &table.kind, ty)?;
                             imports.tables.push(table_addr);
                         }
                         (ExternVal::Memory(memory_addr), ImportKind::Memory(ty)) => {
-                            let mem = store.get_mem(memory_addr)?;
-                            let (size, kind) = { (mem.page_count(), mem.kind) };
+                            let mem = store.get_mem(&memory_addr);
+                            let (size, kind) = { (mem.page_count, mem.kind) };
                             Self::compare_memory_types(import, &kind, ty, Some(size))?;
                             imports.memories.push(memory_addr);
                         }
                         (ExternVal::Func(func_addr), ImportKind::Function(ty)) => {
-                            let func = store.get_func(func_addr)?;
+                            let func = store.get_func(&func_addr);
                             let import_func_type = module
                                 .0
                                 .func_types