Attach ambient function vars to lambda sets

crates/compiler/types/src/pretty_print.rs

@@ -342,6 +342,7 @@ fn find_names_needed(
             solved,
             recursion_var,
             unspecialized,
+            ambient_function: _,
         }) => {
             for slice_index in solved.variables() {
                 let slice = subs[slice_index];
@@ -721,6 +722,7 @@ fn write_content<'a>(
             solved,
             recursion_var,
             unspecialized,
+            ambient_function: _,
         }) => {
             debug_assert!(env.debug.print_lambda_sets);
 

crates/compiler/types/src/solved_types.rs

@@ -187,6 +187,7 @@ pub fn to_type(
             Type::ClosureTag {
                 name: *name,
                 captures: new_args,
+                ambient_function: var_store.fresh(),
             }
         }
         FunctionOrTagUnion(tag_name, symbol, ext) => {

crates/compiler/types/src/subs.rs

@@ -15,8 +15,7 @@ use crate::unification_table::{self, UnificationTable};
 // if your changes cause this number to go down, great!
 // please change it to the lower number.
 // if it went up, maybe check that the change is really required
-roc_error_macros::assert_sizeof_all!(Descriptor, 5 * 8);
-roc_error_macros::assert_sizeof_all!(Content, 3 * 8 + 4);
+roc_error_macros::assert_sizeof_all!(Descriptor, 5 * 8 + 4);
 roc_error_macros::assert_sizeof_all!(FlatType, 3 * 8);
 roc_error_macros::assert_sizeof_all!(UnionTags, 12);
 roc_error_macros::assert_sizeof_all!(RecordFields, 2 * 8);
@@ -809,6 +808,7 @@ fn subs_fmt_content(this: &Content, subs: &Subs, f: &mut fmt::Formatter) -> fmt:
             solved,
             recursion_var,
             unspecialized,
+            ambient_function: ambient_function_var,
         }) => {
             write!(f, "LambdaSet([")?;
 
@@ -839,7 +839,7 @@ fn subs_fmt_content(this: &Content, subs: &Subs, f: &mut fmt::Formatter) -> fmt:
                     region
                 )?;
             }
-            write!(f, ")")
+            write!(f, ", ^<{:?}>)", ambient_function_var)
         }
         Content::RangedNumber(range) => {
             write!(f, "RangedNumber( {:?})", range)
@@ -1920,6 +1920,7 @@ impl Subs {
         recursive: Variable,
         solved_lambdas: UnionLambdas,
         unspecialized_lambdas: SubsSlice<Uls>,
+        ambient_function_var: Variable,
     ) {
         let (rec_var, new_tags) = self.mark_union_recursive_help(recursive, solved_lambdas);
 
@@ -1927,6 +1928,7 @@ impl Subs {
             solved: new_tags,
             recursion_var: OptVariable::from(rec_var),
             unspecialized: unspecialized_lambdas,
+            ambient_function: ambient_function_var,
         });
 
         self.set_content(recursive, new_lambda_set);
@@ -2167,10 +2169,11 @@ impl From<Content> for Descriptor {
     }
 }
 
-roc_error_macros::assert_sizeof_all!(Content, 3 * 8 + 4);
+roc_error_macros::assert_sizeof_all!(Content, 4 * 8);
 roc_error_macros::assert_sizeof_all!((Symbol, AliasVariables, Variable), 2 * 8 + 4);
 roc_error_macros::assert_sizeof_all!(AliasVariables, 8);
 roc_error_macros::assert_sizeof_all!(FlatType, 3 * 8);
+roc_error_macros::assert_sizeof_all!(LambdaSet, 3 * 8 + 4);
 
 roc_error_macros::assert_sizeof_aarch64!((Variable, Option<Lowercase>), 4 * 8);
 roc_error_macros::assert_sizeof_wasm!((Variable, Option<Lowercase>), 4 * 4);
