wasm: support 128-bit numbers as arguments to Roc procedures

compiler/gen_wasm/src/backend.rs

@@ -17,7 +17,7 @@ use roc_std::RocDec;
 
 use crate::layout::{CallConv, ReturnMethod, WasmLayout};
 use crate::low_level::{call_higher_order_lowlevel, LowLevelCall};
-use crate::storage::{Storage, StoredValue, StoredValueKind};
+use crate::storage::{Storage, StoredValue, StoredVarKind};
 use crate::wasm_module::linking::{DataSymbol, WasmObjectSymbol};
 use crate::wasm_module::sections::{
     ConstExpr, DataMode, DataSegment, Export, Global, GlobalType, Import, ImportDesc, Limits,
@@ -414,10 +414,8 @@ impl<'a> WasmBackend<'a> {
         // We never use the `return` instruction. Instead, we break from this block.
         self.start_block();
 
-        for (layout, symbol) in proc.args {
-            self.storage
-                .allocate(*layout, *symbol, StoredValueKind::Parameter);
-        }
+        self.storage
+            .allocate_args(proc.args, &mut self.code_builder, self.env.arena);
 
         if let Some(ty) = ret_type {
             let ret_var = self.storage.create_anonymous_local(ty);
@@ -660,8 +658,8 @@ impl<'a> WasmBackend<'a> {
             }
 
             let kind = match following {
-                Stmt::Ret(ret_sym) if *sym == *ret_sym => StoredValueKind::ReturnValue,
-                _ => StoredValueKind::Variable,
+                Stmt::Ret(ret_sym) if *sym == *ret_sym => StoredVarKind::ReturnValue,
+                _ => StoredVarKind::Variable,
             };
 
             self.stmt_let_store_expr(*sym, layout, expr, kind);
@@ -677,9 +675,9 @@ impl<'a> WasmBackend<'a> {
         sym: Symbol,
         layout: &Layout<'a>,
         expr: &Expr<'a>,
-        kind: StoredValueKind,
+        kind: StoredVarKind,
     ) {
-        let sym_storage = self.storage.allocate(*layout, sym, kind);
+        let sym_storage = self.storage.allocate_var(*layout, sym, kind);
 
         self.expr(sym, expr, layout, &sym_storage);
 
@@ -817,10 +815,10 @@ impl<'a> WasmBackend<'a> {
         // make locals for join pointer parameters
         let mut jp_param_storages = Vec::with_capacity_in(parameters.len(), self.env.arena);
         for parameter in parameters.iter() {
-            let mut param_storage = self.storage.allocate(
+            let mut param_storage = self.storage.allocate_var(
                 parameter.layout,
                 parameter.symbol,
-                StoredValueKind::Variable,
+                StoredVarKind::Variable,
             );
             param_storage = self.storage.ensure_value_has_local(
                 &mut self.code_builder,
@@ -1420,7 +1418,7 @@ impl<'a> WasmBackend<'a> {
                             elem_sym,
                             elem_layout,
                             &expr,
-                            StoredValueKind::Variable,
+                            StoredVarKind::Variable,
                         );
 
                         elem_sym

compiler/gen_wasm/src/storage.rs

@@ -10,8 +10,7 @@ use crate::layout::{CallConv, ReturnMethod, StackMemoryFormat, WasmLayout};
 use crate::wasm_module::{Align, CodeBuilder, LocalId, ValueType, VmSymbolState};
 use crate::{copy_memory, round_up_to_alignment, CopyMemoryConfig, PTR_TYPE};
 
-pub enum StoredValueKind {
-    Parameter,
+pub enum StoredVarKind {
     Variable,
     ReturnValue,
 }
@@ -124,7 +123,7 @@ impl<'a> Storage<'a> {
         id
     }
 
-    /// Allocate storage for a Roc value
+    /// Allocate storage for a Roc variable
     ///
     /// Wasm primitives (i32, i64, f32, f64) are allocated "storage" on the VM stack.
     /// This is really just a way to model how the stack machine works as a sort of
@@ -134,31 +133,21 @@ impl<'a> Storage<'a> {
     ///
     /// Structs and Tags are stored in memory rather than in Wasm primitives.
     /// They are allocated a certain offset and size in the stack frame.
-    pub fn allocate(
+    pub fn allocate_var(
         &mut self,
         layout: Layout<'a>,
         symbol: Symbol,
-        kind: StoredValueKind,
+        kind: StoredVarKind,
     ) -> StoredValue {
         let next_local_id = self.get_next_local_id();
         let wasm_layout = WasmLayout::new(&layout);
         self.symbol_layouts.insert(symbol, layout);
 
         let storage = match wasm_layout {
-            WasmLayout::Primitive(value_type, size) => match kind {
-                StoredValueKind::Parameter => {
-                    self.arg_types.push(value_type);
-                    StoredValue::Local {
-                        local_id: next_local_id,
-                        value_type,
-                        size,
-                    }
-                }
-                _ => StoredValue::VirtualMachineStack {
-                    vm_state: VmSymbolState::NotYetPushed,
-                    value_type,
-                    size,
-                },
+            WasmLayout::Primitive(value_type, size) => StoredValue::VirtualMachineStack {
+                vm_state: VmSymbolState::NotYetPushed,
+                value_type,
+                size,
             },
 
             WasmLayout::StackMemory {
@@ -167,18 +156,7 @@ impl<'a> Storage<'a> {
                 format,
             } => {
                 let location = match kind {
-                    StoredValueKind::Parameter => {
-                        if size > 0 {
-                            self.arg_types.push(PTR_TYPE);
-                            StackMemoryLocation::PointerArg(next_local_id)
-                        } else {
-                            // An argument with zero size is purely conceptual, and will not exist in Wasm.
-                            // However we need to track the symbol, so we treat it like a local variable.
-                            StackMemoryLocation::FrameOffset(0)
-                        }
-                    }
-
-                    StoredValueKind::Variable => {
+                    StoredVarKind::Variable => {
                         if self.stack_frame_pointer.is_none() && size > 0 {
                             self.stack_frame_pointer = Some(next_local_id);
                             self.local_types.push(PTR_TYPE);
@@ -192,7 +170,7 @@ impl<'a> Storage<'a> {
                         StackMemoryLocation::FrameOffset(offset as u32)
                     }
 
-                    StoredValueKind::ReturnValue => StackMemoryLocation::PointerArg(LocalId(0)),
+                    StoredVarKind::ReturnValue => StackMemoryLocation::PointerArg(LocalId(0)),
                 };
 
                 StoredValue::StackMemory {
@@ -209,6 +187,98 @@ impl<'a> Storage<'a> {
         storage
     }
 
+    /// Allocate storage for a Roc procedure argument
+    /// Each argument is also a local variable. Their indices come before other locals.
+    /// Structs and Tags are passed as pointers into the caller's frame
+    /// 128-bit numbers are passed as two i64's, but we immediately store them in the
+    /// stack frame, because it's a lot easier to keep track of the data flow.
+    pub fn allocate_args(
+        &mut self,
+        args: &[(Layout<'a>, Symbol)],
+        code_builder: &mut CodeBuilder,
+        arena: &'a Bump,
+    ) {
+        let mut wide_number_args = Vec::with_capacity_in(args.len(), arena);
+
+        for (layout, symbol) in args {
+            self.symbol_layouts.insert(*symbol, *layout);
+            let wasm_layout = WasmLayout::new(layout);
+            let local_index = self.arg_types.len() as u32;
+
+            let storage = match wasm_layout {
+                WasmLayout::Primitive(value_type, size) => {
+                    self.arg_types.push(value_type);
+                    StoredValue::Local {
+                        local_id: LocalId(local_index),
+                        value_type,
+                        size,
+                    }
+                }
+                WasmLayout::StackMemory {
+                    size,
+                    alignment_bytes,
+                    format,
+                } => {
+                    use StackMemoryFormat::*;
+
+                    self.arg_types
+                        .extend_from_slice(CallConv::C.stack_memory_arg_types(size, format));
+
+                    let location = match format {
+                        Int128 | Float128 | Decimal => {
+                            // passed as two i64's but stored in the stack frame
+                            wide_number_args.push(local_index);
+                            let loc =
+                                StackMemoryLocation::FrameOffset(self.stack_frame_size as u32);
+                            self.stack_frame_size += size as i32;
+                            loc
+                        }
+                        DataStructure => {
+                            if size == 0 {
+                                // An argument with zero size is purely conceptual, and will not exist in Wasm.
+                                // However we need to track the symbol, so we treat it like a local variable.
+                                StackMemoryLocation::FrameOffset(0)
+                            } else {
+                                StackMemoryLocation::PointerArg(LocalId(local_index))
+                            }
+                        }
+                    };
+
+                    StoredValue::StackMemory {
+                        location,
+                        size,
+                        alignment_bytes,
+                        format,
+                    }
+                }
+            };
+
+            self.symbol_storage_map.insert(*symbol, storage.clone());
+        }
+
+        // If any arguments are 128-bit numbers, store them in the stack frame
+        // This makes it easier to keep track of which symbols are on the Wasm value stack
+        // The frame pointer will be the next local after the arguments
+        if self.stack_frame_size > 0 {
+            let frame_ptr = LocalId(self.arg_types.len() as u32);
+            self.stack_frame_pointer = Some(frame_ptr);
+            self.local_types.push(PTR_TYPE);
+
+            let mut offset = 0;
+            for arg_index in wide_number_args.iter().copied() {
+                code_builder.get_local(frame_ptr);
+                code_builder.get_local(LocalId(arg_index));
+                code_builder.i64_store(Align::Bytes8, offset);
+
+                code_builder.get_local(frame_ptr);
+                code_builder.get_local(LocalId(arg_index + 1));
+                code_builder.i64_store(Align::Bytes8, offset + 8);
+
+                offset += 16;
+            }
+        }
+    }
+
     /// Get storage info for a given symbol
     pub fn get(&self, sym: &Symbol) -> &StoredValue {
         self.symbol_storage_map.get(sym).unwrap_or_else(|| {