Supply type information in func_new for components (#11944)

Originally the `Linker::func_new` type for components was modeled after
core wasm where a type was provided when the type was defined. This was
difficult, however, because there's no way to construct types right now
and instead a component had to be created and compiled to acquire type
information from it. Later various refactorings meant that it was
possible to drop the type information entirely from `func_new` meaning
that it was "typeless" in a sense, and the functionality relied on the
tagged nature of `Val` which always knows its type.

This has proven a bit difficult to integrate into Python in some work
I've been doing. Namely in Python component values are (at least IMO)
best represented as native Python values where possible. Native Python
doesn't distinguish, however, between integer bit widths (or
signededness). This means that converting a Python value to a component
value requires type information to guide the conversion process.
Currently when defining a function in a linker there's no way to get at
this type information meaning that the types would need to be supplied
in Python, bringing up the same shortcomings of not being able to create
types.

In lieu of filling out type constructors, which is itself quite
nontrivial, I've opted to do the next-most-easiest thing which is to
supply the type to the callee when an import is invoked. This means that
the callee has all the type information necessary. While this is
somewhat costly in that it clones a few arcs it's expected to be a drop
in the bucket compared to the inefficiencies of `Val` so it's at least
in the same spirit of the cost of the API right now.

My intention is to use this to update the C API for components and
eventually end up with type information when Python is invoked in
wasmtime-py. Also while I'm using wasmtime-py as a motivating case here
that's sort of just one example of a host embedding which doesn't have
the full fidelity of Rust's type system and might benefit from this.

Author: alexcrichton

benches/call.rs

