more wast async support and more tests

This also fixes a bug in `wasmprinter` keeping track of core functions. Signed-off-by: Joel Dice <[email protected]>
bytecodealliance · Nov 15, 2024 · 05d06d1 · 05d06d1
1 parent b91ad7d
commit 05d06d1
Show file tree

Hide file tree

Showing 12 changed files with 246 additions and 10 deletions.
diff --git a/crates/wasmprinter/src/component.rs b/crates/wasmprinter/src/component.rs
@@ -933,7 +933,9 @@ impl Printer<'_, '_> {
                     self.print_idx(&state.component.type_names, type_index)?;
                     self.end_group()?;
                     self.end_group()?;
+                    debug_assert_eq!(state.core.func_to_type.len(), state.core.funcs as usize);
                     state.core.funcs += 1;
+                    state.core.func_to_type.push(None);
                 }
                 CanonicalFunction::TaskWait { async_, memory } => {
                     self.start_group("core func ")?;
@@ -946,7 +948,9 @@ impl Printer<'_, '_> {
                     self.print_idx(&state.component.type_names, memory)?;
                     self.end_group()?;
                     self.end_group()?;
+                    debug_assert_eq!(state.core.func_to_type.len(), state.core.funcs as usize);
                     state.core.funcs += 1;
+                    state.core.func_to_type.push(None);
                 }
                 CanonicalFunction::TaskPoll { async_, memory } => {
                     self.start_group("core func ")?;
@@ -959,7 +963,9 @@ impl Printer<'_, '_> {
                     self.print_idx(&state.component.type_names, memory)?;
                     self.end_group()?;
                     self.end_group()?;
+                    debug_assert_eq!(state.core.func_to_type.len(), state.core.funcs as usize);
                     state.core.funcs += 1;
+                    state.core.func_to_type.push(None);
                 }
                 CanonicalFunction::TaskYield { async_ } => {
                     self.start_group("core func ")?;
@@ -971,7 +977,9 @@ impl Printer<'_, '_> {
                     }
                     self.end_group()?;
                     self.end_group()?;
+                    debug_assert_eq!(state.core.func_to_type.len(), state.core.funcs as usize);
                     state.core.funcs += 1;
+                    state.core.func_to_type.push(None);
                 }
                 CanonicalFunction::SubtaskDrop => {
                     self.start_group("core func ")?;
@@ -980,7 +988,9 @@ impl Printer<'_, '_> {
                     self.start_group("canon subtask.drop")?;
                     self.end_group()?;
                     self.end_group()?;
+                    debug_assert_eq!(state.core.func_to_type.len(), state.core.funcs as usize);
                     state.core.funcs += 1;
+                    state.core.func_to_type.push(None);
                 }
                 CanonicalFunction::StreamNew { ty } => {
                     self.start_group("core func ")?;
@@ -990,7 +1000,9 @@ impl Printer<'_, '_> {
                     self.print_idx(&state.component.type_names, ty)?;
                     self.end_group()?;
                     self.end_group()?;
+                    debug_assert_eq!(state.core.func_to_type.len(), state.core.funcs as usize);
                     state.core.funcs += 1;
+                    state.core.func_to_type.push(None);
                 }
                 CanonicalFunction::StreamRead { ty, options } => {
                     self.start_group("core func ")?;
@@ -1002,7 +1014,9 @@ impl Printer<'_, '_> {
                     self.print_canonical_options(state, &options)?;
                     self.end_group()?;
                     self.end_group()?;
+                    debug_assert_eq!(state.core.func_to_type.len(), state.core.funcs as usize);
                     state.core.funcs += 1;
+                    state.core.func_to_type.push(None);
                 }
                 CanonicalFunction::StreamWrite { ty, options } => {
                     self.start_group("core func ")?;
@@ -1014,7 +1028,9 @@ impl Printer<'_, '_> {
                     self.print_canonical_options(state, &options)?;
                     self.end_group()?;
                     self.end_group()?;
+                    debug_assert_eq!(state.core.func_to_type.len(), state.core.funcs as usize);
                     state.core.funcs += 1;
+                    state.core.func_to_type.push(None);
                 }
                 CanonicalFunction::StreamCancelRead { ty, async_ } => {
                     self.start_group("core func ")?;
@@ -1027,7 +1043,9 @@ impl Printer<'_, '_> {
                     }
                     self.end_group()?;
                     self.end_group()?;
+                    debug_assert_eq!(state.core.func_to_type.len(), state.core.funcs as usize);
                     state.core.funcs += 1;
+                    state.core.func_to_type.push(None);
                 }
                 CanonicalFunction::StreamCancelWrite { ty, async_ } => {
                     self.start_group("core func ")?;
@@ -1040,7 +1058,9 @@ impl Printer<'_, '_> {
                     }
                     self.end_group()?;
                     self.end_group()?;
+                    debug_assert_eq!(state.core.func_to_type.len(), state.core.funcs as usize);
                     state.core.funcs += 1;
+                    state.core.func_to_type.push(None);
                 }
                 CanonicalFunction::StreamCloseReadable { ty } => {
                     self.start_group("core func ")?;
@@ -1050,7 +1070,9 @@ impl Printer<'_, '_> {
                     self.print_idx(&state.component.type_names, ty)?;
                     self.end_group()?;
                     self.end_group()?;
+                    debug_assert_eq!(state.core.func_to_type.len(), state.core.funcs as usize);
                     state.core.funcs += 1;
+                    state.core.func_to_type.push(None);
                 }
                 CanonicalFunction::StreamCloseWritable { ty } => {
                     self.start_group("core func ")?;
@@ -1060,7 +1082,9 @@ impl Printer<'_, '_> {
                     self.print_idx(&state.component.type_names, ty)?;
                     self.end_group()?;
                     self.end_group()?;
+                    debug_assert_eq!(state.core.func_to_type.len(), state.core.funcs as usize);
                     state.core.funcs += 1;
+                    state.core.func_to_type.push(None);
                 }
                 CanonicalFunction::FutureNew { ty } => {
                     self.start_group("core func ")?;
@@ -1070,7 +1094,9 @@ impl Printer<'_, '_> {
                     self.print_idx(&state.component.type_names, ty)?;
                     self.end_group()?;
                     self.end_group()?;
+                    debug_assert_eq!(state.core.func_to_type.len(), state.core.funcs as usize);
                     state.core.funcs += 1;
+                    state.core.func_to_type.push(None);
                 }
                 CanonicalFunction::FutureWrite { ty, options } => {
                     self.start_group("core func ")?;
@@ -1082,7 +1108,9 @@ impl Printer<'_, '_> {
                     self.print_canonical_options(state, &options)?;
                     self.end_group()?;
                     self.end_group()?;
+                    debug_assert_eq!(state.core.func_to_type.len(), state.core.funcs as usize);
                     state.core.funcs += 1;
+                    state.core.func_to_type.push(None);
                 }
                 CanonicalFunction::FutureRead { ty, options } => {
                     self.start_group("core func ")?;
@@ -1094,7 +1122,9 @@ impl Printer<'_, '_> {
                     self.print_canonical_options(state, &options)?;
                     self.end_group()?;
                     self.end_group()?;
+                    debug_assert_eq!(state.core.func_to_type.len(), state.core.funcs as usize);
                     state.core.funcs += 1;
+                    state.core.func_to_type.push(None);
                 }
                 CanonicalFunction::FutureCancelRead { ty, async_ } => {
                     self.start_group("core func ")?;
@@ -1107,7 +1137,9 @@ impl Printer<'_, '_> {
                     }
                     self.end_group()?;
                     self.end_group()?;
+                    debug_assert_eq!(state.core.func_to_type.len(), state.core.funcs as usize);
                     state.core.funcs += 1;
+                    state.core.func_to_type.push(None);
                 }
                 CanonicalFunction::FutureCancelWrite { ty, async_ } => {
                     self.start_group("core func ")?;
@@ -1120,7 +1152,9 @@ impl Printer<'_, '_> {
                     }
                     self.end_group()?;
                     self.end_group()?;
+                    debug_assert_eq!(state.core.func_to_type.len(), state.core.funcs as usize);
                     state.core.funcs += 1;
+                    state.core.func_to_type.push(None);
                 }
                 CanonicalFunction::FutureCloseReadable { ty } => {
                     self.start_group("core func ")?;
@@ -1130,7 +1164,9 @@ impl Printer<'_, '_> {
                     self.print_idx(&state.component.type_names, ty)?;
                     self.end_group()?;
                     self.end_group()?;
+                    debug_assert_eq!(state.core.func_to_type.len(), state.core.funcs as usize);
                     state.core.funcs += 1;
+                    state.core.func_to_type.push(None);
                 }
                 CanonicalFunction::FutureCloseWritable { ty } => {
                     self.start_group("core func ")?;
@@ -1140,7 +1176,9 @@ impl Printer<'_, '_> {
                     self.print_idx(&state.component.type_names, ty)?;
                     self.end_group()?;
                     self.end_group()?;
+                    debug_assert_eq!(state.core.func_to_type.len(), state.core.funcs as usize);
                     state.core.funcs += 1;
+                    state.core.func_to_type.push(None);
                 }
                 CanonicalFunction::ErrorContextNew { options } => {
                     self.start_group("core func ")?;
@@ -1150,7 +1188,9 @@ impl Printer<'_, '_> {
                     self.print_canonical_options(state, &options)?;
                     self.end_group()?;
                     self.end_group()?;
+                    debug_assert_eq!(state.core.func_to_type.len(), state.core.funcs as usize);
                     state.core.funcs += 1;
+                    state.core.func_to_type.push(None);
                 }
                 CanonicalFunction::ErrorContextDebugMessage { options } => {
                     self.start_group("core func ")?;
@@ -1160,7 +1200,9 @@ impl Printer<'_, '_> {
                     self.print_canonical_options(state, &options)?;
                     self.end_group()?;
                     self.end_group()?;
+                    debug_assert_eq!(state.core.func_to_type.len(), state.core.funcs as usize);
                     state.core.funcs += 1;
+                    state.core.func_to_type.push(None);
                 }
                 CanonicalFunction::ErrorContextDrop => {
                     self.start_group("core func ")?;
@@ -1169,7 +1211,9 @@ impl Printer<'_, '_> {
                     self.start_group("canon error-context.drop")?;
                     self.end_group()?;
                     self.end_group()?;
+                    debug_assert_eq!(state.core.func_to_type.len(), state.core.funcs as usize);
                     state.core.funcs += 1;
+                    state.core.func_to_type.push(None);
                 }
             }
         }
@@ -1365,7 +1409,7 @@ impl Printer<'_, '_> {
                         self.start_group("core func ")?;
                         self.print_name(&state.core.func_names, state.core.funcs)?;
                         self.end_group()?;
-                        debug_assert!(state.core.func_to_type.len() == state.core.funcs as usize);
+                        debug_assert_eq!(state.core.func_to_type.len(), state.core.funcs as usize);
                         state.core.funcs += 1;
                         state.core.func_to_type.push(None)
                     }
@@ -1391,7 +1435,7 @@ impl Printer<'_, '_> {
                         self.start_group("core tag ")?;
                         self.print_name(&state.core.tag_names, state.core.tags as u32)?;
                         self.end_group()?;
-                        debug_assert!(state.core.tag_to_type.len() == state.core.tags as usize);
+                        debug_assert_eq!(state.core.tag_to_type.len(), state.core.tags as usize);
                         state.core.tags += 1;
                         state.core.tag_to_type.push(None)
                     }

diff --git a/crates/wast/src/component/binary.rs b/crates/wast/src/component/binary.rs
@@ -355,6 +355,10 @@ impl<'a> Encoder<'a> {
                 self.core_func_names.push(name);
                 self.funcs.thread_hw_concurrency();
             }
+            CanonicalFuncKind::TaskReturn(info) => {
+                self.core_func_names.push(name);
+                self.funcs.task_return(info.ty.into());
+            }
         }
 
         self.flush(Some(self.funcs.id()));

diff --git a/crates/wast/src/component/expand.rs b/crates/wast/src/component/expand.rs
@@ -266,7 +266,8 @@ impl<'a> Expander<'a> {
             | CanonicalFuncKind::ResourceRep(_)
             | CanonicalFuncKind::ResourceDrop(_)
             | CanonicalFuncKind::ThreadSpawn(_)
-            | CanonicalFuncKind::ThreadHwConcurrency(_) => {}
+            | CanonicalFuncKind::ThreadHwConcurrency(_)
+            | CanonicalFuncKind::TaskReturn(_) => {}
         }
     }
 
@@ -322,6 +323,12 @@ impl<'a> Expander<'a> {
                     kind: CanonicalFuncKind::ThreadHwConcurrency(mem::take(info)),
                 }))
             }
+            CoreFuncKind::TaskReturn(info) => Some(ComponentField::CanonicalFunc(CanonicalFunc {
+                span: func.span,
+                id: func.id,
+                name: func.name,
+                kind: CanonicalFuncKind::TaskReturn(mem::take(info)),
+            })),
         }
     }
 

diff --git a/crates/wast/src/component/func.rs b/crates/wast/src/component/func.rs
@@ -53,6 +53,7 @@ pub enum CoreFuncKind<'a> {
     ResourceRep(CanonResourceRep<'a>),
     ThreadSpawn(CanonThreadSpawn<'a>),
     ThreadHwConcurrency(CanonThreadHwConcurrency),
+    TaskReturn(CanonTaskReturn<'a>),
 }
 
 impl<'a> Parse<'a> for CoreFuncKind<'a> {
@@ -79,6 +80,8 @@ impl<'a> Parse<'a> for CoreFuncKind<'a> {
                 Ok(CoreFuncKind::ThreadSpawn(parser.parse()?))
             } else if l.peek::<kw::thread_hw_concurrency>()? {
                 Ok(CoreFuncKind::ThreadHwConcurrency(parser.parse()?))
+            } else if l.peek::<kw::task_return>()? {
+                Ok(CoreFuncKind::TaskReturn(parser.parse()?))
             } else {
                 Err(l.error())
             }
@@ -270,6 +273,8 @@ pub enum CanonicalFuncKind<'a> {
 
     ThreadSpawn(CanonThreadSpawn<'a>),
     ThreadHwConcurrency(CanonThreadHwConcurrency),
+
+    TaskReturn(CanonTaskReturn<'a>),
 }
 
 /// Information relating to lifting a core function.
@@ -460,6 +465,31 @@ impl Default 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>,
+}
+
+impl<'a> Parse<'a> for CanonTaskReturn<'a> {
+    fn parse(parser: Parser<'a>) -> Result<Self> {
+        parser.parse::<kw::task_return>()?;
+
+        Ok(Self {
+            ty: parser.parse()?,
+        })
+    }
+}
+
+impl Default for CanonTaskReturn<'_> {
+    fn default() -> Self {
+        Self {
+            ty: Index::Num(0, Span::from_offset(0)),
+        }
+    }
+}
+
 #[derive(Debug)]
 /// Canonical ABI options.
 pub enum CanonOpt<'a> {

diff --git a/crates/wast/src/component/resolve.rs b/crates/wast/src/component/resolve.rs
@@ -387,6 +387,9 @@ impl<'a> Resolver<'a> {
                 self.resolve_ns(&mut info.ty, Ns::CoreType)?;
             }
             CanonicalFuncKind::ThreadHwConcurrency(_) => {}
+            CanonicalFuncKind::TaskReturn(info) => {
+                self.resolve_ns(&mut info.ty, Ns::CoreType)?;
+            }
         }
 
         Ok(())
@@ -887,7 +890,8 @@ impl<'a> ComponentState<'a> {
                 | CanonicalFuncKind::ResourceRep(_)
                 | CanonicalFuncKind::ResourceDrop(_)
                 | CanonicalFuncKind::ThreadSpawn(_)
-                | CanonicalFuncKind::ThreadHwConcurrency(_) => {
+                | CanonicalFuncKind::ThreadHwConcurrency(_)
+                | CanonicalFuncKind::TaskReturn(_) => {
                     self.core_funcs.register(f.id, "core func")?
                 }
             },

diff --git a/crates/wast/src/lib.rs b/crates/wast/src/lib.rs
@@ -556,6 +556,7 @@ pub mod kw {
     custom_keyword!(thread);
     custom_keyword!(thread_spawn = "thread.spawn");
     custom_keyword!(thread_hw_concurrency = "thread.hw_concurrency");
+    custom_keyword!(task_return = "task.return");
     custom_keyword!(wait);
     custom_keyword!(definition);
     custom_keyword!(r#async = "async");