Inline polymorphic calls at usage sites

This is a bit.. ugly, or at least seems suboptimal, but I can't think of
a better way to do it currently aside from demanding a uniform
representation, which we probably don't want to do.

Another option is something like the defunctionalization we perform
today, except also capturing potential uses of nested functions in the
closure tag of an encompassing lambda. So for example,

```
f = \x -> \y -> 1
```

would now record a lambdaset with the data `[Test.f
[TypeOfInnerClos1]]`, where `TypeOfInnerClos1` is e.g.
`[Test.f.innerClos1 I8, Test.f.innerClos1 I16]`, symbolizing that the
inner closure may be specialized to take an I8 or I16. Then at the time
that we create the capture set for `f`, we create a tag noting what
specialization should be used for the inner closure, and apply the
current defunctionalization algorithm. So effectively, the type of the
inner closure becomes a capture.

I'm not sure if this is any better, or if it has more problems.
@folkertdev any thoughts?

Closes #2322

compiler/mono/src/ir.rs

@@ -4179,9 +4179,6 @@ pub fn with_hole<'a>(
                         LocalFunction(_) => {
                             unreachable!("if this was known to be a function, we would not be here")
                         }
-                        UnspecializedExpr(_) => {
-                            unreachable!("if this was known to be an unspecialized expression, we would not be here")
-                        }
                         Imported(thunk_name) => {
                             debug_assert!(procs.is_imported_module_thunk(thunk_name));
 
@@ -4240,6 +4237,49 @@ pub fn with_hole<'a>(
                                 unreachable!("calling a non-closure layout")
                             }
                         },
+                        UnspecializedExpr(symbol) => match full_layout {
+                            RawFunctionLayout::Function(arg_layouts, lambda_set, ret_layout) => {
+                                let closure_data_symbol = env.unique_symbol();
+
+                                result = match_on_lambda_set(
+                                    env,
+                                    lambda_set,
+                                    closure_data_symbol,
+                                    arg_symbols,
+                                    arg_layouts,
+                                    ret_layout,
+                                    assigned,
+                                    hole,
+                                );
+
+                                let (lambda_expr, lambda_expr_var) =
+                                    procs.partial_exprs.get(symbol).unwrap();
+
+                                let snapshot = env.subs.snapshot();
+                                let cache_snapshot = layout_cache.snapshot();
+                                let _unified = roc_unify::unify::unify(
+                                    env.subs,
+                                    fn_var,
+                                    lambda_expr_var,
+                                    roc_unify::unify::Mode::Eq,
+                                );
+
+                                result = with_hole(
+                                    env,
+                                    lambda_expr.clone(),
+                                    fn_var,
+                                    procs,
+                                    layout_cache,
+                                    closure_data_symbol,
+                                    env.arena.alloc(result),
+                                );
+                                env.subs.rollback_to(snapshot);
+                                layout_cache.rollback_to(cache_snapshot);
+                            }
+                            RawFunctionLayout::ZeroArgumentThunk(_) => {
+                                unreachable!("calling a non-closure layout")
+                            }
+                        },
                         NotASymbol => {
                             // the expression is not a symbol. That means it's an expression
                             // evaluating to a function value.
@@ -5184,6 +5224,7 @@ fn is_literal_like(expr: &roc_can::expr::Expr) -> bool {
             | ZeroArgumentTag { .. }
             | Tag { .. }
             | Record { .. }
+            | Call(..)
     )
 }
 

compiler/test_gen/src/gen_primitives.rs

@@ -3207,3 +3207,20 @@ fn polymophic_expression_captured_inside_closure() {
         u8
     );
 }
+
+#[test]
+#[cfg(any(feature = "gen-llvm"))]
+fn issue_2322() {
+    assert_evals_to!(
+        indoc!(
+            r#"
+            double = \x -> x * 2
+            doubleBind = \x -> (\_ -> double x)
+            doubleThree = doubleBind 3
+            doubleThree {}
+            "#
+        ),
+        6,
+        u8
+    )
+}

compiler/test_mono/generated/aliased_polymorphic_closure.txt

