wasm: generate code for ZigCC wrapper function

compiler/gen_wasm/src/backend.rs

@@ -36,7 +36,7 @@ pub enum ProcSource {
     Helper,
     /// Wrapper function for higher-order calls from Zig to Roc,
     /// to work around Zig's incorrect implementation of C calling convention in Wasm
-    ZigCallConvWrapper,
+    ZigCallConvWrapper(usize),
 }
 
 #[derive(Debug)]
@@ -274,6 +274,98 @@ impl<'a> WasmBackend<'a> {
         self.module.names.append_function(wasm_fn_index, name_bytes);
     }
 
+    /// Procs that are called from higher-order Zig builtins currently need a wrapper,
+    /// because the Zig compiler has bugs in its C calling convention for Wasm.
+    /// Whenever Zig fixes this, we should be able to remove the wrapper entirely.
+    pub fn build_zigcc_wrapper(&mut self, wrapper_idx: usize, inner_idx: usize) {
+        let ProcLookupData {
+            layout: proc_layout,
+            ..
+        } = self.proc_lookup[wrapper_idx];
+        let ProcLayout {
+            arguments: arg_layouts,
+            result,
+        } = proc_layout;
+
+        let ret_sym = self.create_symbol("##ret");
+        let ret_layout = WasmLayout::new(&result);
+        let is_stack_return = match ret_layout.return_method() {
+            ReturnMethod::Primitive(_) => false,
+            ReturnMethod::NoReturnValue => false,
+            ReturnMethod::WriteToPointerArg => {
+                // Return variable must be at index 0
+                self.storage
+                    .allocate(result, ret_sym, StoredValueKind::Parameter);
+                true
+            }
+        };
+
+        let mut arg_symbols = Vec::with_capacity_in(arg_layouts.len(), self.env.arena);
+        let mut frame_writes = Vec::with_capacity_in(arg_layouts.len() * 2, self.env.arena);
+
+        for (i, arg) in arg_layouts.iter().enumerate() {
+            let arg_name = format!("arg{}", i);
+            let symbol = self.create_symbol(&arg_name);
+            arg_symbols.push(symbol);
+            self.storage
+                .allocate_zigcc_arg(arg, symbol, &mut frame_writes);
+        }
+
+        if !is_stack_return {
+            // Local variables must come *after* the arguments
+            self.storage
+                .allocate(result, ret_sym, StoredValueKind::Variable);
+        }
+
+        // Write structs to the stack frame
+        if !frame_writes.is_empty() {
+            let frame_ptr = self.storage.create_anonymous_local(PTR_TYPE);
+            self.storage.stack_frame_pointer = Some(frame_ptr);
+
+            for (zig_arg, value_type, offset) in frame_writes.into_iter() {
+                self.code_builder.get_local(frame_ptr);
+                self.code_builder.get_local(zig_arg);
+                if value_type == ValueType::I32 {
+                    let align = Align::from_stack_offset(Align::Bytes4, offset);
+                    self.code_builder.i32_store(align, offset);
+                } else {
+                    let align = Align::from_stack_offset(Align::Bytes8, offset);
+                    self.code_builder.i64_store(align, offset);
+                }
+            }
+        }
+
+        // Call the wrapped inner function
+        let ProcLookupData { linker_index, .. } = self.proc_lookup[inner_idx];
+        let (param_types, ret_type) = self.storage.load_symbols_for_call(
+            self.env.arena,
+            &mut self.code_builder,
+            arg_symbols.into_bump_slice(),
+            ret_sym,
+            &ret_layout,
+            CallConv::C,
+        );
+        let wasm_fn_index = self.fn_index_offset + inner_idx 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_index, num_wasm_args, has_return_val);
+
+        // Setup & teardown the stack frame
+        self.code_builder.build_fn_header_and_footer(
+            &self.storage.local_types,
+            self.storage.stack_frame_size,
+            self.storage.stack_frame_pointer,
+        );
+
+        self.module.add_function_signature(Signature {
+            param_types: self.storage.arg_types.clone(),
+            ret_type,
+        });
+
+        self.reset();
+    }
+
     /**********************************************************
 
             STATEMENTS

compiler/gen_wasm/src/lib.rs

@@ -170,7 +170,7 @@ pub fn build_module_without_wrapper<'a>(
         env.arena,
     );
     let mut helper_iter = helper_procs.iter();
-    for source in sources {
+    for (idx, source) in sources.iter().enumerate() {
         use ProcSource::*;
         match source {
             Roc => { /* already generated */ }
@@ -179,9 +179,7 @@ pub fn build_module_without_wrapper<'a>(
                     backend.build_proc(proc);
                 }
             }
