wasm: Work around Zig returning small structs as integers

compiler/gen_wasm/src/backend.rs

@@ -200,15 +200,19 @@ impl<'a> WasmBackend<'a> {
     }
 
     fn start_proc(&mut self, proc: &Proc<'a>) {
+        use ReturnMethod::*;
         let ret_layout = WasmLayout::new(&proc.ret_layout);
 
-        let ret_type = match ret_layout.return_method() {
-            ReturnMethod::Primitive(ty) => Some(ty),
-            ReturnMethod::NoReturnValue => None,
-            ReturnMethod::WriteToPointerArg => {
+        let ret_type = match ret_layout.return_method(CallConv::C) {
+            Primitive(ty) => Some(ty),
+            NoReturnValue => None,
+            WriteToPointerArg => {
                 self.storage.arg_types.push(PTR_TYPE);
                 None
             }
+            ZigPackedStruct => {
+                internal_error!("C calling convention does not return Zig packed structs")
+            }
         };
 
         // Create a block so we can exit the function without skipping stack frame "pop" code.
@@ -849,7 +853,8 @@ impl<'a> WasmBackend<'a> {
             return self.expr_call_low_level(lowlevel, arguments, ret_sym, ret_layout, ret_storage);
         }
 
-        let (param_types, ret_type) = self.storage.load_symbols_for_call(
+        let (num_wasm_args, has_return_val, ret_zig_packed_struct) =
+            self.storage.load_symbols_for_call(
                 self.env.arena,
                 &mut self.code_builder,
                 arguments,
@@ -857,13 +862,12 @@ impl<'a> WasmBackend<'a> {
                 &wasm_layout,
                 CallConv::C,
             );
+        debug_assert!(!ret_zig_packed_struct);
 
         let iter = self.proc_lookup.iter().enumerate();
         for (roc_proc_index, (ir_sym, pl, linker_sym_index)) in iter {
             if *ir_sym == func_sym && pl == proc_layout {
                 let wasm_fn_index = self.fn_index_offset + roc_proc_index as u32;
-                let num_wasm_args = param_types.len();
-                let has_return_val = ret_type.is_some();
                 self.code_builder.call(
                     wasm_fn_index,
                     *linker_sym_index,

compiler/gen_wasm/src/layout.rs

@@ -4,9 +4,6 @@ use roc_mono::layout::{Layout, UnionLayout};
 use crate::wasm_module::ValueType;
 use crate::{PTR_SIZE, PTR_TYPE, TARGET_INFO};
 
-/// Manually keep up to date with the Zig version we are using for builtins
-pub const BUILTINS_ZIG_VERSION: ZigVersion = ZigVersion::Zig9;
-
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum ReturnMethod {
     /// This layout is returned from a Wasm function "normally" as a Primitive
@@ -15,6 +12,8 @@ pub enum ReturnMethod {
     WriteToPointerArg,
     /// This layout is empty and requires no return value or argument (e.g. refcount helpers)
     NoReturnValue,
+    /// This layout is returned as a packed struct in an integer. Only used by Zig, not C.
+    ZigPackedStruct,
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -124,34 +123,23 @@ impl WasmLayout {
         }
     }
 
-    pub fn return_method(&self) -> ReturnMethod {
+    pub fn return_method(&self, conv: CallConv) -> ReturnMethod {
         match self {
             Self::Primitive(ty, _) => ReturnMethod::Primitive(*ty),
-            Self::StackMemory { size, .. } => {
-                if *size == 0 {
-                    ReturnMethod::NoReturnValue
-                } else {
-                    ReturnMethod::WriteToPointerArg
+            Self::StackMemory { size, format, .. } => {
+                conv.stack_memory_return_method(*size, *format)
             }
         }
     }
 }
-}
-
-#[allow(dead_code)]
-#[derive(PartialEq, Eq)]
-pub enum ZigVersion {
-    Zig8,
-    Zig9,
-}
 
 #[derive(Debug, Clone, Copy)]
 pub enum CallConv {
     /// The C calling convention, as defined here:
     /// https://github.com/WebAssembly/tool-conventions/blob/main/BasicCABI.md
     C,
-    /// The calling convention that Zig 0.8 or 0.9 generates for Wasm when we *ask* it
-    /// for the .C calling convention, due to bugs in both versions of the Zig compiler.
+    /// The calling convention that Zig 0.9 generates for Wasm when we *ask* it
+    /// for the .C calling convention, due to bugs in the Zig compiler.
     Zig,
 }
 
@@ -183,7 +171,7 @@ impl CallConv {
                             &[I32] // Small struct: pass by value
                         } else if size <= 8 {
                             &[I64] // Small struct: pass by value
-                        } else if size <= 12 && BUILTINS_ZIG_VERSION == ZigVersion::Zig9 {
+                        } else if size <= 12 {
                             &[I64, I32] // Medium struct: pass by value, as two Wasm arguments
                         } else if size <= 16 {
                             &[I64, I64] // Medium struct: pass by value, as two Wasm arguments
@@ -195,4 +183,30 @@ impl CallConv {
             }
         }
     }
+
+    pub fn stack_memory_return_method(&self, size: u32, format: StackMemoryFormat) -> ReturnMethod {
+        use ReturnMethod::*;
+        use StackMemoryFormat::*;
+
+        match format {
+            Int128 | Float128 | Decimal => WriteToPointerArg,
+
+            DataStructure => {
+                if size == 0 {
+                    return NoReturnValue;
+                }
+                match self {
+                    CallConv::C => WriteToPointerArg,
+
+                    CallConv::Zig => {
+                        if size <= 8 {
+                            ZigPackedStruct
+                        } else {
+                            WriteToPointerArg
+                        }
+                    }
+                }
+            }
+        }
+    }
 }

compiler/gen_wasm/src/low_level.rs

@@ -1,4 +1,3 @@
-use bumpalo::collections::Vec;
 use roc_builtins::bitcode::{self, FloatWidth, IntWidth};
 use roc_error_macros::internal_error;
 use roc_module::low_level::{LowLevel, LowLevel::*};
@@ -126,7 +125,7 @@ impl<'a> LowLevelCall<'a> {
     /// For numerical ops, this just pushes the arguments to the Wasm VM's value stack
     /// It implements the calling convention used by Zig for both numbers and structs
     /// Result is the type signature of the call
-    fn load_args(&self, backend: &mut WasmBackend<'a>) -> (Vec<'a, ValueType>, Option<ValueType>) {
+    fn load_args(&self, backend: &mut WasmBackend<'a>) -> (usize, bool, bool) {
         backend.storage.load_symbols_for_call(
             backend.env.arena,
             &mut backend.code_builder,
@@ -138,8 +137,29 @@ impl<'a> LowLevelCall<'a> {
     }
 
     fn load_args_and_call_zig(&self, backend: &mut WasmBackend<'a>, name: &'a str) {
-        let (param_types, ret_type) = self.load_args(backend);
-        backend.call_zig_builtin_after_loading_args(name, param_types.len(), ret_type.is_some());
+        let (num_wasm_args, has_return_val, ret_zig_packed_struct) = self.load_args(backend);
+        backend.call_zig_builtin_after_loading_args(name, num_wasm_args, has_return_val);
+
+        if ret_zig_packed_struct {
+            match self.ret_storage {
+                StoredValue::StackMemory {
+                    size,
+                    alignment_bytes,
+                    ..
+                } => {
+                    // The address of the return value was already loaded before the call
+                    let align = Align::from(alignment_bytes);
+                    if size > 4 {
+                        backend.code_builder.i64_store(align, 0);
+                    } else {
+                        backend.code_builder.i32_store(align, 0);
+                    }
+                }
+                _ => {
+                    internal_error!("Zig packed struct should always be stored to StackMemory")
+                }
+            }
+        }
     }
 
     /// Wrap an integer whose Wasm representation is i32

compiler/gen_wasm/src/storage.rs

@@ -6,9 +6,7 @@ use roc_error_macros::internal_error;
 use roc_module::symbol::Symbol;
 use roc_mono::layout::Layout;
 
-use crate::layout::{
-    CallConv, ReturnMethod, StackMemoryFormat, WasmLayout, ZigVersion, BUILTINS_ZIG_VERSION,
-};
+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};
 
@@ -326,7 +324,7 @@ impl<'a> Storage<'a> {
                     code_builder.i32_load(align, offset);
                 } else if *size <= 8 {
                     code_builder.i64_load(align, offset);
-                } else if *size <= 12 && BUILTINS_ZIG_VERSION == ZigVersion::Zig9 {
+                } else if *size <= 12 {
                     code_builder.i64_load(align, offset);
                     code_builder.get_local(local_id);
                     code_builder.i32_load(align, offset + 8);
@@ -396,37 +394,47 @@ impl<'a> Storage<'a> {
         return_symbol: Symbol,
         return_layout: &WasmLayout,
         call_conv: CallConv,
-    ) -> (Vec<'a, ValueType>, Option<ValueType>) {
-        let mut wasm_arg_types = Vec::with_capacity_in(arguments.len() * 2 + 1, arena);
-        let mut wasm_args = Vec::with_capacity_in(arguments.len() * 2 + 1, arena);
+    ) -> (usize, bool, bool) {
+        use ReturnMethod::*;
 
-        let return_method = return_layout.return_method();
-        let return_type = match return_method {
-            ReturnMethod::Primitive(ty) => Some(ty),
-            ReturnMethod::NoReturnValue => None,
-            ReturnMethod::WriteToPointerArg => {
-                wasm_arg_types.push(PTR_TYPE);
-                wasm_args.push(return_symbol);
-                None
+        let mut num_wasm_args = 0;
+        let mut symbols_to_load = Vec::with_capacity_in(arguments.len() * 2 + 1, arena);
+
+        let return_method = return_layout.return_method(call_conv);
+        let has_return_val = match return_method {
+            Primitive(_) => true,
+            NoReturnValue => false,
+            WriteToPointerArg => {
+                num_wasm_args += 1;
+                symbols_to_load.push(return_symbol);
+                false
+            }
+            ZigPackedStruct => {
+                // Workaround for Zig's incorrect implementation of the C calling convention.
+                // We need to copy the packed struct into the stack frame
+                // Load the address before the call so that afterward, it will be 2nd on the value stack,
+                // ready for the store instruction.
+                symbols_to_load.push(return_symbol);
+                true
             }
         };
 
         for arg in arguments {
             let stored = self.symbol_storage_map.get(arg).unwrap();
             let arg_types = stored.arg_types(call_conv);
-            wasm_arg_types.extend_from_slice(arg_types);
+            num_wasm_args += arg_types.len();
             match arg_types.len() {
                 0 => {}
-                1 => wasm_args.push(*arg),
-                2 => wasm_args.extend_from_slice(&[*arg, *arg]),
+                1 => symbols_to_load.push(*arg),
+                2 => symbols_to_load.extend_from_slice(&[*arg, *arg]),
                 n => internal_error!("Cannot have {} Wasm arguments for 1 Roc argument", n),
             }
         }
 
         // If the symbols were already at the top of the stack, do nothing!
         // Should be common for simple cases, due to the structure of the Mono IR
-        if !code_builder.verify_stack_match(&wasm_args) {
-            if return_method == ReturnMethod::WriteToPointerArg {
+        if !code_builder.verify_stack_match(&symbols_to_load) {
+            if matches!(return_method, WriteToPointerArg | ZigPackedStruct) {
                 self.load_return_address_ccc(code_builder, return_symbol);
             };
 
@@ -438,7 +446,11 @@ impl<'a> Storage<'a> {
             }
         }
 
-        (wasm_arg_types, return_type)
+        (
+            num_wasm_args,
+            has_return_val,
+            return_method == ZigPackedStruct,
+        )
     }
 
     /// Generate code to copy a StoredValue to an arbitrary memory location