Attach ambient function vars to lambda sets

crates/ast/src/solve_type.rs

@@ -1421,6 +1421,7 @@ fn adjust_rank_content(
             recursion_var,
             // TODO: handle unspecialized
             unspecialized: _,
+            ambient_function: _,
         }) => {
             let mut rank = group_rank;
 
@@ -1623,6 +1624,7 @@ fn instantiate_rigids_help(
             recursion_var,
             // TODO: handle unspecialized
             unspecialized: _,
+            ambient_function: _,
         }) => {
             if let Some(rec_var) = recursion_var.into_variable() {
                 instantiate_rigids_help(subs, max_rank, pools, rec_var);
@@ -1903,6 +1905,7 @@ fn deep_copy_var_help(
             solved,
             recursion_var,
             unspecialized,
+            ambient_function,
         }) => {
             let mut new_variable_slices = Vec::with_capacity(solved.len());
 
@@ -1937,6 +1940,7 @@ fn deep_copy_var_help(
                 recursion_var: new_rec_var,
                 // TODO: actually copy
                 unspecialized,
+                ambient_function,
             });
 
             subs.set(copy, make_descriptor(new_content));

crates/compiler/can/src/def.rs

@@ -952,8 +952,11 @@ fn resolve_abilities<'a>(
 
             let signature = {
                 let mut signature = member_annot.typ;
-                signature
+                signature.instantiate_lambda_sets_as_unspecialized(
-                    .instantiate_lambda_sets_as_unspecialized(var_bound_to_ability, member_sym);
+                    var_store,
+                    var_bound_to_ability,
+                    member_sym,
+                );
                 signature
             };
 

crates/compiler/constrain/src/expr.rs

@@ -143,6 +143,7 @@ fn constrain_untyped_closure(
         constraints,
         name,
         region,
+        fn_var,
         captured_symbols,
         closure_var,
         &mut vars,
@@ -911,6 +912,7 @@ pub fn constrain_expr(
             let lambda_set = Type::ClosureTag {
                 name: *closure_name,
                 captures: vec![],
+                ambient_function: *function_var,
             };
 
             let closure_type = Type::Variable(*closure_var);
@@ -1390,6 +1392,7 @@ fn constrain_function_def(
                 constraints,
                 loc_symbol.value,
                 region,
+                expr_var,
                 &function_def.captured_symbols,
                 closure_var,
                 &mut vars,
@@ -2040,6 +2043,7 @@ fn constrain_typed_def(
                 constraints,
                 *name,
                 region,
+                *fn_var,
                 captured_symbols,
                 closure_var,
                 &mut vars,
@@ -2514,6 +2518,7 @@ fn constrain_closure_size(
     constraints: &mut Constraints,
     name: Symbol,
     region: Region,
+    ambient_function: Variable,
     captured_symbols: &[(Symbol, Variable)],
     closure_var: Variable,
     variables: &mut Vec<Variable>,
@@ -2542,6 +2547,7 @@ fn constrain_closure_size(
     let closure_type = Type::ClosureTag {
         name,
         captures: captured_types,
+        ambient_function,
     };
 
     let finalizer = constraints.equal_types_var(
@@ -2811,6 +2817,7 @@ fn constraint_recursive_function(
                 constraints,
                 loc_symbol.value,
                 region,
+                expr_var,
                 &function_def.captured_symbols,
                 closure_var,
                 &mut vars,
@@ -3220,6 +3227,7 @@ fn rec_defs_help(
                             constraints,
                             *name,
                             region,
+                            *fn_var,
                             captured_symbols,
                             closure_var,
                             &mut vars,

crates/compiler/derive/src/encoding.rs

@@ -325,6 +325,9 @@ fn to_encoder_string(env: &mut Env<'_>, fn_name: Symbol) -> (Expr, Variable) {
     let (body, this_encoder_var) =
         wrap_in_encode_custom(env, encode_string_call, encoder_var, s_sym, Variable::STR);
 
+    // Create fn_var for ambient capture; we fix it up below.
+    let fn_var = synth_var(env.subs, Content::Error);
+
     // -[fn_name]->
     let fn_name_labels = UnionLambdas::insert_into_subs(env.subs, once((fn_name, vec![])));
     let fn_clos_var = synth_var(
@@ -333,11 +336,12 @@ fn to_encoder_string(env: &mut Env<'_>, fn_name: Symbol) -> (Expr, Variable) {
             solved: fn_name_labels,
             recursion_var: OptVariable::NONE,
             unspecialized: SubsSlice::default(),
+            ambient_function: fn_var,
         }),
     );
     // Str -[fn_name]-> (typeof Encode.record [ .. ] = Encoder fmt)
-    let fn_var = synth_var(
+    env.subs.set_content(
-        env.subs,
+        fn_var,
         Content::Structure(FlatType::Func(
             string_var_slice,
             fn_clos_var,
@@ -527,6 +531,9 @@ fn to_encoder_record(
     let (body, this_encoder_var) =
         wrap_in_encode_custom(env, encode_record_call, encoder_var, rcd_sym, record_var);
 
+    // Create fn_var for ambient capture; we fix it up below.
+    let fn_var = synth_var(env.subs, Content::Error);
+
     // -[fn_name]->
     let fn_name_labels = UnionLambdas::insert_into_subs(env.subs, once((fn_name, vec![])));
     let fn_clos_var = synth_var(
@@ -535,12 +542,13 @@ fn to_encoder_record(
             solved: fn_name_labels,
             recursion_var: OptVariable::NONE,
             unspecialized: SubsSlice::default(),
+            ambient_function: fn_var,
         }),
     );
     // typeof rcd -[fn_name]-> (typeof Encode.record [ .. ] = Encoder fmt)
     let record_var_slice = SubsSlice::insert_into_subs(env.subs, once(record_var));
-    let fn_var = synth_var(
+    env.subs.set_content(
-        env.subs,
+        fn_var,
         Content::Structure(FlatType::Func(
             record_var_slice,
             fn_clos_var,
@@ -758,6 +766,9 @@ fn to_encoder_tag_union(
         tag_union_var,
     );
 
+    // Create fn_var for ambient capture; we fix it up below.
+    let fn_var = synth_var(env.subs, Content::Error);
+
     // -[fn_name]->
     let fn_name_labels = UnionLambdas::insert_into_subs(env.subs, once((fn_name, vec![])));
     let fn_clos_var = synth_var(
@@ -766,12 +777,13 @@ fn to_encoder_tag_union(
             solved: fn_name_labels,
             recursion_var: OptVariable::NONE,
             unspecialized: SubsSlice::default(),
+            ambient_function: fn_var,
         }),
     );
     // tag_union_var -[fn_name]-> whole_tag_encoders_var
     let tag_union_var_slice = SubsSlice::insert_into_subs(env.subs, once(tag_union_var));
-    let fn_var = synth_var(
+    env.subs.set_content(
-        env.subs,
+        fn_var,
         Content::Structure(FlatType::Func(
             tag_union_var_slice,
             fn_clos_var,
@@ -871,6 +883,9 @@ fn wrap_in_encode_custom(
         CalledVia::Space,
     );
 
+    // Create fn_var for ambient capture; we fix it up below.
+    let fn_var = synth_var(env.subs, Content::Error);
+
     // -[[FN_name captured_var]]->
     let fn_name_labels =
         UnionLambdas::insert_into_subs(env.subs, once((fn_name, vec![captured_var])));
@@ -880,13 +895,14 @@ fn wrap_in_encode_custom(
             solved: fn_name_labels,
             recursion_var: OptVariable::NONE,
             unspecialized: SubsSlice::default(),
+            ambient_function: fn_var,
         }),
     );
 
     // bytes, fmt -[[FN_name captured_var]]-> Encode.appendWith bytes encoder fmt
     let args_slice = SubsSlice::insert_into_subs(env.subs, vec![bytes_var, fmt_var]);
-    let fn_var = synth_var(
+    env.subs.set_content(
-        env.subs,
+        fn_var,
         Content::Structure(FlatType::Func(args_slice, fn_clos_var, Variable::LIST_U8)),
     );
 

crates/compiler/mono/src/copy.rs

@@ -620,6 +620,7 @@ fn deep_copy_type_vars<'a>(
                 solved,
                 recursion_var,
                 unspecialized,
+                ambient_function,
             }) => {
                 let new_rec_var = recursion_var.map(|var| descend_var!(var));
                 for variables_slice_index in solved.variables() {
@@ -630,6 +631,7 @@ fn deep_copy_type_vars<'a>(
                     let Uls(var, _, _) = subs[uls_index];
                     descend_var!(var);
                 }
+                let new_ambient_function = descend_var!(ambient_function);
 
                 perform_clone!({
                     let new_variable_slices =
@@ -657,6 +659,7 @@ fn deep_copy_type_vars<'a>(
                         solved: new_solved,
                         recursion_var: new_rec_var,
                         unspecialized: new_unspecialized,
+                        ambient_function: new_ambient_function,
                     })
                 })
             }

crates/compiler/mono/src/ir.rs

@@ -5114,6 +5114,7 @@ fn late_resolve_ability_specialization<'a>(
             solved,
             unspecialized,
             recursion_var: _,
+            ambient_function: _,
         } = env.subs.get_lambda_set(*lambda_set);
 
         debug_assert!(unspecialized.is_empty());

crates/compiler/mono/src/layout.rs

@@ -1127,6 +1127,7 @@ fn resolve_lambda_set(subs: &Subs, mut var: Variable) -> ResolvedLambdaSet {
                 solved,
                 recursion_var: _,
                 unspecialized,
+                ambient_function: _,
             }) => {
                 debug_assert!(
                     unspecialized.is_empty(),
@@ -1966,6 +1967,7 @@ fn layout_from_lambda_set<'a>(
         solved,
         recursion_var,
         unspecialized,
+        ambient_function: _,
     } = lset;
 
     if !unspecialized.is_empty() {

crates/compiler/mono/src/layout_soa.rs

@@ -277,6 +277,7 @@ impl LambdaSet {
             solved,
             recursion_var: _,
             unspecialized: _,
+            ambient_function: _,
         } = lset;
 
         // TODO: handle unspecialized
@@ -700,6 +701,7 @@ impl Layout {
             solved,
             recursion_var,
             unspecialized: _,
+            ambient_function: _,
         } = lset;
 
         // TODO: handle unspecialized lambda set

crates/compiler/solve/src/solve.rs

@@ -1784,6 +1784,7 @@ fn compact_lambda_set<P: Phase>(
         solved,
         recursion_var,
         unspecialized,
+        ambient_function,
     } = subs.get_lambda_set(this_lambda_set);
     let target_rank = subs.get_rank(this_lambda_set);
 
@@ -1820,11 +1821,14 @@ fn compact_lambda_set<P: Phase>(
                             subs,
                             vec![(specialization_symbol, vec![])],
                         );
+                        // TODO: This is WRONG, fix it!
+                        let ambient_function = subs.fresh_unnamed_flex_var();
                         let lambda_set_for_derived = subs.fresh(Descriptor {
                             content: LambdaSet(subs::LambdaSet {
                                 solved: specialization_symbol_slice,
                                 recursion_var: OptVariable::NONE,
                                 unspecialized: SubsSlice::default(),
+                                ambient_function,
                             }),
                             rank: target_rank,
                             mark: Mark::NONE,
@@ -1928,6 +1932,7 @@ fn compact_lambda_set<P: Phase>(
         solved,
         recursion_var,
         unspecialized: new_unspecialized_slice,
+        ambient_function,
     });
     subs.set_content(this_lambda_set, partial_compacted_lambda_set);
 
@@ -2192,7 +2197,11 @@ fn type_to_variable<'a>(
                 register_with_known_var(subs, destination, rank, pools, content)
             }
 
-            ClosureTag { name, captures } => {
+            ClosureTag {
+                name,
+                captures,
+                ambient_function,
+            } => {
                 let union_lambdas =
                     create_union_lambda(subs, rank, pools, arena, *name, captures, &mut stack);
 
@@ -2202,20 +2211,25 @@ fn type_to_variable<'a>(
                     // is to begin with.
                     recursion_var: OptVariable::NONE,
                     unspecialized: SubsSlice::default(),
+                    ambient_function: *ambient_function,
                 });
 
                 register_with_known_var(subs, destination, rank, pools, content)
             }
-            UnspecializedLambdaSet(uls) => {
+            UnspecializedLambdaSet {
+                unspecialized,
+                ambient_function,
+            } => {
                 let unspecialized = SubsSlice::extend_new(
                     &mut subs.unspecialized_lambda_sets,
-                    std::iter::once(*uls),
+                    std::iter::once(*unspecialized),
                 );
 
                 let content = Content::LambdaSet(subs::LambdaSet {
                     unspecialized,
                     solved: UnionLabels::default(),
                     recursion_var: OptVariable::NONE,
+                    ambient_function: *ambient_function,
                 });
 
                 register_with_known_var(subs, destination, rank, pools, content)
@@ -2877,8 +2891,14 @@ fn check_for_infinite_type(
                 solved,
                 recursion_var: _,
                 unspecialized,
+                ambient_function: ambient_function_var,
             }) => {
-                subs.mark_lambda_set_recursive(recursive, solved, unspecialized);
+                subs.mark_lambda_set_recursive(
+                    recursive,
+                    solved,
+                    unspecialized,
+                    ambient_function_var,
+                );
             }
 
             _other => circular_error(subs, problems, symbol, &loc_var),
@@ -3226,6 +3246,7 @@ fn adjust_rank_content(
             solved,
             recursion_var,
             unspecialized,
+            ambient_function: ambient_function_var,
         }) => {
             let mut rank = group_rank;
 
@@ -3262,6 +3283,11 @@ fn adjust_rank_content(
                 );
             }
 
+            // NEVER TOUCH the ambient function var, it would already have been passed through.
+            {
+                let _ = ambient_function_var;
+            }
+
             rank
         }
 
@@ -3525,6 +3551,7 @@ fn deep_copy_var_help(
                 solved,
                 recursion_var,
                 unspecialized,
+                ambient_function: ambient_function_var,
             }) => {
                 let lambda_set_var = copy;
 
@@ -3545,12 +3572,19 @@ fn deep_copy_var_help(
                     subs.uls_of_var.add(new_var, lambda_set_var);
                 }
 
+                let new_ambient_function_var = work!(ambient_function_var);
+                debug_assert_ne!(
+                    ambient_function_var, new_ambient_function_var,
+                    "lambda set cloned but its ambient function wasn't?"
+                );
+
                 subs.set_content_unchecked(
                     lambda_set_var,
                     LambdaSet(subs::LambdaSet {
                         solved: new_solved,
                         recursion_var: new_rec_var,
                         unspecialized: new_unspecialized,
+                        ambient_function: new_ambient_function_var,
                     }),
                 );
             }

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!(Descriptor, 5 * 8 + 4);
-roc_error_macros::assert_sizeof_all!(Content, 3 * 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) => {
+            Type::UnspecializedLambdaSet {
-                write!(f, "{:?}", uls)
+                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 } => {
+            ClosureTag {
-                captures.iter().any(|t| t.contains_variable(rep_variable))
+                name: _,
-            }
+                captures,
-            UnspecializedLambdaSet(Uls(v, _, _)) => *v == rep_variable,
+                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!(),

crates/compiler/unify/src/unify.rs

@@ -967,6 +967,7 @@ fn unify_lambda_set<M: MetaCollector>(
                     solved: UnionLabels::default(),
                     recursion_var: OptVariable::NONE,
                     unspecialized: SubsSlice::default(),
+                    ambient_function: subs.fresh_unnamed_flex_var(),
                 };
 
                 extract_specialization_lambda_set(subs, ctx, lambda_set, zero_lambda_set)
@@ -1012,6 +1013,7 @@ fn extract_specialization_lambda_set<M: MetaCollector>(
         solved: member_solved,
         recursion_var: member_rec_var,
         unspecialized: member_uls_slice,
+        ambient_function: _,
     } = ability_member_proto_lset;
 
     debug_assert!(
@@ -1220,13 +1222,20 @@ fn unify_lambda_set_help<M: MetaCollector>(
         solved: solved1,
         recursion_var: rec1,
         unspecialized: uls1,
+        ambient_function: ambient_function_var1,
     } = lset1;
     let LambdaSet {
         solved: solved2,
         recursion_var: rec2,
         unspecialized: uls2,
+        ambient_function: ambient_function_var2,
     } = lset2;
 
+    // Assumed precondition: the ambient functions have already been unified, or are in the process
+    // of being unified - otherwise, how could we have reached unification of lambda sets?
+    let _ = ambient_function_var2;
+    let ambient_function_var = ambient_function_var1;
+
     debug_assert!(
         (rec1.into_variable().into_iter())
             .chain(rec2.into_variable().into_iter())
@@ -1293,6 +1302,7 @@ fn unify_lambda_set_help<M: MetaCollector>(
         solved: new_solved,
         recursion_var,
         unspecialized: merged_unspecialized,
+        ambient_function: ambient_function_var,
     });
 
     merge(subs, ctx, new_lambda_set)
@@ -1926,8 +1936,9 @@ fn maybe_mark_union_recursive(subs: &mut Subs, union_var: Variable) {
                     solved,
                     recursion_var: OptVariable::NONE,
                     unspecialized,
+                    ambient_function: ambient_function_var,
                 }) => {
-                    subs.mark_lambda_set_recursive(v, solved, unspecialized);
+                    subs.mark_lambda_set_recursive(v, solved, unspecialized, ambient_function_var);
                     continue 'outer;
                 }
                 _ => { /* fall through */ }
@@ -2685,10 +2696,12 @@ fn unify_function_or_tag_union_and_func<M: MetaCollector>(
 
     {
         let union_tags = UnionLambdas::tag_without_arguments(subs, tag_symbol);
+        let ambient_function_var = if left { ctx.first } else { ctx.second };
         let lambda_set_content = LambdaSet(self::LambdaSet {
             solved: union_tags,
             recursion_var: OptVariable::NONE,
             unspecialized: SubsSlice::default(),
+            ambient_function: ambient_function_var,
         });
 
         let tag_lambda_set = register(