@@ -0,0 +1,17 @@
+procedure Test.2 (Test.6, #Attr.12):
+    let Test.1 : Builtin(Int(U8)) = StructAtIndex 0 #Attr.12;
+    let Test.11 : LambdaSet(LambdaSet { set: [( Test.4, [Builtin(Int(U8))])], representation: Struct([Builtin(Int(U8))]) }) = Struct {Test.1};
+    ret Test.11;
+
+procedure Test.4 (Test.5, #Attr.12):
+    let Test.1 : Builtin(Int(U8)) = StructAtIndex 0 #Attr.12;
+    ret Test.1;
+
+procedure Test.0 ():
+    let Test.1 : Builtin(Int(U8)) = 1i64;
+    let Test.8 : Struct([]) = Struct {};
+    let Test.10 : Struct([]) = Struct {};
+    let Test.14 : LambdaSet(LambdaSet { set: [( Test.2, [Builtin(Int(U8))])], representation: Struct([Builtin(Int(U8))]) }) = Struct {Test.1};
+    let Test.9 : LambdaSet(LambdaSet { set: [( Test.4, [Builtin(Int(U8))])], representation: Struct([Builtin(Int(U8))]) }) = CallByName Test.2 Test.10 Test.14;
+    let Test.7 : Builtin(Int(U8)) = CallByName Test.4 Test.8 Test.9;
+    ret Test.7;

compiler/test_mono/generated/closure_in_list.txt

@@ -1,21 +1,21 @@
 procedure List.7 (#Attr.2):
-    let Test.7 : Builtin(Int(U64)) = lowlevel ListLen #Attr.2;
-    ret Test.7;
+    let Test.8 : Builtin(Int(U64)) = lowlevel ListLen #Attr.2;
+    ret Test.8;
 
 procedure Test.1 (Test.5):
     let Test.2 : Builtin(Int(I64)) = 41i64;
-    let Test.11 : LambdaSet(LambdaSet { set: [( Test.3, [Builtin(Int(I64))])], representation: Struct([Builtin(Int(I64))]) }) = Struct {Test.2};
-    let Test.10 : Builtin(List(LambdaSet(LambdaSet { set: [( Test.3, [Builtin(Int(I64))])], representation: Struct([Builtin(Int(I64))]) }))) = Array [Test.11];
-    ret Test.10;
+    let Test.12 : LambdaSet(LambdaSet { set: [( Test.3, [Builtin(Int(I64))])], representation: Struct([Builtin(Int(I64))]) }) = Struct {Test.2};
+    let Test.11 : Builtin(List(LambdaSet(LambdaSet { set: [( Test.3, [Builtin(Int(I64))])], representation: Struct([Builtin(Int(I64))]) }))) = Array [Test.12];
+    ret Test.11;
 
-procedure Test.3 (Test.9, #Attr.12):
+procedure Test.3 (Test.10, #Attr.12):
     let Test.2 : Builtin(Int(I64)) = StructAtIndex 0 #Attr.12;
     let Test.2 : Builtin(Int(I64)) = 41i64;
     ret Test.2;
 
 procedure Test.0 ():
-    let Test.8 : Struct([]) = Struct {};
-    let Test.4 : Builtin(List(LambdaSet(LambdaSet { set: [( Test.3, [Builtin(Int(I64))])], representation: Struct([Builtin(Int(I64))]) }))) = CallByName Test.1 Test.8;
-    let Test.6 : Builtin(Int(U64)) = CallByName List.7 Test.4;
-    dec Test.4;
+    let Test.9 : Struct([]) = Struct {};
+    let Test.7 : Builtin(List(LambdaSet(LambdaSet { set: [( Test.3, [Builtin(Int(I64))])], representation: Struct([Builtin(Int(I64))]) }))) = CallByName Test.1 Test.9;
+    let Test.6 : Builtin(Int(U64)) = CallByName List.7 Test.7;
+    dec Test.7;
     ret Test.6;

compiler/test_mono/generated/nested_closure.txt

@@ -3,14 +3,14 @@ procedure Test.1 (Test.5):
     let Test.3 : LambdaSet(LambdaSet { set: [( Test.3, [Builtin(Int(I64))])], representation: Struct([Builtin(Int(I64))]) }) = Struct {Test.2};
     ret Test.3;
 
-procedure Test.3 (Test.9, #Attr.12):
+procedure Test.3 (Test.10, #Attr.12):
     let Test.2 : Builtin(Int(I64)) = StructAtIndex 0 #Attr.12;
     let Test.2 : Builtin(Int(I64)) = 42i64;
     ret Test.2;
 
 procedure Test.0 ():
-    let Test.8 : Struct([]) = Struct {};
-    let Test.4 : LambdaSet(LambdaSet { set: [( Test.3, [Builtin(Int(I64))])], representation: Struct([Builtin(Int(I64))]) }) = CallByName Test.1 Test.8;
     let Test.7 : Struct([]) = Struct {};
-    let Test.6 : Builtin(Int(I64)) = CallByName Test.3 Test.7 Test.4;
+    let Test.9 : Struct([]) = Struct {};
+    let Test.8 : LambdaSet(LambdaSet { set: [( Test.3, [Builtin(Int(I64))])], representation: Struct([Builtin(Int(I64))]) }) = CallByName Test.1 Test.9;
+    let Test.6 : Builtin(Int(I64)) = CallByName Test.3 Test.7 Test.8;
     ret Test.6;

compiler/test_mono/generated/optional_when.txt

@@ -1,42 +1,42 @@
 procedure Num.24 (#Attr.2, #Attr.3):
-    let Test.17 : Builtin(Int(I64)) = lowlevel NumMul #Attr.2 #Attr.3;
-    ret Test.17;
+    let Test.18 : Builtin(Int(I64)) = lowlevel NumMul #Attr.2 #Attr.3;
+    ret Test.18;
 
 procedure Test.1 (Test.6):
-    let Test.22 : Builtin(Bool) = StructAtIndex 1 Test.6;
-    let Test.23 : Builtin(Bool) = false;
-    let Test.24 : Builtin(Bool) = lowlevel Eq Test.23 Test.22;
-    if Test.24 then
-        let Test.8 : Builtin(Int(I64)) = StructAtIndex 0 Test.6;
-        ret Test.8;
-    else
-        let Test.10 : Builtin(Int(I64)) = StructAtIndex 0 Test.6;
-        ret Test.10;
-
-procedure Test.1 (Test.6):
-    let Test.33 : Builtin(Bool) = false;
-    let Test.34 : Builtin(Bool) = lowlevel Eq Test.33 Test.6;
-    if Test.34 then
+    let Test.22 : Builtin(Bool) = false;
+    let Test.23 : Builtin(Bool) = lowlevel Eq Test.22 Test.6;
+    if Test.23 then
         let Test.8 : Builtin(Int(I64)) = 3i64;
         ret Test.8;
     else
         let Test.10 : Builtin(Int(I64)) = 5i64;
         ret Test.10;
 
+procedure Test.1 (Test.6):
+    let Test.30 : Builtin(Bool) = StructAtIndex 1 Test.6;
+    let Test.31 : Builtin(Bool) = false;
+    let Test.32 : Builtin(Bool) = lowlevel Eq Test.31 Test.30;
+    if Test.32 then
+        let Test.8 : Builtin(Int(I64)) = StructAtIndex 0 Test.6;
+        ret Test.8;
+    else
+        let Test.10 : Builtin(Int(I64)) = StructAtIndex 0 Test.6;
+        ret Test.10;
+
 procedure Test.0 ():
-    let Test.37 : Builtin(Bool) = true;
-    let Test.5 : Builtin(Int(I64)) = CallByName Test.1 Test.37;
-    let Test.35 : Builtin(Bool) = false;
-    let Test.3 : Builtin(Int(I64)) = CallByName Test.1 Test.35;
-    let Test.28 : Builtin(Int(I64)) = 11i64;
-    let Test.29 : Builtin(Bool) = true;
-    let Test.27 : Struct([Builtin(Int(I64)), Builtin(Bool)]) = Struct {Test.28, Test.29};
-    let Test.4 : Builtin(Int(I64)) = CallByName Test.1 Test.27;
-    let Test.25 : Builtin(Int(I64)) = 7i64;
-    let Test.26 : Builtin(Bool) = false;
-    let Test.19 : Struct([Builtin(Int(I64)), Builtin(Bool)]) = Struct {Test.25, Test.26};
-    let Test.2 : Builtin(Int(I64)) = CallByName Test.1 Test.19;
-    let Test.18 : Builtin(Int(I64)) = CallByName Num.24 Test.2 Test.3;
-    let Test.16 : Builtin(Int(I64)) = CallByName Num.24 Test.18 Test.4;
-    let Test.15 : Builtin(Int(I64)) = CallByName Num.24 Test.16 Test.5;
+    let Test.40 : Builtin(Int(I64)) = 7i64;
+    let Test.41 : Builtin(Bool) = false;
+    let Test.39 : Struct([Builtin(Int(I64)), Builtin(Bool)]) = Struct {Test.40, Test.41};
+    let Test.35 : Builtin(Int(I64)) = CallByName Test.1 Test.39;
+    let Test.38 : Builtin(Bool) = false;
+    let Test.36 : Builtin(Int(I64)) = CallByName Test.1 Test.38;
+    let Test.25 : Builtin(Int(I64)) = CallByName Num.24 Test.35 Test.36;
+    let Test.33 : Builtin(Int(I64)) = 11i64;
+    let Test.34 : Builtin(Bool) = true;
+    let Test.27 : Struct([Builtin(Int(I64)), Builtin(Bool)]) = Struct {Test.33, Test.34};
+    let Test.26 : Builtin(Int(I64)) = CallByName Test.1 Test.27;
+    let Test.16 : Builtin(Int(I64)) = CallByName Num.24 Test.25 Test.26;
+    let Test.24 : Builtin(Bool) = true;
+    let Test.17 : Builtin(Int(I64)) = CallByName Test.1 Test.24;
+    let Test.15 : Builtin(Int(I64)) = CallByName Num.24 Test.16 Test.17;
     ret Test.15;

compiler/test_mono/src/tests.rs

@@ -1238,6 +1238,19 @@ fn monomorphized_applied_tag() {
     )
 }
 
+#[mono_test]
+fn aliased_polymorphic_closure() {
+    indoc!(
+        r#"
+        n : U8
+        n = 1
+        f = \{} -> (\a -> n)
+        g = f {}
+        g {}
+        "#
+    )
+}
+
 // #[ignore]
 // #[mono_test]
 // fn static_str_closure() {