Remove unspecialized following unification links

compiler/solve/src/solve.rs

@@ -1735,14 +1735,12 @@ fn find_specializaton_lambda_sets(
                 uls_before.into_iter().chain(uls_after.into_iter()).copied(),
             );
 
-            subs.set_content(
+            let new_lambda_set_content = Content::LambdaSet(LambdaSet {
-                lambda_set,
-                Content::LambdaSet(LambdaSet {
                 solved,
                 recursion_var,
                 unspecialized: new_unspecialized,
-                }),
+            });
-            );
+            subs.set_content(lambda_set, new_lambda_set_content);
 
             let old_specialized =
                 specialization_lambda_sets.insert(specialized_lset_region, lambda_set);
@@ -1812,6 +1810,7 @@ fn compact_lambda_set(
                 //
                 //   {a, b} = default {}
                 //   #        ^^^^^^^ {} -[{a: t1, b: t2}:default:1]
+                new_unspecialized.push(uls);
                 continue;
             }
             Alias(opaque, _, _, AliasKind::Opaque) => opaque,

compiler/solve/tests/solve_expr.rs

@@ -277,7 +277,7 @@ mod solve_expr {
             let var = find_type_at(region, &decls)
                 .unwrap_or_else(|| panic!("No type for {} ({:?})!", &text, region));
 
-            let actual_str = name_and_print_var(var, subs, home, &interns, PrintLambdaSets::No);
+            let actual_str = name_and_print_var(var, subs, home, &interns, PrintLambdaSets::Yes);
 
             let elaborated = match find_ability_member_at(region, &decls) {
                 Some((member, specialization_id)) => {
@@ -6318,8 +6318,8 @@ mod solve_expr {
                 "#
             ),
             &[
-                "u8 : U8 -> Encoder Linear",
+                "u8 : U8 -[[u8(22)]]-> Encoder Linear",
-                "toEncoder : MyU8 -> Encoder fmt | fmt has Format",
+                "toEncoder : MyU8 -[[toEncoder(23)]]-> Encoder fmt | fmt has Format",
                 "myU8Bytes : List U8",
             ],
         )
@@ -6426,7 +6426,7 @@ mod solve_expr {
                     a
                 "#
             ),
-            &["A#default : {} -> A"],
+            &["A#default : {} -[[default(5)]]-> A"],
         )
     }
 
@@ -6469,7 +6469,9 @@ mod solve_expr {
                      # ^^^^^^^^^
                 "#
             ),
-            &["Encoding#toEncoder : { a : Str } -> Encoder fmt | fmt has EncoderFormatting"],
+            &[
+                "Encoding#toEncoder : { a : Str } -[[] + { a : Str }:toEncoder(2):1]-> Encoder fmt | fmt has EncoderFormatting",
+            ],
         )
     }
 
@@ -6489,7 +6491,7 @@ mod solve_expr {
                      # ^^^^^^^^^
                 "#
             ),
-            &["Encoding#toEncoder : { a : A } -> Encoder fmt | fmt has EncoderFormatting"],
+            &["Encoding#toEncoder : { a : A } -[[] + { a : A }:toEncoder(2):1]-> Encoder fmt | fmt has EncoderFormatting"],
         )
     }
 }

compiler/types/src/pretty_print.rs

@@ -585,13 +585,21 @@ fn write_content<'a>(
         }) => {
             buf.push_str("[[");
 
+            let print_symbol = |symbol: &Symbol| {
+                format!(
+                    "{}({})",
+                    symbol.as_str(&env.interns),
+                    symbol.ident_id().index(),
+                )
+            };
+
             write_sorted_tags2(
                 env,
                 ctx,
                 subs,
                 buf,
                 solved.unsorted_lambdas(subs),
-                |symbol| symbol.as_str(&env.interns),
+                print_symbol,
             );
 
             buf.push(']');
@@ -619,7 +627,7 @@ fn write_content<'a>(
                     Parens::Unnecessary,
                 );
                 buf.push(':');
-                buf.push_str(member.as_str(&env.interns));
+                buf.push_str(&print_symbol(member));
                 buf.push(':');
                 buf.push_str(&region.to_string());
             }
@@ -743,13 +751,13 @@ fn write_ext_content<'a>(
     }
 }
 
