Add support for Zig/LLVM "fast calling convention"

compiler/gen_wasm/src/backend.rs

@@ -444,6 +444,13 @@ impl<'a> WasmBackend<'a> {
 
                 Ok(())
             }
+
+            Stmt::Refcounting(_modify, following) => {
+                // TODO: actually deal with refcounting. For hello world, we just skipped it.
+                self.build_stmt(following, ret_layout)?;
+                Ok(())
+            }
+
             x => Err(format!("statement not yet implemented: {:?}", x)),
         }
     }
@@ -558,7 +565,8 @@ impl<'a> WasmBackend<'a> {
         arguments: &'a [Symbol],
         return_layout: WasmLayout,
     ) -> Result<(), String> {
-        self.storage.load_symbols(&mut self.code_builder, arguments);
+        self.storage
+            .load_symbols_fastcc(&mut self.code_builder, arguments, &return_layout);
 
         let build_result = decode_low_level(
             &mut self.code_builder,
@@ -764,11 +772,27 @@ impl<'a> WasmBackend<'a> {
             },
 
             None => {
-                let mut param_types = Vec::with_capacity_in(arguments.len(), self.env.arena);
-                param_types.extend(arguments.iter().map(|a| self.storage.get(a).value_type()));
+                let mut param_types = Vec::with_capacity_in(1 + arguments.len(), self.env.arena);
+
+                let ret_type = if ret_layout.is_stack_memory() {
+                    param_types.push(ValueType::I32);
+                    None
+                } else {
+                    Some(ret_layout.value_type())
+                };
+
+                for arg in arguments {
+                    param_types.push(match self.storage.get(arg) {
+                        StoredValue::StackMemory { size, .. } if *size > 4 && *size <= 8 => {
+                            ValueType::I64
+                        }
+                        _ => ValueType::I32,
+                    });
+                }
+
                 let signature_index = self.module.types.insert(Signature {
                     param_types,
-                    ret_type: Some(ret_layout.value_type()), // TODO: handle builtins with no return value
+                    ret_type,
                 });
 
                 let import_index = self.module.import.entries.len() as u32;

compiler/gen_wasm/src/lib.rs

@@ -7,6 +7,7 @@ pub mod wasm_module;
 use bumpalo::{self, collections::Vec, Bump};
 
 use roc_collections::all::{MutMap, MutSet};
+use roc_module::low_level::LowLevel;
 use roc_module::symbol::{Interns, Symbol};
 use roc_mono::ir::{Proc, ProcLayout};
 use roc_mono::layout::LayoutIds;
@@ -47,34 +48,50 @@ pub fn build_module_help<'a>(
     procedures: MutMap<(Symbol, ProcLayout<'a>), Proc<'a>>,
 ) -> Result<WasmModule<'a>, String> {
     let mut layout_ids = LayoutIds::default();
-    let mut proc_symbols = Vec::with_capacity_in(procedures.len(), env.arena);
+    let mut generated_procs = Vec::with_capacity_in(procedures.len(), env.arena);
+    let mut generated_symbols = Vec::with_capacity_in(procedures.len(), env.arena);
     let mut linker_symbols = Vec::with_capacity_in(procedures.len() * 2, env.arena);
-    let mut exports = Vec::with_capacity_in(procedures.len(), env.arena);
+    let mut exports = Vec::with_capacity_in(4, env.arena);
 
-    // Collect the symbols & names for the procedures
-    for (i, (sym, layout)) in procedures.keys().enumerate() {
-        proc_symbols.push(*sym);
+    // Collect the symbols & names for the procedures,
+    // and filter out procs we're going to inline
+    let mut fn_index = 0;
+    for ((sym, layout), proc) in procedures.into_iter() {
+        if LowLevel::from_inlined_wrapper(sym).is_some() {
+            continue;
+        }
+        generated_procs.push(proc);
+        generated_symbols.push(sym);
 
         let fn_name = layout_ids
-            .get_toplevel(*sym, layout)
-            .to_symbol_string(*sym, &env.interns);
+            .get_toplevel(sym, &layout)
+            .to_symbol_string(sym, &env.interns);
 
-        if env.exposed_to_host.contains(sym) {
+        if env.exposed_to_host.contains(&sym) {
             exports.push(Export {
                 name: fn_name.clone(),
                 ty: ExportType::Func,
-                index: i as u32,
+                index: fn_index as u32,
             });
         }
 
-        let linker_sym = SymInfo::for_function(i as u32, fn_name);
+        let linker_sym = SymInfo::for_function(fn_index as u32, fn_name);
         linker_symbols.push(linker_sym);
+
+        fn_index += 1;
     }
 
-    // Main loop: Build the Wasm module
+    // Build the Wasm module
     let (mut module, linker_symbols) = {
-        let mut backend = WasmBackend::new(env, layout_ids, proc_symbols, linker_symbols, exports);
-        for ((sym, _), proc) in procedures.into_iter() {
+        let mut backend = WasmBackend::new(
+            env,
+            layout_ids,
+            generated_symbols.clone(),
+            linker_symbols,
+            exports,
+        );
+
+        for (proc, sym) in generated_procs.into_iter().zip(generated_symbols) {
             backend.build_proc(proc, sym)?;
         }
         (backend.module, backend.linker_symbols)

compiler/gen_wasm/src/storage.rs

@@ -5,7 +5,7 @@ use roc_collections::all::MutMap;
 use roc_module::symbol::Symbol;
 
 use crate::layout::WasmLayout;
-use crate::wasm_module::{CodeBuilder, LocalId, ValueType, VmSymbolState};
+use crate::wasm_module::{Align, CodeBuilder, LocalId, ValueType, VmSymbolState};
 use crate::{copy_memory, round_up_to_alignment, CopyMemoryConfig, PTR_SIZE, PTR_TYPE};
 
 pub enum StoredValueKind {
@@ -194,6 +194,67 @@ impl<'a> Storage<'a> {
         })
     }
 
+    /// Load a symbol using the C Calling Convention
+    fn load_symbol_ccc(&mut self, code_builder: &mut CodeBuilder, sym: Symbol) {
+        let storage = self.get(&sym).to_owned();
+        match storage {
+            StoredValue::VirtualMachineStack {
+                vm_state,
+                value_type,
+                size,
+            } => {
+                let next_local_id = self.get_next_local_id();
+                let maybe_next_vm_state = code_builder.load_symbol(sym, vm_state, next_local_id);
+                match maybe_next_vm_state {
+                    // The act of loading the value changed the VM state, so update it
+                    Some(next_vm_state) => {
+                        self.symbol_storage_map.insert(
+                            sym,
+                            StoredValue::VirtualMachineStack {
+                                vm_state: next_vm_state,
+                                value_type,
+                                size,
+                            },
+                        );
+                    }
+                    None => {
+                        // Loading the value required creating a new local, because
+                        // it was not in a convenient position in the VM stack.
+                        self.local_types.push(value_type);
+                        self.symbol_storage_map.insert(
+                            sym,
+                            StoredValue::Local {
+                                local_id: next_local_id,
+                                value_type,
+                                size,
+                            },
+                        );
+                    }
+                }
+            }
+            StoredValue::Local { local_id, .. }
+            | StoredValue::StackMemory {
+                location: StackMemoryLocation::PointerArg(local_id),
+                ..
+            } => {
+                code_builder.get_local(local_id);
+                code_builder.set_top_symbol(sym);
+            }
+
+            StoredValue::StackMemory {
+                location: StackMemoryLocation::FrameOffset(offset),
+                ..
+            } => {
+                code_builder.get_local(self.stack_frame_pointer.unwrap());
+                if offset != 0 {
+                    code_builder.i32_const(offset as i32);
+                    code_builder.i32_add();
+                }
+                code_builder.set_top_symbol(sym);
+            }
+        }
+    }
+
     /// Load symbols to the top of the VM stack
     /// Avoid calling this method in a loop with one symbol at a time! It will work,
     /// but it generates very inefficient Wasm code.
@@ -204,61 +265,55 @@ impl<'a> Storage<'a> {
             return;
         }
         for sym in symbols.iter() {
-            let storage = self.get(sym).to_owned();
-            match storage {
-                StoredValue::VirtualMachineStack {
-                    vm_state,
-                    value_type,
-                    size,
-                } => {
-                    let next_local_id = self.get_next_local_id();
-                    let maybe_next_vm_state =
-                        code_builder.load_symbol(*sym, vm_state, next_local_id);
-                    match maybe_next_vm_state {
-                        // The act of loading the value changed the VM state, so update it
-                        Some(next_vm_state) => {
-                            self.symbol_storage_map.insert(
-                                *sym,
-                                StoredValue::VirtualMachineStack {
-                                    vm_state: next_vm_state,
-                                    value_type,
-                                    size,
-                                },
-                            );
-                        }
-                        None => {
-                            // Loading the value required creating a new local, because
-                            // it was not in a convenient position in the VM stack.
-                            self.local_types.push(value_type);
-                            self.symbol_storage_map.insert(
-                                *sym,
-                                StoredValue::Local {
-                                    local_id: next_local_id,
-                                    value_type,
-                                    size,
-                                },
-                            );
-                        }
-                    }
-                }
-                StoredValue::Local { local_id, .. }
-                | StoredValue::StackMemory {
-                    location: StackMemoryLocation::PointerArg(local_id),
-                    ..
-                } => {
-                    code_builder.get_local(local_id);
-                    code_builder.set_top_symbol(*sym);
+            self.load_symbol_ccc(code_builder, *sym);
+        }
+    }
+
+    /// Load symbols in a way compatible with LLVM's "fast calling convention"
+    /// A bug in Zig means it always uses this for Wasm even when we specify C calling convention.
+    /// It squashes small structs into primitive values where possible, avoiding stack memory
+    /// in favour of CPU registers (or VM stack values, which eventually become CPU registers).
+    /// We need to convert some of our structs from our internal C-like representation to work with Zig.
+    /// Why not just always use the fastcc representation? Because of non-Zig platforms.
+    pub fn load_symbols_fastcc(
+        &mut self,
+        code_builder: &mut CodeBuilder,
+        symbols: &[Symbol],
+        return_layout: &WasmLayout,
+    ) {
+        if return_layout.is_stack_memory() {
+            // Load the address where the return value should be written
+            self.load_symbol_ccc(code_builder, symbols[0]);
+        };
+
+        for sym in symbols {
+            if let StoredValue::StackMemory {
+                location,
+                size,
+                alignment_bytes,
+            } = self.get(sym)
+            {
+                if *size == 0 {
+                    unimplemented!("Passing zero-sized values is not implemented yet");
+                } else if *size > 8 {
+                    return self.load_symbol_ccc(code_builder, *sym);
                 }
 
-                StoredValue::StackMemory {
-                    location: StackMemoryLocation::FrameOffset(offset),
-                    ..
-                } => {
-                    code_builder.get_local(self.stack_frame_pointer.unwrap());
-                    code_builder.i32_const(offset as i32);
-                    code_builder.i32_add();
-                    code_builder.set_top_symbol(*sym);
+                let (local_id, offset) = location.local_and_offset(self.stack_frame_pointer);
+                code_builder.get_local(local_id);
+                let align = Align::from(*alignment_bytes);
+
+                if *size == 1 {
+                    code_builder.i32_load8_u(align, offset);
+                } else if *size == 2 {
+                    code_builder.i32_load16_u(align, offset);
+                } else if *size <= 4 {
+                    code_builder.i32_load(align, offset);
+                } else {
+                    code_builder.i64_load(align, offset);
                 }
+            } else {
+                self.load_symbol_ccc(code_builder, *sym);
             }
         }
     }