drop borrows earlier for async exports in Rust generator (#1124)

While testing runtime borrow enforcement in Wasmtime, I discovered that the
guest code generated by `wit-bindgen` was not dropping borrows early enough for
async exports.  Specifically, all borrows need to be dropped before
`task.return` is called, which is what this patch does.

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

Author: dicej

crates/rust/src/bindgen.rs

@@ -20,6 +20,7 @@ pub(super) struct FunctionBindgen<'a, 'b> {
     pub import_return_pointer_area_align: usize,
     pub handle_decls: Vec<String>,
     always_owned: bool,
+    pub async_result_name: Option<String>,
 }
 
 impl<'a, 'b> FunctionBindgen<'a, 'b> {
@@ -45,6 +46,7 @@ impl<'a, 'b> FunctionBindgen<'a, 'b> {
             import_return_pointer_area_align: 0,
             handle_decls: Vec::new(),
             always_owned,
+            async_result_name: None,
         }
     }
 
@@ -877,10 +879,8 @@ impl Bindgen for FunctionBindgen<'_, '_> {
 
             Instruction::CallInterface { func, .. } => {
                 if self.async_ {
-                    let tmp = self.tmp();
-                    let result = format!("result{tmp}");
-                    self.push_str(&format!("let {result} = "));
-                    results.push(result);
+                    self.push_str(&format!("let result = "));
+                    results.push("result".into());
                 } else {
                     self.let_results(func.results.len(), results);
                 };
@@ -950,6 +950,7 @@ impl Bindgen for FunctionBindgen<'_, '_> {
 
             Instruction::AsyncPostCallInterface { func } => {
                 let result = &operands[0];
+                self.async_result_name = Some(result.clone());
                 results.push("result".into());
                 let params = (0..func.results.len())
                     .map(|_| {
@@ -973,12 +974,18 @@ impl Bindgen for FunctionBindgen<'_, '_> {
                 // `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}, |{params}| {{
                     "
                 );

crates/rust/src/interface.rs

@@ -1085,9 +1085,22 @@ pub mod vtable{ordinal} {{
             needs_cleanup_list,
             src,
             handle_decls,
+            async_result_name,
             ..
         } = f;
         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
+            // are dropped before we call `async_support::first_poll`, which may
+            // call `task.return` at which point any borrow handles must already
+            // have been dropped.
+            //
+            // The corresponding close bracket is emitted in the
+            // `Instruction::AsyncPostCallInterface` case inside
+            // `FunctionBindgen::emit`.
+            uwriteln!(self.src, "let {name} = {{");
+        }
         for decl in handle_decls {
             self.src.push_str(&decl);
             self.src.push_str("\n");