@@ -2244,6 +2247,12 @@ pub struct LambdaSet {
     pub recursion_var: OptVariable,
     /// Lambdas we won't know until an ability specialization is resolved.
     pub unspecialized: SubsSlice<Uls>,
+
+    /// Backlink to the function wrapping this lambda set.
+    /// This should never be unified against when unifying a lambda set; that would evidently
+    /// introduce an infinite unification.
+    /// This is used for the ambient lambda set unification algorithm.
+    pub ambient_function: Variable,
 }
 
 #[derive(Clone, Copy, Debug, Default)]
@@ -3135,6 +3144,7 @@ fn occurs(
                 solved,
                 recursion_var,
                 unspecialized: _,
+                ambient_function: _,
             }) => {
                 let mut new_seen = seen.to_owned();
                 new_seen.push(root_var);
@@ -3318,6 +3328,7 @@ fn explicit_substitute(
                 solved,
                 recursion_var,
                 unspecialized,
+                ambient_function: ambient_function_var,
             }) => {
                 // NOTE recursion_var is not substituted, verify that this is correct!
                 let new_solved = explicit_substitute_union(subs, from, to, solved, seen);
@@ -3332,6 +3343,7 @@ fn explicit_substitute(
                         solved: new_solved,
                         recursion_var,
                         unspecialized,
+                        ambient_function: ambient_function_var,
                     }),
                 );
 
@@ -3439,6 +3451,7 @@ fn get_var_names(
                 solved,
                 recursion_var,
                 unspecialized,
+                ambient_function: _,
             }) => {
                 let taken_names = get_var_names_union(subs, solved, taken_names);
                 let mut taken_names = match recursion_var.into_variable() {
@@ -4156,10 +4169,12 @@ impl StorageSubs {
                 solved,
                 recursion_var,
                 unspecialized,
+                ambient_function: ambient_function_var,
             }) => LambdaSet(self::LambdaSet {
                 solved: Self::offset_lambda_set(offsets, *solved),
                 recursion_var: recursion_var.map(|v| Self::offset_variable(offsets, v)),
                 unspecialized: Self::offset_uls_slice(offsets, *unspecialized),
+                ambient_function: Self::offset_variable(offsets, *ambient_function_var),
             }),
             RangedNumber(range) => RangedNumber(*range),
             Error => Content::Error,
@@ -4578,6 +4593,7 @@ fn storage_copy_var_to_help(env: &mut StorageCopyVarToEnv<'_>, var: Variable) ->
             solved,
             recursion_var,
             unspecialized,
+            ambient_function: ambient_function_var,
         }) => {
             let new_solved = storage_copy_union(env, solved);
             let new_rec_var = recursion_var.map(|v| storage_copy_var_to_help(env, v));
@@ -4593,10 +4609,13 @@ fn storage_copy_var_to_help(env: &mut StorageCopyVarToEnv<'_>, var: Variable) ->
                 env.target[target_index] = Uls(new_var, sym, region);
             }
 
+            let new_ambient_function_var = storage_copy_var_to_help(env, ambient_function_var);
+
             let new_content = LambdaSet(self::LambdaSet {
                 solved: new_solved,
                 recursion_var: new_rec_var,
                 unspecialized: new_unspecialized,
+                ambient_function: new_ambient_function_var,
             });
             env.target.set(copy, make_descriptor(new_content));
             copy
@@ -5036,6 +5055,7 @@ fn copy_import_to_help(env: &mut CopyImportEnv<'_>, max_rank: Rank, var: Variabl
             solved,
             recursion_var,
             unspecialized,
+            ambient_function: ambient_function_var,
         }) => {
             let new_solved = copy_union(env, max_rank, solved);
             let new_rec_var =
@@ -5054,10 +5074,13 @@ fn copy_import_to_help(env: &mut CopyImportEnv<'_>, max_rank: Rank, var: Variabl
                 }
             }
 
