gen-wasm: box captured data before passing it to a higher-order lowlevel

Previously, when it existed, captured data would always be represented
as a struct, and hence implicitly boxed. That meant that passing
captured data would always pass a pointer. However, now, captured data
can live unwrapped. This poses a challenge for the higher-order
lowlevels, which always expect captures data to be passed as an opaque
pointer. As such, always box captured data.

It's possible to optimize this so that only unwrapped captures data
needs to be boxed; however, that is not attempted in this patch.

crates/compiler/gen_wasm/src/backend.rs

@@ -517,7 +517,7 @@ impl<'a> WasmBackend<'a> {
 
         // Load all the arguments for the inner function
         for (i, wrapper_arg) in wrapper_arg_layouts.iter().enumerate() {
-            let is_closure_data = i == 0; // Skip closure data (first for wrapper, last for inner)
+            let is_closure_data = i == 0; // Skip closure data (first for wrapper, last for inner). We'll handle it below.
             let is_return_pointer = i == wrapper_arg_layouts.len() - 1; // Skip return pointer (may not be an arg for inner. And if it is, swaps from end to start)
             if is_closure_data || is_return_pointer {
                 continue;
@@ -540,7 +540,16 @@ impl<'a> WasmBackend<'a> {
         // If the inner function has closure data, it's the last arg of the inner fn
         let closure_data_layout = wrapper_arg_layouts[0];
         if closure_data_layout.stack_size(TARGET_INFO) > 0 {
+            // The closure data exists, and will have been passed in to the wrapper as boxed.
+            let inner_closure_data_layout = match closure_data_layout {
+                Layout::Boxed(inner) => inner,
+                other => internal_error!(
+                    "Expected a boxed layout for wrapped closure data, got {:?}",
+                    other
+                ),
+            };
             self.code_builder.get_local(LocalId(0));
+            self.dereference_boxed_value(inner_closure_data_layout);
         }
 
         // Call the wrapped inner function
@@ -1870,7 +1879,7 @@ impl<'a> WasmBackend<'a> {
     /// If the data size is known at compile time, pass it in comptime_data_size.
     /// If size is only known at runtime, push *data* size to the VM stack first.
     /// Leaves the *data* address on the VM stack
-    fn allocate_with_refcount(
+    pub fn allocate_with_refcount(
         &mut self,
         comptime_data_size: Option<u32>,
         alignment_bytes: u32,

crates/compiler/gen_wasm/src/low_level.rs

@@ -2117,6 +2117,17 @@ pub fn call_higher_order_lowlevel<'a>(
         ..
     } = passed_function;
 
+    // The zig lowlevel builtins expect the passed functions' closure data to always
+    // be sent as an opaque pointer. On the Roc side, however, we need to call the passed function
+    // with the Roc representation of the closure data. There are three possible cases for that
+    // representation:
+    //
+    // 1. The closure data is a struct
+    // 2. The closure data is an unwrapped value
+    // 3. There is no closure data
+    //
+    // To uniformly deal with the first two cases, always box these layouts before calling the
+    // builtin; the wrapper around the passed function will unbox them.
     let (closure_data_layout, closure_data_exists) =
         match backend.storage.symbol_layouts[captured_environment] {
             Layout::LambdaSet(lambda_set) => {
@@ -2135,6 +2146,49 @@ pub fn call_higher_order_lowlevel<'a>(
             x => internal_error!("Closure data has an invalid layout\n{:?}", x),
         };
 
+    let (wrapped_captured_environment, wrapped_captures_layout) = if closure_data_exists {
+        // If there is closure data, make sure we box it before passing it to the external builtin
+        // impl.
+        let boxed_sym = backend.create_symbol("boxed_captures");
+        let boxed_captures_layout = Layout::Boxed(backend.env.arena.alloc(closure_data_layout));
+
+        // create a local variable for the box
+        let boxed_storage = backend.storage.allocate_var(
+            boxed_captures_layout,
+            boxed_sym,
+            crate::storage::StoredVarKind::Variable,
+        );
+        let boxed_local_id = match backend.storage.ensure_value_has_local(
+            &mut backend.code_builder,
+            boxed_sym,
+            boxed_storage.clone(),
+        ) {
+            StoredValue::Local { local_id, .. } => local_id,
+            _ => internal_error!("Expected a local"),
+        };
+
+        // allocate heap memory and load its data address onto the value stack
+        let (size, alignment) = closure_data_layout.stack_size_and_alignment(TARGET_INFO);
+        backend.allocate_with_refcount(Some(size), alignment, 1);
+
+        // store the pointer value from the value stack into the local variable
+        backend.code_builder.set_local(boxed_local_id);
+
+        // copy the argument to the pointer address
+        backend.storage.copy_value_to_memory(
+            &mut backend.code_builder,
+            boxed_local_id,
+            0,
+            *captured_environment,
+        );
+
+        (boxed_sym, boxed_captures_layout)
+    } else {
+        // If we don't capture anything, pass along the captured environment as-is - the wrapper
+        // function will take care not to unbox this.
+        (*captured_environment, closure_data_layout)
+    };
+
     // 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 helper_proc_source = {
@@ -2168,7 +2222,7 @@ pub fn call_higher_order_lowlevel<'a>(
             argument_layouts.len()
         };
 
-        wrapper_arg_layouts.push(closure_data_layout);
+        wrapper_arg_layouts.push(wrapped_captures_layout);
         wrapper_arg_layouts.extend(
             argument_layouts
                 .iter()
@@ -2208,7 +2262,7 @@ pub fn call_higher_order_lowlevel<'a>(
             .get_refcount_fn_index(Layout::Builtin(Builtin::Int(IntWidth::I32)), HelperOp::Inc);
         backend.get_fn_ptr(inc_fn)
     } else {
-        let inc_fn = backend.get_refcount_fn_index(closure_data_layout, HelperOp::Inc);
+        let inc_fn = backend.get_refcount_fn_index(wrapped_captures_layout, HelperOp::Inc);
         backend.get_fn_ptr(inc_fn)
     };
 
@@ -2222,7 +2276,7 @@ pub fn call_higher_order_lowlevel<'a>(
             wrapper_fn_ptr,
             inc_fn_ptr,
             closure_data_exists,
-            *captured_environment,
+            wrapped_captured_environment,
             *owns_captured_environment,
         ),
 
@@ -2235,7 +2289,7 @@ pub fn call_higher_order_lowlevel<'a>(
             wrapper_fn_ptr,
             inc_fn_ptr,
             closure_data_exists,
-            *captured_environment,
+            wrapped_captured_environment,
             *owns_captured_environment,
         ),
 
@@ -2248,7 +2302,7 @@ pub fn call_higher_order_lowlevel<'a>(
             wrapper_fn_ptr,
             inc_fn_ptr,
             closure_data_exists,
-            *captured_environment,
+            wrapped_captured_environment,
             *owns_captured_environment,
         ),
 
@@ -2261,7 +2315,7 @@ pub fn call_higher_order_lowlevel<'a>(
             wrapper_fn_ptr,
             inc_fn_ptr,
             closure_data_exists,
-            *captured_environment,
+            wrapped_captured_environment,
             *owns_captured_environment,
         ),
 
@@ -2284,7 +2338,9 @@ pub fn call_higher_order_lowlevel<'a>(
             backend.storage.load_symbol_zig(cb, *xs);
             cb.i32_const(wrapper_fn_ptr);
             if closure_data_exists {
-                backend.storage.load_symbols(cb, &[*captured_environment]);
+                backend
+                    .storage
+                    .load_symbols(cb, &[wrapped_captured_environment]);
             } else {
                 // load_symbols assumes that a zero-size arg should be eliminated in code gen,
                 // but that's a specialization that our Zig code doesn't have! Pass a null pointer.