add component-model-async/task-builtins.wast test

This is another piece of bytecodealliance#9582 which I'm splitting out to make review easier.

This test includes components which use the `task.backpressure`, `task.return`,
`task.wait`, `task.poll`, `task.yield`, and `subtask.drop` builtins.

The rest of the changes fill in some TODOs to make the test pass.

Signed-off-by: Joel Dice <[email protected]>

Author: dicej

crates/cranelift/src/compiler/component.rs

@@ -95,34 +95,25 @@ impl<'a> TrampolineCompiler<'a> {
                 self.translate_always_trap();
             }
             Trampoline::TaskBackpressure { instance } => {
-                _ = instance;
-                todo!()
+                self.translate_task_backpressure_call(*instance)
             }
             Trampoline::TaskReturn { results } => self.translate_task_return_call(*results),
             Trampoline::TaskWait {
                 instance,
                 async_,
                 memory,
             } => {
-                _ = (instance, async_, memory);
-                todo!()
+                self.translate_task_wait_or_poll_call(*instance, *async_, *memory, host::task_wait)
             }
             Trampoline::TaskPoll {
                 instance,
                 async_,
                 memory,
             } => {
-                _ = (instance, async_, memory);
-                todo!()
-            }
-            Trampoline::TaskYield { async_ } => {
-                _ = async_;
-                todo!()
-            }
-            Trampoline::SubtaskDrop { instance } => {
-                _ = instance;
-                todo!()
+                self.translate_task_wait_or_poll_call(*instance, *async_, *memory, host::task_poll)
             }
+            Trampoline::TaskYield { async_ } => self.translate_task_yield_call(*async_),
+            Trampoline::SubtaskDrop { instance } => self.translate_subtask_drop_call(*instance),
             Trampoline::StreamNew { ty } => {
                 _ = ty;
                 todo!()
@@ -261,7 +252,16 @@ impl<'a> TrampolineCompiler<'a> {
         }
     }
 