-            ZigCallConvWrapper => {
-                todo!("Generate Wasm wrapper to convert from Zig CC to CCC");
-            }
+            ZigCallConvWrapper(inner_idx) => backend.build_zigcc_wrapper(idx, *inner_idx),
         }
     }
 

compiler/gen_wasm/src/storage.rs

@@ -211,6 +211,86 @@ impl<'a> Storage<'a> {
         storage
     }
 
+    /// Allocate storage for an argument that will be passed from Zig code
+    /// (Zig *should* implement the C calling convention here, but it has bugs we need to work around)
+    pub fn allocate_zigcc_arg(
+        &mut self,
+        layout: &Layout<'a>,
+        symbol: Symbol,
+        frame_writes: &mut Vec<(LocalId, ValueType, u32)>,
+    ) {
+        let wasm_layout = WasmLayout::new(layout);
+        self.symbol_layouts.insert(symbol, *layout);
+
+        match wasm_layout {
+            WasmLayout::Primitive(value_type, size) => {
+                self.arg_types.push(value_type);
+                let storage = StoredValue::Local {
+                    local_id: self.get_next_local_id(),
+                    value_type,
+                    size,
+                };
+                self.symbol_storage_map.insert(symbol, storage);
+            }
+
+            WasmLayout::StackMemory {
+                size,
+                alignment_bytes,
+                format,
+            } => {
+                // Stack frame offset where we'll write the Zig argument value
+                let location = 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 rather than an argument.
+                    StackMemoryLocation::FrameOffset(0)
+                } else if size > 16 {
+                    // For larger structs, Zig passes a pointer to stack memory in the Zig caller. That suits us. Just pass it through.
+                    self.arg_types.push(PTR_TYPE);
+                    StackMemoryLocation::PointerArg(self.get_next_local_id())
+                } else {
+                    // Zig passes small structs as primitive values, but Roc expects a pointer to stack memory
+
+                    // Generate the Zig-compatible argument(s)
+                    let types: &[ValueType] = if size <= 4 {
+                        &[ValueType::I32]
+                    } else if size <= 8 {
+                        &[ValueType::I64]
+                    } else if size <= 12 {
+                        &[ValueType::I64, ValueType::I32]
+                    } else {
+                        &[ValueType::I64, ValueType::I64]
+                    };
+
+                    // Allocate space in the stack frame, so we can pass a pointer to Roc
+                    let mut offset: u32 =
+                        round_up_to_alignment!(self.stack_frame_size as u32, alignment_bytes);
+                    let loc = StackMemoryLocation::FrameOffset(offset);
+                    self.stack_frame_size = (offset + size) as i32;
+
+                    // Make a note of which writes we need to do
+                    // Note: We can't do the writes until after we know how many Wasm args we have.
+                    // We need a LocalId for the stack frame pointer and it must come after the args.
+                    for ty in types.iter() {
+                        let local_id = LocalId(self.arg_types.len() as u32);
+                        frame_writes.push((local_id, *ty, offset));
+                        offset += if *ty == ValueType::I32 { 4 } else { 8 };
+                    }
+
+                    loc
+                };
+
+                let storage = StoredValue::StackMemory {
+                    location,
+                    size,
+                    alignment_bytes,
+                    format,
+                };
+
+                self.symbol_storage_map.insert(symbol, storage);
+            }
+        }
+    }
+
     /// Get storage info for a given symbol
     pub fn get(&self, sym: &Symbol) -> &StoredValue {
         self.symbol_storage_map.get(sym).unwrap_or_else(|| {

compiler/gen_wasm/src/wasm_module/code_builder.rs

@@ -70,7 +70,7 @@ impl std::fmt::Debug for VmBlock<'_> {
 /// Rust representation matches Wasm encoding.
 /// It's an error to specify alignment higher than the "natural" alignment of the instruction
 #[repr(u8)]
-#[derive(Clone, Copy, Debug)]
+#[derive(Clone, Copy, Debug, PartialEq, PartialOrd)]
 pub enum Align {
     Bytes1 = 0,
     Bytes2 = 1,
@@ -78,6 +78,23 @@ pub enum Align {
     Bytes8 = 3,
 }
 
+impl Align {
+    /// Calculate the largest possible alignment for a load/store at a given stack frame offset
+    /// Assumes the stack frame is aligned to at least 8 bytes
+    pub fn from_stack_offset(max_align: Align, offset: u32) -> Align {
+        if (max_align == Align::Bytes8) && (offset & 7 == 0) {
+            return Align::Bytes8;
+        }
+        if (max_align >= Align::Bytes4) && (offset & 3 == 0) {
+            return Align::Bytes4;
+        }
+        if (max_align >= Align::Bytes2) && (offset & 1 == 0) {
+            return Align::Bytes2;
+        }
+        return Align::Bytes1;
+    }
+}
+
 impl From<u32> for Align {
     fn from(x: u32) -> Align {
         match x {