From 02409730f25a27a3c09abfb6abfbbf9fb611bffb Mon Sep 17 00:00:00 2001
From: Dan Gohman <dev@sunfishcode.online>
Date: Fri, 23 Sep 2022 15:57:03 -0700
Subject: [PATCH] Move `RET_AREA` variables onto the stack for import bindings.

Move `RET_AREA` variables out of static memory and onto the stack, for
import bindings.

This makes it easier to use wit-bindgen to generate polyfill modules
that import their linear memory.
---
 crates/gen-guest-c/src/lib.rs    | 26 ++++++++++++++++++++++++--
 crates/gen-guest-rust/src/lib.rs | 26 ++++++++++++++++++++++++--
 crates/wit-parser/src/abi.rs     |  2 +-
 3 files changed, 49 insertions(+), 5 deletions(-)

diff --git a/crates/gen-guest-c/src/lib.rs b/crates/gen-guest-c/src/lib.rs
index 1c8cd9eb9..c2509d22c 100644
--- a/crates/gen-guest-c/src/lib.rs
+++ b/crates/gen-guest-c/src/lib.rs
@@ -1347,7 +1347,10 @@ impl Generator for C {
             }
         }
 
-        if self.return_pointer_area_size > 0 {
+        // Declare a statically-allocated return area, if needed. We only do
+        // this for export bindings, because import bindings allocate their
+        // return-area on the stack.
+        if !self.in_import && self.return_pointer_area_size > 0 {
             uwrite!(
                 self.src.c,
                 "
@@ -1472,7 +1475,26 @@ impl Bindgen for FunctionBindgen<'_> {
         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);
         let ptr = self.locals.tmp("ptr");
-        uwriteln!(self.src, "int32_t {} = (int32_t) &RET_AREA;", ptr);
+
+        if self.gen.in_import {
+            // Declare a stack-allocated return area. We only do this for
+            // imports, because exports need their return area to be live until
+            // the post-return call.
+            uwrite!(
+                self.src,
+                "
+                    __attribute__((aligned({})))
+                    uint8_t ret_area[{}];
+                ",
+                align,
+                size,
+            );
+            uwriteln!(self.src, "int32_t {} = (int32_t) &ret_area;", ptr);
+        } else {
+            // Declare a statically-allocated return area.
+            uwriteln!(self.src, "int32_t {} = (int32_t) &RET_AREA;", ptr);
+        }
+
         ptr
     }
 
diff --git a/crates/gen-guest-rust/src/lib.rs b/crates/gen-guest-rust/src/lib.rs
index c574ac0c2..e13a94124 100644
--- a/crates/gen-guest-rust/src/lib.rs
+++ b/crates/gen-guest-rust/src/lib.rs
@@ -653,7 +653,7 @@ impl Generator for RustWasm {
     }
 
     fn finish_functions(&mut self, iface: &Interface, dir: Direction) {
-        if self.return_pointer_area_align > 0 {
+        if !self.in_import && self.return_pointer_area_align > 0 {
             self.src.push_str(&format!(
                 "
                     #[repr(align({align}))]
@@ -810,6 +810,16 @@ impl FunctionBindgen<'_> {
         self.push_str(";\n}\n");
         "wit_import".to_string()
     }
+
+    fn ret_area_name(&self, iface: &Interface) -> String {
+        // For imports, we allocate the return area on the stack; for exports,
+        // we statically allocate it.
+        if self.gen.in_import {
+            format!("__{}_ret_area", iface.name.to_snake_case())
+        } else {
+            format!("__{}_RET_AREA", iface.name.to_shouty_snake_case())
+        }
+    }
 }
 
 impl RustFunctionGenerator for FunctionBindgen<'_> {
@@ -880,10 +890,22 @@ impl Bindgen for FunctionBindgen<'_> {
         self.gen.return_pointer_area_align = self.gen.return_pointer_area_align.max(align);
         let tmp = self.tmp();
 
+        if self.gen.in_import {
+            self.push_str(&format!(
+                "
+                    #[repr(align({align}))]
+                    struct {ty}([u8; {size}]);
+                    let mut {name}: {ty} = {ty}([0; {size}]);
+                ",
+                ty = RustWasm::ret_area_type_name(iface),
+                name = self.ret_area_name(iface),
+            ));
+        }
+
         self.push_str(&format!(
             "let ptr{} = {}.0.as_mut_ptr() as i32;\n",
             tmp,
-            RustWasm::ret_area_name(iface),
+            self.ret_area_name(iface),
         ));
         format!("ptr{}", tmp)
     }
diff --git a/crates/wit-parser/src/abi.rs b/crates/wit-parser/src/abi.rs
index 7ffbbd18f..2770cb5b8 100644
--- a/crates/wit-parser/src/abi.rs
+++ b/crates/wit-parser/src/abi.rs
@@ -1114,7 +1114,7 @@ impl<'a, B: Bindgen> Generator<'a, B> {
                         .record(func.params.iter().map(|t| &t.1));
                     let ptr = match self.variant {
                         // When a wasm module calls an import it will provide
-                        // static space that isn't dynamically allocated.
+                        // space that isn't explicitly deallocated.
                         AbiVariant::GuestImport => {
                             self.bindgen.return_pointer(self.iface, size, align)
                         }