feat: typed host funcs

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

8 files changed, 248 insertions, 78 deletions

diff --git a/crates/parser/src/lib.rs b/crates/parser/src/lib.rs
index 3667f1d..0618683 100644
--- a/crates/parser/src/lib.rs
+++ b/crates/parser/src/lib.rs
@@ -23,7 +23,7 @@ mod module;
 use alloc::vec::Vec;
 pub use error::*;
 use module::ModuleReader;
-use tinywasm_types::{Function, WasmFunction};
+use tinywasm_types::WasmFunction;
 use wasmparser::Validator;
 
 pub use tinywasm_types::TinyWasmModule;
@@ -108,12 +108,10 @@ impl TryFrom<ModuleReader> for TinyWasmModule {
             .code
             .into_iter()
             .zip(func_types)
-            .map(|(f, ty)| {
-                Function::WasmFunction(WasmFunction {
-                    instructions: f.body,
-                    locals: f.locals,
-                    ty,
-                })
+            .map(|(f, ty)| WasmFunction {
+                instructions: f.body,
+                locals: f.locals,
+                ty_addr: ty,
             })
             .collect::<Vec<_>>();
 
diff --git a/crates/tinywasm/src/func.rs b/crates/tinywasm/src/func.rs
index ef07bf0..74043a9 100644
--- a/crates/tinywasm/src/func.rs
+++ b/crates/tinywasm/src/func.rs
@@ -1,6 +1,6 @@
-use alloc::{format, string::String, string::ToString, vec, vec::Vec};
+use alloc::{boxed::Box, format, string::String, string::ToString, vec, vec::Vec};
 use log::{debug, info};
-use tinywasm_types::{FuncAddr, FuncType, Function, WasmValue};
+use tinywasm_types::{FuncAddr, FuncType, ValType, WasmValue};
 
 use crate::{
     runtime::{CallFrame, Stack},
@@ -130,6 +130,21 @@ macro_rules! impl_into_wasm_value_tuple {
     }
 }
 
+macro_rules! impl_into_wasm_value_tuple_single {
+    ($T:ident) => {
+        impl IntoWasmValueTuple for $T {
+            fn into_wasm_value_tuple(self) -> Vec<WasmValue> {
+                vec![self.into()]
+            }
+        }
+    };
+}
+
+impl_into_wasm_value_tuple_single!(i32);
+impl_into_wasm_value_tuple_single!(i64);
+impl_into_wasm_value_tuple_single!(f32);
+impl_into_wasm_value_tuple_single!(f64);
+
 impl_into_wasm_value_tuple!();
 impl_into_wasm_value_tuple!(T1);
 impl_into_wasm_value_tuple!(T1, T2);
@@ -163,6 +178,27 @@ macro_rules! impl_from_wasm_value_tuple {
     }
 }
 
+macro_rules! impl_from_wasm_value_tuple_single {
+    ($T:ident) => {
+        impl FromWasmValueTuple for $T {
+            fn from_wasm_value_tuple(values: Vec<WasmValue>) -> Result<Self> {
+                #[allow(unused_variables, unused_mut)]
+                let mut iter = values.into_iter();
+                Ok($T::try_from(
+                    iter.next()
+                        .ok_or(Error::Other("Not enough values in WasmValue vector".to_string()))?,
+                )
+                .map_err(|_| Error::Other("Could not convert WasmValue to expected type".to_string()))?)
+            }
+        }
+    };
+}
+
+impl_from_wasm_value_tuple_single!(i32);
+impl_from_wasm_value_tuple_single!(i64);
+impl_from_wasm_value_tuple_single!(f32);
+impl_from_wasm_value_tuple_single!(f64);
+
 impl_from_wasm_value_tuple!();
 impl_from_wasm_value_tuple!(T1);
 impl_from_wasm_value_tuple!(T1, T2);
@@ -172,3 +208,70 @@ impl_from_wasm_value_tuple!(T1, T2, T3, T4, T5);
 impl_from_wasm_value_tuple!(T1, T2, T3, T4, T5, T6);
 impl_from_wasm_value_tuple!(T1, T2, T3, T4, T5, T6, T7);
 impl_from_wasm_value_tuple!(T1, T2, T3, T4, T5, T6, T7, T8);
+
+pub trait ValTypesFromTuple {
+    fn val_types() -> Box<[ValType]>;
+}
+
+pub trait ToValType {
+    fn to_val_type() -> ValType;
+}
+
+impl ToValType for i32 {
+    fn to_val_type() -> ValType {
+        ValType::I32
+    }
+}
+
+impl ToValType for i64 {
+    fn to_val_type() -> ValType {
+        ValType::I64
+    }
+}
+
+impl ToValType for f32 {
+    fn to_val_type() -> ValType {
+        ValType::F32
+    }
+}
+
+impl ToValType for f64 {
+    fn to_val_type() -> ValType {
+        ValType::F64
+    }
+}
+
+macro_rules! impl_val_types_from_tuple {
+    ($($t:ident),+) => {
+        impl<$($t),+> ValTypesFromTuple for ($($t,)+)
+        where
+            $($t: ToValType,)+
+        {
+            fn val_types() -> Box<[ValType]> {
+                Box::new([$($t::to_val_type(),)+])
+            }
+        }
+    };
+}
+
+impl ValTypesFromTuple for () {
+    fn val_types() -> Box<[ValType]> {
+        Box::new([])
+    }
+}
+
+impl<T1> ValTypesFromTuple for T1
+where
+    T1: ToValType,
+{
+    fn val_types() -> Box<[ValType]> {
+        Box::new([T1::to_val_type()])
+    }
+}
+
+impl_val_types_from_tuple!(T1);
+impl_val_types_from_tuple!(T1, T2);
+impl_val_types_from_tuple!(T1, T2, T3);
+impl_val_types_from_tuple!(T1, T2, T3, T4);
+impl_val_types_from_tuple!(T1, T2, T3, T4, T5);
+impl_val_types_from_tuple!(T1, T2, T3, T4, T5, T6);
diff --git a/crates/tinywasm/src/imports.rs b/crates/tinywasm/src/imports.rs
index aa9add1..41dd1a2 100644
--- a/crates/tinywasm/src/imports.rs
+++ b/crates/tinywasm/src/imports.rs
@@ -1,13 +1,41 @@
-use crate::Result;
+use core::fmt::Debug;
+
+use crate::{
+    func::{FromWasmValueTuple, IntoWasmValueTuple, ValTypesFromTuple},
+    Result,
+};
 use alloc::{
     collections::BTreeMap,
     string::{String, ToString},
+    sync::Arc,
+    vec::Vec,
 };
 use tinywasm_types::{
-    ExternVal, ExternalKind, FuncAddr, GlobalType, MemoryType, ModuleInstanceAddr, TableType, WasmValue,
+    ExternVal, ExternalKind, GlobalType, MemoryType, ModuleInstanceAddr, TableType, WasmFunction, WasmValue,
 };
 
 #[derive(Debug)]
+pub(crate) enum Function {
+    Host(HostFunction),
+    Wasm(WasmFunction),
+}
+
+/// A host function
+pub struct HostFunction {
+    pub(crate) ty: tinywasm_types::FuncType,
+    pub(crate) func: Arc<dyn Fn(&mut crate::Store, &[WasmValue]) -> Result<Vec<WasmValue>> + 'static + Send + Sync>,
+}
+
+impl Debug for HostFunction {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        f.debug_struct("HostFunction")
+            .field("ty", &self.ty)
+            .field("func", &"...")
+            .finish()
+    }
+}
+
+#[derive(Debug)]
 #[non_exhaustive]
 /// An external value
 pub enum Extern {
@@ -21,13 +49,13 @@ pub enum Extern {
     Memory(ExternMemory),
 
     /// A function
-    Func(ExternFunc),
+    Func(HostFunction),
 }
 
 /// A function
 #[derive(Debug)]
 pub struct ExternFunc {
-    pub(crate) addr: FuncAddr,
+    pub(crate) inner: HostFunction,
 }
 
 /// A global value
@@ -72,6 +100,45 @@ impl Extern {
         Self::Memory(ExternMemory { ty })
     }
 
+    /// Create a new function import
+    pub fn func(
+        ty: &tinywasm_types::FuncType,
+        func: impl Fn(&mut crate::Store, &[WasmValue]) -> Result<Vec<WasmValue>> + 'static + Send + Sync,
+    ) -> Self {
+        let inner_func = move |store: &mut crate::Store, args: &[WasmValue]| {
+            let args = args.to_vec();
+            func(store, &args)
+        };
+
+        Self::Func(HostFunction {
+            func: Arc::new(inner_func),
+            ty: ty.clone(),
+        })
+    }
+
+    /// Create a new typed function import
+    pub fn typed_func<P, R>(func: impl Fn(&mut crate::Store, P) -> Result<R> + 'static + Send + Sync) -> Self
+    where
+        P: FromWasmValueTuple + ValTypesFromTuple,
+        R: IntoWasmValueTuple + ValTypesFromTuple,
+    {
+        let inner_func = move |store: &mut crate::Store, args: &[WasmValue]| -> Result<Vec<WasmValue>> {
+            let args = P::from_wasm_value_tuple(args.to_vec())?;
+            let result = func(store, args)?;
+            Ok(result.into_wasm_value_tuple())
+        };
+
+        let ty = tinywasm_types::FuncType {
+            params: P::val_types(),
+            results: R::val_types(),
+        };
+
+        Self::Func(HostFunction {
+            func: Arc::new(inner_func),
+            ty: ty.clone(),
+        })
+    }
+
     pub(crate) fn kind(&self) -> ExternalKind {
         match self {
             Self::Global(_) => ExternalKind::Global,
diff --git a/crates/tinywasm/src/instance.rs b/crates/tinywasm/src/instance.rs
index 7fd30aa..ad1b53c 100644
--- a/crates/tinywasm/src/instance.rs
+++ b/crates/tinywasm/src/instance.rs
@@ -153,8 +153,9 @@ impl ModuleInstance {
             .ok_or_else(|| Error::Other(format!("Export not found: {}", name)))?;
 
         let func_addr = self.0.func_addrs[export.index as usize];
-        let func = store.get_func(func_addr as usize)?;
-        let ty = self.0.types[func.ty_addr() as usize].clone();
+        let func_inst = store.get_func(func_addr as usize)?;
+        let func = func_inst.assert_wasm()?;
+        let ty = self.0.types[func.ty_addr as usize].clone();
 
         Ok(FuncHandle {
             addr: export.index,
@@ -207,8 +208,9 @@ impl ModuleInstance {
             .get(func_index as usize)
             .expect("No func addr for start func, this is a bug");
 
-        let func = store.get_func(*func_addr as usize)?;
-        let ty = self.0.types[func.ty_addr() as usize].clone();
+        let func_inst = store.get_func(*func_addr as usize)?;
+        let func = func_inst.assert_wasm()?;
+        let ty = self.0.types[func.ty_addr as usize].clone();
 
         Ok(Some(FuncHandle {
             module: self.clone(),
diff --git a/crates/tinywasm/src/runtime/executor/mod.rs b/crates/tinywasm/src/runtime/executor/mod.rs
index e494a27..5e191cc 100644
--- a/crates/tinywasm/src/runtime/executor/mod.rs
+++ b/crates/tinywasm/src/runtime/executor/mod.rs
@@ -7,7 +7,7 @@ use crate::{
     CallFrame, Error, LabelArgs, ModuleInstance, Result, Store, Trap,
 };
 use alloc::{string::ToString, vec::Vec};
-use tinywasm_types::{Function, Instruction};
+use tinywasm_types::Instruction;
 
 mod macros;
 mod traits;
@@ -129,17 +129,13 @@ fn exec_one(
             // prepare the call frame
             let func_idx = module.resolve_func_addr(*v);
             let func_inst = store.get_func(func_idx as usize)?;
-            let func_ty = module.func_ty(func_inst.ty_addr());
+            let func = func_inst.assert_wasm()?;
+            let func_ty = module.func_ty(func.ty_addr);
 
             debug!("params: {:?}", func_ty.params);
             debug!("stack: {:?}", stack.values);
             let params = stack.values.pop_n(func_ty.params.len())?;
 
-            let func = match &func_inst.func {
-                Function::WasmFunction(wasm_func) => wasm_func,
-                _ => return Err(Error::UnsupportedFeature("Host functions cannot be called".to_string())),
-            };
-
             let call_frame = CallFrame::new_raw(*v as usize, &params, func.locals.to_vec());
 
             // push the call frame
@@ -162,7 +158,8 @@ fn exec_one(
 
             // prepare the call frame
             let func_inst = store.get_func(func_addr as usize)?;
-            let func_ty = module.func_ty(func_inst.ty_addr());
+            let func = func_inst.assert_wasm()?;
+            let func_ty = module.func_ty(func.ty_addr);
 
             if func_ty != call_ty {
                 return Err(Trap::IndirectCallTypeMismatch {
@@ -174,15 +171,6 @@ fn exec_one(
 
             let params = stack.values.pop_n(func_ty.params.len())?;
 
-            let func = match &func_inst.func {
-                Function::WasmFunction(wasm_func) => wasm_func,
-                _ => {
-                    return Err(Error::UnsupportedFeature(
-                        "Host functions cannot be called indirectly".to_string(),
-                    ))
-                }
-            };
-
             let call_frame = CallFrame::new_raw(func_addr as usize, &params, func.locals.to_vec());
 
             // push the call frame
diff --git a/crates/tinywasm/src/store.rs b/crates/tinywasm/src/store.rs
index 9c8a75e..6293cca 100644
--- a/crates/tinywasm/src/store.rs
+++ b/crates/tinywasm/src/store.rs
@@ -7,13 +7,13 @@ use core::{
 
 use alloc::{format, rc::Rc, string::ToString, vec, vec::Vec};
 use tinywasm_types::{
-    Addr, Data, Element, ElementKind, FuncAddr, Function, Global, GlobalType, Import, Instruction, MemAddr, MemoryArch,
+    Addr, Data, Element, ElementKind, FuncAddr, Global, GlobalType, Import, Instruction, MemAddr, MemoryArch,
     MemoryType, ModuleInstanceAddr, TableAddr, TableType, TypeAddr, ValType, WasmFunction,
 };
 
 use crate::{
     runtime::{self, DefaultRuntime},
-    Error, Extern, LinkedImports, ModuleInstance, RawWasmValue, Result, Trap,
+    Error, Extern, Function, LinkedImports, ModuleInstance, RawWasmValue, Result, Trap,
 };
 
 // global store id counter
@@ -114,11 +114,14 @@ impl Store {
     }
 
     /// Add functions to the store, returning their addresses in the store
-    pub(crate) fn add_funcs(&mut self, funcs: Vec<Function>, idx: ModuleInstanceAddr) -> Vec<FuncAddr> {
+    pub(crate) fn add_funcs(&mut self, funcs: Vec<WasmFunction>, idx: ModuleInstanceAddr) -> Vec<FuncAddr> {
         let func_count = self.data.funcs.len();
         let mut func_addrs = Vec::with_capacity(func_count);
         for (i, func) in funcs.into_iter().enumerate() {
-            self.data.funcs.push(Rc::new(FunctionInstance { func, owner: idx }));
+            self.data.funcs.push(Rc::new(FunctionInstance {
+                func: Function::Wasm(func),
+                owner: idx,
+            }));
             func_addrs.push((i + func_count) as FuncAddr);
         }
         func_addrs
@@ -404,17 +407,13 @@ const fn cold() {}
 impl FunctionInstance {
     pub(crate) fn assert_wasm(&self) -> Result<&WasmFunction> {
         match &self.func {
-            Function::WasmFunction(w) => Ok(w),
-            Function::HostFunction(_) => {
+            Function::Wasm(w) => Ok(w),
+            Function::Host(_) => {
                 cold();
                 Err(Error::Other("expected wasm function".to_string()))
             }
         }
     }
-
-    pub(crate) fn ty_addr(&self) -> TypeAddr {
-        self.func.ty()
-    }
 }
 
 /// A WebAssembly Table Instance
diff --git a/crates/tinywasm/tests/testsuite/run.rs b/crates/tinywasm/tests/testsuite/run.rs
index 462149e..5dd42ee 100644
--- a/crates/tinywasm/tests/testsuite/run.rs
+++ b/crates/tinywasm/tests/testsuite/run.rs
@@ -22,19 +22,60 @@ impl TestSuite {
     fn imports(registered_modules: Vec<(String, ModuleInstanceAddr)>) -> Result<Imports> {
         let mut imports = Imports::new();
 
-        let memory = Extern::memory(MemoryType::new_32(1, Some(2)));
         let table = Extern::table(
             TableType::new(ValType::FuncRef, 10, Some(20)),
             WasmValue::default_for(ValType::FuncRef),
         );
 
+        let print = Extern::typed_func(|_: &mut tinywasm::Store, _: ()| {
+            log::debug!("print");
+            Ok(())
+        });
+
+        let print_i32 = Extern::typed_func(|_: &mut tinywasm::Store, arg: i32| {
+            log::debug!("print_i32: {}", arg);
+            Ok(())
+        });
+
+        let print_i64 = Extern::typed_func(|_: &mut tinywasm::Store, arg: i64| {
+            log::debug!("print_i64: {}", arg);
+            Ok(())
+        });
+
+        let print_f32 = Extern::typed_func(|_: &mut tinywasm::Store, arg: f32| {
+            log::debug!("print_f32: {}", arg);
+            Ok(())
+        });
+
+        let print_f64 = Extern::typed_func(|_: &mut tinywasm::Store, arg: f64| {
+            log::debug!("print_f64: {}", arg);
+            Ok(())
+        });
+
+        let print_i32_f32 = Extern::typed_func(|_: &mut tinywasm::Store, args: (i32, f32)| {
+            log::debug!("print_i32_f32: {}, {}", args.0, args.1);
+            Ok(())
+        });
+
+        let print_i64_f64 = Extern::typed_func(|_: &mut tinywasm::Store, args: (i64, f64)| {
+            log::debug!("print_i64_f64: {}, {}", args.0, args.1);
+            Ok(())
+        });
+
         imports
-            .define("spectest", "memory", memory)?
+            .define("spectest", "memory", Extern::memory(MemoryType::new_32(1, Some(2))))?
             .define("spectest", "table", table)?
             .define("spectest", "global_i32", Extern::global(WasmValue::I32(666), false))?
             .define("spectest", "global_i64", Extern::global(WasmValue::I64(666), false))?
             .define("spectest", "global_f32", Extern::global(WasmValue::F32(666.0), false))?
-            .define("spectest", "global_f64", Extern::global(WasmValue::F64(666.0), false))?;
+            .define("spectest", "global_f64", Extern::global(WasmValue::F64(666.0), false))?
+            .define("spectest", "print", print)?
+            .define("spectest", "print_i32", print_i32)?
+            .define("spectest", "print_i64", print_i64)?
+            .define("spectest", "print_f32", print_f32)?
+            .define("spectest", "print_f64", print_f64)?
+            .define("spectest", "print_i32_f32", print_i32_f32)?
+            .define("spectest", "print_i64_f64", print_i64_f64)?;
 
         for (name, addr) in registered_modules {
             imports.link_module(&name, addr)?;
diff --git a/crates/types/src/lib.rs b/crates/types/src/lib.rs
index 5b290ca..069f134 100644
--- a/crates/types/src/lib.rs
+++ b/crates/types/src/lib.rs
@@ -28,7 +28,7 @@ extern crate alloc;
 mod instructions;
 use core::{fmt::Debug, ops::Range};
 
-use alloc::boxed::Box;
+use alloc::{boxed::Box, sync::Arc, vec::Vec};
 pub use instructions::*;
 
 /// A TinyWasm WebAssembly Module
@@ -45,7 +45,7 @@ pub struct TinyWasmModule {
     pub start_func: Option<FuncAddr>,
 
     /// The functions of the WebAssembly module.
-    pub funcs: Box<[Function]>,
+    pub funcs: Box<[WasmFunction]>,
 
     /// The types of the WebAssembly module.
     pub func_types: Box<[FuncType]>,
@@ -169,12 +169,6 @@ impl From<f64> for WasmValue {
     }
 }
 
-// impl From<i128> for WasmValue {
-//     fn from(i: i128) -> Self {
-//         Self::V128(i)
-//     }
-// }
-
 impl TryFrom<WasmValue> for i32 {
     type Error = ();
 
@@ -339,35 +333,13 @@ impl FuncType {
     }
 }
 
-// A WebAssembly Function
-#[derive(Debug, Clone)]
-pub enum Function {
-    WasmFunction(WasmFunction),
-    HostFunction(HostFunction),
-}
-
-impl Function {
-    pub fn ty(&self) -> TypeAddr {
-        match self {
-            Self::WasmFunction(f) => f.ty,
-            Self::HostFunction(f) => f.ty,
-        }
-    }
-}
-
 #[derive(Debug, Clone)]
 pub struct WasmFunction {
-    pub ty: TypeAddr,
+    pub ty_addr: TypeAddr,
     pub instructions: Box<[Instruction]>,
     pub locals: Box<[ValType]>,
 }
 
-#[derive(Debug, Clone)]
-pub struct HostFunction {
-    pub ty: TypeAddr,
-    pub func: fn(&mut [WasmValue]) -> Result<(), ()>,
-}
-
 /// A WebAssembly Module Export
 #[derive(Debug, Clone)]
 pub struct Export {