Pass actual stack pointers around instead of address 8

This commit updates the implementation of passing return pointers from
JS to wasm to actually modify the wasm's shadow stack pointer instead of
manufacturing the arbitrary address of 8. The purpose of this is to
correctly handle threaded scenarios where each thread will write to its
own area instead of everyone trying to compete at address 8.

The implementation here will lazily, if necessary, export the stack
pointer we find the module and modify it as necessary.

Closes #2218

crates/cli-support/src/js/binding.rs

@@ -463,13 +463,6 @@ impl<'a, 'b> JsBuilder<'a, 'b> {
 }
 
 fn instruction(js: &mut JsBuilder, instr: &Instruction, log_error: &mut bool) -> Result<(), Error> {
-    // Here first properly aligned nonzero address is chosen to be the
-    // out-pointer. We use the address for a BigInt64Array sometimes which
-    // means it needs to be 8-byte aligned. Otherwise valid code is
-    // unlikely to ever be working around address 8, so this should be a
-    // safe address to use for returning data through.
-    let retptr_val = 8;
-
     match instr {
         Instruction::Standard(wit_walrus::Instruction::ArgGet(n)) => {
             let arg = js.arg(*n).to_string();
@@ -566,7 +559,20 @@ fn instruction(js: &mut JsBuilder, instr: &Instruction, log_error: &mut bool) ->
             js.string_to_memory(*mem, *malloc, *realloc)?;
         }
 
-        Instruction::Retptr => js.stack.push(retptr_val.to_string()),
+        Instruction::Retptr { size } => {
+            let sp = match js.cx.aux.shadow_stack_pointer {
+                Some(s) => js.cx.export_name_of(s),
+                // In theory this shouldn't happen since malloc is included in
+                // most wasm binaries (and may be gc'd out) and that almost
+                // always pulls in a stack pointer. We can try to synthesize
+                // something here later if necessary.
+                None => bail!("failed to find shadow stack pointer"),
+            };
+            js.prelude(&format!("const retptr = wasm.{}.value - {};", sp, size));
+            js.prelude(&format!("wasm.{}.value = retptr;", sp));
+            js.finally(&format!("wasm.{}.value += {};", sp, size));
+            js.stack.push("retptr".to_string());
+        }
 
         Instruction::StoreRetptr { ty, offset, mem } => {
             let (mem, size) = match ty {
@@ -599,7 +605,7 @@ fn instruction(js: &mut JsBuilder, instr: &Instruction, log_error: &mut bool) ->
             // If we're loading from the return pointer then we must have pushed
             // it earlier, and we always push the same value, so load that value
             // here
-            let expr = format!("{}()[{} / {} + {}]", mem, retptr_val, size, offset);
+            let expr = format!("{}()[retptr / {} + {}]", mem, size, offset);
             js.prelude(&format!("var r{} = {};", offset, expr));
             js.push(format!("r{}", offset));
         }

crates/cli-support/src/multivalue.rs

@@ -64,7 +64,7 @@ fn extract_xform<'a>(
     // If the first instruction is a `Retptr`, then this must be an exported
     // adapter which calls a wasm-defined function. Something we'd like to
     // adapt to multi-value!
-    if let Some(Instruction::Retptr) = instructions.first().map(|e| &e.instr) {
+    if let Some(Instruction::Retptr { .. }) = instructions.first().map(|e| &e.instr) {
         instructions.remove(0);
         let mut types = Vec::new();
         instructions.retain(|instruction| match &instruction.instr {

crates/cli-support/src/wit/mod.rs

@@ -1292,8 +1292,23 @@ impl<'a> Context<'a> {
         // everything into the outgoing arguments.
         let mut instructions = Vec::new();
         if uses_retptr {
+            let size = ret.input.iter().fold(0, |sum, ty| {
+                let size = match ty {
+                    AdapterType::I32 => 4,
+                    AdapterType::I64 => 8,
+                    AdapterType::F32 => 4,
+                    AdapterType::F64 => 8,
+                    _ => panic!("unsupported type in retptr {:?}", ty),
+                };
+                let sum_rounded_up = (sum + (size - 1)) & (!(size - 1));
+                sum_rounded_up + size
+            });
+            // Round the number of bytes up to a 16-byte alignment to ensure the
+            // stack pointer is always 16-byte aligned (which LLVM currently
+            // requires).
+            let size = (size + 15) & (!15);
             instructions.push(InstructionData {
-                instr: Instruction::Retptr,
+                instr: Instruction::Retptr { size },
                 stack_change: StackChange::Modified {
                     pushed: 1,
                     popped: 0,

crates/cli-support/src/wit/section.rs

@@ -234,7 +234,7 @@ fn translate_instruction(
             mem: *mem,
             malloc: *malloc,
         }),
-        StoreRetptr { .. } | LoadRetptr { .. } | Retptr => {
+        StoreRetptr { .. } | LoadRetptr { .. } | Retptr { .. } => {
             bail!("return pointers aren't supported in wasm interface types");
         }
         I32FromBool | BoolFromI32 => {

crates/cli-support/src/wit/standard.rs

@@ -113,8 +113,11 @@ pub enum Instruction {
         offset: usize,
         mem: walrus::MemoryId,
     },
-    /// An instruction which pushes the return pointer onto the stack.
-    Retptr,
+    /// An instruction which pushes the return pointer onto the stack, reserving
+    /// `size` bytes of space.
+    Retptr {
+        size: u32,
+    },
 
     /// Pops a `bool` from the stack and pushes an `i32` equivalent
     I32FromBool,