Merge pull request #4754 from roc-lang/i4733

Properly handle imported top levels that appear in a non-unary lambda set

crates/compiler/mono/src/ir.rs

@@ -8218,49 +8218,71 @@ fn specialize_symbol<'a>(
                         Err(e) => return_on_layout_error_help!(env, e, "specialize_symbol"),
                     };
 
-                    if procs.is_imported_module_thunk(original) {
-                        let layout = match raw {
-                            RawFunctionLayout::ZeroArgumentThunk(layout) => layout,
-                            RawFunctionLayout::Function(_, lambda_set, _) => {
-                                Layout::LambdaSet(lambda_set)
-                            }
-                        };
+                    match raw {
+                        RawFunctionLayout::Function(_, lambda_set, _)
+                            if !procs.is_imported_module_thunk(original) =>
+                        {
+                            let lambda_name =
+                                find_lambda_name(env, layout_cache, lambda_set, original, &[]);
 
-                        let raw = RawFunctionLayout::ZeroArgumentThunk(layout);
-                        let top_level = ProcLayout::from_raw(
-                            env.arena,
-                            &layout_cache.interner,
-                            raw,
-                            CapturesNiche::no_niche(),
-                        );
+                            debug_assert!(
+                                lambda_name.no_captures(),
+                                "imported functions are top-level and should never capture"
+                            );
 
-                        procs.insert_passed_by_name(
-                            env,
-                            arg_var,
-                            LambdaName::no_niche(original),
-                            top_level,
-                            layout_cache,
-                        );
+                            let function_ptr_layout = ProcLayout::from_raw(
+                                env.arena,
+                                &layout_cache.interner,
+                                raw,
+                                lambda_name.captures_niche(),
+                            );
+                            procs.insert_passed_by_name(
+                                env,
+                                arg_var,
+                                lambda_name,
+                                function_ptr_layout,
+                                layout_cache,
+                            );
 
-                        force_thunk(env, original, layout, assign_to, env.arena.alloc(result))
-                    } else {
-                        // Imported symbol, so it must have no captures niche (since
-                        // top-levels can't capture)
-                        let top_level = ProcLayout::from_raw(
-                            env.arena,
-                            &layout_cache.interner,
-                            raw,
-                            CapturesNiche::no_niche(),
-                        );
-                        procs.insert_passed_by_name(
-                            env,
-                            arg_var,
-                            LambdaName::no_niche(original),
-                            top_level,
-                            layout_cache,
-                        );
+                            construct_closure_data(
+                                env,
+                                procs,
+                                layout_cache,
+                                lambda_set,
+                                lambda_name,
+                                &[],
+                                assign_to,
+                                env.arena.alloc(result),
+                            )
+                        }
+                        _ => {
+                            // This is an imported ZAT that returns either a value, or the closure
+                            // data for a lambda set.
+                            let layout = match raw {
+                                RawFunctionLayout::ZeroArgumentThunk(layout) => layout,
+                                RawFunctionLayout::Function(_, lambda_set, _) => {
+                                    Layout::LambdaSet(lambda_set)
+                                }
+                            };
 
-                        let_empty_struct(assign_to, env.arena.alloc(result))
+                            let raw = RawFunctionLayout::ZeroArgumentThunk(layout);
+                            let top_level = ProcLayout::from_raw(
+                                env.arena,
+                                &layout_cache.interner,
+                                raw,
+                                CapturesNiche::no_niche(),
+                            );
+
+                            procs.insert_passed_by_name(
+                                env,
+                                arg_var,
+                                LambdaName::no_niche(original),
+                                top_level,
+                                layout_cache,
+                            );
+
+                            force_thunk(env, original, layout, assign_to, env.arena.alloc(result))
+                        }
                     }
                 }
 

crates/compiler/test_gen/src/gen_tags.rs

@@ -2083,3 +2083,53 @@ fn unify_types_with_fixed_fixpoints_outside_fixing_region() {
         RocStr
     );
 }
