Repository: wasmi-labs/wasmi (Rust)

Issue:
Implement Wasmtime's `Func::new` constructor for dynamic function arguments and results
Context is me trying to update to a recent version of wasmi in smoldot.

In smoldot, I don't know at compile-time what the types of the functions imported by a wasm module are.
In other words, in smoldot, I don't actually inject strongly-typed functions into a `Linker` or similar. Instead, every host function has the same implementation that manually checks whether the number and the type of the parameters provided by the user corresponds to the function being called.

I don't think it is possible to do that with a recent version of wasmi.
Let me know if I'm wrong, but in order to create a `Func` I need to provide a closure whose parameters and return type match the parameters and return type of the Wasm function. Because I don't know these parameters/return type at compilation time, this isn't possible for me.



Discussion:
One way to circumvent the issue is to have in the source code one closure for each possible combination of parameter and return type, and then have a huge `match` block that picks the correct one at runtime.
This would however be very hacky.
I guess [this constructor](https://docs.rs/wasmtime/5.0.0/wasmtime/struct.Func.html#method.new) from the Wasmtime API would suffice your needs.
If so this should be fairly straight forward to implement. I left it out so far from the `wasmi` API since I thought the `wrap` constructor is the better API and strictly superior to the `new` constructor. But this use case might proof me wrong.
> I guess [this constructor](https://docs.rs/wasmtime/5.0.0/wasmtime/struct.Func.html#method.new) from the Wasmtime API would suffice your needs.

Yes :+1: 



Patch:
diff --git a/crates/wasmi/src/engine/func_args.rs b/crates/wasmi/src/engine/func_args.rs
index 37af254533..52c0ccff94 100644
--- a/crates/wasmi/src/engine/func_args.rs
+++ b/crates/wasmi/src/engine/func_args.rs
@@ -1,5 +1,6 @@
+use crate::{value::WithType, Value};
 use core::cmp;
-use wasmi_core::{DecodeUntypedSlice, EncodeUntypedSlice, UntypedValue};
+use wasmi_core::{DecodeUntypedSlice, EncodeUntypedSlice, UntypedError, UntypedValue};
 
 #[derive(Debug)]
 pub struct FuncParams<'a> {
@@ -17,8 +18,9 @@ pub struct FuncParams<'a> {
     len_results: usize,
 }
 
-/// Utility type to ensure at compile time that host functions always
-/// call [`FuncParams::write_results`] at the end of their execution.
+/// Utility type to ensure at compile time that host functions always call
+/// [`FuncParams::encode_results`] or [`FuncParams::encode_results_from_slice`]
+/// at the end of their execution.
 #[derive(Debug)]
 pub struct FuncResults {}
 
@@ -42,32 +44,72 @@ impl<'a> FuncParams<'a> {
         }
     }
 
-    /// Returns the host function parameters as `T`.
+    /// Returns a slice over the untyped function parameters.
+    fn params(&self) -> &[UntypedValue] {
+        &self.params_results[..self.len_params]
+    }
+
+    /// Returns an exclusive reference to the slice of function results.
+    fn results(&mut self) -> &mut [UntypedValue] {
+        &mut self.params_results[..self.len_results]
+    }
+
+    /// Decodes and returns the executed host function parameters as `T`.
     ///
     /// # Panics
     ///
     /// If the number of function parameters dictated by `T` does not match.
-    pub fn read_params<T>(&self) -> T
+    pub fn decode_params<T>(&self) -> T
     where
         T: DecodeUntypedSlice,
     {
-        let params_buffer = &self.params_results[..self.len_params];
-        UntypedValue::decode_slice::<T>(params_buffer)
+        UntypedValue::decode_slice::<T>(self.params())
             .unwrap_or_else(|error| panic!("encountered unexpected invalid tuple length: {error}"))
     }
 
-    /// Sets the results of the function invocation.
+    /// Decodes and stores the executed host functions parameters into `values`.
+    ///
+    /// # Panics
+    ///
+    /// If the number of host function parameters and items in `values` does not match.
+    pub fn decode_params_into_slice(&self, values: &mut [Value]) -> Result<(), UntypedError> {
+        assert_eq!(self.params().len(), values.len());
+        self.params().iter().zip(values).for_each(|(src, dst)| {
+            *dst = src.with_type(dst.ty());
+        });
+        Ok(())
+    }
+
+    /// Encodes the results of the host function invocation as `T`.
     ///
     /// # Panics
     ///
