fix lowering results for async exports (#1129)

Previously, we were either dropping the result too early
(e.g. `wit_bindgen_rt::async_support::ErrorContext`) or not at all
(e.g. `String` and `Vec`) due to the lowering code expecting to pass ownership
to the caller, which would later free memory using a post-return function.  But
async exports don't have post-return functions; they use `task.return` instead,
and they should free any memory (and/or drop handles) after `task.return`
returns.  So now we do that!

Signed-off-by: Joel Dice <joel.dice@fermyon.com>

Author: dicej

crates/core/src/abi.rs

@@ -1474,11 +1474,12 @@ impl<'a, B: Bindgen> Generator<'a, B> {
     }
 
     fn list_realloc(&self) -> Option<&'static str> {
-        // Lowering parameters calling a wasm import means
-        // we don't need to pass ownership, but we pass
-        // ownership in all other cases.
-        match (self.variant, self.lift_lower) {
-            (AbiVariant::GuestImport, LiftLower::LowerArgsLiftResults) => None,
+        // Lowering parameters calling a wasm import _or_ returning a result
+        // from an async-lifted wasm export means we don't need to pass
+        // ownership, but we pass ownership in all other cases.
+        match (self.variant, self.lift_lower, self.async_) {
+            (AbiVariant::GuestImport, LiftLower::LowerArgsLiftResults, _)
+            | (AbiVariant::GuestExport, LiftLower::LiftArgsLowerResults, true) => None,
             _ => Some("cabi_realloc"),
         }
     }

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

@@ -122,12 +122,12 @@ unsafe fn poll(state: *mut FutureState) -> Poll<()> {
 #[doc(hidden)]
 pub fn first_poll<T: 'static>(
     future: impl Future<Output = T> + 'static,
-    fun: impl FnOnce(T) + 'static,
+    fun: impl FnOnce(&T) + 'static,
 ) -> *mut u8 {
     let state = Box::into_raw(Box::new(FutureState {
         todo: 0,
         tasks: Some(
-            [Box::pin(future.map(fun)) as BoxFuture]
+            [Box::pin(future.map(|v| fun(&v))) as BoxFuture]
                 .into_iter()
                 .collect(),
         ),

crates/rust/src/bindgen.rs

@@ -21,6 +21,7 @@ pub(super) struct FunctionBindgen<'a, 'b> {
     pub handle_decls: Vec<String>,
     always_owned: bool,
     pub async_result_name: Option<String>,
+    emitted_cleanup: bool,
 }
 
 impl<'a, 'b> FunctionBindgen<'a, 'b> {
@@ -47,10 +48,15 @@ impl<'a, 'b> FunctionBindgen<'a, 'b> {
             handle_decls: Vec::new(),
             always_owned,
             async_result_name: None,
+            emitted_cleanup: false,
         }
     }
 
     fn emit_cleanup(&mut self) {
+        if self.emitted_cleanup {
+            return;
+        }
+        self.emitted_cleanup = true;
         for (ptr, layout) in mem::take(&mut self.cleanup) {
             let alloc = self.gen.path_to_std_alloc_module();
             self.push_str(&format!(
@@ -1001,10 +1007,11 @@ impl Bindgen for FunctionBindgen<'_, '_> {
                     self.src,
                     "\
                             {func}({});
-                        }});
                     ",
                     operands.join(", ")
                 );
+                self.emit_cleanup();
+                self.src.push_str("});\n");
             }
 
             Instruction::Flush { amt } => {

crates/rust/src/interface.rs

@@ -1088,7 +1088,14 @@ pub mod vtable{ordinal} {{
             async_result_name,
             ..
         } = f;
-        assert!(!needs_cleanup_list);
+        if async_ {
+            if needs_cleanup_list {
+                let vec = self.path_to_vec();
+                uwriteln!(self.src, "let mut cleanup_list = {vec}::new();");
+            }
+        } else {
+            assert!(!needs_cleanup_list);
+        }
         if let Some(name) = async_result_name {
             // When `async_result_name` is `Some(_)`, we wrap the call and any
             // `handle_decls` in a block scope to ensure any resource handles