@@ -797,10 +797,10 @@ mod component {
             bench_instance(group, store, &instance, "typed", is_async);
 
             let mut untyped = component::Linker::new(&engine);
-            untyped.root().func_new("nop", |_, _, _| Ok(())).unwrap();
+            untyped.root().func_new("nop", |_, _, _, _| Ok(())).unwrap();
             untyped
                 .root()
-                .func_new("nop-params-and-results", |_caller, params, results| {
+                .func_new("nop-params-and-results", |_caller, _ty, params, results| {
                     assert_eq!(params.len(), 2);
                     match params[0] {
                         component::Val::U32(0) => {}

crates/c-api/src/component/linker.rs

@@ -123,7 +123,7 @@ pub unsafe extern "C" fn wasmtime_component_linker_instance_add_func(
 
     let result = linker_instance
         .linker_instance
-        .func_new(&name, move |ctx, args, rets| {
+        .func_new(&name, move |ctx, _ty, args, rets| {
             let _ = &foreign;
 
             let args = args

crates/fuzzing/src/oracles.rs

@@ -1090,6 +1090,7 @@ pub fn dynamic_component_api_target(input: &mut arbitrary::Unstructured) -> arbi
         .root()
         .func_new(IMPORT_FUNCTION, {
             move |mut cx: StoreContextMut<'_, (Vec<Val>, Option<Vec<Val>>)>,
+                  _,
                   params: &[Val],
                   results: &mut [Val]|
                   -> Result<()> {

crates/misc/component-async-tests/tests/scenario/round_trip.rs

@@ -352,7 +352,7 @@ pub async fn test_round_trip(
         linker
             .root()
             .instance("local:local/baz")?
-            .func_new_concurrent("[async]foo", |_, params, results| {
+            .func_new_concurrent("[async]foo", |_, _, params, results| {
                 Box::pin(async move {
                     sleep(Duration::from_millis(10)).await;
                     let Some(Val::String(s)) = params.into_iter().next() else {

crates/misc/component-async-tests/tests/scenario/round_trip_direct.rs

@@ -94,7 +94,7 @@ async fn test_round_trip_direct(
         wasmtime_wasi::p2::add_to_linker_async(&mut linker)?;
         linker
             .root()
-            .func_new_concurrent("[async]foo", |_, params, results| {
+            .func_new_concurrent("[async]foo", |_, _, params, results| {
                 Box::pin(async move {
                     sleep(Duration::from_millis(10)).await;
                     let Some(Val::String(s)) = params.into_iter().next() else {

crates/misc/component-async-tests/tests/scenario/round_trip_many.rs

@@ -341,7 +341,7 @@ async fn test_round_trip_many(
         linker
             .root()
             .instance("local:local/many")?
-            .func_new_concurrent("[async]foo", |_, params, results| {
+            .func_new_concurrent("[async]foo", |_, _, params, results| {
                 Box::pin(async move {
                     sleep(Duration::from_millis(10)).await;
                     let mut params = params.into_iter();

crates/wasmtime/src/runtime/component/func/host.rs

@@ -4,6 +4,7 @@ use crate::component::concurrent::{Accessor, Status};
 use crate::component::func::{LiftContext, LowerContext, Options};
 use crate::component::matching::InstanceType;
 use crate::component::storage::slice_to_storage_mut;
+use crate::component::types::ComponentFunc;
 use crate::component::{ComponentNamedList, ComponentType, Instance, Lift, Lower, Val};
 use crate::prelude::*;
 use crate::runtime::vm::component::{
@@ -121,6 +122,7 @@ impl HostFunc {
     where
         F: Fn(
                 StoreContextMut<'_, T>,
+                ComponentFunc,
                 Vec<Val>,
                 usize,
             ) -> Pin<Box<dyn Future<Output = Result<Vec<Val>>> + Send + 'static>>
@@ -141,13 +143,18 @@ impl HostFunc {
 
     pub(crate) fn new_dynamic<T: 'static, F>(func: F) -> Arc<HostFunc>
     where
-        F: Fn(StoreContextMut<'_, T>, &[Val], &mut [Val]) -> Result<()> + Send + Sync + 'static,
+        F: Fn(StoreContextMut<'_, T>, ComponentFunc, &[Val], &mut [Val]) -> Result<()>
+            + Send
+            + Sync
+            + 'static,
     {
-        Self::new_dynamic_canonical::<T, _>(move |store, mut params_and_results, result_start| {
-            let (params, results) = params_and_results.split_at_mut(result_start);
-            let result = func(store, params, results).map(move |()| params_and_results);
-            Box::pin(async move { result })
-        })
+        Self::new_dynamic_canonical::<T, _>(
+            move |store, ty, mut params_and_results, result_start| {
+                let (params, results) = params_and_results.split_at_mut(result_start);
+                let result = func(store, ty, params, results).map(move |()| params_and_results);
+                Box::pin(async move { result })
+            },
+        )
     }
 
     #[cfg(feature = "component-model-async")]
@@ -156,6 +163,7 @@ impl HostFunc {
         T: 'static,
         F: for<'a> Fn(
                 &'a Accessor<T>,
+                ComponentFunc,
                 &'a [Val],
                 &'a mut [Val],
             ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>>
@@ -164,16 +172,18 @@ impl HostFunc {
             + 'static,
     {
         let func = Arc::new(func);
-        Self::new_dynamic_canonical::<T, _>(move |store, mut params_and_results, result_start| {
-            let func = func.clone();
-            Box::pin(store.wrap_call(move |accessor| {
-                Box::pin(async move {
-                    let (params, results) = params_and_results.split_at_mut(result_start);
-                    func(accessor, params, results).await?;
-                    Ok(params_and_results)
-                })
-            }))
-        })
+        Self::new_dynamic_canonical::<T, _>(
+            move |store, ty, mut params_and_results, result_start| {
+                let func = func.clone();
+                Box::pin(store.wrap_call(move |accessor| {
+                    Box::pin(async move {
+                        let (params, results) = params_and_results.split_at_mut(result_start);
+                        func(accessor, ty, params, results).await?;
+                        Ok(params_and_results)
+                    })
+                }))
+            },
+        )
     }
 
     pub fn typecheck(&self, ty: TypeFuncIndex, types: &InstanceType<'_>) -> Result<()> {
@@ -700,6 +710,7 @@ unsafe fn call_host_dynamic<T, F>(
 where
     F: Fn(
             StoreContextMut<'_, T>,
+            ComponentFunc,
             Vec<Val>,
             usize,
         ) -> Pin<Box<dyn Future<Output = Result<Vec<Val>>> + Send + 'static>>
@@ -735,6 +746,7 @@ where
     } else {
         MAX_FLAT_PARAMS
     };
+    let ty = ComponentFunc::from(ty, &lift.instance_type());
 
     let ret_index = unsafe {
         dynamic_params_load(
@@ -763,7 +775,7 @@ where
                 validate_inbounds_dynamic(&result_tys.abi, lower.as_slice_mut(), &retptr)?
             };
 
-            let future = closure(store.as_context_mut(), params_and_results, result_start);
+            let future = closure(store.as_context_mut(), ty, params_and_results, result_start);
 
             let task = instance.first_poll(store, future, caller_instance, {
                 let types = types.clone();
@@ -810,7 +822,7 @@ where
             );
         }
     } else {
-        let future = closure(store.as_context_mut(), params_and_results, result_start);
+        let future = closure(store.as_context_mut(), ty, params_and_results, result_start);
         let result_vals = concurrent::poll_and_block(store.0, future, caller_instance)?;
         let result_vals = &result_vals[result_start..];
 
@@ -915,6 +927,7 @@ extern "C" fn dynamic_entrypoint<T, F>(
 where
     F: Fn(
             StoreContextMut<'_, T>,
+            ComponentFunc,
             Vec<Val>,
             usize,
         ) -> Pin<Box<dyn Future<Output = Result<Vec<Val>>> + Send + 'static>>
@@ -932,7 +945,7 @@ where
                 TypeFuncIndex::from_u32(ty),
                 OptionsIndex::from_u32(options),
                 NonNull::slice_from_raw_parts(storage, storage_len).as_mut(),
-                move |store, params, results| (*data.as_ptr())(store, params, results),
+                &*data.as_ptr(),
             )
         })
     }

crates/wasmtime/src/runtime/component/linker.rs

@@ -337,7 +337,7 @@ impl<T: 'static> Linker<T> {
                     let fully_qualified_name = parent_instance
                         .map(|parent| format!("{parent}#{item_name}"))
                         .unwrap_or_else(|| item_name.to_owned());
-                    linker.func_new(&item_name, move |_, _, _| {
+                    linker.func_new(&item_name, move |_, _, _, _| {
                         bail!("unknown import: `{fully_qualified_name}` has not been defined")
                     })?;
                 }
@@ -635,15 +635,15 @@ impl<T: 'static> LinkerInstance<'_, T> {
     /// let mut linker = Linker::<()>::new(&engine);
     ///
     /// // Sample function that takes no arguments.
-    /// linker.root().func_new("thunk", |_store, params, results| {
+    /// linker.root().func_new("thunk", |_store, _ty, params, results| {
     ///     assert!(params.is_empty());
     ///     assert!(results.is_empty());
     ///     println!("Look ma, host hands!");
     ///     Ok(())
     /// })?;
     ///
     /// // This function takes one argument and returns one result.
-    /// linker.root().func_new("is-even", |_store, params, results| {
+    /// linker.root().func_new("is-even", |_store, _ty, params, results| {
     ///     assert_eq!(params.len(), 1);
     ///     let param = match params[0] {
     ///         Val::U32(n) => n,
@@ -665,7 +665,10 @@ impl<T: 'static> LinkerInstance<'_, T> {
     pub fn func_new(
         &mut self,
         name: &str,
-        func: impl Fn(StoreContextMut<'_, T>, &[Val], &mut [Val]) -> Result<()> + Send + Sync + 'static,
+        func: impl Fn(StoreContextMut<'_, T>, types::ComponentFunc, &[Val], &mut [Val]) -> Result<()>
+        + Send
+        + Sync
+        + 'static,
     ) -> Result<()> {
         self.insert(name, Definition::Func(HostFunc::new_dynamic(func)))?;
         Ok(())
@@ -682,6 +685,7 @@ impl<T: 'static> LinkerInstance<'_, T> {
     where
         F: for<'a> Fn(
                 StoreContextMut<'a, T>,
+                types::ComponentFunc,
                 &'a [Val],
                 &'a mut [Val],
             ) -> Box<dyn Future<Output = Result<()>> + Send + 'a>
@@ -693,8 +697,9 @@ impl<T: 'static> LinkerInstance<'_, T> {
             self.engine.config().async_support,
             "cannot use `func_new_async` without enabling async support in the config"
         );
-        let ff = move |store: StoreContextMut<'_, T>, params: &[Val], results: &mut [Val]| {
-            store.with_blocking(|store, cx| cx.block_on(Pin::from(f(store, params, results)))?)
+        let ff = move |store: StoreContextMut<'_, T>, ty, params: &[Val], results: &mut [Val]| {
+            store
+                .with_blocking(|store, cx| cx.block_on(Pin::from(f(store, ty, params, results)))?)
         };
         return self.func_new(name, ff);
     }
@@ -712,6 +717,7 @@ impl<T: 'static> LinkerInstance<'_, T> {
         T: 'static,
         F: for<'a> Fn(
                 &'a Accessor<T>,
+                types::ComponentFunc,
                 &'a [Val],
                 &'a mut [Val],
             ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>>

tests/all/component_model/import.rs

@@ -148,7 +148,7 @@ fn simple() -> Result<()> {
     let mut linker = Linker::new(&engine);
     linker.root().func_new(
         "a",
-        |mut store: StoreContextMut<'_, Option<String>>, args, _results| {
+        |mut store: StoreContextMut<'_, Option<String>>, _, args, _results| {
             if let Val::String(s) = &args[0] {
                 assert!(store.data().is_none());
                 *store.data_mut() = Some(s.to_string());
@@ -249,7 +249,7 @@ fn functions_in_instances() -> Result<()> {
     let mut linker = Linker::new(&engine);
     linker.instance("test:test/foo")?.func_new(
         "a",
-        |mut store: StoreContextMut<'_, Option<String>>, args, _results| {
+        |mut store: StoreContextMut<'_, Option<String>>, _, args, _results| {
             if let Val::String(s) = &args[0] {
                 assert!(store.data().is_none());
                 *store.data_mut() = Some(s.to_string());
@@ -462,7 +462,7 @@ fn attempt_to_reenter_during_host() -> Result<()> {
     let mut linker = Linker::new(&engine);
     linker.root().func_new(
         "thunk",
-        |mut store: StoreContextMut<'_, DynamicState>, _, _| {
+        |mut store: StoreContextMut<'_, DynamicState>, _, _, _| {
             let func = store.data_mut().func.take().unwrap();
             let trap = func.call(&mut store, &[], &mut []).unwrap_err();
             assert_eq!(
@@ -848,7 +848,7 @@ async fn test_stack_and_heap_args_and_rets(concurrent: bool) -> Result<()> {
     if concurrent {
         linker
             .root()
-            .func_new_concurrent("f1", |_, args, results| {
+            .func_new_concurrent("f1", |_, _, args, results| {
                 if let Val::U32(x) = &args[0] {
                     assert_eq!(*x, 1);
                     Box::pin(async {
@@ -861,7 +861,7 @@ async fn test_stack_and_heap_args_and_rets(concurrent: bool) -> Result<()> {
             })?;
         linker
             .root()
-            .func_new_concurrent("f2", |_, args, results| {
+            .func_new_concurrent("f2", |_, _, args, results| {
                 if let Val::Tuple(tuple) = &args[0] {
                     if let Val::String(s) = &tuple[0] {
                         assert_eq!(s.deref(), "abc");
@@ -878,7 +878,7 @@ async fn test_stack_and_heap_args_and_rets(concurrent: bool) -> Result<()> {
             })?;
         linker
             .root()
-            .func_new_concurrent("f3", |_, args, results| {
+            .func_new_concurrent("f3", |_, _, args, results| {
                 if let Val::U32(x) = &args[0] {
                     assert_eq!(*x, 8);
                     Box::pin(async {
@@ -891,7 +891,7 @@ async fn test_stack_and_heap_args_and_rets(concurrent: bool) -> Result<()> {
             })?;
         linker
             .root()
-            .func_new_concurrent("f4", |_, args, results| {
+            .func_new_concurrent("f4", |_, _, args, results| {
                 if let Val::Tuple(tuple) = &args[0] {
                     if let Val::String(s) = &tuple[0] {
                         assert_eq!(s.deref(), "abc");
@@ -907,7 +907,7 @@ async fn test_stack_and_heap_args_and_rets(concurrent: bool) -> Result<()> {
                 }
             })?;
     } else {
-        linker.root().func_new("f1", |_, args, results| {
+        linker.root().func_new("f1", |_, _, args, results| {
             if let Val::U32(x) = &args[0] {
                 assert_eq!(*x, 1);
                 results[0] = Val::U32(2);
@@ -916,7 +916,7 @@ async fn test_stack_and_heap_args_and_rets(concurrent: bool) -> Result<()> {
                 panic!()
             }
         })?;
-        linker.root().func_new("f2", |_, args, results| {
+        linker.root().func_new("f2", |_, _, args, results| {
             if let Val::Tuple(tuple) = &args[0] {
                 if let Val::String(s) = &tuple[0] {
                     assert_eq!(s.deref(), "abc");
@@ -929,7 +929,7 @@ async fn test_stack_and_heap_args_and_rets(concurrent: bool) -> Result<()> {
                 panic!()
             }
         })?;
-        linker.root().func_new("f3", |_, args, results| {
+        linker.root().func_new("f3", |_, _, args, results| {
             if let Val::U32(x) = &args[0] {
                 assert_eq!(*x, 8);
                 results[0] = Val::String("xyz".into());
@@ -938,7 +938,7 @@ async fn test_stack_and_heap_args_and_rets(concurrent: bool) -> Result<()> {
                 panic!();
             }
         })?;
-        linker.root().func_new("f4", |_, args, results| {
+        linker.root().func_new("f4", |_, _, args, results| {
             if let Val::Tuple(tuple) = &args[0] {
                 if let Val::String(s) = &tuple[0] {
                     assert_eq!(s.deref(), "abc");
@@ -1126,7 +1126,7 @@ fn no_actual_wasm_code() -> Result<()> {
     let mut linker = Linker::new(&engine);
     linker
         .root()
-        .func_new("f", |mut store: StoreContextMut<'_, u32>, _, _| {
+        .func_new("f", |mut store: StoreContextMut<'_, u32>, _, _, _| {
             *store.data_mut() += 1;
             Ok(())
         })?;

tests/all/intrinsics.rs

@@ -93,7 +93,7 @@ fn native_loads_and_stores() -> Result<()> {
         linker.instance("host")?.func_new("get-pointer", {
             let ptr = ptr as usize;
             let ptr = u64::try_from(ptr).unwrap();
-            move |_cx, _args, results| {
+            move |_cx, _, _args, results| {
                 results[0] = component::Val::U64(ptr);
                 Ok(())
             }