Refactor SymbolStorage to emphasise stack/non-stack params

We don't have much (any) code that cares about the difference between
heap pointers and other primitives, but we have a _lot_ of code that cares
if it's stack memory. So let's encode it that way.

compiler/gen_wasm/src/backend.rs

@@ -13,7 +13,8 @@ use roc_mono::layout::{Builtin, Layout};
 use crate::layout::WasmLayout;
 use crate::storage::SymbolStorage;
 use crate::{
-    allocate_stack_frame, copy_memory, free_stack_frame, round_up_to_alignment, LocalId, PTR_TYPE,
+    allocate_stack_frame, copy_memory, free_stack_frame, round_up_to_alignment, LocalId, PTR_SIZE,
+    PTR_TYPE,
 };
 
 // Don't allocate any constant data at address zero or near it. Would be valid, but bug-prone.
@@ -172,9 +173,20 @@ impl<'a> WasmBackend<'a> {
                         value_type,
                         size,
                     },
-                    _ => SymbolStorage::ParamPointer {
+
+                    WasmLayout::HeapMemory => SymbolStorage::ParamPrimitive {
                         local_id,
-                        wasm_layout,
+                        value_type: PTR_TYPE,
+                        size: PTR_SIZE,
+                    },
+
+                    WasmLayout::StackMemory {
+                        size,
+                        alignment_bytes,
+                    } => SymbolStorage::ParamStackMemory {
+                        local_id,
+                        size,
+                        alignment_bytes,
                     },
                 }
             }
