wasm: ProcLayout for higher order wrapper fns

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(usize),
+    HigherOrderWrapper(usize),
 }
 
 #[derive(Debug)]
@@ -121,21 +121,25 @@ impl<'a> WasmBackend<'a> {
         self.helper_proc_gen.take_procs()
     }
 
-    fn register_helper_proc(&mut self, new_proc_info: (Symbol, ProcLayout<'a>)) {
+    pub fn register_helper_proc(
-        let (new_proc_sym, new_proc_layout) = new_proc_info;
+        &mut self,
+        symbol: Symbol,
+        layout: ProcLayout<'a>,
+        source: ProcSource,
+    ) -> u32 {
         let wasm_fn_index = self.proc_lookup.len() as u32;
         let linker_sym_index = self.module.linking.symbol_table.len() as u32;
 
         let name = self
             .layout_ids
-            .get_toplevel(new_proc_sym, &new_proc_layout)
+            .get_toplevel(symbol, &layout)
-            .to_symbol_string(new_proc_sym, self.interns);
+            .to_symbol_string(symbol, self.interns);
 
         self.proc_lookup.push(ProcLookupData {
-            name: new_proc_sym,
+            name: symbol,
-            layout: new_proc_layout,
+            layout,
             linker_index: linker_sym_index,
-            source: ProcSource::Helper,
+            source,
         });
 
         let linker_symbol = SymInfo::Function(WasmObjectSymbol::Defined {
@@ -144,6 +148,8 @@ impl<'a> WasmBackend<'a> {
             name,
         });
         self.module.linking.symbol_table.push(linker_symbol);
+
+        wasm_fn_index
     }
 
     pub fn finalize(self) -> (WasmModule<'a>, Vec<'a, u32>) {
@@ -164,7 +170,7 @@ impl<'a> WasmBackend<'a> {
     pub fn register_symbol_debug_names(&self) {}
 
     /// Create an IR Symbol for an anonymous value (such as ListLiteral)
-    fn create_symbol(&mut self, debug_name: &str) -> Symbol {
+    pub fn create_symbol(&mut self, debug_name: &str) -> Symbol {
         let ident_ids = self
             .interns
             .all_ident_ids
@@ -652,8 +658,8 @@ impl<'a> WasmBackend<'a> {
         }
 
         // If any new specializations were created, register their symbol data
-        for spec in new_specializations.into_iter() {
+        for (spec_sym, spec_layout) in new_specializations.into_iter() {
-            self.register_helper_proc(spec);
+            self.register_helper_proc(spec_sym, spec_layout, ProcSource::Helper);
         }
 
         self.stmt(rc_stmt);
@@ -1044,8 +1050,8 @@ impl<'a> WasmBackend<'a> {
             .call_specialized_equals(ident_ids, arg_layout, arguments);
 
         // If any new specializations were created, register their symbol data
-        for spec in new_specializations.into_iter() {
+        for (spec_sym, spec_layout) in new_specializations.into_iter() {
-            self.register_helper_proc(spec);
+            self.register_helper_proc(spec_sym, spec_layout, ProcSource::Helper);
         }
 
         // Generate Wasm code for the IR call expression
@@ -1529,8 +1535,8 @@ impl<'a> WasmBackend<'a> {
             .call_reset_refcount(ident_ids, layout, argument);
 
         // If any new specializations were created, register their symbol data
-        for spec in new_specializations.into_iter() {
+        for (spec_sym, spec_layout) in new_specializations.into_iter() {
-            self.register_helper_proc(spec);
+            self.register_helper_proc(spec_sym, spec_layout, ProcSource::Helper);
         }
 
         // Generate Wasm code for the IR call expression

compiler/gen_wasm/src/lib.rs

@@ -179,7 +179,7 @@ pub fn build_module_without_wrapper<'a>(
                     backend.build_proc(proc);
                 }
             }
-            ZigCallConvWrapper(inner_idx) => backend.build_zigcc_wrapper(idx, *inner_idx),
+            HigherOrderWrapper(inner_idx) => backend.build_zigcc_wrapper(idx, *inner_idx),
         }
     }
 

compiler/gen_wasm/src/low_level.rs

@@ -3,11 +3,11 @@ use roc_builtins::bitcode::{self, FloatWidth, IntWidth};
 use roc_error_macros::internal_error;
 use roc_module::low_level::LowLevel;
 use roc_module::symbol::Symbol;
-use roc_mono::ir::{HigherOrderLowLevel, PassedFunction};
+use roc_mono::ir::{HigherOrderLowLevel, PassedFunction, ProcLayout};
 use roc_mono::layout::{Builtin, Layout, UnionLayout};
 use roc_mono::low_level::HigherOrder;
 
-use crate::backend::WasmBackend;
+use crate::backend::{ProcLookupData, ProcSource, WasmBackend};
 use crate::layout::CallConv;
 use crate::layout::{StackMemoryFormat, WasmLayout};
 use crate::storage::{StackMemoryLocation, Storage, StoredValue};
@@ -937,7 +937,7 @@ pub fn call_higher_order_lowlevel<'a>(
     backend: &mut WasmBackend<'a>,
     return_sym: Symbol,
     return_layout: &Layout<'a>,
-    higher_order: &HigherOrderLowLevel<'a>,
+    higher_order: &'a HigherOrderLowLevel<'a>,
 ) {
     use HigherOrder::*;
 
@@ -948,6 +948,7 @@ pub fn call_higher_order_lowlevel<'a>(
     } = higher_order;
 
     let PassedFunction {
+        name: fn_name,
         argument_layouts,
         return_layout: result_layout,
         owns_captured_environment,
@@ -955,6 +956,56 @@ pub fn call_higher_order_lowlevel<'a>(
         ..
     } = passed_function;
 
+    // We create a wrapper around the passed function, which just unboxes the arguments.
+    // This allows Zig builtins to have a generic pointer-based interface.
+    let source = {
+        let passed_proc_layout = ProcLayout {
+            arguments: argument_layouts,
+            result: *result_layout,
+        };
+        let passed_proc_index = backend
+            .proc_lookup
+            .iter()
+            .position(|ProcLookupData { name, layout, .. }| {
+                name == fn_name && layout == &passed_proc_layout
+            })
+            .unwrap();
+        ProcSource::HigherOrderWrapper(passed_proc_index)
+    };
+    let wrapper_sym = backend.create_symbol(&format!("#wrap#{:?}", fn_name));
+    let wrapper_layout = {
+        let mut wrapper_arg_layouts: Vec<Layout<'a>> =
+            Vec::with_capacity_in(argument_layouts.len() + 1, backend.env.arena);
+
+        let closure_data_layout = backend.storage.symbol_layouts[captured_environment];
+
+        let last_arg_is_closure_data: bool = match closure_data_layout {
+            Layout::LambdaSet(lambda_set) => lambda_set.runtime_representation() != Layout::UNIT,
+            x => internal_error!("Closure data has an invalid layout\n{:?}", x),
+        };
+        let n_non_closure_args = if last_arg_is_closure_data {
+            argument_layouts.len() - 1
+        } else {
+            argument_layouts.len()
+        };
+
+        wrapper_arg_layouts.push(closure_data_layout);
+        wrapper_arg_layouts.extend(
+            argument_layouts
+                .iter()
+                .take(n_non_closure_args)
+                .map(Layout::Boxed),
+        );
+        wrapper_arg_layouts.push(Layout::Boxed(result_layout));
+
+        ProcLayout {
+            arguments: wrapper_arg_layouts.into_bump_slice(),
+            result: Layout::UNIT,
+        }
+    };
+
+    let wrapper_fn_idx = backend.register_helper_proc(wrapper_sym, wrapper_layout, source);
+
     /*
         TODO indirect calls
             passed_function
@@ -969,10 +1020,20 @@ pub fn call_higher_order_lowlevel<'a>(
                 append to (or find in) backend.proc_lookup
                 calculate its index
                 insert index into elements table
+
+        questions about the closure data
+        - is the first arg always closure data? no, last
+        - Do I need a special case where if the closure layout is Unit then no argument exists on the Proc?
+            (Otherwise I suppose there would always an arg with Unit layout and I haven't seen that)
+        - if the closure exists then does it reliably appear first or last in the argument list on the Proc?
+            - or is it absent from that
+        - if there is no closure data then I need a dummy incrementor for Unit?
+            yes
+
+
     */
 
-    let caller_fn_idx: i32 = todo!();
+    let inc_fn_idx: i32 = 123;
-    let inc_fn_idx: i32 = todo!();
 
     match op {
         // List.map : List elem_old, (elem_old -> elem_new) -> List elem_new
@@ -994,7 +1055,7 @@ pub fn call_higher_order_lowlevel<'a>(
             // Load return pointer & argument values
             backend.storage.load_symbols(cb, &[return_sym]);
             backend.storage.load_symbol_zig(cb, *xs);
-            cb.i32_const(caller_fn_idx);
+            cb.i32_const(wrapper_fn_idx as i32);
             backend.storage.load_symbols(cb, &[*captured_environment]);
             cb.i32_const(inc_fn_idx);
             cb.i32_const(*owns_captured_environment as i32);