-    /// If the number of results does not match the expected amount.
-    pub fn write_results<T>(self, results: T) -> FuncResults
+    /// If the number of results dictated by `T` does not match the expected amount.
+    pub fn encode_results<T>(mut self, values: T) -> FuncResults
     where
         T: EncodeUntypedSlice,
     {
-        let results_buffer = &mut self.params_results[..self.len_results];
-        UntypedValue::encode_slice::<T>(results_buffer, results)
+        UntypedValue::encode_slice::<T>(self.results(), values)
             .unwrap_or_else(|error| panic!("encountered unexpected invalid tuple length: {error}"));
         FuncResults {}
     }
+
+    /// Encodes the results of the host function invocation given the `values` slice.
+    ///
+    /// # Panics
+    ///
+    /// If the number of expected results does not match the length of `values`.
+    pub fn encode_results_from_slice(
+        mut self,
+        values: &[Value],
+    ) -> Result<FuncResults, UntypedError> {
+        assert_eq!(self.results().len(), values.len());
+        self.results()
+            .iter_mut()
+            .zip(values)
+            .for_each(|(dst, src)| {
+                *dst = src.clone().into();
+            });
+        Ok(FuncResults {})
+    }
 }
diff --git a/crates/wasmi/src/func/into_func.rs b/crates/wasmi/src/func/into_func.rs
index a5e6b8eaa0..8150431b96 100644
--- a/crates/wasmi/src/func/into_func.rs
+++ b/crates/wasmi/src/func/into_func.rs
@@ -76,10 +76,10 @@ macro_rules! impl_into_func {
                 );
                 let trampoline = HostFuncTrampoline::new(
                     move |caller: Caller<T>, params_results: FuncParams| -> Result<FuncResults, Trap> {
-                        let ($($tuple,)*): Self::Params = params_results.read_params();
+                        let ($($tuple,)*): Self::Params = params_results.decode_params();
                         let results: Self::Results =
                             (self)(caller, $($tuple),*).into_fallible()?;
-                        Ok(params_results.write_results(results))
+                        Ok(params_results.encode_results(results))
                     },
                 );
                 (signature, trampoline)
diff --git a/crates/wasmi/src/func/mod.rs b/crates/wasmi/src/func/mod.rs
index 6fc749bab1..4b0df90976 100644
--- a/crates/wasmi/src/func/mod.rs
+++ b/crates/wasmi/src/func/mod.rs
@@ -23,7 +23,7 @@ use super::{
     Stored,
 };
 use crate::{core::Trap, engine::ResumableCall, Error, Value};
-use alloc::sync::Arc;
+use alloc::{boxed::Box, sync::Arc};
 use core::{fmt, fmt::Debug, num::NonZeroU32};
 use wasmi_arena::ArenaIndex;
 
