Bookkeeping for unspecialized lambda sets

compiler/solve/src/solve.rs

@@ -3382,6 +3382,8 @@ fn deep_copy_var_help(
                 recursion_var,
                 unspecialized,
             }) => {
+                let lambda_set_var = copy;
+
                 let new_solved = copy_union!(solved);
                 let new_rec_var = recursion_var.map(|v| work!(v));
                 let new_unspecialized = SubsSlice::reserve_uls_slice(subs, unspecialized.len());
@@ -3396,11 +3398,11 @@ fn deep_copy_var_help(
 
                     subs[new_uls_index] = Uls(new_var, sym, region);
 
-                    // TODO: bookkeeping of new_var -> lambda set
+                    subs.uls_of_var.add(new_var, lambda_set_var);
                 }
 
                 subs.set_content_unchecked(
-                    copy,
+                    lambda_set_var,
                     LambdaSet(subs::LambdaSet {
                         solved: new_solved,
                         recursion_var: new_rec_var,

compiler/types/src/subs.rs

@@ -3,6 +3,7 @@ use crate::types::{
     name_type_var, AliasKind, ErrorType, Problem, RecordField, RecordFieldsError, TypeExt, Uls,
 };
 use roc_collections::all::{ImMap, ImSet, MutSet, SendMap};
+use roc_collections::{VecMap, VecSet};
 use roc_error_macros::internal_error;
 use roc_module::ident::{Lowercase, TagName, Uppercase};
 use roc_module::symbol::Symbol;
@@ -241,6 +242,7 @@ impl Subs {
                 unspecialized_lambda_sets: unspecialized_lambda_sets.to_vec(),
                 tag_name_cache: Default::default(),
                 problems: Default::default(),
+                uls_of_var: Default::default(),
             },
             exposed_vars_by_symbol,
         )
@@ -306,6 +308,25 @@ impl Subs {
     }
 }
 
+/// Mapping of variables to [Content::LambdaSet]s containing unspecialized lambda sets depending on
+/// that variable.
+#[derive(Clone, Default)]
+pub struct UlsOfVar(VecMap<Variable, VecSet<Variable>>);
+
+impl UlsOfVar {
+    pub fn add(&mut self, var: Variable, dependent_lambda_set: Variable) -> bool {
+        let set = self.0.get_or_insert(var, Default::default);
+        set.insert(dependent_lambda_set)
+    }
+
+    pub fn remove_dependents(
+        &mut self,
+        var: Variable,
+    ) -> Option<impl IntoIterator<Item = Variable>> {
+        self.0.remove(&var).map(|(_, v)| v)
+    }
+}
+
 #[derive(Clone)]
 pub struct Subs {
     utable: UnificationTable,
@@ -318,6 +339,7 @@ pub struct Subs {
     pub unspecialized_lambda_sets: Vec<Uls>,
     pub tag_name_cache: TagNameCache,
     pub problems: Vec<Problem>,
+    pub uls_of_var: UlsOfVar,
 }
 
 #[derive(Debug, Clone, Default)]
@@ -1567,6 +1589,7 @@ impl Subs {
             unspecialized_lambda_sets: Vec::new(),
             tag_name_cache: Default::default(),
             problems: Vec::new(),
+            uls_of_var: Default::default(),
         };
 
         subs.utable.reserve(capacity);