Handle recursive variables in building lambda set representations

crates/compiler/mono/src/layout.rs

@@ -10,7 +10,8 @@ use roc_problem::can::RuntimeError;
 use roc_target::{PtrWidth, TargetInfo};
 use roc_types::num::NumericRange;
 use roc_types::subs::{
-    self, Content, FlatType, Label, RecordFields, Subs, UnionTags, UnsortedUnionLabels, Variable,
+    self, Content, FlatType, Label, OptVariable, RecordFields, Subs, UnionTags,
+    UnsortedUnionLabels, Variable,
 };
 use roc_types::types::{gather_fields_unsorted_iter, RecordField, RecordFieldsError};
 use std::cmp::Ordering;
@@ -866,7 +867,14 @@ impl<'a> LambdaSet<'a> {
             .set
             .iter()
             .find(|(name, layouts)| comparator(*name, layouts))
-            .expect("no lambda set found");
+            .unwrap_or_else(|| {
+                internal_error!(
+                    "no lambda set found for ({:?}, {:#?}): {:#?}",
+                    function_symbol,
+                    captures_layouts,
+                    self
+                )
+            });
 
         LambdaName {
             name: *name,
@@ -971,7 +979,7 @@ impl<'a> LambdaSet<'a> {
         target_info: TargetInfo,
     ) -> Result<Self, LayoutProblem> {
         match resolve_lambda_set(subs, closure_var) {
-            ResolvedLambdaSet::Set(mut lambdas) => {
+            ResolvedLambdaSet::Set(mut lambdas, opt_recursion_var) => {
                 // sort the tags; make sure ordering stays intact!
                 lambdas.sort_by_key(|(sym, _)| *sym);
 
@@ -992,6 +1000,9 @@ impl<'a> LambdaSet<'a> {
                         seen: Vec::new_in(arena),
                         target_info,
                     };
+                    if let Some(rec_var) = opt_recursion_var.into_variable() {
+                        env.insert_seen(rec_var);
+                    }
 
                     for var in variables {
                         arguments.push(Layout::from_var(&mut env, *var)?);
@@ -1048,6 +1059,7 @@ impl<'a> LambdaSet<'a> {
                     arena,
                     subs,
                     set_with_variables,
+                    opt_recursion_var.into_variable(),
                     target_info,
                 ));
 
@@ -1071,10 +1083,11 @@ impl<'a> LambdaSet<'a> {
         arena: &'a Bump,
         subs: &Subs,
         tags: std::vec::Vec<(Symbol, std::vec::Vec<Variable>)>,
+        opt_rec_var: Option<Variable>,
         target_info: TargetInfo,
     ) -> Layout<'a> {
         // otherwise, this is a closure with a payload
-        let variant = union_sorted_tags_help(arena, tags, None, subs, target_info);
+        let variant = union_sorted_tags_help(arena, tags, opt_rec_var, subs, target_info);
 
         use UnionVariant::*;
         match variant {
@@ -1108,7 +1121,20 @@ impl<'a> LambdaSet<'a> {
                         Layout::Union(UnionLayout::NonRecursive(tag_arguments.into_bump_slice()))
                     }
 
-                    _ => panic!("handle recursive layouts"),
+                    Recursive {
+                        sorted_tag_layouts: tags,
+                    } => {
+                        debug_assert!(tags.len() > 1);
+
+                        let mut tag_arguments = Vec::with_capacity_in(tags.len(), arena);
+
+                        for (_, tag_args) in tags.iter() {
+                            tag_arguments.push(&tag_args[0..]);
+                        }
+                        Layout::Union(UnionLayout::Recursive(tag_arguments.into_bump_slice()))
+                    }
+
+                    layout => panic!("handle recursive layout: {:?}", layout),
                 }
             }
         }
@@ -1130,7 +1156,10 @@ impl<'a> LambdaSet<'a> {
 }
 
 enum ResolvedLambdaSet {
-    Set(std::vec::Vec<(Symbol, std::vec::Vec<Variable>)>),
+    Set(
+        std::vec::Vec<(Symbol, std::vec::Vec<Variable>)>,
+        OptVariable,
+    ),
     /// TODO: figure out if this can happen in a correct program, or is the result of a bug in our
     /// compiler. See https://github.com/rtfeldman/roc/issues/3163.
     Unbound,
@@ -1142,7 +1171,7 @@ fn resolve_lambda_set(subs: &Subs, mut var: Variable) -> ResolvedLambdaSet {
         match subs.get_content_without_compacting(var) {
             Content::LambdaSet(subs::LambdaSet {
                 solved,
-                recursion_var: _,
+                recursion_var,
                 unspecialized,
                 ambient_function: _,
             }) => {
@@ -1153,7 +1182,7 @@ fn resolve_lambda_set(subs: &Subs, mut var: Variable) -> ResolvedLambdaSet {
                     subs.uls_of_var
                 );
                 roc_types::pretty_print::push_union(subs, solved, &mut set);
-                return ResolvedLambdaSet::Set(set);
+                return ResolvedLambdaSet::Set(set, *recursion_var);
             }
             Content::RecursionVar { structure, .. } => {
                 var = *structure;
@@ -2130,10 +2159,14 @@ fn layout_from_flat_type<'a>(
             }
         }
         Func(_, closure_var, _) => {
-            let lambda_set =
-                LambdaSet::from_var(env.arena, env.subs, closure_var, env.target_info)?;
+            if env.is_seen(closure_var) {
+                Ok(Layout::RecursivePointer)
+            } else {
+                let lambda_set =
+                    LambdaSet::from_var(env.arena, env.subs, closure_var, env.target_info)?;
 
-            Ok(Layout::LambdaSet(lambda_set))
+                Ok(Layout::LambdaSet(lambda_set))
+            }
         }
         Record(fields, ext_var) => {
             // extract any values from the ext_var
@@ -2686,6 +2719,9 @@ where
         seen: Vec::new_in(arena),
         target_info,
     };
+    if let Some(rec_var) = opt_rec_var {
+        env.insert_seen(rec_var);
+    }
 
     match tags_vec.len() {
         0 => {

crates/compiler/test_mono/src/tests.rs

@@ -1922,3 +1922,22 @@ fn issue_3669() {
         "#
     )
 }
+
+#[mono_test]
+fn issue_3444() {
+    indoc!(
+        r#"
+        app "test" provides [main] to "./platform"
+
+        combine = \a, b -> (\x -> x |> a |> b )
+        const = \x -> (\_y -> x)
+
+        list = []
+
+        res : Str -> Str
+        res = List.walk list (const "z") (\c1, c2 -> combine c1 c2)
+
+        main = res "hello"
+        "#
+    )
+}