Admit duplicate lambdas in lambda sets when their captures don't unify

crates/compiler/solve/tests/solve_expr.rs

@@ -6953,6 +6953,68 @@ mod solve_expr {
                 "#
             ),
             "Str",
+        )
+    }
+
+    #[test]
+    fn lambda_sets_collide_with_captured_var() {
+        infer_queries!(
+            indoc!(
+                r#"
+                capture : a -> ({} -> Str)
+                capture = \val ->
+                    thunk =
+                        \{} ->
+                            when val is
+                                _ -> ""
+                    thunk
+
+                x : [True, False]
+
+                fun =
+                    when x is
+                        True -> capture ""
+                        False -> capture {}
+                fun
+                #^^^{-1}
+                "#
+            ),
+            &["fun : {} -[[thunk(5) {}, thunk(5) Str]]-> Str"],
+        );
+    }
+
+    #[test]
+    fn lambda_sets_collide_with_captured_function() {
+        infer_queries!(
+            indoc!(
+                r#"
+                Lazy a : {} -> a
+
+                after : Lazy a, (a -> Lazy b) -> Lazy b
+                after = \effect, map ->
+                    thunk = \{} ->
+                        when map (effect {}) is
+                            b -> b {}
+                    thunk
+
+                f = \_ -> \_ -> ""
+                g = \{ s1 } -> \_ -> s1
+
+                x : [True, False]
+
+                fun =
+                    when x is
+                        True -> after (\{} -> "") f
+                        False -> after (\{} -> {s1: "s1"}) g
+                fun
+                #^^^{-1}
+                "#
+            ),
+            &[
+                "fun : {} -[[thunk(9) (({} -[[15(15)]]-> { s1 : Str })) ({ s1 : Str } -[[g(4)]]-> ({} -[[13(13) Str]]-> Str)), \
+                             thunk(9) (({} -[[14(14)]]-> Str)) (Str -[[f(3)]]-> ({} -[[11(11)]]-> Str))]]-> Str",
+            ],
+            print_only_under_alias = true,
         );
     }
 }

crates/compiler/types/src/subs.rs

@@ -1872,6 +1872,10 @@ impl Subs {
         self.utable.unioned(left, right)
     }
 
+    pub fn equivalent_without_compacting(&self, left: Variable, right: Variable) -> bool {
+        self.utable.unioned_without_compacting(left, right)
+    }
+
     pub fn redundant(&self, var: Variable) -> bool {
         self.utable.is_redirect(var)
     }

crates/compiler/types/src/unification_table.rs

@@ -323,6 +323,10 @@ impl UnificationTable {
         self.root_key(a) == self.root_key(b)
     }
 
+    pub fn unioned_without_compacting(&self, a: Variable, b: Variable) -> bool {
+        self.root_key_without_compacting(a) == self.root_key_without_compacting(b)
+    }
+
     // custom very specific helpers
     #[inline(always)]
     pub fn get_rank_set_mark(&mut self, key: Variable, mark: Mark) -> Rank {

crates/compiler/unify/src/unify.rs

@@ -1000,17 +1000,11 @@ fn unify_lambda_set_help<M: MetaCollector>(
         in_both,
     } = separate_union_lambdas(subs, solved1, solved2);
 
-    let num_shared = in_both.len();
+    let mut new_lambdas = vec![];
+    for (lambda_name, (vars1, vars2)) in in_both {
+        let mut captures_unify = vars1.len() == vars2.len();
 
-    let mut joined_lambdas = vec![];
+        if captures_unify {
-    for (tag_name, (vars1, vars2)) in in_both {
-        let mut matching_vars = vec![];
-
-        if vars1.len() != vars2.len() {
-            continue; // this is a type mismatch; not adding the tag will trigger it below.
-        }
-
-        let num_vars = vars1.len();
             for (var1, var2) in (vars1.into_iter()).zip(vars2.into_iter()) {
                 let (var1, var2) = (subs[var1], subs[var2]);
 
@@ -1024,20 +1018,35 @@ fn unify_lambda_set_help<M: MetaCollector>(
                 maybe_mark_union_recursive(subs, var1);
                 maybe_mark_union_recursive(subs, var2);
 
+                let snapshot = subs.snapshot();
                 let outcome = unify_pool::<M>(subs, pool, var1, var2, ctx.mode);
 
-            if outcome.mismatches.is_empty() {
+                if !outcome.mismatches.is_empty() {
-                matching_vars.push(var1);
+                    captures_unify = false;
+                    subs.rollback_to(snapshot);
+                    // Continue so the other variables can unify if possible, allowing us to re-use
+                    // shared variables.
+                }
             }
         }
 
-        if matching_vars.len() == num_vars {
+        if captures_unify {
-            joined_lambdas.push((tag_name, matching_vars));
+            debug_assert!((subs.get_subs_slice(vars1).iter())
+                .zip(subs.get_subs_slice(vars2).iter())
+                .all(|(v1, v2)| subs.equivalent_without_compacting(*v1, *v2)));
+
+            new_lambdas.push((lambda_name, subs.get_subs_slice(vars1).to_vec()));
+        } else {
+            debug_assert!((subs.get_subs_slice(vars1).iter())
+                .zip(subs.get_subs_slice(vars2).iter())
+                .any(|(v1, v2)| !subs.equivalent_without_compacting(*v1, *v2)));
+
+            new_lambdas.push((lambda_name, subs.get_subs_slice(vars1).to_vec()));
+            new_lambdas.push((lambda_name, subs.get_subs_slice(vars2).to_vec()));
         }
     }
 
-    if joined_lambdas.len() == num_shared {
+    let all_lambdas = new_lambdas;
-        let all_lambdas = joined_lambdas;
     let all_lambdas = merge_sorted(
         all_lambdas,
         only_in_1.into_iter().map(|(name, subs_slice)| {
@@ -1090,13 +1099,6 @@ fn unify_lambda_set_help<M: MetaCollector>(
     });
 
     merge(subs, ctx, new_lambda_set)
-    } else {
-        mismatch!(
-            "Problem with lambda sets: there should be {:?} matching lambda, but only found {:?}",
-            num_shared,
-            &joined_lambdas
-        )
-    }
 }
 
 /// Ensures that a non-recursive tag union, when unified with a recursion var to become a recursive