-fn write_sorted_tags2<'a, 'b, L: 'b>(
+fn write_sorted_tags2<'a, L>(
-    env: &'b Env,
+    env: &Env,
     ctx: &mut Context<'a>,
     subs: &'a Subs,
     buf: &mut String,
-    tags: UnsortedUnionLabels<'b, L>,
+    tags: UnsortedUnionLabels<L>,
-    label_to_string: impl Fn(&'b L) -> &'b str,
+    label_to_string: impl Fn(&L) -> String,
 ) where
     L: Label + Ord,
 {
@@ -766,7 +774,7 @@ fn write_sorted_tags2<'a, 'b, L: 'b>(
             any_written_yet = true;
         }
 
-        buf.push_str(label_to_string(label));
+        buf.push_str(&label_to_string(label));
 
         for var in vars {
             buf.push(' ');
@@ -935,7 +943,7 @@ fn write_flat_type<'a>(
 
             // Sort the fields so they always end up in the same order.
             let (tags, new_ext_var) = tags.unsorted_tags_and_ext(subs, *ext_var);
-            write_sorted_tags2(env, ctx, subs, buf, tags, |tag| tag.0.as_str());
+            write_sorted_tags2(env, ctx, subs, buf, tags, |tag| tag.0.as_str().to_string());
 
             buf.push(']');
 
@@ -965,7 +973,7 @@ fn write_flat_type<'a>(
             buf.push('[');
 
             let (tags, new_ext_var) = tags.unsorted_tags_and_ext(subs, *ext_var);
-            write_sorted_tags2(env, ctx, subs, buf, tags, |tag| tag.0.as_str());
+            write_sorted_tags2(env, ctx, subs, buf, tags, |tag| tag.0.as_str().to_string());
 
             buf.push(']');
 

compiler/types/src/subs.rs

@@ -315,7 +315,8 @@ pub struct UlsOfVar(VecMap<Variable, VecSet<Variable>>);
 
 impl UlsOfVar {
     pub fn add(&mut self, var: Variable, dependent_lambda_set: Variable) -> bool {
-        // TODO: should we be checking root key here?
+        // NOTE: this adds the var directly without following unification links.
+        // [Subs::remove_dependent_unspecialized_lambda_sets] follows unifications when removing.
         let set = self.0.get_or_insert(var, Default::default);
         set.insert(dependent_lambda_set)
     }
@@ -325,7 +326,8 @@ impl UlsOfVar {
         var: Variable,
         dependent_lambda_sets: impl IntoIterator<Item = Variable>,
     ) {
-        // TODO: should we be checking root key here?
+        // NOTE: this adds the var directly without following unification links.
+        // [Subs::remove_dependent_unspecialized_lambda_sets] follows unifications when removing.
         let set = self.0.get_or_insert(var, Default::default);
         set.extend(dependent_lambda_sets);
     }
@@ -344,6 +346,7 @@ impl UlsOfVar {
         self.0.remove(&var).map(|(_, v)| v)
     }
 
+    /// NOTE: this does not follow unification links.
     pub fn drain(self) -> impl Iterator<Item = (Variable, impl Iterator<Item = Variable>)> {
         self.0
             .into_iter()
@@ -1991,6 +1994,21 @@ impl Subs {
             _ => internal_error!("not a lambda set"),
         }
     }
+
+    pub fn remove_dependent_unspecialized_lambda_sets(
+        &mut self,
+        var: Variable,
+    ) -> impl Iterator<Item = Variable> + '_ {
+        let utable = &self.utable;
+        let root_var = utable.root_key_without_compacting(var);
+
+        self.uls_of_var
+            .0
+            .drain_filter(move |cand_var, _| {
+                utable.root_key_without_compacting(*cand_var) == root_var
+            })
+            .flat_map(|(_, lambda_set_vars)| lambda_set_vars.into_iter())
+    }
 }
 
 #[inline(always)]

compiler/unify/src/unify.rs

@@ -2162,11 +2162,10 @@ fn merge_flex_able_with_concrete(
     //
     // If we ever organize ability implementations so that they are well-known before any other
     // unification is done, they can be solved in-band here!
-    if let Some(uls_of_concrete) = subs.uls_of_var.remove_dependents(flex_able_var) {
+    let uls_of_concrete = subs.remove_dependent_unspecialized_lambda_sets(flex_able_var);
     outcome
         .lambda_sets_to_specialize
         .extend(flex_able_var, uls_of_concrete);
-    }
 
     outcome
 }