update task.return to match spec (#1989)

dicej · web-flow · commit 4e37f73e2690 · 2025-01-28T21:43:27.000Z
* update `task.return` to match spec As of WebAssembly/component-model#431, `task.return` takes component-level result types rather than a core function type representing the flattening of those types. This updates `wasm-tools` to match. Signed-off-by: Joel Dice <joel.dice@fermyon.com> * remove `task.return` support for named results Signed-off-by: Joel Dice <joel.dice@fermyon.com> * add assertion to `ComponentState::task_return` Signed-off-by: Joel Dice <joel.dice@fermyon.com> --------- Signed-off-by: Joel Dice <joel.dice@fermyon.com>
diff --git a/crates/wasm-encoder/src/component/builder.rs b/crates/wasm-encoder/src/component/builder.rs
@@ -392,7 +392,7 @@ impl ComponentBuilder {
     }
 
     /// Declares a new `task.return` intrinsic.
-    pub fn task_return(&mut self, ty: u32) -> u32 {
+    pub fn task_return(&mut self, ty: Option<impl Into<ComponentValType>>) -> u32 {
         self.canonical_functions().task_return(ty);
         inc(&mut self.core_funcs)
     }
diff --git a/crates/wasm-encoder/src/component/canonicals.rs b/crates/wasm-encoder/src/component/canonicals.rs
@@ -1,4 +1,4 @@
-use crate::{encode_section, ComponentSection, ComponentSectionId, Encode};
+use crate::{encode_section, ComponentSection, ComponentSectionId, ComponentValType, Encode};
 use alloc::vec::Vec;
 
 /// Represents options for canonical function definitions.
@@ -188,9 +188,15 @@ impl CanonicalFunctionSection {
     /// Defines a function which returns a result to the caller of a lifted
     /// export function.  This allows the callee to continue executing after
     /// returning a result.
-    pub fn task_return(&mut self, ty: u32) -> &mut Self {
+    pub fn task_return(&mut self, ty: Option<impl Into<ComponentValType>>) -> &mut Self {
         self.bytes.push(0x09);
-        ty.encode(&mut self.bytes);
+        if let Some(ty) = ty {
+            self.bytes.push(0x00);
+            ty.into().encode(&mut self.bytes);
+        } else {
+            self.bytes.push(0x01);
+            0_usize.encode(&mut self.bytes);
+        }
         self.num_added += 1;
         self
     }
diff --git a/crates/wasm-encoder/src/reencode/component.rs b/crates/wasm-encoder/src/reencode/component.rs
@@ -970,8 +970,8 @@ pub mod component_utils {
             wasmparser::CanonicalFunction::TaskBackpressure => {
                 section.task_backpressure();
             }
-            wasmparser::CanonicalFunction::TaskReturn { type_index } => {
-                section.task_return(reencoder.type_index(type_index));
+            wasmparser::CanonicalFunction::TaskReturn { result } => {
+                section.task_return(result.map(|ty| reencoder.component_val_type(ty)));
             }
             wasmparser::CanonicalFunction::TaskWait { async_, memory } => {
                 section.task_wait(async_, reencoder.memory_index(memory));
diff --git a/crates/wasmparser/src/readers/component/canonicals.rs b/crates/wasmparser/src/readers/component/canonicals.rs
@@ -1,6 +1,8 @@
 use crate::limits::MAX_WASM_CANONICAL_OPTIONS;
 use crate::prelude::*;
-use crate::{BinaryReader, FromReader, Result, SectionLimited};
+use crate::{
+    BinaryReader, BinaryReaderError, ComponentValType, FromReader, Result, SectionLimited,
+};
 
 /// Represents options for component functions.
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -80,10 +82,8 @@ pub enum CanonicalFunction {
     /// function.  This allows the callee to continue executing after returning
     /// a result.
     TaskReturn {
-        /// Core function type whose parameters represent the flattened
-        /// representation of the component-level results to be returned by the
-        /// currently executing task.
-        type_index: u32,
+        /// The result type, if any.
+        result: Option<ComponentValType>,
     },
     /// A function which waits for at least one outstanding async
     /// task/stream/future to make progress, returning the first such event.
@@ -275,7 +275,20 @@ impl<'a> FromReader<'a> for CanonicalFunction {
             0x06 => CanonicalFunction::ThreadHwConcurrency,
             0x08 => CanonicalFunction::TaskBackpressure,
             0x09 => CanonicalFunction::TaskReturn {
-                type_index: reader.read()?,
+                result: match reader.read_u8()? {
+                    0x00 => Some(reader.read()?),
+                    0x01 => {
+                        if reader.read_u8()? == 0 {
+                            None
+                        } else {
+                            return Err(BinaryReaderError::new(
+                                "named results not allowed for `task.return` intrinsic",
+                                reader.original_position() - 2,
+                            ));
+                        }
+                    }
+                    x => return reader.invalid_leading_byte(x, "`task.return` result"),
+                },
             },
             0x0a => CanonicalFunction::TaskWait {
                 async_: reader.read()?,
diff --git a/crates/wasmparser/src/validator.rs b/crates/wasmparser/src/validator.rs
@@ -1313,8 +1313,8 @@ impl Validator {
                     crate::CanonicalFunction::TaskBackpressure => {
                         current.task_backpressure(types, offset, features)
                     }
-                    crate::CanonicalFunction::TaskReturn { type_index } => {
-                        current.task_return(type_index, types, offset, features)
+                    crate::CanonicalFunction::TaskReturn { result } => {
+                        current.task_return(&result, types, offset, features)
                     }
                     crate::CanonicalFunction::TaskWait { async_, memory } => {
                         current.task_wait(async_, memory, types, offset, features)
diff --git a/crates/wasmparser/src/validator/component.rs b/crates/wasmparser/src/validator/component.rs
@@ -19,9 +19,9 @@ use crate::prelude::*;
 use crate::validator::names::{ComponentName, ComponentNameKind, KebabStr, KebabString};
 use crate::{
     BinaryReaderError, CanonicalOption, ComponentExportName, ComponentExternalKind,
-    ComponentOuterAliasKind, ComponentTypeRef, CompositeInnerType, CompositeType, ExternalKind,
-    FuncType, GlobalType, InstantiationArgKind, MemoryType, PackedIndex, RefType, Result, SubType,
-    TableType, TypeBounds, ValType, WasmFeatures,
+    ComponentOuterAliasKind, ComponentTypeRef, CompositeInnerType, ExternalKind, FuncType,
+    GlobalType, InstantiationArgKind, MemoryType, PackedIndex, RefType, Result, SubType, TableType,
+    TypeBounds, ValType, WasmFeatures,
 };
 use core::mem;
 
@@ -1102,7 +1102,7 @@ impl ComponentState {
 
     pub fn task_return(
         &mut self,
-        type_index: u32,
+        result: &Option<crate::ComponentValType>,
         types: &mut TypeAlloc,
         offset: usize,
         features: &WasmFeatures,
@@ -1114,20 +1114,32 @@ impl ComponentState {
             )
         }
 
-        let id = self.type_id_at(type_index, offset)?;
-        let Some(SubType {
-            composite_type:
-                CompositeType {
-                    inner: CompositeInnerType::Func(_),
-                    ..
-                },
-            ..
-        }) = types.get(id)
-        else {
-            bail!(offset, "invalid `task.return` type index");
-        };
+        let info = ComponentFuncType {
+            info: TypeInfo::new(),
+            params: result
+                .iter()
+                .map(|ty| {
+                    Ok((
+                        KebabString::new("v").unwrap(),
+                        match ty {
+                            crate::ComponentValType::Primitive(ty) => {
+                                ComponentValType::Primitive(*ty)
+                            }
+                            crate::ComponentValType::Type(index) => {
+                                ComponentValType::Type(self.defined_type_at(*index, offset)?)
+                            }
+                        },
+                    ))
+                })
+                .collect::<Result<_>>()?,
+            results: Box::new([]),
+        }
+        .lower(types, Abi::LiftSync);
 
-        self.core_funcs.push(id);
+        assert!(info.results.iter().next().is_none());
+
+        self.core_funcs
+            .push(types.intern_func_type(FuncType::new(info.params.iter(), []), offset));
         Ok(())
     }
 
diff --git a/crates/wasmprinter/src/component.rs b/crates/wasmprinter/src/component.rs
@@ -914,9 +914,15 @@ impl Printer<'_, '_> {
                 CanonicalFunction::TaskBackpressure => {
                     self.print_intrinsic(state, "canon task.backpressure", &|_, _| Ok(()))?;
                 }
-                CanonicalFunction::TaskReturn { type_index } => {
-                    self.print_intrinsic(state, "canon task.return ", &|me, state| {
-                        me.print_idx(&state.component.type_names, type_index)
+                CanonicalFunction::TaskReturn { result } => {
+                    self.print_intrinsic(state, "canon task.return", &|me, state| {
+                        if let Some(ty) = result {
+                            me.result.write_str(" ")?;
+                            me.start_group("result ")?;
+                            me.print_component_val_type(state, &ty)?;
+                            me.end_group()?;
+                        }
+                        Ok(())
                     })?;
                 }
                 CanonicalFunction::TaskWait { async_, memory } => {
diff --git a/crates/wast/src/component/binary.rs b/crates/wast/src/component/binary.rs
@@ -363,7 +363,11 @@ impl<'a> Encoder<'a> {
             }
             CanonicalFuncKind::TaskReturn(info) => {
                 self.core_func_names.push(name);
-                self.funcs.task_return(info.ty.into());
+                self.funcs.task_return(
+                    info.result
+                        .as_ref()
+                        .map(|ty| wasm_encoder::ComponentValType::from(ty)),
+                );
             }
             CanonicalFuncKind::TaskWait(info) => {
                 self.core_func_names.push(name);
diff --git a/crates/wast/src/component/func.rs b/crates/wast/src/component/func.rs
@@ -521,17 +521,24 @@ impl<'a> Parse<'a> for CanonThreadHwConcurrency {
 /// Information relating to the `task.return` intrinsic.
 #[derive(Debug)]
 pub struct CanonTaskReturn<'a> {
-    /// The core function type representing the signature of this intrinsic.
-    pub ty: Index<'a>,
+    /// The type of the result which may be returned with this intrinsic.
+    pub result: Option<ComponentValType<'a>>,
 }
 
 impl<'a> Parse<'a> for CanonTaskReturn<'a> {
     fn parse(parser: Parser<'a>) -> Result<Self> {
         parser.parse::<kw::task_return>()?;
 
-        Ok(Self {
-            ty: parser.parse()?,
-        })
+        let result = if parser.peek2::<kw::result>()? {
+            Some(parser.parens(|p| {
+                p.parse::<kw::result>()?.0;
+                p.parse()
+            })?)
+        } else {
+            None
+        };
+
+        Ok(Self { result })
     }
 }
 
diff --git a/crates/wast/src/component/resolve.rs b/crates/wast/src/component/resolve.rs
@@ -392,7 +392,9 @@ impl<'a> Resolver<'a> {
             | CanonicalFuncKind::SubtaskDrop
             | CanonicalFuncKind::ErrorContextDrop => {}
             CanonicalFuncKind::TaskReturn(info) => {
-                self.resolve_ns(&mut info.ty, Ns::CoreType)?;
+                if let Some(ty) = &mut info.result {
+                    self.component_val_type(ty)?;
+                }
             }
             CanonicalFuncKind::TaskWait(info) => {
                 self.core_item_ref(&mut info.memory)?;
diff --git a/crates/wit-component/src/encoding.rs b/crates/wit-component/src/encoding.rs
@@ -84,8 +84,8 @@ use wasm_encoder::*;
 use wasmparser::Validator;
 use wit_parser::{
     abi::{AbiVariant, WasmSignature, WasmType},
-    Docs, Function, FunctionKind, InterfaceId, LiveTypes, Resolve, Results, Stability, Type,
-    TypeDefKind, TypeId, TypeOwner, WorldItem, WorldKey,
+    Function, FunctionKind, InterfaceId, LiveTypes, Resolve, Results, Stability, Type, TypeDefKind,
+    TypeId, TypeOwner, WorldItem, WorldKey,
 };
 
 const INDIRECT_TABLE_NAME: &str = "$imports";
@@ -1718,38 +1718,22 @@ impl<'a> EncodingState<'a> {
                     AbiVariant::GuestImport,
                 )
             }
-            Import::ExportedTaskReturn(function) => {
-                let signature = resolve.wasm_signature(
-                    AbiVariant::GuestImport,
-                    &Function {
-                        name: String::new(),
-                        kind: FunctionKind::Freestanding,
-                        params: match &function.results {
-                            Results::Named(params) => params.clone(),
-                            Results::Anon(ty) => vec![("v".to_string(), *ty)],
-                        },
-                        results: Results::Named(Vec::new()),
-                        docs: Docs::default(),
-                        stability: Stability::Unknown,
-                    },
-                );
-                let (type_index, encoder) = self.component.core_type();
-                encoder.core().function(
-                    signature.params.into_iter().map(into_val_type),
-                    signature.results.into_iter().map(into_val_type),
-                );
-
-                let index = self.component.task_return(type_index);
-                return Ok((ExportKind::Func, index));
+            Import::ExportedTaskReturn(interface, function) => {
+                let mut encoder = self.root_export_type_encoder(*interface);
 
-                fn into_val_type(ty: WasmType) -> ValType {
-                    match ty {
-                        WasmType::I32 | WasmType::Pointer | WasmType::Length => ValType::I32,
-                        WasmType::I64 | WasmType::PointerOrI64 => ValType::I64,
-                        WasmType::F32 => ValType::F32,
-                        WasmType::F64 => ValType::F64,
+                let result = match &function.results {
+                    Results::Named(rs) => {
+                        if rs.is_empty() {
+                            None
+                        } else {
+                            bail!("named results not supported for `task.return` intrinsic")
+                        }
                     }
-                }
+                    Results::Anon(ty) => Some(encoder.encode_valtype(resolve, ty)?),
+                };
+
+                let index = self.component.task_return(result);
+                return Ok((ExportKind::Func, index));
             }
             Import::TaskBackpressure => {
                 let index = self.component.task_backpressure();
diff --git a/crates/wit-component/src/validation.rs b/crates/wit-component/src/validation.rs
@@ -243,7 +243,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(Function),
+    ExportedTaskReturn(Option<InterfaceId>, Function),
 
     /// A `canon task.backpressure` intrinsic.
     ///
@@ -551,7 +551,7 @@ impl ImportMap {
                 // it's associated with in general.  Instead, the host will
                 // compare it with the expected type at runtime and trap if
                 // necessary.
-                Some(Import::ExportedTaskReturn(func))
+                Some(Import::ExportedTaskReturn(interface_id, func))
             } else {
                 None
             })
diff --git a/crates/wit-component/tests/components/async-builtins/component.wat b/crates/wit-component/tests/components/async-builtins/component.wat
@@ -104,13 +104,11 @@
     (export "[error-context-debug-message;encoding=utf8;realloc=cabi_realloc]" (func 6))
     (export "[error-context-drop]" (func 7))
   )
-  (core type (;0;) (func (param i32 i32)))
-  (core func (;8;) (canon task.return 0))
+  (core func (;8;) (canon task.return (result string)))
   (core instance (;2;)
     (export "[task-return]foo" (func 8))
   )
-  (core type (;1;) (func (param i32 i32)))
-  (core func (;9;) (canon task.return 1))
+  (core func (;9;) (canon task.return (result string)))
   (core instance (;3;)
     (export "[task-return]foo" (func 9))
   )
diff --git a/tests/local/component-model-async/task-builtins.wast b/tests/local/component-model-async/task-builtins.wast
@@ -24,8 +24,7 @@
   (core module $m
     (import "" "task.return" (func $task-return (param i32)))
   )
-  (core type $task-return-type (func (param i32)))
-  (core func $task-return (canon task.return $task-return-type))
+  (core func $task-return (canon task.return (result u32)))
   (core instance $i (instantiate $m (with "" (instance (export "task.return" (func $task-return))))))
 )
 
diff --git a/tests/local/missing-features/component-model/async.wast b/tests/local/missing-features/component-model/async.wast
@@ -46,8 +46,7 @@
     (core module $m
       (import "" "task.return" (func $task-return (param i32)))
     )
-    (core type $task-return-type (func (param i32)))
-    (core func $task-return (canon task.return $task-return-type))
+    (core func $task-return (canon task.return (result u32)))
     (core instance $i (instantiate $m (with "" (instance (export "task.return" (func $task-return))))))
   )
   "`task.return` requires the component model async feature"
diff --git a/tests/snapshots/local/component-model-async/task-builtins.wast.json b/tests/snapshots/local/component-model-async/task-builtins.wast.json
diff --git a/tests/snapshots/local/component-model-async/task-builtins.wast/2.print b/tests/snapshots/local/component-model-async/task-builtins.wast/2.print
diff --git a/tests/snapshots/local/missing-features/component-model/async.wast.json b/tests/snapshots/local/missing-features/component-model/async.wast.json

Original file line number	Diff line number	Diff line change
`@@ -392,7 +392,7 @@ impl ComponentBuilder {`
`392`	`392`	`}`
`393`	`393`
`394`	`394`	/// Declares a new `task.return` intrinsic.
`395`		`- pub fn task_return(&mut self, ty: u32) -> u32 {`
	`395`	`+ pub fn task_return(&mut self, ty: Option<impl Into<ComponentValType>>) -> u32 {`
`396`	`396`	`self.canonical_functions().task_return(ty);`
`397`	`397`	`inc(&mut self.core_funcs)`
`398`	`398`	`}`
Original file line number	Diff line number	Diff line change
`@@ -970,8 +970,8 @@ pub mod component_utils {`
`970`	`970`	`wasmparser::CanonicalFunction::TaskBackpressure => {`
`971`	`971`	`section.task_backpressure();`
`972`	`972`	`}`
`973`		`- wasmparser::CanonicalFunction::TaskReturn { type_index } => {`
`974`		`- section.task_return(reencoder.type_index(type_index));`
	`973`	`+ wasmparser::CanonicalFunction::TaskReturn { result } => {`
	`974`	`+ section.task_return(result.map(\|ty\| reencoder.component_val_type(ty)));`
`975`	`975`	`}`
`976`	`976`	`wasmparser::CanonicalFunction::TaskWait { async_, memory } => {`
`977`	`977`	`section.task_wait(async_, reencoder.memory_index(memory));`
Original file line number	Diff line number	Diff line change
`@@ -1313,8 +1313,8 @@ impl Validator {`
`1313`	`1313`	`crate::CanonicalFunction::TaskBackpressure => {`
`1314`	`1314`	`current.task_backpressure(types, offset, features)`
`1315`	`1315`	`}`
`1316`		`- crate::CanonicalFunction::TaskReturn { type_index } => {`
`1317`		`- current.task_return(type_index, types, offset, features)`
	`1316`	`+ crate::CanonicalFunction::TaskReturn { result } => {`
	`1317`	`+ current.task_return(&result, types, offset, features)`
`1318`	`1318`	`}`
`1319`	`1319`	`crate::CanonicalFunction::TaskWait { async_, memory } => {`
`1320`	`1320`	`current.task_wait(async_, memory, types, offset, features)`
Original file line number	Diff line number	Diff line change
`@@ -363,7 +363,11 @@ impl<'a> Encoder<'a> {`
`363`	`363`	`}`
`364`	`364`	`CanonicalFuncKind::TaskReturn(info) => {`
`365`	`365`	`self.core_func_names.push(name);`
`366`		`- self.funcs.task_return(info.ty.into());`
	`366`	`+ self.funcs.task_return(`
	`367`	`+ info.result`
	`368`	`+ .as_ref()`
	`369`	`+ .map(\|ty\| wasm_encoder::ComponentValType::from(ty)),`
	`370`	`+ );`
`367`	`371`	`}`
`368`	`372`	`CanonicalFuncKind::TaskWait(info) => {`
`369`	`373`	`self.core_func_names.push(name);`