Refactor how guest exported tasks work (#1268)

* Refactor how guest exported tasks work

This commit refactors some of the code generation and ABI instructions
used for exported tasks. Instead of a future + closure they're
represented now only as a single future which internally handles the
invocation of `task.return`. No major functional change here, just some
refactoring.

* Review comments

Author: alexcrichton

crates/core/src/abi.rs

@@ -1,3 +1,4 @@
+use std::fmt;
 pub use wit_parser::abi::{AbiVariant, WasmSignature, WasmType};
 use wit_parser::{
     align_to_arch, Alignment, ArchitectureSize, ElementInfo, Enum, Flags, FlagsRepr, Function,
@@ -497,11 +498,13 @@ def_instruction! {
         /// Same as `CallWasm`, except the dual where an interface is being
         /// called rather than a raw wasm function.
         ///
-        /// Note that this will be used for async functions.
+        /// Note that this will be used for async functions, and `async_`
+        /// indicates whether the function should be invoked in an async
+        /// fashion.
         CallInterface {
             func: &'a Function,
             async_: bool,
-        } : [func.params.len()] => [if *async_ { 1 } else { usize::from(func.result.is_some()) }],
+        } : [func.params.len()] => [usize::from(func.result.is_some())],
 
         /// Returns `amt` values on the stack. This is always the last
         /// instruction.
@@ -551,16 +554,13 @@ def_instruction! {
             blocks: usize,
         } : [1] => [0],
 
-        /// Generate code to run after `CallInterface` for an async-lifted export.
+        /// Call `task.return` for an async-lifted export.
         ///
-        /// For example, this might include task management for the
-        /// future/promise/task returned by the call made for `CallInterface`.
-        AsyncPostCallInterface { func: &'a Function } : [1] => [usize::from(func.result.is_some()) + 1],
-
-        /// Call `task.return` for an async-lifted export once the task returned
-        /// by `CallInterface` and managed by `AsyncPostCallInterface`
-        /// yields a value.
-        AsyncCallReturn { name: &'a str, params: &'a [WasmType] } : [params.len()] => [0],
+        /// This will call core wasm import `name` which will be mapped to
+        /// `task.return` later on. The function given has `params` as its
+        /// parameters and it will return no results. This is used to pass the
+        /// lowered representation of a function's results to `task.return`.
+        AsyncTaskReturn { name: &'a str, params: &'a [WasmType] } : [params.len()] => [0],
 
         /// Force the evaluation of the specified number of expressions and push
         /// the results to the stack.
@@ -654,7 +654,7 @@ pub trait Bindgen {
     /// The intermediate type for fragments of code for this type.
     ///
     /// For most languages `String` is a suitable intermediate type.
-    type Operand: Clone;
+    type Operand: Clone + fmt::Debug;
 
     /// Emit code to implement the given instruction.
     ///
@@ -1044,14 +1044,13 @@ impl<'a, B: Bindgen> Generator<'a, B> {
                 }
 
                 // ... and that allows us to call the interface types function
-                self.emit(&Instruction::CallInterface {
-                    func,
-                    async_: async_,
-                });
-
-                let (lower_to_memory, async_results) = if async_ {
-                    self.emit(&Instruction::AsyncPostCallInterface { func });
+                self.emit(&Instruction::CallInterface { func, async_ });
 
+                // Asynchronous functions will call `task.return` after the
+                // interface function completes, so lowering is conditional
+                // based on slightly different logic for the `task.return`
+                // intrinsic.
+                let (lower_to_memory, async_flat_results) = if async_ {
                     let mut results = Vec::new();
                     if let Some(ty) = &func.result {
                         self.resolve.push_flat(ty, &mut results);
@@ -1124,18 +1123,15 @@ impl<'a, B: Bindgen> Generator<'a, B> {
                     }
                 }
 
-                if let Some(results) = async_results {
+                if let Some(results) = async_flat_results {
                     let name = &format!("[task-return]{}", func.name);
+                    let params = if lower_to_memory {
+                        &[WasmType::Pointer]
+                    } else {
+                        &results[..]
+                    };
 
-                    self.emit(&Instruction::AsyncCallReturn {
-                        name,
-                        params: &if results.len() > MAX_FLAT_PARAMS {
-                            vec![WasmType::Pointer]
-                        } else {
-                            results
-                        },
-                    });
-                    self.emit(&Instruction::Return { func, amt: 1 });
+                    self.emit(&Instruction::AsyncTaskReturn { name, params });
                 } else {
                     self.emit(&Instruction::Return {
                         func,
@@ -1150,8 +1146,9 @@ impl<'a, B: Bindgen> Generator<'a, B> {
 
         assert!(
             self.stack.is_empty(),
-            "stack has {} items remaining",
-            self.stack.len()
+            "stack has {} items remaining: {:?}",
+            self.stack.len(),
+            self.stack,
         );
     }
 

crates/csharp/src/function.rs

@@ -1256,8 +1256,7 @@ impl Bindgen for FunctionBindgen<'_, '_> {
                 results.extend(operands.iter().take(*amt).map(|v| v.clone()));
             }
 
-            Instruction::AsyncPostCallInterface { .. }
-            | Instruction::AsyncCallReturn { .. }
+            Instruction::AsyncTaskReturn { .. }
             | Instruction::FutureLower { .. }
             | Instruction::FutureLift { .. }
             | Instruction::StreamLower { .. }

crates/guest-rust/rt/src/async_support.rs

@@ -54,7 +54,7 @@ pub use futures;
 type BoxFuture = Pin<Box<dyn Future<Output = ()> + 'static>>;
 
 /// Represents a task created by either a call to an async-lifted export or a
-/// future run using `block_on` or `first_poll`.
+/// future run using `block_on` or `start_task`.
 struct FutureState {
     /// Remaining work to do (if any) before this task can be considered "done".
     ///
@@ -287,22 +287,18 @@ const STATUS_STARTING: u32 = 1;
 const STATUS_STARTED: u32 = 2;
 const STATUS_RETURNED: u32 = 3;
 
-/// Poll the future generated by a call to an async-lifted export once, calling
-/// the specified closure (presumably backed by a call to `task.return`) when it
-/// generates a value.
+/// Starts execution of the `task` provided, an asynchronous computation.
 ///
-/// This will return an appropriate status code to be returned from the
-/// exported function.
+/// This is used for async-lifted exports at their definition site. The
+/// representation of the export is `task` and this function is called from the
+/// entrypoint. The code returned here is the same as the callback associated
+/// with this export, and the callback will be used if this task doesn't exit
+/// immediately with its result.
 #[doc(hidden)]
-pub fn first_poll<T: 'static>(
-    future: impl Future<Output = T> + 'static,
-    fun: impl FnOnce(&T) + 'static,
-) -> u32 {
+pub fn start_task(task: impl Future<Output = ()> + 'static) -> u32 {
     // Allocate a new `FutureState` which will track all state necessary for
     // our exported task.
-    let state = Box::into_raw(Box::new(FutureState::new(Box::pin(
-        future.map(|v| fun(&v)),
-    ))));
+    let state = Box::into_raw(Box::new(FutureState::new(Box::pin(task))));
 
     // Store our `FutureState` into our context-local-storage slot and then
     // pretend we got EVENT_NONE to kick off everything.

crates/moonbit/src/lib.rs

@@ -2660,8 +2660,7 @@ impl Bindgen for FunctionBindgen<'_, '_> {
                 results.extend(operands.iter().take(*amt).map(|v| v.clone()));
             }
 
-            Instruction::AsyncPostCallInterface { .. }
-            | Instruction::AsyncCallReturn { .. }
+            Instruction::AsyncTaskReturn { .. }
             | Instruction::FutureLower { .. }
             | Instruction::FutureLift { .. }
             | Instruction::StreamLower { .. }

crates/rust/src/bindgen.rs

@@ -8,7 +8,6 @@ use wit_bindgen_core::{dealias, uwrite, uwriteln, wit_parser::*, Source};
 pub(super) struct FunctionBindgen<'a, 'b> {
     pub r#gen: &'b mut InterfaceGenerator<'a>,
     params: Vec<String>,
-    async_: bool,
     wasm_import_module: &'b str,
     pub src: Source,
     blocks: Vec<String>,
@@ -19,7 +18,6 @@ pub(super) struct FunctionBindgen<'a, 'b> {
     pub import_return_pointer_area_align: Alignment,
     pub handle_decls: Vec<String>,
     always_owned: bool,
-    pub async_result_name: Option<String>,
 }
 
 pub const POINTER_SIZE_EXPRESSION: &str = "::core::mem::size_of::<*const u8>()";
@@ -28,14 +26,12 @@ impl<'a, 'b> FunctionBindgen<'a, 'b> {
     pub(super) fn new(
         r#gen: &'b mut InterfaceGenerator<'a>,
         params: Vec<String>,
-        async_: bool,
         wasm_import_module: &'b str,
         always_owned: bool,
     ) -> FunctionBindgen<'a, 'b> {
         FunctionBindgen {
             r#gen,
             params,
-            async_,
             wasm_import_module,
             src: Default::default(),
             blocks: Vec::new(),
@@ -46,7 +42,6 @@ impl<'a, 'b> FunctionBindgen<'a, 'b> {
             import_return_pointer_area_align: Default::default(),
             handle_decls: Vec::new(),
             always_owned,
-            async_result_name: None,
         }
     }
 
@@ -859,13 +854,35 @@ impl Bindgen for FunctionBindgen<'_, '_> {
                 self.push_str(");\n");
             }
 
-            Instruction::CallInterface { func, .. } => {
-                if self.async_ {
-                    self.push_str(&format!("let result = "));
-                    results.push("result".into());
-                } else {
-                    self.let_results(usize::from(func.result.is_some()), results);
-                };
+            Instruction::CallInterface { func, async_ } => {
+                let prev_src = mem::replace(self.src.as_mut_string(), String::new());
+
+                self.let_results(usize::from(func.result.is_some()), results);
+                // If this function has a result and it's async, then after
+                // this instruction the result is going to be lowered. This
+                // lowering must happen in terms of `&T`, so force the result
+                // of this expression to have `&` in front.
+                if func.result.is_some() && *async_ {
+                    self.push_str("&");
+                }
+
+                // Force evaluation of all arguments to happen in a sub-block
+                // for this call. This is where `handle_decls` are pushed to
+                // ensure that they're scoped to just this block. Additionally
+                // this is where `prev_src` is pushed, just before the call.
+                //
+                // The purpose of this is to force `borrow<T>` arguments to get
+                // dropped before the call to `task.return` that will happen
+                // after an async call. For normal calls this happens because
+                // the result of this function is returned directly, but for
+                // async calls the return happens as part of this function as
+                // a call to `task.return` via the `AsyncTaskReturn`
+                // instruction.
+                self.push_str("{\n");
+                for decl in self.handle_decls.drain(..) {
+                    self.src.push_str(&decl);
+                }
+                self.push_str(&prev_src);
                 let constructor_type = match &func.kind {
                     FunctionKind::Freestanding | FunctionKind::AsyncFreestanding => {
                         self.push_str(&format!("T::{}", to_rust_ident(&func.name)));
@@ -884,12 +901,7 @@ impl Bindgen for FunctionBindgen<'_, '_> {
                             .as_deref()
                             .unwrap()
                             .to_upper_camel_case();
-                        let call = if self.async_ {
-                            let async_support = self.r#gen.r#gen.async_support_path();
-                            format!("{async_support}::futures::FutureExt::map(T::new")
-                        } else {
-                            format!("{ty}::new(T::new",)
-                        };
+                        let call = format!("{ty}::new(T::new");
                         self.push_str(&call);
                         Some(ty)
                     }
@@ -910,67 +922,19 @@ impl Bindgen for FunctionBindgen<'_, '_> {
                     }
                 }
                 self.push_str(")");
-                if let Some(ty) = constructor_type {
-                    self.push_str(&if self.async_ {
-                        format!(", {ty}::new)")
-                    } else {
-                        ")".into()
-                    });
-                }
-                if self.async_ {
+                if *async_ {
                     self.push_str(".await");
                 }
-                self.push_str(";\n");
-            }
-
-            Instruction::AsyncPostCallInterface { func } => {
-                let result = &operands[0];
-                self.async_result_name = Some(result.clone());
-                results.push("result".into());
-                let params = (0..usize::from(func.result.is_some()))
-                    .map(|_| {
-                        let tmp = self.tmp();
-                        let param = format!("result{}", tmp);
-                        results.push(param.clone());
-                        param
-                    })
-                    .collect::<Vec<_>>()
-                    .join(", ");
-                let params = if func.result.is_none() {
-                    format!("({params})")
-                } else {
-                    params
-                };
-                let async_support = self.r#gen.r#gen.async_support_path();
-                // TODO: This relies on `abi::Generator` emitting
-                // `AsyncCallReturn` immediately after this instruction to
-                // complete the incomplete expression we generate here.  We
-                // should refactor this so it's less fragile (e.g. have
-                // `abi::Generator` emit a `AsyncCallReturn` first, which would
-                // push a closure expression we can consume here).
-                //
-                // Also, see `InterfaceGenerator::generate_guest_export` for
-                // where we emit the open bracket corresponding to the `};` we
-                // emit here.
-                //
-                // The async-specific `Instruction`s will probably need to be
-                // refactored anyway once we start implementing support for
-                // other languages besides Rust.
-                uwriteln!(
-                    self.src,
-                    "\
-                            {result}
-                        }};
-                        let result = {async_support}::first_poll({result}, move |{params}| {{
-                    "
-                );
+                if constructor_type.is_some() {
+                    self.push_str(")");
+                }
+                self.push_str("\n};\n");
             }
 
-            Instruction::AsyncCallReturn { name, params } => {
+            Instruction::AsyncTaskReturn { name, params } => {
                 let func = self.declare_import(name, params, &[]);
 
                 uwriteln!(self.src, "{func}({});", operands.join(", "));
-                self.src.push_str("});\n");
             }
 
             Instruction::Flush { amt } => {