wit-component: fix async task.return options for export signatures

crates/wit-component/src/encoding.rs

@@ -203,7 +203,7 @@ impl RequiredOptions {
         }
         if let AbiVariant::GuestExportAsync | AbiVariant::GuestExportAsyncStackful = abi {
             ret |= RequiredOptions::ASYNC;
-            ret |= task_return_options_and_type(resolve, func.result).0;
+            ret |= task_return_options_and_type(resolve, func).0;
         }
         ret
     }
@@ -1791,8 +1791,9 @@ impl<'a> EncodingState<'a> {
                     AbiVariant::GuestImport,
                 )
             }
-            Import::ExportedTaskReturn(key, interface, func, result) => {
-                let (options, _sig) = task_return_options_and_type(resolve, *result);
+            Import::ExportedTaskReturn(key, interface, func) => {
+                let (options, _sig) = task_return_options_and_type(resolve, func);
+                let result_ty = func.result;
                 if options.is_empty() {
                     // Note that an "import type encoder" is used here despite
                     // this being for an exported function if the `interface`
@@ -1805,22 +1806,25 @@ impl<'a> EncodingState<'a> {
                         self.root_export_type_encoder(*interface)
                     };
 
-                    let result = match result {
+                    let result = match result_ty.as_ref() {
                         Some(ty) => Some(encoder.encode_valtype(resolve, ty)?),
                         None => None,
                     };
                     let index = self.component.task_return(result, []);
                     Ok((ExportKind::Func, index))
                 } else {
                     let metadata = &self.info.module_metadata_for(for_module);
-                    let encoding = metadata.export_encodings.get(resolve, key, func).unwrap();
+                    let encoding = metadata
+                        .export_encodings
+                        .get(resolve, key, &func.name)
+                        .unwrap();
                     Ok(self.materialize_shim_import(
                         shims,
                         &ShimKind::TaskReturn {
                             for_module,
                             interface: *interface,
-                            func,
-                            result: *result,
+                            func: &func.name,
+                            result: result_ty,
                             encoding,
                         },
                     ))
@@ -2638,16 +2642,16 @@ impl<'a> Shims<'a> {
                 // If `task.return` needs to be indirect then generate a shim
                 // for it, otherwise skip the shim and let it get materialized
                 // naturally later.
-                Import::ExportedTaskReturn(key, interface, func, ty) => {
-                    let (options, sig) = task_return_options_and_type(resolve, *ty);
+                Import::ExportedTaskReturn(key, interface, func) => {
+                    let (options, sig) = task_return_options_and_type(resolve, func);
                     if options.is_empty() {
                         continue;
                     }
                     let name = self.shims.len().to_string();
                     let encoding = world
                         .module_metadata_for(for_module)
                         .export_encodings
-                        .get(resolve, key, func)
+                        .get(resolve, key, &func.name)
                         .ok_or_else(|| {
                             anyhow::anyhow!(
                                 "missing component metadata for export of \
@@ -2656,12 +2660,12 @@ impl<'a> Shims<'a> {
                         })?;
                     self.push(Shim {
                         name,
-                        debug_name: format!("task-return-{func}"),
+                        debug_name: format!("task-return-{}", func.name),
                         options,
                         kind: ShimKind::TaskReturn {
                             interface: *interface,
-                            func,
-                            result: *ty,
+                            func: &func.name,
+                            result: func.result,
                             for_module,
                             encoding,
                         },
@@ -3027,15 +3031,15 @@ impl<'a> Shims<'a> {
 
 fn task_return_options_and_type(
     resolve: &Resolve,
-    ty: Option<Type>,
+    func: &Function,
 ) -> (RequiredOptions, WasmSignature) {
     let func_tmp = Function {
         name: String::new(),
         kind: FunctionKind::Freestanding,
-        params: match ty {
+        params: match &func.result {
             Some(ty) => vec![Param {
                 name: "a".to_string(),
-                ty,
+                ty: *ty,
                 span: Default::default(),
             }],
             None => Vec::new(),
@@ -3046,7 +3050,9 @@ fn task_return_options_and_type(
         span: Default::default(),
     };
     let abi = AbiVariant::GuestImport;
-    let options = RequiredOptions::for_import(resolve, &func_tmp, abi);
+    let mut options = RequiredOptions::for_import(resolve, func, abi);
+    // `task.return` does not support a `realloc` canonical option.
+    options.remove(RequiredOptions::REALLOC);
     let sig = resolve.wasm_signature(abi, &func_tmp);
     (options, sig)
 }

crates/wit-component/src/encoding/world.rs

@@ -403,9 +403,9 @@ impl<'a> ComponentWorld<'a> {
 
                 // The `task.return` intrinsic needs to be able to refer to the
                 // type that is being returned.
-                Import::ExportedTaskReturn(.., ty) => {
-                    if let Some(ty) = ty {
-                        live.add_type(resolve, ty);
+                Import::ExportedTaskReturn(.., func) => {
+                    if let Some(ty) = func.result {
+                        live.add_type(resolve, &ty);
                     }
                 }
 

crates/wit-component/src/validation.rs

@@ -268,7 +268,7 @@ pub enum Import {
     /// As of this writing, only async-lifted exports use `task.return`, but the
     /// plan is to also support it for sync-lifted exports in the future as
     /// well.
-    ExportedTaskReturn(WorldKey, Option<InterfaceId>, String, Option<Type>),
+    ExportedTaskReturn(WorldKey, Option<InterfaceId>, Function),
 
     /// A `canon task.cancel` intrinsic for an exported function.
     ///
@@ -901,12 +901,7 @@ impl ImportMap {
                 let key = key.unwrap_or_else(|| WorldKey::Name(name.to_string()));
                 // TODO: should call `validate_func_sig` but would require
                 // calculating the expected signature based of `func.result`.
-                return Ok(Some(Import::ExportedTaskReturn(
-                    key,
-                    id,
-                    func.name.clone(),
-                    func.result,
-                )));
+                return Ok(Some(Import::ExportedTaskReturn(key, id, func.clone())));
             }
             if names.task_cancel(name) {
                 let expected = FuncType::new([], []);

crates/wit-component/tests/components/async-task-return-param-string/component.wat

@@ -0,0 +1,81 @@
+(component
+  (core module $main (;0;)
+    (type (;0;) (func))
+    (type (;1;) (func (param i32 i32) (result i32)))
+    (type (;2;) (func (param i32 i32 i32) (result i32)))
+    (type (;3;) (func (param i32 i32 i32 i32) (result i32)))
+    (import "[export]$root" "[task-cancel]" (func (;0;) (type 0)))
+    (import "[export]$root" "[task-return]foo" (func (;1;) (type 0)))
+    (memory (;0;) 1)
+    (export "[async-lift]foo" (func 2))
+    (export "[callback][async-lift]foo" (func 3))
+    (export "memory" (memory 0))
+    (export "cabi_realloc" (func 4))
+    (func (;2;) (type 1) (param i32 i32) (result i32)
+      unreachable
+    )
+    (func (;3;) (type 2) (param i32 i32 i32) (result i32)
+      unreachable
+    )
+    (func (;4;) (type 3) (param i32 i32 i32 i32) (result i32)
+      unreachable
+    )
+    (@producers
+      (processed-by "wit-component" "$CARGO_PKG_VERSION")
+      (processed-by "my-fake-bindgen" "123.45")
+    )
+  )
+  (core module $wit-component-shim-module (;1;)
+    (type (;0;) (func))
+    (table (;0;) 1 1 funcref)
+    (export "0" (func $task-return-foo))
+    (export "$imports" (table 0))
+    (func $task-return-foo (;0;) (type 0)
+      i32.const 0
+      call_indirect (type 0)
+    )
+    (@producers
+      (processed-by "wit-component" "$CARGO_PKG_VERSION")
+    )
+  )
+  (core module $wit-component-fixup (;2;)
+    (type (;0;) (func))
+    (import "" "0" (func (;0;) (type 0)))
+    (import "" "$imports" (table (;0;) 1 1 funcref))
+    (elem (;0;) (i32.const 0) func 0)
+    (@producers
+      (processed-by "wit-component" "$CARGO_PKG_VERSION")
+    )
+  )
+  (core instance $wit-component-shim-instance (;0;) (instantiate $wit-component-shim-module))
+  (core func $task.cancel (;0;) (canon task.cancel))
+  (alias core export $wit-component-shim-instance "0" (core func $task-return-foo (;1;)))
+  (core instance $"[export]$root" (;1;)
+    (export "[task-cancel]" (func $task.cancel))
+    (export "[task-return]foo" (func $task-return-foo))
+  )
+  (core instance $main (;2;) (instantiate $main
+      (with "[export]$root" (instance $"[export]$root"))
+    )
+  )
+  (alias core export $main "memory" (core memory $memory (;0;)))
+  (alias core export $wit-component-shim-instance "$imports" (core table $"shim table" (;0;)))
+  (alias core export $main "cabi_realloc" (core func $realloc (;2;)))
+  (core func $task.return (;3;) (canon task.return (memory $memory) string-encoding=utf8))
+  (core instance $fixup-args (;3;)
+    (export "$imports" (table $"shim table"))
+    (export "0" (func $task.return))
+  )
+  (core instance $fixup (;4;) (instantiate $wit-component-fixup
+      (with "" (instance $fixup-args))
+    )
+  )
+  (type (;0;) (func (param "s" string)))
+  (alias core export $main "[async-lift]foo" (core func $"[async-lift]foo" (;4;)))
+  (alias core export $main "[callback][async-lift]foo" (core func $"[callback][async-lift]foo" (;5;)))
+  (func $foo (;0;) (type 0) (canon lift (core func $"[async-lift]foo") (memory $memory) (realloc $realloc) string-encoding=utf8 async (callback $"[callback][async-lift]foo")))
+  (export $"#func1 foo" (@name "foo") (;1;) "foo" (func $foo))
+  (@producers
+    (processed-by "wit-component" "$CARGO_PKG_VERSION")
+  )
+)

crates/wit-component/tests/components/async-task-return-param-string/component.wit.print

@@ -0,0 +1,5 @@
+package root:component;
+
+world root {
+  export foo: func(s: string);
+}

crates/wit-component/tests/components/async-task-return-param-string/module.wat

@@ -0,0 +1,8 @@
+(module
+  (import "[export]$root" "[task-cancel]" (func))
+  (import "[export]$root" "[task-return]foo" (func))
+  (func (export "[async-lift]foo") (param i32 i32) (result i32) unreachable)
+  (func (export "[callback][async-lift]foo") (param i32 i32 i32) (result i32) unreachable)
+  (memory (export "memory") 1)
+  (func (export "cabi_realloc") (param i32 i32 i32 i32) (result i32) unreachable)
+)

crates/wit-component/tests/components/async-task-return-param-string/module.wit

@@ -0,0 +1,5 @@
+package foo:foo;
+
+world module {
+  export foo: func(s: string);
+}