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,17 +194,9 @@ impl<'a> Storage<'a> {
         })
     }
 
-    /// 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.
-    pub fn load_symbols(&mut self, code_builder: &mut CodeBuilder, symbols: &[Symbol]) {
-        if code_builder.verify_stack_match(symbols) {
-            // The symbols were already at the top of the stack, do nothing!
-            // This should be quite common due to the structure of the Mono IR
-            return;
-        }
-        for sym in symbols.iter() {
-            let storage = self.get(sym).to_owned();
+    /// 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,
@@ -212,13 +204,12 @@ impl<'a> Storage<'a> {
                 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);
+                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,
+                            sym,
                             StoredValue::VirtualMachineStack {
                                 vm_state: next_vm_state,
                                 value_type,
@@ -231,7 +222,7 @@ impl<'a> Storage<'a> {
                         // it was not in a convenient position in the VM stack.
                         self.local_types.push(value_type);
                         self.symbol_storage_map.insert(
-                                *sym,
+                            sym,
                             StoredValue::Local {
                                 local_id: next_local_id,
                                 value_type,
@@ -247,7 +238,7 @@ impl<'a> Storage<'a> {
                 ..
             } => {
                 code_builder.get_local(local_id);
-                    code_builder.set_top_symbol(*sym);
+                code_builder.set_top_symbol(sym);
             }
 
             StoredValue::StackMemory {
@@ -255,10 +246,74 @@ impl<'a> Storage<'a> {
                 ..
             } => {
                 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);
                 }
+                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.
+    pub fn load_symbols(&mut self, code_builder: &mut CodeBuilder, symbols: &[Symbol]) {
+        if code_builder.verify_stack_match(symbols) {
+            // The symbols were already at the top of the stack, do nothing!
+            // This should be quite common due to the structure of the Mono IR
+            return;
+        }
+        for sym in symbols.iter() {
+            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);
+                }
+
+                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);
             }
         }
     }

compiler/test_gen/src/helpers/wasm.rs

@@ -9,7 +9,7 @@ use roc_can::builtins::builtin_defs_map;
 use roc_collections::all::{MutMap, MutSet};
 use roc_gen_wasm::MEMORY_NAME;
 
-use tempfile::tempdir;
+use tempfile::{TempDir, tempdir};
 
 const TEST_WRAPPER_NAME: &str = "test_wrapper";
 
@@ -118,39 +118,34 @@ pub fn helper_wasm<'a, T: Wasm32TestResult>(
     let mut module_bytes = std::vec::Vec::with_capacity(4096);
     wasm_module.serialize_mut(&mut module_bytes);
 
-    // for debugging (e.g. with wasm2wat or wasm-objdump)
-    if false {
-        use std::io::Write;
-
-        let mut hash_state = DefaultHasher::new();
-        src.hash(&mut hash_state);
-        let src_hash = hash_state.finish();
-
-        // Filename contains a hash of the Roc test source code. Helpful when comparing across commits.
-        let dir = "/tmp/roc/gen_wasm";
-        std::fs::create_dir_all(dir).unwrap();
-        let path = format!("{}/test-{:016x}.wasm", dir, src_hash);
-
-        // Print out filename (appears just after test name)
-        println!("dumping file {:?}", path);
-
-        match std::fs::File::create(path) {
-            Err(e) => eprintln!("Problem creating wasm debug file: {:?}", e),
-            Ok(mut file) => {
-                file.write_all(&module_bytes).unwrap();
-            }
-        }
-    }
-
     // now, do wasmer stuff
 
     use wasmer::{Instance, Module, Store};
 
     let store = Store::default();
 
+    // Keep the final .wasm file for debugging with wasm-objdump or wasm2wat
+    const DEBUG_WASM_FILE: bool = true;
+
     let wasmer_module = {
-        let dir = tempdir().unwrap();
-        let dirpath = dir.path();
+        let tmp_dir: TempDir; // directory for normal test runs, deleted when dropped
+        let debug_dir: String; // persistent directory for debugging
+
+        let dirpath: &Path = 
+            if DEBUG_WASM_FILE {
+                // Directory name based on a hash of the Roc source
+                let mut hash_state = DefaultHasher::new();
+                src.hash(&mut hash_state);
+                let src_hash = hash_state.finish();
+                debug_dir = format!("/tmp/roc/gen_wasm/{:016x}", src_hash);
+                std::fs::create_dir_all(&debug_dir).unwrap();
+                println!("Debug command:\n\twasm-objdump -sdx {}/final.wasm", &debug_dir);
+                Path::new(&debug_dir)
+            } else {
+                tmp_dir = tempdir().unwrap();
+                tmp_dir.path()
+            };
+
         let final_wasm_file = dirpath.join("final.wasm");
         let app_o_file = dirpath.join("app.o");
 

compiler/test_gen/src/wasm_str.rs

@@ -296,31 +296,23 @@ fn long_str_literal() {
     );
 }
 
-// #[test]
-// fn small_str_concat_empty_first_arg() {
-//     assert_llvm_evals_to!(
-//         r#"Str.concat "" "JJJJJJJJJJJJJJJ""#,
-//         [
-//             0x4a,
-//             0x4a,
-//             0x4a,
-//             0x4a,
-//             0x4a,
-//             0x4a,
-//             0x4a,
-//             0x4a,
-//             0x4a,
-//             0x4a,
-//             0x4a,
-//             0x4a,
-//             0x4a,
-//             0x4a,
-//             0x4a,
-//             0b1000_1111
-//         ],
-//         [u8; 16]
-//     );
-// }
+#[test]
+fn small_str_concat_empty_first_arg() {
+    assert_evals_to!(
+        r#"Str.concat "" "JJJJJJJ""#,
+        [
+            0x4a,
+            0x4a,
+            0x4a,
+            0x4a,
+            0x4a,
+            0x4a,
+            0x4a,
+            0b1000_0111
+        ],
+        [u8; 8]
+    );
+}
 
 // #[test]
 // fn small_str_concat_empty_second_arg() {