@@ -287,12 +299,17 @@ impl<'a> WasmBackend<'a> {
             // Simple optimisation: if we are just returning the expression, we don't need a local
             Stmt::Let(let_sym, expr, layout, Stmt::Ret(ret_sym)) if let_sym == ret_sym => {
                 let wasm_layout = WasmLayout::new(layout);
-                if let WasmLayout::StackMemory { .. } = wasm_layout {
+                if let WasmLayout::StackMemory {
+                    size,
+                    alignment_bytes,
+                } = wasm_layout
+                {
                     // Map this symbol to the first argument (pointer into caller's stack)
                     // Saves us from having to copy it later
-                    let storage = SymbolStorage::ParamPointer {
+                    let storage = SymbolStorage::ParamStackMemory {
                         local_id: LocalId(0),
-                        wasm_layout,
+                        size,
+                        alignment_bytes,
                     };
                     self.symbol_storage_map.insert(*let_sym, storage);
                 }
@@ -326,14 +343,10 @@ impl<'a> WasmBackend<'a> {
                         alignment_bytes,
                         ..
                     }
-                    | ParamPointer {
+                    | ParamStackMemory {
                         local_id,
-                        wasm_layout:
-                            WasmLayout::StackMemory {
                         size,
                         alignment_bytes,
-                                ..
-                            },
                     } => {
                         let from = *local_id;
                         let to = LocalId(0);
@@ -342,7 +355,6 @@ impl<'a> WasmBackend<'a> {
 
                     ParamPrimitive { local_id, .. }
                     | VarPrimitive { local_id, .. }
-                    | ParamPointer { local_id, .. }
                     | VarHeapMemory { local_id, .. } => {
                         self.instructions.push(GetLocal(local_id.0));
                         self.instructions.push(Return); // TODO: branch instead of return so we can clean up stack
@@ -537,10 +549,7 @@ impl<'a> WasmBackend<'a> {
         if let Layout::Struct(field_layouts) = layout {
             match storage {
                 SymbolStorage::VarStackMemory { local_id, size, .. }
-                | SymbolStorage::ParamPointer {
+                | SymbolStorage::ParamStackMemory { local_id, size, .. } => {
-                    local_id,
-                    wasm_layout: WasmLayout::StackMemory { size, .. },
-                } => {
                     if size > 0 {
                         let mut relative_offset = 0;
                         for (field, _) in fields.iter().zip(field_layouts.iter()) {

compiler/gen_wasm/src/storage.rs

@@ -1,18 +1,14 @@
-use crate::{copy_memory, layout::WasmLayout, LocalId, ALIGN_1, ALIGN_2, ALIGN_4, ALIGN_8};
+use crate::{copy_memory, LocalId, ALIGN_1, ALIGN_2, ALIGN_4, ALIGN_8};
 use parity_wasm::elements::{Instruction, Instruction::*, ValueType};
 
 #[derive(Debug, Clone)]
 pub enum SymbolStorage {
-    ParamPrimitive {
+    VarPrimitive {
         local_id: LocalId,
         value_type: ValueType,
         size: u32,
     },
-    ParamPointer {
+    ParamPrimitive {
-        local_id: LocalId,
-        wasm_layout: WasmLayout,
-    },
-    VarPrimitive {
         local_id: LocalId,
         value_type: ValueType,
         size: u32,
@@ -23,6 +19,11 @@ pub enum SymbolStorage {
         offset: u32,
         alignment_bytes: u32,
     },
+    ParamStackMemory {
+        local_id: LocalId,
+        size: u32,
+        alignment_bytes: u32,
+    },
     VarHeapMemory {
         local_id: LocalId,
     },
@@ -32,7 +33,7 @@ impl SymbolStorage {
     pub fn local_id(&self) -> LocalId {
         match self {
             Self::ParamPrimitive { local_id, .. } => *local_id,
-            Self::ParamPointer { local_id, .. } => *local_id,
+            Self::ParamStackMemory { local_id, .. } => *local_id,
             Self::VarPrimitive { local_id, .. } => *local_id,
             Self::VarStackMemory { local_id, .. } => *local_id,
             Self::VarHeapMemory { local_id, .. } => *local_id,
@@ -43,7 +44,7 @@ impl SymbolStorage {
         match self {
             Self::ParamPrimitive { value_type, .. } => *value_type,
             Self::VarPrimitive { value_type, .. } => *value_type,
-            Self::ParamPointer { .. } => ValueType::I32,
+            Self::ParamStackMemory { .. } => ValueType::I32,
             Self::VarStackMemory { .. } => ValueType::I32,
             Self::VarHeapMemory { .. } => ValueType::I32,
         }
@@ -51,11 +52,7 @@ impl SymbolStorage {
 
     pub fn has_stack_memory(&self) -> bool {
         match self {
-            Self::ParamPointer {
+            Self::ParamStackMemory { .. } => true,
-                wasm_layout: WasmLayout::StackMemory { .. },
-                ..
-            } => true,
-            Self::ParamPointer { .. } => false,
             Self::VarStackMemory { .. } => true,
             Self::ParamPrimitive { .. } => false,
             Self::VarPrimitive { .. } => false,
@@ -70,14 +67,10 @@ impl SymbolStorage {
                 alignment_bytes,
                 ..
             }
-            | Self::ParamPointer {
+            | Self::ParamStackMemory {
-                wasm_layout:
-                    WasmLayout::StackMemory {
                 size,
                 alignment_bytes,
                 ..
-                    },
-                ..
             } => (*size, *alignment_bytes),
 
             _ => (0, 0),
@@ -123,13 +116,10 @@ impl SymbolStorage {
                 Ok(*size)
             }
 
-            Self::ParamPointer {
+            Self::ParamStackMemory {
                 local_id,
-                wasm_layout:
-                    WasmLayout::StackMemory {
                 size,
                 alignment_bytes,
-                    },
             }
             | Self::VarStackMemory {
                 local_id,
@@ -148,7 +138,7 @@ impl SymbolStorage {
                 Ok(*size)
             }
 
-            Self::ParamPointer { local_id, .. } | Self::VarHeapMemory { local_id, .. } => {
+            Self::VarHeapMemory { local_id, .. } => {
                 instructions.push(GetLocal(to_pointer.0));
                 instructions.push(GetLocal(local_id.0));
                 instructions.push(I32Store(ALIGN_4, to_offset));