+            let new_ambient_function_var = copy_import_to_help(env, max_rank, ambient_function_var);
+
             let new_content = LambdaSet(self::LambdaSet {
                 solved: new_solved,
                 recursion_var: new_rec_var,
                 unspecialized: new_unspecialized,
+                ambient_function: new_ambient_function_var,
             });
 
             env.target.set(copy, make_descriptor(new_content));
@@ -5214,6 +5237,7 @@ fn instantiate_rigids_help(subs: &mut Subs, max_rank: Rank, initial: Variable) {
                 solved,
                 recursion_var,
                 unspecialized,
+                ambient_function: _,
             }) => {
                 for slice_index in solved.variables() {
                     let slice = subs.variable_slices[slice_index.index as usize];

crates/compiler/types/src/types.rs

@@ -240,8 +240,12 @@ pub enum Type {
     ClosureTag {
         name: Symbol,
         captures: Vec<Type>,
+        ambient_function: Variable,
+    },
+    UnspecializedLambdaSet {
+        unspecialized: Uls,
+        ambient_function: Variable,
     },
-    UnspecializedLambdaSet(Uls),
     DelayedAlias(AliasCommon),
     Alias {
         symbol: Symbol,
@@ -312,11 +316,22 @@ impl Clone for Type {
             Self::FunctionOrTagUnion(arg0, arg1, arg2) => {
                 Self::FunctionOrTagUnion(arg0.clone(), *arg1, arg2.clone())
             }
-            Self::ClosureTag { name, captures } => Self::ClosureTag {
+            Self::ClosureTag {
+                name,
+                captures,
+                ambient_function,
+            } => Self::ClosureTag {
                 name: *name,
                 captures: captures.clone(),
+                ambient_function: *ambient_function,
+            },
+            Self::UnspecializedLambdaSet {
+                unspecialized,
+                ambient_function,
+            } => Self::UnspecializedLambdaSet {
+                unspecialized: *unspecialized,
+                ambient_function: *ambient_function,
             },
-            Self::UnspecializedLambdaSet(uls) => Self::UnspecializedLambdaSet(*uls),
             Self::DelayedAlias(arg0) => Self::DelayedAlias(arg0.clone()),
             Self::Alias {
                 symbol,
@@ -622,7 +637,11 @@ impl fmt::Debug for Type {
                     }
                 }
             }
-            Type::ClosureTag { name, captures } => {
+            Type::ClosureTag {
+                name,
+                captures,
+                ambient_function: _,
+            } => {
                 write!(f, "ClosureTag(")?;
 
                 write!(f, "{:?}, ", name)?;
@@ -655,8 +674,11 @@ impl fmt::Debug for Type {
             Type::RangedNumber(range_vars) => {
                 write!(f, "Ranged({:?})", range_vars)
             }
-            Type::UnspecializedLambdaSet(uls) => {
-                write!(f, "{:?}", uls)
+            Type::UnspecializedLambdaSet {
+                unspecialized,
+                ambient_function: _,
+            } => {
+                write!(f, "{:?}", unspecialized)
             }
         }
     }
@@ -713,7 +735,11 @@ impl Type {
                     stack.push(closure);
                     stack.push(ret);
                 }
-                ClosureTag { name: _, captures } => stack.extend(captures),
+                ClosureTag {
+                    name: _,
+                    captures,
+                    ambient_function: _,
+                } => stack.extend(captures),
                 TagUnion(tags, ext) => {
                     for (_, args) in tags {
                         stack.extend(args.iter_mut());
@@ -795,7 +821,10 @@ impl Type {
                     stack.extend(args);
                 }
                 RangedNumber(_) => {}
-                UnspecializedLambdaSet(Uls(v, _, _)) => {
+                UnspecializedLambdaSet {
+                    unspecialized: Uls(v, _, _),
+                    ambient_function: _,
+                } => {
                     debug_assert!(
                         substitutions.get(v).is_none(),
                         "unspecialized lambda sets should never be substituted before solving"
@@ -824,7 +853,11 @@ impl Type {
                     stack.push(closure);
                     stack.push(ret);
                 }
-                ClosureTag { name: _, captures } => {
+                ClosureTag {
+                    name: _,
+                    captures,
+                    ambient_function: _,
+                } => {
                     stack.extend(captures);
                 }
                 TagUnion(tags, ext) => {
@@ -910,7 +943,10 @@ impl Type {
                     stack.extend(args);
                 }
                 RangedNumber(_) => {}
-                UnspecializedLambdaSet(Uls(v, _, _)) => {
+                UnspecializedLambdaSet {
+                    unspecialized: Uls(v, _, _),
+                    ambient_function: _,
+                } => {
                     debug_assert!(
                         substitutions.get(v).is_none(),
                         "unspecialized lambda sets should never be substituted before solving"
@@ -1013,7 +1049,7 @@ impl Type {
                 Ok(())
             }
             RangedNumber(_) => Ok(()),
-            UnspecializedLambdaSet(..) => Ok(()),
+            UnspecializedLambdaSet { .. } => Ok(()),
             EmptyRec | EmptyTagUnion | ClosureTag { .. } | Erroneous(_) | Variable(_) => Ok(()),
         }
     }
@@ -1070,7 +1106,10 @@ impl Type {
             Apply(symbol, _, _) if *symbol == rep_symbol => true,
             Apply(_, args, _) => args.iter().any(|arg| arg.contains_symbol(rep_symbol)),
             RangedNumber(_) => false,
-            UnspecializedLambdaSet(Uls(_, sym, _)) => *sym == rep_symbol,
+            UnspecializedLambdaSet {
+                unspecialized: Uls(_, sym, _),
+                ambient_function: _,
+            } => *sym == rep_symbol,
             EmptyRec | EmptyTagUnion | ClosureTag { .. } | Erroneous(_) | Variable(_) => false,
         }
     }
@@ -1093,10 +1132,15 @@ impl Type {
                     || args.iter().any(|arg| arg.contains_variable(rep_variable))
             }
             FunctionOrTagUnion(_, _, ext) => Self::contains_variable_ext(ext, rep_variable),
-            ClosureTag { name: _, captures } => {
-                captures.iter().any(|t| t.contains_variable(rep_variable))
-            }
-            UnspecializedLambdaSet(Uls(v, _, _)) => *v == rep_variable,
+            ClosureTag {
+                name: _,
+                captures,
+                ambient_function: _,
+            } => captures.iter().any(|t| t.contains_variable(rep_variable)),
+            UnspecializedLambdaSet {
+                unspecialized: Uls(v, _, _),
+                ambient_function: _,
+            } => *v == rep_variable,
             RecursiveTagUnion(_, tags, ext) | TagUnion(tags, ext) => {
                 Self::contains_variable_ext(ext, rep_variable)
                     || tags
@@ -1384,17 +1428,18 @@ impl Type {
                 }
             }
             RangedNumber(_) => {}
-            UnspecializedLambdaSet(..) => {}
+            UnspecializedLambdaSet { .. } => {}
             EmptyRec | EmptyTagUnion | ClosureTag { .. } | Erroneous(_) | Variable(_) => {}
         }
     }
 
     pub fn instantiate_lambda_sets_as_unspecialized(
         &mut self,
+        var_store: &mut VarStore,
         able_var: Variable,
         ability_member: Symbol,
     ) {
-        instantiate_lambda_sets_as_unspecialized(self, able_var, ability_member)
+        instantiate_lambda_sets_as_unspecialized(var_store, self, able_var, ability_member)
     }
 
     pub fn is_tag_union_like(&self) -> bool {
@@ -1520,7 +1565,10 @@ fn symbols_help(initial: &Type) -> Vec<Symbol> {
                 output.push(*alias);
             }
             RangedNumber(_) => {}
-            UnspecializedLambdaSet(Uls(_, _sym, _)) => {
+            UnspecializedLambdaSet {
+                unspecialized: Uls(_, _sym, _),
+                ambient_function: _,
+            } => {
                 // ignore the member symbol because unspecialized lambda sets are internal-only
             }
             EmptyRec | EmptyTagUnion | ClosureTag { .. } | Erroneous(_) | Variable(_) => {}
@@ -1565,12 +1613,19 @@ fn variables_help(tipe: &Type, accum: &mut ImSet<Variable>) {
                 variables_help(ext, accum);
             }
         }
-        ClosureTag { name: _, captures } => {
+        ClosureTag {
+            name: _,
+            captures,
+            ambient_function: _,
+        } => {
             for t in captures {
                 variables_help(t, accum);
             }
         }
-        UnspecializedLambdaSet(Uls(v, _, _)) => {
+        UnspecializedLambdaSet {
+            unspecialized: Uls(v, _, _),
+            ambient_function: _,
+        } => {
             accum.insert(*v);
         }
         TagUnion(tags, ext) => {
@@ -1700,7 +1755,11 @@ fn variables_help_detailed(tipe: &Type, accum: &mut VariableDetail) {
                 variables_help_detailed(ext, accum);
             }
         }
-        ClosureTag { name: _, captures } => {
+        ClosureTag {
+            name: _,
+            captures,
+            ambient_function: _,
+        } => {
             for t in captures {
                 variables_help_detailed(t, accum);
             }
@@ -1721,7 +1780,10 @@ fn variables_help_detailed(tipe: &Type, accum: &mut VariableDetail) {
                 variables_help_detailed(ext, accum);
             }
         }
-        UnspecializedLambdaSet(Uls(var, _, _)) => {
+        UnspecializedLambdaSet {
+            unspecialized: Uls(var, _, _),
+            ambient_function: _,
+        } => {
             accum.type_variables.insert(*var);
         }
         RecursiveTagUnion(rec, tags, ext) => {
@@ -2717,6 +2779,7 @@ pub fn gather_tags(subs: &Subs, other_fields: UnionTags, var: Variable) -> TagUn
 }
 
 fn instantiate_lambda_sets_as_unspecialized(
+    var_store: &mut VarStore,
     typ: &mut Type,
     able_var: Variable,
     ability_member: Symbol,
@@ -2728,7 +2791,10 @@ fn instantiate_lambda_sets_as_unspecialized(
 
     let mut new_uls = || {
         region += 1;
-        Type::UnspecializedLambdaSet(Uls(able_var, ability_member, region))
+        Type::UnspecializedLambdaSet {
+            unspecialized: Uls(able_var, ability_member, region),
+            ambient_function: var_store.fresh(),
+        }
     };
 
     while let Some(typ) = stack.pop() {
@@ -2762,10 +2828,14 @@ fn instantiate_lambda_sets_as_unspecialized(
             Type::FunctionOrTagUnion(_, _, ext) => {
                 stack.extend(ext.iter_mut());
             }
-            Type::ClosureTag { name: _, captures } => {
+            Type::ClosureTag {
+                name: _,
+                captures,
+                ambient_function: _,
+            } => {
                 stack.extend(captures.iter_mut().rev());
             }
-            Type::UnspecializedLambdaSet(..) => {
+            Type::UnspecializedLambdaSet { .. } => {
                 internal_error!("attempting to re-instantiate ULS")
             }
             Type::DelayedAlias(AliasCommon {
@@ -2841,12 +2911,15 @@ mod test {
 
         let able_var = var_store.fresh();
         let member = Symbol::UNDERSCORE;
-        typ.instantiate_lambda_sets_as_unspecialized(able_var, member);
+        typ.instantiate_lambda_sets_as_unspecialized(&mut var_store, able_var, member);
 
         macro_rules! check_uls {
             ($typ:expr, $region:literal) => {{
                 match $typ {
-                    Type::UnspecializedLambdaSet(Uls(var1, member1, $region)) => {
+                    Type::UnspecializedLambdaSet {
+                        unspecialized: Uls(var1, member1, $region),
+                        ambient_function: _,
+                    } => {
                         assert!(var1 == able_var && member1 == member)
                     }
                     _ => panic!(),