update task.return to match spec

dicej · dicej · commit 42b2673cbf89 · 2025-01-28T11:38:48.000-07:00
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>
diff --git a/crates/wasm-encoder/src/component/builder.rs b/crates/wasm-encoder/src/component/builder.rs
@@ -391,9 +391,20 @@ impl ComponentBuilder {
         inc(&mut self.core_funcs)
     }
 
-    /// Declares a new `task.return` intrinsic.
-    pub fn task_return(&mut self, ty: u32) -> u32 {
-        self.canonical_functions().task_return(ty);
+    /// Declares a new `task.return` intrinsic using named results.
+    pub fn task_return_named<'a, R, T>(&mut self, results: R) -> u32
+    where
+        R: IntoIterator<Item = (&'a str, T)>,
+        R::IntoIter: ExactSizeIterator,
+        T: Into<ComponentValType>,
+    {
+        self.canonical_functions().task_return_named(results);
+        inc(&mut self.core_funcs)
+    }
+
+    /// Declares a new `task.return` intrinsic using an anonymous result.
+    pub fn task_return_anon(&mut self, ty: impl Into<ComponentValType>) -> u32 {
+        self.canonical_functions().task_return_anon(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,39 @@ 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 {
+    ///
+    /// This version accepts a list of named result types.  See
+    /// [Self::task_return_anon] for a version that accepts a single, anonymous result
+    /// type.
+    pub fn task_return_named<'a, R, T>(&mut self, results: R) -> &mut Self
+    where
+        R: IntoIterator<Item = (&'a str, T)>,
+        R::IntoIter: ExactSizeIterator,
+        T: Into<ComponentValType>,
+    {
         self.bytes.push(0x09);
-        ty.encode(&mut self.bytes);
+        self.bytes.push(0x01);
+        let results = results.into_iter();
+        results.len().encode(&mut self.bytes);
+        for (name, ty) in results {
+            name.encode(&mut self.bytes);
+            ty.into().encode(&mut self.bytes);
+        }
+        self.num_added += 1;
+        self
+    }
+
+    /// 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.
+    ///
+    /// This version accepts a single, anonymous result type.  See
+    /// [Self::task_return_named] for a version that accepts a list of named result
+    /// types.
+    pub fn task_return_anon(&mut self, ty: impl Into<ComponentValType>) -> &mut Self {
+        self.bytes.push(0x09);
+        self.bytes.push(0x00);
+        ty.into().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,17 @@ 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 { results } => {
+                match results {
+                    wasmparser::ComponentFuncResult::Unnamed(ty) => {
+                        section.task_return_anon(reencoder.component_val_type(ty))
+                    }
+                    wasmparser::ComponentFuncResult::Named(results) => section.task_return_named(
+                        results
+                            .iter()
+                            .map(|(name, ty)| (*name, 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,6 @@
 use crate::limits::MAX_WASM_CANONICAL_OPTIONS;
 use crate::prelude::*;
-use crate::{BinaryReader, FromReader, Result, SectionLimited};
+use crate::{BinaryReader, ComponentFuncResult, FromReader, Result, SectionLimited};
 
 /// Represents options for component functions.
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -32,7 +32,7 @@ pub enum CanonicalOption {
 
 /// Represents a canonical function in a WebAssembly component.
 #[derive(Debug, Clone, Eq, PartialEq)]
-pub enum CanonicalFunction {
+pub enum CanonicalFunction<'a> {
     /// The function lifts a core WebAssembly function to the canonical ABI.
     Lift {
         /// The index of the core WebAssembly function to lift.
@@ -80,10 +80,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(s).
+        results: ComponentFuncResult<'a>,
     },
     /// A function which waits for at least one outstanding async
     /// task/stream/future to make progress, returning the first such event.
@@ -231,10 +229,10 @@ pub enum CanonicalFunction {
 }
 
 /// A reader for the canonical section of a WebAssembly component.
-pub type ComponentCanonicalSectionReader<'a> = SectionLimited<'a, CanonicalFunction>;
+pub type ComponentCanonicalSectionReader<'a> = SectionLimited<'a, CanonicalFunction<'a>>;
 
-impl<'a> FromReader<'a> for CanonicalFunction {
-    fn from_reader(reader: &mut BinaryReader<'a>) -> Result<CanonicalFunction> {
+impl<'a> FromReader<'a> for CanonicalFunction<'a> {
+    fn from_reader(reader: &mut BinaryReader<'a>) -> Result<CanonicalFunction<'a>> {
         Ok(match reader.read_u8()? {
             0x00 => match reader.read_u8()? {
                 0x00 => {
@@ -275,7 +273,7 @@ impl<'a> FromReader<'a> for CanonicalFunction {
             0x06 => CanonicalFunction::ThreadHwConcurrency,
             0x08 => CanonicalFunction::TaskBackpressure,
             0x09 => CanonicalFunction::TaskReturn {
-                type_index: reader.read()?,
+                results: reader.read()?,
             },
             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 { results } => {
+                        current.task_return(&results, 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,
+    ComponentFuncResult, 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: &ComponentFuncResult,
         types: &mut TypeAlloc,
         offset: usize,
         features: &WasmFeatures,
@@ -1114,20 +1114,31 @@ 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(|(name, ty)| {
+                    let ty = match ty {
+                        crate::ComponentValType::Primitive(ty) => ComponentValType::Primitive(*ty),
+                        crate::ComponentValType::Type(index) => {
+                            ComponentValType::Type(self.defined_type_at(*index, offset)?)
+                        }
+                    };
 
-        self.core_funcs.push(id);
+                    let Some(name) = KebabString::new(name.unwrap_or("v")) else {
+                        bail!(offset, "non-kebab name found in `task.return` result list")
+                    };
+
+                    Ok((name, ty))
+                })
+                .collect::<Result<_>>()?,
+            results: Box::new([]),
+        }
+        .lower(types, Abi::LiftSync);
+
+        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,19 @@ 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 { results } => {
+                    self.print_intrinsic(state, "canon task.return", &|me, state| {
+                        for (name, ty) in results.iter() {
+                            me.result.write_str(" ")?;
+                            me.start_group("result ")?;
+                            if let Some(name) = name {
+                                me.print_str(name)?;
+                                me.result.write_str(" ")?;
+                            }
+                            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,30 @@ impl<'a> Encoder<'a> {
             }
             CanonicalFuncKind::TaskReturn(info) => {
                 self.core_func_names.push(name);
-                self.funcs.task_return(info.ty.into());
+                match &info.results[..] {
+                    [ComponentFunctionResult { name: None, ty }] => {
+                        self.funcs.task_return_anon(ty);
+                    }
+                    results => {
+                        // Note that we synthesize names for any unnamed results
+                        // we encounter.
+                        let results = results
+                            .iter()
+                            .enumerate()
+                            .map(|(index, result)| {
+                                (
+                                    result
+                                        .name
+                                        .map(|s| s.to_string())
+                                        .unwrap_or_else(|| format!("v{index}")),
+                                    &result.ty,
+                                )
+                            })
+                            .collect::<Vec<_>>();
+                        self.funcs
+                            .task_return_named(results.iter().map(|(n, v)| (n.as_str(), *v)));
+                    }
+                }
             }
             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,16 +521,22 @@ 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 types (and optionally, names) of the results which may be returned
+    /// with this intrinsic.
+    pub results: Box<[ComponentFunctionResult<'a>]>,
 }
 
 impl<'a> Parse<'a> for CanonTaskReturn<'a> {
     fn parse(parser: Parser<'a>) -> Result<Self> {
         parser.parse::<kw::task_return>()?;
 
+        let mut results: Vec<ComponentFunctionResult> = Vec::new();
+        while parser.peek2::<kw::result>()? {
+            results.push(parser.parens(|p| p.parse())?);
+        }
+
         Ok(Self {
-            ty: parser.parse()?,
+            results: results.into(),
         })
     }
 }
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)?;
+                for result in info.results.iter_mut() {
+                    self.component_val_type(&mut result.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,26 @@ 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),
-                );
+            Import::ExportedTaskReturn(interface, function) => {
+                let mut encoder = self.root_export_type_encoder(*interface);
 
-                let index = self.component.task_return(type_index);
-                return Ok((ExportKind::Func, index));
+                enum EncodedResults<'a> {
+                    Named(Vec<(&'a str, ComponentValType)>),
+                    Anon(ComponentValType),
+                }
 
-                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 results = match &function.results {
+                    Results::Named(rs) => {
+                        EncodedResults::Named(encoder.encode_params(resolve, rs)?)
                     }
-                }
+                    Results::Anon(ty) => EncodedResults::Anon(encoder.encode_valtype(resolve, ty)?),
+                };
+
+                let index = match results {
+                    EncodedResults::Named(rs) => self.component.task_return_named(rs),
+                    EncodedResults::Anon(ty) => self.component.task_return_anon(ty),
+                };
+                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
diff --git a/tests/local/component-model-async/task-builtins.wast b/tests/local/component-model-async/task-builtins.wast
diff --git a/tests/local/missing-features/component-model/async.wast b/tests/local/missing-features/component-model/async.wast
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
`@@ -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 { results } => {`
	`1317`	`+ current.task_return(&results, types, offset, features)`
`1318`	`1318`	`}`
`1319`	`1319`	`crate::CanonicalFunction::TaskWait { async_, memory } => {`
`1320`	`1320`	`current.task_wait(async_, memory, types, offset, features)`