Store arguments for async params/results on the stack (#1185)

After #1176 there's no longer any need to store these values on the
heap, so store them on the stack instead. This also updates to store
params/results in an overlapping allocation which should be reasonable
as it's local per-import and results are never written before parameters
are read. (and params are never read after results are written).

Author: alexcrichton

crates/core/src/abi.rs

@@ -551,19 +551,8 @@ def_instruction! {
             blocks: usize,
         } : [1] => [0],
 
-        /// Allocate the parameter and/or return areas to use for an
-        /// async-lowered import call.
-        ///
-        /// This cannot be allocated on the (shadow-)stack since it needs to
-        /// remain valid until the callee has finished using the buffers, which
-        /// may be after we pop the current stack frame.
-        AsyncMalloc { size: usize, align: usize } : [0] => [1],
-
         /// Call an async-lowered import.
-        ///
-        /// `size` and `align` are used to deallocate the parameter area
-        /// allocated using `AsyncMalloc` after the callee task returns a value.
-        AsyncCallWasm { name: &'a str, size: usize, align: usize } : [2] => [0],
+        AsyncCallWasm { name: &'a str } : [2] => [0],
 
         /// Generate code to run after `CallInterface` for an async-lifted export.
         ///
@@ -913,18 +902,15 @@ impl<'a, B: Bindgen> Generator<'a, B> {
                     self_.stack.push(ptr);
                 };
 
-                let params_size_align = if self.async_ {
+                if self.async_ {
                     let ElementInfo { size, align } = self
                         .bindgen
                         .sizes()
                         .record(func.params.iter().map(|(_, ty)| ty));
-                    self.emit(&Instruction::AsyncMalloc {
-                        size: size.size_wasm32(),
-                        align: align.align_wasm32(),
-                    });
-                    let ptr = self.stack.pop().unwrap();
+                    let ptr = self
+                        .bindgen
+                        .return_pointer(size.size_wasm32(), align.align_wasm32());
                     lower_to_memory(self, ptr);
-                    Some((size, align))
                 } else {
                     if !sig.indirect_params {
                         // If the parameters for this function aren't indirect
@@ -966,47 +952,39 @@ impl<'a, B: Bindgen> Generator<'a, B> {
                         };
                         lower_to_memory(self, ptr);
                     }
-                    None
-                };
+                }
 
-                // If necessary we may need to prepare a return pointer for
-                // this ABI.
-                let dealloc_size_align =
-                    if let Some((params_size, params_align)) = params_size_align {
-                        let ElementInfo { size, align } =
-                            self.bindgen.sizes().record(func.result.iter());
-                        self.emit(&Instruction::AsyncMalloc {
-                            size: size.size_wasm32(),
-                            align: align.align_wasm32(),
-                        });
-                        let ptr = self.stack.pop().unwrap();
+                if self.async_ {
+                    let ElementInfo { size, align } =
+                        self.bindgen.sizes().record(func.result.iter());
+                    let ptr = self
+                        .bindgen
+                        .return_pointer(size.size_wasm32(), align.align_wasm32());
+                    self.return_pointer = Some(ptr.clone());
+                    self.stack.push(ptr);
+
+                    assert_eq!(self.stack.len(), 2);
+                    self.emit(&Instruction::AsyncCallWasm {
+                        name: &format!("[async-lower]{}", func.name),
+                    });
+                } else {
+                    // If necessary we may need to prepare a return pointer for
+                    // this ABI.
+                    if self.variant == AbiVariant::GuestImport && sig.retptr {
+                        let info = self.bindgen.sizes().params(&func.result);
+                        let ptr = self
+                            .bindgen
+                            .return_pointer(info.size.size_wasm32(), info.align.align_wasm32());
                         self.return_pointer = Some(ptr.clone());
                         self.stack.push(ptr);
+                    }
 
-                        assert_eq!(self.stack.len(), 2);
-                        self.emit(&Instruction::AsyncCallWasm {
-                            name: &format!("[async-lower]{}", func.name),
-                            size: params_size.size_wasm32(),
-                            align: params_align.align_wasm32(),
-                        });
-                        Some((size, align))
-                    } else {
-                        if self.variant == AbiVariant::GuestImport && sig.retptr {
-                            let info = self.bindgen.sizes().params(&func.result);
-                            let ptr = self
-                                .bindgen
-                                .return_pointer(info.size.size_wasm32(), info.align.align_wasm32());
-                            self.return_pointer = Some(ptr.clone());
-                            self.stack.push(ptr);
-                        }
-
-                        assert_eq!(self.stack.len(), sig.params.len());
-                        self.emit(&Instruction::CallWasm {
-                            name: &func.name,
-                            sig: &sig,
-                        });
-                        None
-                    };
+                    assert_eq!(self.stack.len(), sig.params.len());
+                    self.emit(&Instruction::CallWasm {
+                        name: &func.name,
+                        sig: &sig,
+                    });
+                }
 
                 if !(sig.retptr || self.async_) {
                     // With no return pointer in use we can simply lift the
@@ -1043,14 +1021,6 @@ impl<'a, B: Bindgen> Generator<'a, B> {
                     self.emit(&Instruction::Flush {
                         amt: usize::from(func.result.is_some()),
                     });
-
-                    if let Some((size, align)) = dealloc_size_align {
-                        self.stack.push(ptr);
-                        self.emit(&Instruction::GuestDeallocate {
-                            size: size.size_wasm32(),
-                            align: align.align_wasm32(),
-                        });
-                    }
                 }
 
                 self.emit(&Instruction::Return {

crates/csharp/src/function.rs

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

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

@@ -3,7 +3,6 @@
 
 extern crate std;
 use core::sync::atomic::{AtomicBool, Ordering};
-use std::alloc::{self, Layout};
 use std::any::Any;
 use std::boxed::Box;
 use std::collections::{hash_map, HashMap};
@@ -157,7 +156,6 @@ pub fn first_poll<T: 'static>(
 #[doc(hidden)]
 pub async unsafe fn await_result(
     import: unsafe extern "C" fn(*mut u8, *mut u8) -> i32,
-    params_layout: Layout,
     params: *mut u8,
     results: *mut u8,
 ) {
@@ -182,17 +180,13 @@ pub async unsafe fn await_result(
             let (tx, rx) = oneshot::channel();
             CALLS.insert(call, tx);
             rx.await.unwrap();
-            alloc::dealloc(params, params_layout);
         }
         STATUS_STARTED => {
-            alloc::dealloc(params, params_layout);
             let (tx, rx) = oneshot::channel();
             CALLS.insert(call, tx);
             rx.await.unwrap();
         }
-        STATUS_RETURNED | STATUS_DONE => {
-            alloc::dealloc(params, params_layout);
-        }
+        STATUS_RETURNED | STATUS_DONE => {}
         _ => unreachable!("unrecognized async call status"),
     }
 

crates/moonbit/src/lib.rs

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

crates/rust/src/bindgen.rs

@@ -278,7 +278,14 @@ impl Bindgen for FunctionBindgen<'_, '_> {
         // stack whereas exports use a per-module return area to cut down on
         // stack usage. Note that for imports this also facilitates "adapter
         // modules" for components to not have data segments.
-        if self.gen.in_import {
+        if size == 0 {
+            // If the size requested is 0 then we know it won't be written to so
+            // hand out a null pointer. This can happen with async for example
+            // when the params or results are zero-sized.
+            uwrite!(self.src, "let ptr{tmp} = core::ptr::null_mut::<u8>();");
+        } else if self.gen.in_import {
+            // Import return areas are stored on the stack since this stack
+            // frame will be around for the entire function call.
             self.import_return_pointer_area_size = self.import_return_pointer_area_size.max(size);
             self.import_return_pointer_area_align =
                 self.import_return_pointer_area_align.max(align);
@@ -287,6 +294,9 @@ impl Bindgen for FunctionBindgen<'_, '_> {
                 "let ptr{tmp} = ret_area.0.as_mut_ptr().cast::<u8>();"
             );
         } else {
+            // Export return areas are stored in `static` memory as they need to
+            // persist beyond the function call itself (and are cleaned-up in
+            // `post-return`).
             self.gen.return_pointer_area_size = self.gen.return_pointer_area_size.max(size);
             self.gen.return_pointer_area_align = self.gen.return_pointer_area_align.max(align);
             uwriteln!(
@@ -870,20 +880,14 @@ impl Bindgen for FunctionBindgen<'_, '_> {
                 self.push_str(");\n");
             }
 
-            Instruction::AsyncCallWasm { name, size, align } => {
+            Instruction::AsyncCallWasm { name, .. } => {
                 let func = self.declare_import(name, &[WasmType::Pointer; 2], &[WasmType::I32]);
 
                 let async_support = self.gen.gen.async_support_path();
-                let tmp = self.tmp();
-                let layout = format!("layout{tmp}");
-                let alloc = self.gen.path_to_std_alloc_module();
-                self.push_str(&format!(
-                    "let {layout} = {alloc}::Layout::from_size_align_unchecked({size}, {align});\n",
-                ));
                 let operands = operands.join(", ");
                 uwriteln!(
                     self.src,
-                    "{async_support}::await_result({func}, {layout}, {operands}).await;"
+                    "{async_support}::await_result({func}, {operands}).await;"
                 );
             }
 
@@ -951,19 +955,6 @@ impl Bindgen for FunctionBindgen<'_, '_> {
                 self.push_str(";\n");
             }
 
-            Instruction::AsyncMalloc { size, align } => {
-                let alloc = self.gen.path_to_std_alloc_module();
-                let tmp = self.tmp();
-                let ptr = format!("ptr{tmp}");
-                let layout = format!("layout{tmp}");
-                uwriteln!(
-                    self.src,
-                    "let {layout} = {alloc}::Layout::from_size_align_unchecked({size}, {align});
-                     let {ptr} = {alloc}::alloc({layout});"
-                );
-                results.push(ptr);
-            }
-
             Instruction::AsyncPostCallInterface { func } => {
                 let result = &operands[0];
                 self.async_result_name = Some(result.clone());

crates/teavm-java/src/lib.rs

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