Invalidate layout-cached variables correctly after merging

crates/compiler/late_solve/src/lib.rs

@@ -16,7 +16,8 @@ use roc_solve::solve::Pools;
 use roc_solve::specialize::{compact_lambda_sets_of_vars, DerivedEnv, Phase};
 use roc_types::subs::{get_member_lambda_sets_at_region, Content, FlatType, LambdaSet};
 use roc_types::subs::{ExposedTypesStorageSubs, Subs, Variable};
-use roc_unify::unify::{unify as unify_unify, Env, Mode, Unified};
+use roc_unify::unify::MetaCollector;
+use roc_unify::unify::{Env, Mode, Unified};
 
 pub use roc_solve::ability::Resolved;
 pub use roc_types::subs::instantiate_rigids;
@@ -313,6 +314,28 @@ impl Phase for LatePhase<'_> {
     }
 }
 
+#[derive(Debug, Default)]
+struct ChangedVariableCollector {
+    changed: Vec<Variable>,
+}
+
+impl MetaCollector for ChangedVariableCollector {
+    const UNIFYING_SPECIALIZATION: bool = false;
+
+    #[inline(always)]
+    fn record_specialization_lambda_set(&mut self, _member: Symbol, _region: u8, _var: Variable) {}
+
+    #[inline(always)]
+    fn record_changed_variable(&mut self, subs: &Subs, var: Variable) {
+        self.changed.push(subs.get_root_key_without_compacting(var))
+    }
+
+    #[inline(always)]
+    fn union(&mut self, other: Self) {
+        self.changed.extend(other.changed)
+    }
+}
+
 /// Unifies two variables and performs lambda set compaction.
 /// Ranks and other ability demands are disregarded.
 #[allow(clippy::too_many_arguments)]
