Allocate and free stack frames

compiler/gen_wasm/src/backend.rs

@@ -10,8 +10,8 @@ use roc_module::symbol::Symbol;
 use roc_mono::ir::{CallType, Expr, JoinPointId, Literal, Proc, Stmt};
 use roc_mono::layout::{Builtin, Layout};
 
-use crate::layout::{WasmLayout};
-use crate::{PTR_TYPE, copy_memory, LocalId};
+use crate::layout::WasmLayout;
+use crate::{allocate_stack_frame, copy_memory, free_stack_frame, LocalId, PTR_TYPE};
 
 // Don't allocate any constant data at address zero or near it. Would be valid, but bug-prone.
 // Follow Emscripten's example by using 1kB (4 bytes would probably do)
@@ -110,14 +110,17 @@ impl<'a> WasmBackend<'a> {
 
         self.build_stmt(&proc.body, &proc.ret_layout)?;
 
-        // functions must end with an End instruction/opcode
-        self.instructions.push(Instruction::End);
+        let mut final_instructions = Vec::with_capacity(self.instructions.len() + 10);
+        allocate_stack_frame(&mut final_instructions, self.stack_memory as i32);
+        final_instructions.extend(self.instructions.clone());
+        free_stack_frame(&mut final_instructions, self.stack_memory as i32);
+        final_instructions.push(Instruction::End);
 
         let function_def = builder::function()
             .with_signature(signature)
             .body()
             .with_locals(self.locals.clone())
-            .with_instructions(Instructions::new(self.instructions.clone()))
+            .with_instructions(Instructions::new(final_instructions))
             .build() // body
             .build(); // function
 
@@ -130,13 +133,13 @@ impl<'a> WasmBackend<'a> {
     }
 
     fn insert_local(&mut self, layout: WasmLayout, symbol: Symbol, kind: LocalKind) -> LocalId {
-        self.stack_memory += layout.stack_memory();
-
         match kind {
             LocalKind::Parameter => {
+                // Don't increment stack_memory! Structs are allocated in caller's stack memory and passed as pointers.
                 self.arg_types.push(layout.value_type());
             }
             LocalKind::Variable => {
+                self.stack_memory += layout.stack_memory();
                 self.locals.push(Local::new(1, layout.value_type()));
             }
         }
@@ -228,7 +231,13 @@ impl<'a> WasmBackend<'a> {
                         let to = LocalId(0);
                         let copy_size: u32 = *size;
                         let copy_alignment_bytes: u32 = *alignment_bytes;
-                        copy_memory(&mut self.instructions, from, to, copy_size, copy_alignment_bytes)?;
+                        copy_memory(
+                            &mut self.instructions,
+                            from,
+                            to,
+                            copy_size,
+                            copy_alignment_bytes,
+                        )?;
                     }
                 }
 

compiler/gen_wasm/src/lib.rs

@@ -25,7 +25,7 @@ pub const ALIGN_8: u32 = 3;
 pub const STACK_POINTER_GLOBAL_ID: u32 = 0;
 
 #[derive(Clone, Copy, Debug)]
-struct LocalId(u32);
+pub struct LocalId(u32);
 
 pub struct Env<'a> {
     pub arena: &'a Bump, // not really using this much, parity_wasm works with std::vec a lot
@@ -152,3 +152,23 @@ fn copy_memory(
     }
     Ok(())
 }
+
+pub fn allocate_stack_frame(instructions: &mut Vec<Instruction>, size: i32) {
+    if size == 0 {
+        return;
+    }
+    instructions.push(GetGlobal(STACK_POINTER_GLOBAL_ID));
+    instructions.push(I32Const(size));
+    instructions.push(I32Sub);
+    instructions.push(SetGlobal(STACK_POINTER_GLOBAL_ID));
+}
+
+pub fn free_stack_frame(instructions: &mut Vec<Instruction>, size: i32) {
+    if size == 0 {
+        return;
+    }
+    instructions.push(GetGlobal(STACK_POINTER_GLOBAL_ID));
+    instructions.push(I32Const(size));
+    instructions.push(I32Add);
+    instructions.push(SetGlobal(STACK_POINTER_GLOBAL_ID));
+}

compiler/gen_wasm/tests/helpers/wasm32_test_result.rs

@@ -35,20 +35,12 @@ pub trait Wasm32TestResult {
     fn build_wrapper_body(main_function_index: u32) -> Vec<Instruction>;
 }
 
-fn build_wrapper_body_prelude(stack_memory_size: usize) -> Vec<Instruction> {
-    vec![
-        GetGlobal(STACK_POINTER_GLOBAL_ID),
-        I32Const(stack_memory_size as i32),
-        I32Sub,
-        SetGlobal(STACK_POINTER_GLOBAL_ID),
-    ]
-}
-
 macro_rules! build_wrapper_body_primitive {
     ($store_instruction: expr, $align: expr) => {
         fn build_wrapper_body(main_function_index: u32) -> Vec<Instruction> {
-            const MAX_ALIGNED_SIZE: usize = 16;
-            let mut instructions = build_wrapper_body_prelude(MAX_ALIGNED_SIZE);
+            const MAX_ALIGNED_SIZE: i32 = 16;
+            let mut instructions = Vec::with_capacity(9);
+            allocate_stack_frame(&mut instructions, MAX_ALIGNED_SIZE);
             instructions.extend([
                 GetGlobal(STACK_POINTER_GLOBAL_ID),
                 //
@@ -76,7 +68,8 @@ macro_rules! wasm_test_result_primitive {
 }
 
 fn build_wrapper_body_stack_memory(main_function_index: u32, size: usize) -> Vec<Instruction> {
-    let mut instructions = build_wrapper_body_prelude(size);
+    let mut instructions = Vec::with_capacity(8);
+    allocate_stack_frame(&mut instructions, size as i32);
     instructions.extend([
         //
         // Call the main function with the allocated address to write the result.