-    fn translate_task_return_call(&mut self, results: TypeTupleIndex) {
+    fn translate_intrinsic_libcall(
+        &mut self,
+        vmctx: ir::Value,
+        get_libcall: fn(
+            &dyn TargetIsa,
+            &mut ir::Function,
+        ) -> (ir::SigRef, ComponentBuiltinFunctionIndex),
+        args: &[ir::Value],
+        result: ir::types::Type,
+    ) {
         match self.abi {
             Abi::Wasm => {}
 
@@ -274,29 +274,39 @@ impl<'a> TrampolineCompiler<'a> {
             }
         }
 
+        let call = self.call_libcall(vmctx, get_libcall, args);
+
+        if result == ir::types::I64 {
+            let result = self.builder.func.dfg.inst_results(call)[0];
+            let result = self.raise_if_negative_one(result);
+            self.abi_store_results(&[result]);
+        } else {
+            if result != ir::types::I8 {
+                todo!("support additional intrinsic return types")
+            }
+            let succeeded = self.builder.func.dfg.inst_results(call)[0];
+            self.raise_if_host_trapped(succeeded);
+            self.builder.ins().return_(&[]);
+        }
+    }
+
+    fn translate_task_return_call(&mut self, results: TypeTupleIndex) {
         let args = self.builder.func.dfg.block_params(self.block0).to_vec();
         let vmctx = args[0];
 
         let (values_vec_ptr, values_vec_len) = self.store_wasm_arguments(&args[2..]);
 
-        let (host_sig, index) = host::task_return(self.isa, &mut self.builder.func);
-        let host_fn = self.load_libcall(vmctx, index);
-
         let ty = self
             .builder
             .ins()
             .iconst(ir::types::I32, i64::from(results.as_u32()));
 
-        let call = self.compiler.call_indirect_host(
-            &mut self.builder,
-            index,
-            host_sig,
-            host_fn,
+        self.translate_intrinsic_libcall(
+            vmctx,
+            host::task_return,
             &[vmctx, ty, values_vec_ptr, values_vec_len],
+            ir::types::I8,
         );
-        let succeeded = self.builder.func.dfg.inst_results(call)[0];
-        self.raise_if_host_trapped(succeeded);
-        self.builder.ins().return_(&[]);
     }
 
     fn translate_async_enter_or_exit(
@@ -311,23 +321,9 @@ impl<'a> TrampolineCompiler<'a> {
         )>,
         result: ir::types::Type,
     ) {
-        match self.abi {
-            Abi::Wasm => {}
-
-            // These trampolines can only actually be called by Wasm, so
-            // let's assert that here.
-            Abi::Array => {
-                self.builder.ins().trap(TRAP_INTERNAL_ASSERT);
-                return;
-            }
-        }
-
         let args = self.builder.func.dfg.block_params(self.block0).to_vec();
         let vmctx = args[0];
 
-        let (host_sig, index) = get_libcall(self.isa, &mut self.builder.func);
-        let host_fn = self.load_libcall(vmctx, index);
-
         let mut callee_args = vec![vmctx];
 
         if let Some((callback, post_return)) = callback_and_post_return {
@@ -361,24 +357,93 @@ impl<'a> TrampolineCompiler<'a> {
         // remaining parameters
         callee_args.extend(args[2..].iter().copied());
 
-        let call = self.compiler.call_indirect_host(
-            &mut self.builder,
-            index,
-            host_sig,
-            host_fn,
+        self.translate_intrinsic_libcall(vmctx, get_libcall, &callee_args, result);
+    }
+
+    fn translate_task_backpressure_call(&mut self, caller_instance: RuntimeComponentInstanceIndex) {
+        let args = self.builder.func.dfg.block_params(self.block0).to_vec();
+        let vmctx = args[0];
+
+        let mut callee_args = vec![
+            vmctx,
+            self.builder
+                .ins()
+                .iconst(ir::types::I32, i64::from(caller_instance.as_u32())),
+        ];
+
+        callee_args.extend(args[2..].iter().copied());
+
+        self.translate_intrinsic_libcall(
+            vmctx,
+            host::task_backpressure,
             &callee_args,
+            ir::types::I8,
         );
+    }
 
-        if result == ir::types::I64 {
-            let result = self.builder.func.dfg.inst_results(call)[0];
-            let result = self.raise_if_negative_one(result);
-            self.abi_store_results(&[result]);
-        } else {
-            assert!(result == ir::types::I8);
-            let succeeded = self.builder.func.dfg.inst_results(call)[0];
-            self.raise_if_host_trapped(succeeded);
-            self.builder.ins().return_(&[]);
-        }
+    fn translate_task_wait_or_poll_call(
+        &mut self,
+        caller_instance: RuntimeComponentInstanceIndex,
+        async_: bool,
+        memory: RuntimeMemoryIndex,
+        get_libcall: fn(
+            &dyn TargetIsa,
+            &mut ir::Function,
+        ) -> (ir::SigRef, ComponentBuiltinFunctionIndex),
+    ) {
+        let pointer_type = self.isa.pointer_type();
+        let args = self.builder.func.dfg.block_params(self.block0).to_vec();
+        let vmctx = args[0];
+
+        let mut callee_args = vec![
+            vmctx,
+            self.builder
+                .ins()
+                .iconst(ir::types::I32, i64::from(caller_instance.as_u32())),
+            self.builder
+                .ins()
+                .iconst(ir::types::I8, if async_ { 1 } else { 0 }),
+            self.builder.ins().load(
+                pointer_type,
+                MemFlags::trusted(),
+                vmctx,
+                i32::try_from(self.offsets.runtime_memory(memory)).unwrap(),
+            ),
+        ];
+
+        callee_args.extend(args[2..].iter().copied());
+
+        self.translate_intrinsic_libcall(vmctx, get_libcall, &callee_args, ir::types::I64);
+    }
+
+    fn translate_task_yield_call(&mut self, async_: bool) {
+        let args = self.builder.func.dfg.block_params(self.block0).to_vec();
+        let vmctx = args[0];
+
+        let callee_args = [
+            vmctx,
+            self.builder
+                .ins()
+                .iconst(ir::types::I8, if async_ { 1 } else { 0 }),
+        ];
+
+        self.translate_intrinsic_libcall(vmctx, host::task_yield, &callee_args, ir::types::I8);
+    }
+
+    fn translate_subtask_drop_call(&mut self, caller_instance: RuntimeComponentInstanceIndex) {
+        let args = self.builder.func.dfg.block_params(self.block0).to_vec();
+        let vmctx = args[0];
+
+        let mut callee_args = vec![
+            vmctx,
+            self.builder
+                .ins()
+                .iconst(ir::types::I32, i64::from(caller_instance.as_u32())),
+        ];
+
+        callee_args.extend(args[2..].iter().copied());
+
+        self.translate_intrinsic_libcall(vmctx, host::subtask_drop, &callee_args, ir::types::I8);
     }
 
     fn translate_lower_import(

crates/environ/src/component.rs

@@ -83,9 +83,19 @@ macro_rules! foreach_builtin_component_function {
             resource_enter_call(vmctx: vmctx);
             resource_exit_call(vmctx: vmctx) -> bool;
 
+            #[cfg(feature = "component-model-async")]
+            task_backpressure(vmctx: vmctx, caller_instance: u32, enabled: u32) -> bool;
             #[cfg(feature = "component-model-async")]
             task_return(vmctx: vmctx, ty: u32, storage: ptr_u8, storage_len: size) -> bool;
             #[cfg(feature = "component-model-async")]
+            task_wait(vmctx: vmctx, caller_instance: u32, async_: u8, memory: ptr_u8, payload: u32) -> u64;
+            #[cfg(feature = "component-model-async")]
+            task_poll(vmctx: vmctx, caller_instance: u32, async_: u8, memory: ptr_u8, payload: u32) -> u64;
+            #[cfg(feature = "component-model-async")]
+            task_yield(vmctx: vmctx, async_: u8) -> bool;
+            #[cfg(feature = "component-model-async")]
+            subtask_drop(vmctx: vmctx, caller_instance: u32, task_id: u32) -> bool;
+            #[cfg(feature = "component-model-async")]
             async_enter(vmctx: vmctx, start: ptr_u8, return_: ptr_u8, caller_instance: u32, task_return_type: u32, params: u32, results: u32) -> bool;
             #[cfg(feature = "component-model-async")]
             async_exit(vmctx: vmctx, callback: ptr_u8, post_return: ptr_u8, caller_instance: u32, callee: ptr_u8, callee_instance: u32, param_count: u32, result_count: u32, flags: u32) -> u64;

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

@@ -1,5 +1,9 @@
 use {
-    crate::{store::StoreInner, vm::VMFuncRef, AsContextMut, ValRaw},
+    crate::{
+        store::StoreInner,
+        vm::{VMFuncRef, VMMemoryDefinition},
+        AsContextMut, ValRaw,
+    },
     anyhow::Result,
     futures::{stream::FuturesUnordered, FutureExt},
     std::{boxed::Box, future::Future, pin::Pin},
@@ -77,6 +81,13 @@ impl<T: 'static> PromisesUnordered<T> {
 /// Trait representing component model ABI async intrinsics and fused adapter
 /// helper functions.
 pub unsafe trait VMComponentAsyncStore {
+    /// The `task.backpressure` intrinsic.
+    fn task_backpressure(
+        &mut self,
+        caller_instance: RuntimeComponentInstanceIndex,
+        enabled: u32,
+    ) -> Result<()>;
+
     /// The `task.return` intrinsic.
     fn task_return(
         &mut self,
@@ -85,6 +96,34 @@ pub unsafe trait VMComponentAsyncStore {
         storage_len: usize,
     ) -> Result<()>;
 
+    /// The `task.wait` intrinsic.
+    fn task_wait(
+        &mut self,
+        caller_instance: RuntimeComponentInstanceIndex,
+        async_: bool,
+        memory: *mut VMMemoryDefinition,
+        payload: u32,
+    ) -> Result<u32>;
+
+    /// The `task.poll` intrinsic.
+    fn task_poll(
+        &mut self,
+        caller_instance: RuntimeComponentInstanceIndex,
+        async_: bool,
+        memory: *mut VMMemoryDefinition,
+        payload: u32,
+    ) -> Result<u32>;
+
+    /// The `task.yield` intrinsic.
+    fn task_yield(&mut self, async_: bool) -> Result<()>;
+
+    /// The `subtask.drop` intrinsic.
+    fn subtask_drop(
+        &mut self,
+        caller_instance: RuntimeComponentInstanceIndex,
+        task_id: u32,
+    ) -> Result<()>;
+
     /// A helper function for fused adapter modules involving calls where one or
     /// both of the functions involved are async functions.
     fn async_enter(
@@ -113,6 +152,15 @@ pub unsafe trait VMComponentAsyncStore {
 }
 
 unsafe impl<T> VMComponentAsyncStore for StoreInner<T> {
+    fn task_backpressure(
+        &mut self,
+        caller_instance: RuntimeComponentInstanceIndex,
+        enabled: u32,
+    ) -> Result<()> {
+        _ = (caller_instance, enabled);
+        todo!()
+    }
+
     fn task_return(
         &mut self,
         ty: TypeTupleIndex,
@@ -123,6 +171,42 @@ unsafe impl<T> VMComponentAsyncStore for StoreInner<T> {
         todo!()
     }
 
+    fn task_wait(
+        &mut self,
+        caller_instance: RuntimeComponentInstanceIndex,
+        async_: bool,
+        memory: *mut VMMemoryDefinition,
+        payload: u32,
+    ) -> Result<u32> {
+        _ = (caller_instance, async_, memory, payload);
+        todo!()
+    }
+
+    fn task_poll(
+        &mut self,
+        caller_instance: RuntimeComponentInstanceIndex,
+        async_: bool,
+        memory: *mut VMMemoryDefinition,
+        payload: u32,
+    ) -> Result<u32> {
+        _ = (caller_instance, async_, memory, payload);
+        todo!()
+    }
+
+    fn task_yield(&mut self, async_: bool) -> Result<()> {
+        _ = async_;
+        todo!()
+    }
+
+    fn subtask_drop(
+        &mut self,
+        caller_instance: RuntimeComponentInstanceIndex,
+        task_id: u32,
+    ) -> Result<()> {
+        _ = (caller_instance, task_id);
+        todo!()
+    }
+
     fn async_enter(
         &mut self,
         start: *mut VMFuncRef,