@@ -325,20 +348,25 @@ pub fn unify(
     exposed_by_module: &ExposedByModule,
     left: Variable,
     right: Variable,
-) -> Result<(), UnificationFailed> {
+) -> Result<Vec<Variable>, UnificationFailed> {
     debug_assert_ne!(
         home,
         ModuleId::DERIVED_SYNTH,
         "derived module can only unify its subs in its own context!"
     );
-    let unified = unify_unify(&mut Env::new(subs), left, right, Mode::EQ);
+    let unified = roc_unify::unify::unify_with_collector::<ChangedVariableCollector>(
+        &mut Env::new(subs),
+        left,
+        right,
+        Mode::EQ,
+    );
 
     match unified {
         Unified::Success {
             vars: _,
             must_implement_ability: _,
             lambda_sets_to_specialize,
-            extra_metadata: _,
+            extra_metadata,
         } => {
             let mut pools = Pools::default();
 
@@ -366,7 +394,7 @@ pub fn unify(
             // here. We only need it for `compact_lambda_sets_of_vars`, which is also used in a
             // solving context where pools are relevant.
 
-            Ok(())
+            Ok(extra_metadata.changed)
         }
         Unified::Failure(..) | Unified::BadType(..) => Err(UnificationFailed),
     }

crates/compiler/mono/src/ir.rs

@@ -1429,14 +1429,19 @@ impl<'a, 'i> Env<'a, 'i> {
 
     /// Unifies two variables and performs lambda set compaction.
     /// Use this rather than [roc_unify::unify] directly!
-    fn unify(&mut self, left: Variable, right: Variable) -> Result<(), UnificationFailed> {
+    fn unify(
+        &mut self,
+        layout_cache: &mut LayoutCache,
+        left: Variable,
+        right: Variable,
+    ) -> Result<(), UnificationFailed> {
         debug_assert_ne!(
             self.home,
             ModuleId::DERIVED_SYNTH,
             "should never be monomorphizing the derived synth module!"
         );
 
-        roc_late_solve::unify(
+        let changed_variables = roc_late_solve::unify(
             self.home,
             self.arena,
             self.subs,
@@ -1445,7 +1450,11 @@ impl<'a, 'i> Env<'a, 'i> {
             self.exposed_by_module,
             left,
             right,
-        )
+        )?;
+
+        layout_cache.invalidate(changed_variables);
+
+        Ok(())
     }
 }
 
@@ -2455,7 +2464,7 @@ fn from_can_let<'a>(
                         // Make sure rigid variables in the annotation are converted to flex variables.
                         instantiate_rigids(env.subs, def.expr_var);
                         // Unify the expr_var with the requested specialization once.
-                        let _res = env.unify(var, def.expr_var);
+                        let _res = env.unify(layout_cache, var, def.expr_var);
 
                         with_hole(
                             env,
@@ -2489,7 +2498,7 @@ fn from_can_let<'a>(
                             "expr marked as having specializations, but it has no type variables!",
                         );
 
-                            let _res = env.unify(var, new_def_expr_var);
+                            let _res = env.unify(layout_cache, var, new_def_expr_var);
 
                             stmt = with_hole(
                                 env,
@@ -3076,7 +3085,7 @@ fn specialize_proc_help<'a>(
     let partial_proc = procs.partial_procs.get_id(partial_proc_id);
     let captured_symbols = partial_proc.captured_symbols;
 
-    let _unified = env.unify(partial_proc.annotation, fn_var);
+    let _unified = env.unify(layout_cache, partial_proc.annotation, fn_var);
 
     // This will not hold for programs with type errors
     // let is_valid = matches!(unified, roc_unify::unify::Unified::Success(_));
@@ -5661,9 +5670,28 @@ fn convert_tag_union<'a>(
                 "Wrapped"
             ) {
                 Layout::Union(ul) => ul,
-                _ => unreachable!(),
+                other => internal_error!(
+                    "unexpected layout {:?} for {:?} ({:?})",
+                    other,
+                    roc_types::subs::SubsFmtContent(
+                        env.subs.get_content_without_compacting(variant_var),
+                        env.subs
+                    ),
+                    (
+                        variant_var,
+                        env.subs.get_root_key_without_compacting(variant_var)
+                    )
+                ),
             };
 
+            dbg!((
+                &union_layout,
+                roc_types::subs::SubsFmtContent(
+                    env.subs.get_content_without_compacting(variant_var),
+                    env.subs
+                )
+            ));
+
             use WrappedVariant::*;
             let (tag, union_layout) = match variant {
                 Recursive { sorted_tag_layouts } => {

crates/compiler/mono/src/layout.rs

@@ -290,6 +290,22 @@ impl<'a> LayoutCache<'a> {
         self.raw_function_cache.pop();
     }
 
+    /// Invalidates the list of given root variables.
+    /// Usually called after unification, when merged variables with changed contents need to be
+    /// invalidated.
+    pub fn invalidate(&mut self, vars: impl IntoIterator<Item = Variable>) {
+        for var in vars.into_iter() {
+            for layer in self.cache.iter_mut().rev() {
+                layer.0.remove(&var);
+                roc_tracing::debug!(?var, "invalidating cached layout");
+            }
+            for layer in self.raw_function_cache.iter_mut().rev() {
+                layer.0.remove(&var);
+                roc_tracing::debug!(?var, "invalidating cached layout");
+            }
+        }
+    }
+
     #[cfg(debug_assertions)]
     pub fn statistics(&self) -> (CacheStatistics, CacheStatistics) {
         (self.stats, self.raw_function_stats)
@@ -1970,6 +1986,13 @@ macro_rules! cached_or_impl {
 
         if criteria.is_cacheable() {
             // The computed layout is cacheable; insert it.
+            dbg!((
+                $var,
+                roc_types::subs::SubsFmtContent(
+                    $self.subs.get_content_without_compacting($var),
+                    $self.subs
+                )
+            ));
             $self
                 .cache
                 .$insert($self.subs, $var, result, criteria.cache_metadata());

crates/compiler/unify/src/unify.rs

@@ -159,6 +159,8 @@ pub trait MetaCollector: Default + std::fmt::Debug {
 
     fn record_specialization_lambda_set(&mut self, member: Symbol, region: u8, var: Variable);
 
+    fn record_changed_variable(&mut self, subs: &Subs, var: Variable);
+
     fn union(&mut self, other: Self);
 }
 
@@ -167,8 +169,13 @@ pub struct NoCollector;
 impl MetaCollector for NoCollector {
     const UNIFYING_SPECIALIZATION: bool = false;
 
+    #[inline(always)]
     fn record_specialization_lambda_set(&mut self, _member: Symbol, _region: u8, _var: Variable) {}
 
+    #[inline(always)]
+    fn record_changed_variable(&mut self, _subs: &Subs, _var: Variable) {}
+
+    #[inline(always)]
     fn union(&mut self, _other: Self) {}
 }
 
@@ -178,10 +185,15 @@ pub struct SpecializationLsetCollector(pub VecMap<(Symbol, u8), Variable>);
 impl MetaCollector for SpecializationLsetCollector {
     const UNIFYING_SPECIALIZATION: bool = true;
 
+    #[inline(always)]
     fn record_specialization_lambda_set(&mut self, member: Symbol, region: u8, var: Variable) {
         self.0.insert((member, region), var);
     }
 
+    #[inline(always)]
+    fn record_changed_variable(&mut self, _subs: &Subs, _var: Variable) {}
+
+    #[inline(always)]
     fn union(&mut self, other: Self) {
         for (k, v) in other.0.into_iter() {
             let _old = self.0.insert(k, v);
@@ -334,6 +346,16 @@ pub fn unify_introduced_ability_specialization(
     unify_help(env, ability_member_signature, specialization_var, mode)
 }
 
+#[inline(always)]
+pub fn unify_with_collector<M: MetaCollector>(
+    env: &mut Env,
+    var1: Variable,
+    var2: Variable,
+    mode: Mode,
+) -> Unified<M> {
+    unify_help(env, var1, var2, mode)
+}
+
 #[inline(always)]
 fn unify_help<M: MetaCollector>(
     env: &mut Env,
@@ -2982,6 +3004,8 @@ fn unify_recursion<M: MetaCollector>(
 }
 
 pub fn merge<M: MetaCollector>(env: &mut Env, ctx: &Context, content: Content) -> Outcome<M> {
+    let mut outcome: Outcome<M> = Outcome::default();
+
     if !env.compute_outcome_only {
         let rank = ctx.first_desc.rank.min(ctx.second_desc.rank);
         let desc = Descriptor {
@@ -2991,10 +3015,17 @@ pub fn merge<M: MetaCollector>(env: &mut Env, ctx: &Context, content: Content) -
             copy: OptVariable::NONE,
         };
 
+        outcome
+            .extra_metadata
+            .record_changed_variable(env.subs, ctx.first);
+        outcome
+            .extra_metadata
+            .record_changed_variable(env.subs, ctx.second);
+
         env.subs.union(ctx.first, ctx.second, desc);
     }
 
-    Outcome::default()
+    outcome
 }
 
 fn register(env: &mut Env, desc: Descriptor, pool: &mut Pool) -> Variable {