feat: add FuncContext

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

4 files changed, 89 insertions, 79 deletions

diff --git a/crates/tinywasm/src/func.rs b/crates/tinywasm/src/func.rs
index 8546b9f..87259f5 100644
--- a/crates/tinywasm/src/func.rs
+++ b/crates/tinywasm/src/func.rs
@@ -4,7 +4,7 @@ use tinywasm_types::{FuncAddr, FuncType, ValType, WasmValue};
 
 use crate::{
     runtime::{CallFrame, Stack},
-    Error, ModuleInstance, Result, Store,
+    Error, FuncContext, ModuleInstance, Result, Store,
 };
 
 #[derive(Debug)]
@@ -55,7 +55,8 @@ impl FuncHandle {
         let wasm_func = match &func_inst.func {
             crate::Function::Host(h) => {
                 let func = h.func.clone();
-                return (func)(store, params);
+                let ctx = FuncContext { store, module: &self.module };
+                return (func)(ctx, params);
             }
             crate::Function::Wasm(ref f) => f,
         };
diff --git a/crates/tinywasm/src/imports.rs b/crates/tinywasm/src/imports.rs
index ff9d061..683f65a 100644
--- a/crates/tinywasm/src/imports.rs
+++ b/crates/tinywasm/src/imports.rs
@@ -41,8 +41,44 @@ pub struct HostFunction {
     pub(crate) func: HostFuncInner,
 }
 
+impl HostFunction {
+    /// Get the function's type
+    pub fn ty(&self) -> &tinywasm_types::FuncType {
+        &self.ty
+    }
+
+    /// Call the function
+    pub fn call(&self, ctx: FuncContext<'_>, args: &[WasmValue]) -> Result<Vec<WasmValue>> {
+        (self.func)(ctx, args)
+    }
+}
+
 pub(crate) type HostFuncInner =
-    Arc<dyn Fn(&mut crate::Store, &[WasmValue]) -> Result<Vec<WasmValue>> + 'static + Send + Sync>;
+    Arc<dyn Fn(FuncContext<'_>, &[WasmValue]) -> Result<Vec<WasmValue>> + 'static + Send + Sync>;
+
+/// The context of a host-function call
+#[derive(Debug)]
+pub struct FuncContext<'a> {
+    pub(crate) store: &'a mut crate::Store,
+    pub(crate) module: &'a crate::ModuleInstance,
+}
+
+impl FuncContext<'_> {
+    /// Get a mutable reference to the store
+    pub fn store_mut(&mut self) -> &mut crate::Store {
+        self.store
+    }
+
+    /// Get a reference to the store
+    pub fn store(&self) -> &crate::Store {
+        self.store
+    }
+
+    /// Get a reference to the module instance
+    pub fn module(&self) -> &crate::ModuleInstance {
+        self.module
+    }
+}
 
 impl Debug for HostFunction {
     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
@@ -110,25 +146,25 @@ impl Extern {
     /// 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,
+        func: impl Fn(FuncContext<'_>, &[WasmValue]) -> Result<Vec<WasmValue>> + 'static + Send + Sync,
     ) -> Self {
-        let inner_func = move |store: &mut crate::Store, args: &[WasmValue]| {
+        let inner_func = move |ctx: FuncContext<'_>, args: &[WasmValue]| {
             let args = args.to_vec();
-            func(store, &args)
+            func(ctx, &args)
         };
 
         Self::Func(Function::Host(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
+    pub fn typed_func<P, R>(func: impl Fn(FuncContext<'_>, 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 inner_func = move |ctx: FuncContext<'_>, args: &[WasmValue]| -> Result<Vec<WasmValue>> {
             let args = P::from_wasm_value_tuple(args.to_vec())?;
-            let result = func(store, args)?;
+            let result = func(ctx, args)?;
             Ok(result.into_wasm_value_tuple())
         };
 
diff --git a/crates/tinywasm/src/runtime/executor/mod.rs b/crates/tinywasm/src/runtime/executor/mod.rs
index 623f857..80ad4b4 100644
--- a/crates/tinywasm/src/runtime/executor/mod.rs
+++ b/crates/tinywasm/src/runtime/executor/mod.rs
@@ -4,7 +4,7 @@ use super::{DefaultRuntime, Stack};
 use crate::{
     log::debug,
     runtime::{BlockType, LabelFrame},
-    CallFrame, Error, LabelArgs, ModuleInstance, Result, Store, Trap,
+    CallFrame, Error, FuncContext, LabelArgs, ModuleInstance, Result, Store, Trap,
 };
 use alloc::{string::ToString, vec::Vec};
 use tinywasm_types::{ElementKind, Instruction};
@@ -136,7 +136,7 @@ fn exec_one(
                 crate::Function::Host(host_func) => {
                     let func = host_func.func.clone();
                     let params = stack.values.pop_params(&host_func.ty.params)?;
-                    let res = (func)(store, &params)?;
+                    let res = (func)(FuncContext { store, module }, &params)?;
                     stack.values.extend_from_typed(&res);
                     return Ok(ExecResult::Ok);
                 }
@@ -175,7 +175,7 @@ fn exec_one(
                 crate::Function::Host(host_func) => {
                     let func = host_func.func.clone();
                     let params = stack.values.pop_params(&host_func.ty.params)?;
-                    let res = (func)(store, &params)?;
+                    let res = (func)(FuncContext { store, module }, &params)?;
                     stack.values.extend_from_typed(&res);
                     return Ok(ExecResult::Ok);
                 }
diff --git a/crates/tinywasm/tests/testsuite/run.rs b/crates/tinywasm/tests/testsuite/run.rs
index d2eab6c..d41ddca 100644
--- a/crates/tinywasm/tests/testsuite/run.rs
+++ b/crates/tinywasm/tests/testsuite/run.rs
@@ -22,42 +22,40 @@ impl TestSuite {
     fn imports(registered_modules: Vec<(String, ModuleInstanceAddr)>) -> Result<Imports> {
         let mut imports = Imports::new();
 
-        let table = Extern::table(
-            TableType::new(ValType::FuncRef, 10, Some(20)),
-            WasmValue::default_for(ValType::FuncRef),
-        );
+        let table =
+            Extern::table(TableType::new(ValType::FuncRef, 10, Some(20)), WasmValue::default_for(ValType::FuncRef));
 
-        let print = Extern::typed_func(|_: &mut tinywasm::Store, _: ()| {
+        let print = Extern::typed_func(|_ctx: tinywasm::FuncContext, _: ()| {
             log::debug!("print");
             Ok(())
         });
 
-        let print_i32 = Extern::typed_func(|_: &mut tinywasm::Store, arg: i32| {
+        let print_i32 = Extern::typed_func(|_ctx: tinywasm::FuncContext, arg: i32| {
             log::debug!("print_i32: {}", arg);
             Ok(())
         });
 
-        let print_i64 = Extern::typed_func(|_: &mut tinywasm::Store, arg: i64| {
+        let print_i64 = Extern::typed_func(|_ctx: tinywasm::FuncContext, arg: i64| {
             log::debug!("print_i64: {}", arg);
             Ok(())
         });
 
-        let print_f32 = Extern::typed_func(|_: &mut tinywasm::Store, arg: f32| {
+        let print_f32 = Extern::typed_func(|_ctx: tinywasm::FuncContext, arg: f32| {
             log::debug!("print_f32: {}", arg);
             Ok(())
         });
 
-        let print_f64 = Extern::typed_func(|_: &mut tinywasm::Store, arg: f64| {
+        let print_f64 = Extern::typed_func(|_ctx: tinywasm::FuncContext, arg: f64| {
             log::debug!("print_f64: {}", arg);
             Ok(())
         });
 
-        let print_i32_f32 = Extern::typed_func(|_: &mut tinywasm::Store, args: (i32, f32)| {
+        let print_i32_f32 = Extern::typed_func(|_ctx: tinywasm::FuncContext, args: (i32, f32)| {
             log::debug!("print_i32_f32: {}, {}", args.0, args.1);
             Ok(())
         });
 
-        let print_f64_f64 = Extern::typed_func(|_: &mut tinywasm::Store, args: (f64, f64)| {
+        let print_f64_f64 = Extern::typed_func(|_ctx: tinywasm::FuncContext, args: (f64, f64)| {
             log::debug!("print_f64_f64: {}, {}", args.0, args.1);
             Ok(())
         });
@@ -165,11 +163,7 @@ impl TestSuite {
                     test_group.add_result(&format!("Wat({})", i), span.linecol_in(wast), result.map(|_| ()));
                 }
 
-                AssertMalformed {
-                    span,
-                    mut module,
-                    message: _,
-                } => {
+                AssertMalformed { span, mut module, message: _ } => {
                     let Ok(module) = module.encode() else {
                         test_group.add_result(&format!("AssertMalformed({})", i), span.linecol_in(wast), Ok(()));
                         continue;
@@ -189,11 +183,7 @@ impl TestSuite {
                     );
                 }
 
-                AssertInvalid {
-                    span,
-                    mut module,
-                    message: _,
-                } => {
+                AssertInvalid { span, mut module, message: _ } => {
                     let res = catch_unwind_silent(move || parse_module_bytes(&module.encode().unwrap()))
                         .map_err(|e| eyre!("failed to parse module (invalid): {:?}", try_downcast_panic(e)))
                         .and_then(|res| res);
@@ -303,15 +293,13 @@ impl TestSuite {
                 Invoke(invoke) => {
                     let name = invoke.name;
                     let res: Result<Result<()>, _> = catch_unwind_silent(|| {
-                        let args = invoke
-                            .args
-                            .into_iter()
-                            .map(wastarg2tinywasmvalue)
-                            .collect::<Result<Vec<_>>>()
-                            .map_err(|e| {
-                                error!("failed to convert args: {:?}", e);
-                                e
-                            })?;
+                        let args =
+                            invoke.args.into_iter().map(wastarg2tinywasmvalue).collect::<Result<Vec<_>>>().map_err(
+                                |e| {
+                                    error!("failed to convert args: {:?}", e);
+                                    e
+                                },
+                            )?;
 
                         exec_fn_instance(last_module.as_ref(), &mut store, invoke.name, &args).map_err(|e| {
                             error!("failed to execute function: {:?}", e);
@@ -320,9 +308,7 @@ impl TestSuite {
                         Ok(())
                     });
 
-                    let res = res
-                        .map_err(|e| eyre!("test panicked: {:?}", try_downcast_panic(e)))
-                        .and_then(|r| r);
+                    let res = res.map_err(|e| eyre!("test panicked: {:?}", try_downcast_panic(e))).and_then(|r| r);
 
                     test_group.add_result(&format!("Invoke({}-{})", name, i), span.linecol_in(wast), res);
                 }
@@ -348,15 +334,13 @@ impl TestSuite {
 
                     let res: Result<Result<()>, _> = catch_unwind_silent(|| {
                         debug!("invoke: {:?}", invoke);
-                        let args = invoke
-                            .args
-                            .into_iter()
-                            .map(wastarg2tinywasmvalue)
-                            .collect::<Result<Vec<_>>>()
-                            .map_err(|e| {
-                                error!("failed to convert args: {:?}", e);
-                                e
-                            })?;
+                        let args =
+                            invoke.args.into_iter().map(wastarg2tinywasmvalue).collect::<Result<Vec<_>>>().map_err(
+                                |e| {
+                                    error!("failed to convert args: {:?}", e);
+                                    e
+                                },
+                            )?;
 
                         let outcomes =
                             exec_fn_instance(last_module.as_ref(), &mut store, invoke.name, &args).map_err(|e| {
@@ -366,14 +350,13 @@ impl TestSuite {
 
                         debug!("outcomes: {:?}", outcomes);
 
-                        let expected = results
-                            .into_iter()
-                            .map(wastret2tinywasmvalue)
-                            .collect::<Result<Vec<_>>>()
-                            .map_err(|e| {
-                                error!("failed to convert expected results: {:?}", e);
-                                e
-                            })?;
+                        let expected =
+                            results.into_iter().map(wastret2tinywasmvalue).collect::<Result<Vec<_>>>().map_err(
+                                |e| {
+                                    error!("failed to convert expected results: {:?}", e);
+                                    e
+                                },
+                            )?;
 
                         debug!("expected: {:?}", expected);
 
@@ -386,26 +369,16 @@ impl TestSuite {
                             ));
                         }
 
-                        outcomes
-                            .iter()
-                            .zip(expected)
-                            .enumerate()
-                            .try_for_each(|(i, (outcome, exp))| {
-                                (outcome.eq_loose(&exp))
-                                    .then_some(())
-                                    .ok_or_else(|| eyre!(" result {} did not match: {:?} != {:?}", i, outcome, exp))
-                            })
+                        outcomes.iter().zip(expected).enumerate().try_for_each(|(i, (outcome, exp))| {
+                            (outcome.eq_loose(&exp))
+                                .then_some(())
+                                .ok_or_else(|| eyre!(" result {} did not match: {:?} != {:?}", i, outcome, exp))
+                        })
                     });
 
-                    let res = res
-                        .map_err(|e| eyre!("test panicked: {:?}", try_downcast_panic(e)))
-                        .and_then(|r| r);
+                    let res = res.map_err(|e| eyre!("test panicked: {:?}", try_downcast_panic(e))).and_then(|r| r);
 
-                    test_group.add_result(
-                        &format!("AssertReturn({}-{})", invoke_name, i),
-                        span.linecol_in(wast),
-                        res,
-                    );
+                    test_group.add_result(&format!("AssertReturn({}-{})", invoke_name, i), span.linecol_in(wast), res);
                 }
                 _ => test_group.add_result(
                     &format!("Unknown({})", i),