fix bug in closure argument unpacking

compiler/gen/src/llvm/build.rs

@@ -1477,10 +1477,9 @@ pub fn build_exp_stmt<'a, 'ctx, 'env>(
         }
         Dec(symbol, cont) => {
             let (value, layout) = load_symbol_and_layout(env, scope, symbol);
-            let layout = layout.clone();
 
             if layout.contains_refcounted() {
-                decrement_refcount_layout(env, parent, layout_ids, value, &layout);
+                decrement_refcount_layout(env, parent, layout_ids, value, layout);
             }
 
             build_exp_stmt(env, layout_ids, scope, parent, cont)
@@ -2323,33 +2322,14 @@ pub fn build_proc<'a, 'ctx, 'env>(
     for (arg_val, (layout, arg_symbol)) in fn_val.get_param_iter().zip(args) {
         set_name(arg_val, arg_symbol.ident_string(&env.interns));
 
-        // the closure argument (if any) comes in as an opaque sequence of bytes.
-        // we need to cast that to the specific closure data layout that the body expects
-        let value = if let Symbol::ARG_CLOSURE = *arg_symbol {
-            // generate a caller function (to be used by the host)
-            // build_closure_caller(env, fn_val);
-            // builder.position_at_end(entry);
-
-            // blindly trust that there is a layout available for the closure data
-            let layout = proc.closure_data_layout.clone().unwrap();
-
-            // cast the input into the type that the body expects
-            let closure_data_type =
-                basic_type_from_layout(env.arena, env.context, &layout, env.ptr_bytes);
-
-            cast_basic_basic(env.builder, arg_val, closure_data_type)
-        } else {
-            arg_val
-        };
-
         let alloca = create_entry_block_alloca(
             env,
             fn_val,
-            value.get_type(),
+            arg_val.get_type(),
             arg_symbol.ident_string(&env.interns),
         );
 
-        builder.build_store(alloca, value);
+        builder.build_store(alloca, arg_val);
 
         scope.insert(*arg_symbol, (layout.clone(), alloca));
     }

compiler/mono/src/ir.rs

@@ -1618,46 +1618,109 @@ fn specialize_external<'a>(
             ret_layout,
         } => {
             // unpack the closure symbols, if any
-            if let CapturedSymbols::Captured(captured) = captured_symbols {
+            match (opt_closure_layout.clone(), captured_symbols) {
-                let mut layouts = Vec::with_capacity_in(captured.len(), env.arena);
+                (Some(closure_layout), CapturedSymbols::Captured(captured)) => {
+                    debug_assert!(!captured.is_empty());
 
-                for (_, variable) in captured.iter() {
+                    let wrapped = closure_layout.get_wrapped();
-                    let layout = layout_cache.from_var(env.arena, *variable, env.subs)?;
-                    layouts.push(layout);
-                }
 
-                let field_layouts = layouts.into_bump_slice();
+                    let internal_layout = closure_layout.internal_layout();
 
-                let wrapped = match &opt_closure_layout {
+                    match internal_layout {
-                    Some(x) => x.get_wrapped(),
+                        Layout::Union(_) => {
-                    None => unreachable!("symbols are captured, so this must be a closure"),
+                            // here we rely on the fact that a union in a closure would be stored in a one-element record
-                };
+                            // a closure layout that is a union must be a of the shape `Closure1 ... | Closure2 ...`
+                            let tag_layout = closure_layout.as_named_layout(proc_name);
 
-                for (index, (symbol, variable)) in captured.iter().enumerate() {
+                            match tag_layout {
-                    // layout is cached anyway, re-using the one found above leads to
+                                Layout::Struct(field_layouts) => {
-                    // issues (combining by-ref and by-move in pattern match
+                                    // NOTE closure unions do not store the tag!
-                    let layout = layout_cache.from_var(env.arena, *variable, env.subs)?;
+                                    let field_layouts = &field_layouts[1..];
-
+                                    for (index, (symbol, _variable)) in captured.iter().enumerate()
-                    // if the symbol has a layout that is dropped from data structures (e.g. `{}`)
+                                    {
-                    // then regenerate the symbol here. The value may not be present in the closure
+                                        let expr = Expr::AccessAtIndex {
-                    // data struct
-                    let expr = {
-                        if layout.is_dropped_because_empty() {
-                            Expr::Struct(&[])
-                        } else {
-                            Expr::AccessAtIndex {
                                             index: index as _,
                                             field_layouts,
                                             structure: Symbol::ARG_CLOSURE,
                                             wrapped,
-                            }
-                        }
                                         };
 
-                    specialized_body =
+                                        let layout = field_layouts[index].clone();
-                        Stmt::Let(*symbol, expr, layout, env.arena.alloc(specialized_body));
+
+                                        specialized_body = Stmt::Let(
+                                            *symbol,
+                                            expr,
+                                            layout,
+                                            env.arena.alloc(specialized_body),
+                                        );
                                     }
                                 }
+                                _ => unreachable!("closure tag must store a struct"),
+                            }
+                        }
+                        Layout::Struct(field_layouts) => {
+                            debug_assert_eq!(
+                                captured.len(),
+                                field_layouts.len(),
+                                "{:?} captures {:?} but has layout {:?}",
+                                proc_name,
+                                &captured,
+                                &field_layouts
+                            );
+
+                            if field_layouts.len() > 1 {
+                                for (index, (symbol, _variable)) in captured.iter().enumerate() {
+                                    let expr = Expr::AccessAtIndex {
+                                        index: index as _,
+                                        field_layouts,
+                                        structure: Symbol::ARG_CLOSURE,
+                                        wrapped,
+                                    };
+
+                                    let layout = field_layouts[index].clone();
+
+                                    specialized_body = Stmt::Let(
+                                        *symbol,
+                                        expr,
+                                        layout,
+                                        env.arena.alloc(specialized_body),
+                                    );
+                                }
+                            } else {
+                                let symbol = captured[0].0;
+
+                                substitute_in_exprs(
+                                    env.arena,
+                                    &mut specialized_body,
+                                    symbol,
+                                    Symbol::ARG_CLOSURE,
+                                );
+                            }
+                        }
+                        _other => {
+                            // the closure argument is not a union or a record
+                            debug_assert_eq!(
+                                captured.len(),
+                                1,
+                                "mismatch {:?} captures {:?}",
+                                proc_name,
+                                &captured
+                            );
+
+                            let symbol = captured[0].0;
+
+                            substitute_in_exprs(
+                                env.arena,
+                                &mut specialized_body,
+                                symbol,
+                                Symbol::ARG_CLOSURE,
+                            );
+                        }
+                    }
+                }
+                (None, CapturedSymbols::None) => {}
+                _ => unreachable!("to closure or not to closure?"),
+            }
 
             // reset subs, so we don't get type errors when specializing for a different signature
             layout_cache.rollback_to(cache_snapshot);

compiler/mono/src/layout.rs

@@ -227,6 +227,10 @@ impl<'a> ClosureLayout<'a> {
         }
     }
 
+    pub fn internal_layout(&self) -> Layout<'a> {
+        self.layout.clone()
+    }
+
     pub fn build_closure_data(
         &self,
         original: Symbol,

compiler/mono/tests/test_mono.rs

@@ -2278,8 +2278,7 @@ mod test_mono {
                     ret Test.3;
 
                 procedure Test.3 (Test.11, #Attr.12):
-                    let Test.2 = Index 0 #Attr.12;
+                    ret #Attr.12;
-                    ret Test.2;
 
                 procedure Test.0 ():
                     let Test.10 = Struct {};
@@ -2328,8 +2327,7 @@ mod test_mono {
                     ret Test.10;
 
                 procedure Test.3 (Test.9, #Attr.12):
-                    let Test.2 = Index 0 #Attr.12;
+                    ret #Attr.12;
-                    ret Test.2;
 
                 procedure Test.0 ():
                     let Test.8 = Struct {};

examples/effect/Main.roc

@@ -3,7 +3,7 @@ app "effect-example"
     imports [ base.Task.{ Task, after } ]
     provides [ main ] to base
 
-main : Task.Task {} I64
+main : Task.Task {} F64
 main =
     Task.after (Task.putLine "foo") \{} ->
         Task.putLine "bar"

examples/effect/thing/platform-dir/Pkg-Config.roc

@@ -11,5 +11,5 @@ platform folkertdev/foo
             getLine : Effect Str
         }
 
-mainForHost : Task.Task {} I64 as Fx
+mainForHost : Task.Task {} F64 as Fx
 mainForHost = main