@@ -66,13 +66,24 @@ impl<T> Clone for FuncEntity<T> {
 
 impl<T> FuncEntity<T> {
     /// Creates a new Wasm function from the given raw parts.
-    pub(crate) fn new_wasm(signature: DedupFuncType, body: FuncBody, instance: Instance) -> Self {
+    pub fn new_wasm(signature: DedupFuncType, body: FuncBody, instance: Instance) -> Self {
         Self {
             internal: FuncEntityInternal::Wasm(WasmFuncEntity::new(signature, body, instance)),
         }
     }
 
-    /// Creates a new host function from the given closure.
+    /// Creates a new host function from the given dynamically typed closure.
+    pub fn new(
+        ctx: impl AsContextMut<UserState = T>,
+        ty: FuncType,
+        func: impl Fn(Caller<'_, T>, &[Value], &mut [Value]) -> Result<(), Trap> + Send + Sync + 'static,
+    ) -> Self {
+        Self {
+            internal: FuncEntityInternal::Host(HostFuncEntity::new(ctx, ty, func)),
+        }
+    }
+
+    /// Creates a new host function from the given dynamically typed closure.
     pub fn wrap<Params, Results>(
         ctx: impl AsContextMut<UserState = T>,
         func: impl IntoFunc<T, Params, Results>,
@@ -204,7 +215,39 @@ impl<T> Debug for HostFuncEntity<T> {
 }
 
 impl<T> HostFuncEntity<T> {
-    /// Creates a new host function from the given closure.
+    /// Creates a new host function from the given dynamically typed closure.
+    pub fn new(
+        mut ctx: impl AsContextMut<UserState = T>,
+        ty: FuncType,
+        func: impl Fn(Caller<'_, T>, &[Value], &mut [Value]) -> Result<(), Trap> + Send + Sync + 'static,
+    ) -> Self {
+        // Preprocess parameters and results buffers so that we can reuse those
+        // computations within the closure implementation. We put both parameters
+        // and results into a single buffer which we can split to minimize the
+        // amount of allocations per trampoline invokation.
+        let params_iter = ty.params().iter().copied().map(Value::default);
+        let results_iter = ty.results().iter().copied().map(Value::default);
+        let len_params = ty.params().len();
+        let params_results: Box<[Value]> = params_iter.chain(results_iter).collect();
+        let trampoline = <HostFuncTrampoline<T>>::new(move |caller, args| {
+            // We are required to clone the buffer because we are operating within a `Fn`.
+            // This way the trampoline closure only has to own a single slice buffer.
+            // Note: An alternative solution is to use interior mutability but that solution
+            //       comes with its own downsides.
+            let mut params_results = params_results.clone();
+            let (params, results) = params_results.split_at_mut(len_params);
+            args.decode_params_into_slice(params).unwrap();
+            func(caller, params, results)?;
+            Ok(args.encode_results_from_slice(results).unwrap())
+        });
+        let signature = ctx.as_context_mut().store.inner.alloc_func_type(ty.clone());
+        Self {
+            signature,
+            trampoline,
+        }
+    }
+
+    /// Creates a new host function from the given statically typed closure.
     pub fn wrap<Params, Results>(
         mut ctx: impl AsContextMut,
         func: impl IntoFunc<T, Params, Results>,
@@ -252,6 +295,38 @@ impl Func {
         &self.0
     }
 
+    /// Creates a new [`Func`] with the given arguments.
+    ///
+    /// This is typically used to create a host-defined function to pass as an import to a Wasm module.
+    ///
+    /// - `ty`: the signature that the given closure adheres to,
+    ///         used to indicate what the inputs and outputs are.
+    /// - `func`: the native code invoked whenever this Func will be called.
+    ///           The closure is provided a [`Caller`] as its first argument
+    ///           which allows it to query information about the [`Instance`]
+    ///           that is assocaited to the call.
+    ///
+    /// # Note
+    ///
+    /// - The given [`FuncType`] `ty` must match the parameters and results otherwise
+    ///   the resulting host [`Func`] might trap during execution.
+    /// - It is the responsibility of the caller of [`Func::new`] to guarantee that
+    ///   the correct amount and types of results are written into the results buffer
+    ///   from the `func` closure. If an incorrect amount of results or types of results
+    ///   is written into the buffer then the remaining computation may fail in unexpected
+    ///   ways. This footgun can be avoided by using the typed [`Func::wrap`] method instead.
+    /// - Prefer using [`Func::wrap`] over this method if possible since [`Func`] instances
+    ///   created using this constructor have runtime overhead for every invokation that
+    ///   can be avoided by using [`Func::wrap`].
+    pub fn new<T>(
+        mut ctx: impl AsContextMut<UserState = T>,
+        ty: FuncType,
+        func: impl Fn(Caller<'_, T>, &[Value], &mut [Value]) -> Result<(), Trap> + Send + Sync + 'static,
+    ) -> Self {
+        let func = FuncEntity::new(ctx.as_context_mut(), ty, func);
+        ctx.as_context_mut().store.alloc_func(func)
+    }
+
     /// Creates a new host function from the given closure.
     pub fn wrap<T, Params, Results>(
         mut ctx: impl AsContextMut<UserState = T>,


Test patch:
diff --git a/crates/wasmi/tests/e2e/v1/func.rs b/crates/wasmi/tests/e2e/v1/func.rs
index 055ef57353..65010433c0 100644
--- a/crates/wasmi/tests/e2e/v1/func.rs
+++ b/crates/wasmi/tests/e2e/v1/func.rs
@@ -3,8 +3,8 @@
 use core::slice;
 
 use assert_matches::assert_matches;
-use wasmi::{errors::FuncError, Engine, Error, Func, Store, Value};
-use wasmi_core::{F32, F64};
+use wasmi::{errors::FuncError, Engine, Error, Func, FuncType, Store, Value};
+use wasmi_core::{ValueType, F32, F64};
 
 fn test_setup() -> Store<()> {
     let engine = Engine::default();
@@ -24,85 +24,178 @@ macro_rules! assert_eq_tuple {
     }
 }
 
-fn setup_add2() -> (Store<()>, Func) {
+// Returns a Wasm store and two binary addition [`Func`] instances.
+fn setup_add2() -> (Store<()>, Func, Func) {
     let mut store = test_setup();
-    // This host function represents a simple binary addition.
     let add2 = Func::wrap(&mut store, |lhs: i32, rhs: i32| lhs + rhs);
-    (store, add2)
+    let add2_dyn = Func::new(
+        &mut store,
+        FuncType::new([ValueType::I32, ValueType::I32], [ValueType::I32]),
+        |_caller, inputs: &[Value], results: &mut [Value]| {
+            assert_eq!(inputs.len(), 2);
+            assert_eq!(results.len(), 1);
+            let lhs = &inputs[0].i32().unwrap();
+            let rhs = &inputs[1].i32().unwrap();
+            results[0] = (lhs + rhs).into();
+            Ok(())
+        },
+    );
+    (store, add2, add2_dyn)
 }
 
 #[test]
 fn dynamic_add2_works() {
-    let (mut store, add2) = setup_add2();
-    let mut result = Value::I32(0);
-    add2.call(
-        &mut store,
-        &[Value::I32(1), Value::I32(2)],
-        slice::from_mut(&mut result),
-    )
-    .unwrap();
-    assert_eq!(result.i32(), Some(3));
+    let (mut store, add2, add2_dyn) = setup_add2();
+    for a in 0..10 {
+        for b in 0..10 {
+            let params = [Value::I32(a), Value::I32(b)];
+            let expected = a + b;
+            let mut result = Value::I32(0);
+            // Call to Func with statically typed closure.
+            add2.call(&mut store, &params, slice::from_mut(&mut result))
+                .unwrap();
+            // Reset result before execution.
+            result = Value::I32(0);
+            // Call to Func with dynamically typed closure.
+            add2_dyn
+                .call(&mut store, &params, slice::from_mut(&mut result))
+                .unwrap();
+            assert_eq!(result.i32(), Some(expected));
+        }
+    }
 }
 
 #[test]
 fn static_add2_works() {
-    let (mut store, add2) = setup_add2();
-    let typed_add2 = add2.typed::<(i32, i32), i32>(&mut store).unwrap();
-    let result = typed_add2.call(&mut store, (1, 2)).unwrap();
-    assert_eq!(result, 3);
+    let (mut store, add2, add2_dyn) = setup_add2();
+    let add2 = add2.typed::<(i32, i32), i32>(&mut store).unwrap();
+    let add2_dyn = add2_dyn.typed::<(i32, i32), i32>(&mut store).unwrap();
+    for a in 0..10 {
+        for b in 0..10 {
+            let expected = a + b;
+            assert_eq!(add2.call(&mut store, (a, b)).unwrap(), expected);
+            assert_eq!(add2_dyn.call(&mut store, (a, b)).unwrap(), expected);
+        }
+    }
 }
 
-fn setup_add3() -> (Store<()>, Func) {
+// Returns a Wasm store and two three-way addition [`Func`] instances.
+fn setup_add3() -> (Store<()>, Func, Func) {
     let mut store = test_setup();
-    // This host function performance a three-way addition.
     let add3 = Func::wrap(&mut store, |v0: i32, v1: i32, v2: i32| v0 + v1 + v2);
-    (store, add3)
+    let add3_dyn = Func::new(
+        &mut store,
+        FuncType::new(
+            [ValueType::I32, ValueType::I32, ValueType::I32],
+            [ValueType::I32],
+        ),
+        |_caller, inputs: &[Value], results: &mut [Value]| {
+            assert_eq!(inputs.len(), 3);
+            assert_eq!(results.len(), 1);
+            let a = &inputs[0].i32().unwrap();
+            let b = &inputs[1].i32().unwrap();
+            let c = &inputs[2].i32().unwrap();
+            results[0] = (a + b + c).into();
+            Ok(())
+        },
+    );
+    (store, add3, add3_dyn)
 }
 
 #[test]
 fn dynamic_add3_works() {
-    let (mut store, add3) = setup_add3();
-    let mut result = Value::I32(0);
-    add3.call(
-        &mut store,
-        &[Value::I32(1), Value::I32(2), Value::I32(3)],
-        slice::from_mut(&mut result),
-    )
-    .unwrap();
-    assert_eq!(result.i32(), Some(6));
+    let (mut store, add3, add3_dyn) = setup_add3();
+    for a in 0..5 {
+        for b in 0..5 {
+            for c in 0..5 {
+                let params = [Value::I32(a), Value::I32(b), Value::I32(c)];
+                let expected = a + b + c;
+                let mut result = Value::I32(0);
+                // Call to Func with statically typed closure.
+                add3.call(&mut store, &params, slice::from_mut(&mut result))
+                    .unwrap();
+                assert_eq!(result.i32(), Some(expected));
+                // Reset result before execution.
+                result = Value::I32(0);
+                // Call to Func with dynamically typed closure.
+                add3_dyn
+                    .call(&mut store, &params, slice::from_mut(&mut result))
+                    .unwrap();
+                assert_eq!(result.i32(), Some(expected));
+            }
+        }
+    }
 }
 
 #[test]
 fn static_add3_works() {
-    let (mut store, add3) = setup_add3();
-    let typed_add3 = add3.typed::<(i32, i32, i32), i32>(&mut store).unwrap();
-    let result = typed_add3.call(&mut store, (1, 2, 3)).unwrap();
-    assert_eq!(result, 6);
+    let (mut store, add3, add3_dyn) = setup_add3();
+    let add3 = add3.typed::<(i32, i32, i32), i32>(&mut store).unwrap();
+    let add3_dyn = add3_dyn.typed::<(i32, i32, i32), i32>(&mut store).unwrap();
+    for a in 0..5 {
+        for b in 0..5 {
+            for c in 0..5 {
+                let expected = a + b + c;
+                assert_eq!(add3.call(&mut store, (a, b, c)).unwrap(), expected);
+                assert_eq!(add3_dyn.call(&mut store, (a, b, c)).unwrap(), expected);
+            }
+        }
+    }
 }
 
-fn setup_duplicate() -> (Store<()>, Func) {
+// Returns a `Store` and two Wasm host functions that duplicate their inputs.
+fn setup_duplicate() -> (Store<()>, Func, Func) {
     let mut store = test_setup();
-    // This host function takes one `i32` argument and returns it twice.
     let duplicate = Func::wrap(&mut store, |value: i32| (value, value));
-    (store, duplicate)
+    let duplicate_dyn = Func::new(
+        &mut store,
+        FuncType::new([ValueType::I32], [ValueType::I32, ValueType::I32]),
+        |_caller, inputs: &[Value], results: &mut [Value]| {
+            assert_eq!(inputs.len(), 1);
+            assert_eq!(results.len(), 2);
+            let input = inputs[0].i32().unwrap();
+            results[0] = input.into();
+            results[1] = input.into();
+            Ok(())
+        },
+    );
+    (store, duplicate, duplicate_dyn)
 }
 
 #[test]
 fn dynamic_duplicate_works() {
-    let (mut store, duplicate) = setup_duplicate();
-    let mut result = [Value::I32(0), Value::I32(0)];
-    duplicate
-        .call(&mut store, &[Value::I32(10)], &mut result)
-        .unwrap();
-    assert_eq!((result[0].i32(), result[1].i32()), (Some(10), Some(10)));
+    let (mut store, duplicate, duplicate_dyn) = setup_duplicate();
+    for input in 0..10 {
+        let params = [Value::I32(input)];
+        let expected = [Value::I32(input), Value::I32(input)];
+        let mut results = [Value::I32(0), Value::I32(0)];
+        // Call to Func with statically typed closure.
+        duplicate.call(&mut store, &params, &mut results).unwrap();
+        assert_eq!(results[0].i32(), expected[0].i32());
+        assert_eq!(results[1].i32(), expected[1].i32());
+        // Reset result before execution.
+        results = [Value::I32(0), Value::I32(0)];
+        // Call to Func with dynamically typed closure.
+        duplicate_dyn
+            .call(&mut store, &params, &mut results)
+            .unwrap();
+        assert_eq!(results[0].i32(), expected[0].i32());
+        assert_eq!(results[1].i32(), expected[1].i32());
+    }
 }
 
 #[test]
 fn static_duplicate_works() {
-    let (mut store, duplicate) = setup_duplicate();
-    let typed_duplicate = duplicate.typed::<i32, (i32, i32)>(&mut store).unwrap();
-    let result = typed_duplicate.call(&mut store, 10).unwrap();
-    assert_eq!(result, (10, 10));
+    let (mut store, duplicate, duplicate_dyn) = setup_duplicate();
+    let duplicate = duplicate.typed::<i32, (i32, i32)>(&mut store).unwrap();
+    let duplicate_dyn = duplicate_dyn.typed::<i32, (i32, i32)>(&mut store).unwrap();
+    for input in 0..10 {
+        assert_eq!(duplicate.call(&mut store, input).unwrap(), (input, input));
+        assert_eq!(
+            duplicate_dyn.call(&mut store, input).unwrap(),
+            (input, input)
+        );
+    }
 }
 
 fn setup_many_params() -> (Store<()>, Func) {