+
+#[test]
+#[cfg(any(feature = "gen-llvm", feature = "gen-wasm"))]
+fn lambda_set_with_imported_toplevels_issue_4733() {
+    assert_evals_to!(
+        indoc!(
+            r#"
+            app "test" provides [main] to "./platform"
+
+            fn = \s ->
+                instr = if s == "*" then (Op Num.mul) else (Op Num.add)
+
+                Op op = instr
+
+                \a -> op a a
+
+            main = ((fn "*") 3) * ((fn "+") 5)
+            "#
+        ),
+        90,
+        i64
+    );
+}
+
+#[test]
+#[cfg(any(feature = "gen-llvm", feature = "gen-wasm"))]
+fn non_unary_union_with_lambda_set_with_imported_toplevels_issue_4733() {
+    assert_evals_to!(
+        indoc!(
+            r#"
+            app "test" provides [main] to "./platform"
+
+            fn = \s ->
+                instr =
+                    if s == "*" then (Op Num.mul)
+                    else if s == "+" then (Op Num.add)
+                    else Noop
+
+                when instr is
+                    Op op -> (\a -> op a a)
+                    _ -> (\a -> a)
+
+
+            main = ((fn "*") 3) * ((fn "+") 5)
+            "#
+        ),
+        90,
+        i64
+    );
+}

crates/compiler/test_mono/generated/issue_4749.txt

@@ -28,7 +28,7 @@ procedure Decode.24 (Decode.95, Decode.114, Decode.97):
     ret Decode.127;
 
 procedure Decode.25 (Decode.98, Decode.99):
-    let Decode.126 : {} = Struct {};
+    let Decode.126 : {} = CallByName Json.41;
     let Decode.125 : {List U8, [C {}, C Str]} = CallByName Decode.24 Decode.98 Decode.126 Decode.99;
     ret Decode.125;
 

crates/compiler/test_mono/generated/lambda_set_with_imported_toplevels_issue_4733.txt

@@ -0,0 +1,49 @@
+procedure Bool.11 (#Attr.2, #Attr.3):
+    let Bool.23 : Int1 = lowlevel Eq #Attr.2 #Attr.3;
+    ret Bool.23;
+
+procedure Bool.2 ():
+    let Bool.24 : Int1 = true;
+    ret Bool.24;
+
+procedure Num.19 (#Attr.2, #Attr.3):
+    let Num.256 : U64 = lowlevel NumAdd #Attr.2 #Attr.3;
+    ret Num.256;
+
+procedure Num.21 (#Attr.2, #Attr.3):
+    let Num.257 : U64 = lowlevel NumMul #Attr.2 #Attr.3;
+    ret Num.257;
+
+procedure Test.0 (Test.8):
+    let Test.23 : Int1 = CallByName Bool.2;
+    if Test.23 then
+        let Test.24 : Int1 = true;
+        ret Test.24;
+    else
+        let Test.22 : Int1 = false;
+        ret Test.22;
+
+procedure Test.5 (Test.6, Test.2):
+    joinpoint Test.19 Test.18:
+        ret Test.18;
+    in
+    switch Test.2:
+        case 0:
+            let Test.20 : U64 = CallByName Num.19 Test.6 Test.6;
+            jump Test.19 Test.20;
+    
+        default:
+            let Test.21 : U64 = CallByName Num.21 Test.6 Test.6;
+            jump Test.19 Test.21;
+    
+
+procedure Test.7 ():
+    let Test.13 : U64 = 3i64;
+    let Test.15 : {} = Struct {};
+    let Test.14 : Int1 = CallByName Test.0 Test.15;
+    let Test.11 : U64 = CallByName Test.5 Test.13 Test.14;
+    let Test.12 : U64 = 9i64;
+    let Test.10 : Int1 = CallByName Bool.11 Test.11 Test.12;
+    expect Test.10;
+    let Test.9 : {} = Struct {};
+    ret Test.9;

crates/compiler/test_mono/src/tests.rs

@@ -2173,3 +2173,22 @@ fn issue_4749() {
         "###
     )
 }
+
+#[mono_test(mode = "test", no_check)]
+fn lambda_set_with_imported_toplevels_issue_4733() {
+    indoc!(
+        r###"
+        interface Test exposes [] imports []
+
+        fn = \{} ->
+            instr : [ Op (U64, U64 -> U64) ]
+            instr = if Bool.true then (Op Num.mul) else (Op Num.add)
+
+            Op op = instr
+
+            \a -> op a a
+
+        expect ((fn {}) 3) == 9
+        "###
+    )
+}