consider the `sub_unification_table` in `stalled_on`

2 files changed, 33 insertions, 24 deletions

diff --git a/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/mod.rs b/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/mod.rs
index 5cd597d4fe6..f1c1801dd3f 100644
--- a/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/mod.rs
+++ b/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/mod.rs
@@ -407,20 +407,21 @@ where
         // If we have run this goal before, and it was stalled, check that any of the goal's
         // args have changed. Otherwise, we don't need to re-run the goal because it'll remain
         // stalled, since it'll canonicalize the same way and evaluation is pure.
-        if let Some(stalled_on) = stalled_on
-            && !stalled_on.stalled_vars.iter().any(|value| self.delegate.is_changed_arg(*value))
-            && !self
-                .delegate
-                .opaque_types_storage_num_entries()
-                .needs_reevaluation(stalled_on.num_opaques)
+        if let Some(GoalStalledOn { num_opaques, ref stalled_vars, ref sub_roots, stalled_cause }) =
+            stalled_on
+            && !stalled_vars.iter().any(|value| self.delegate.is_changed_arg(*value))
+            && !sub_roots
+                .iter()
+                .any(|&vid| self.delegate.sub_unification_table_root_var(vid) != vid)
+            && !self.delegate.opaque_types_storage_num_entries().needs_reevaluation(num_opaques)
         {
             return Ok((
                 NestedNormalizationGoals::empty(),
                 GoalEvaluation {
                     goal,
-                    certainty: Certainty::Maybe(stalled_on.stalled_cause),
+                    certainty: Certainty::Maybe(stalled_cause),
                     has_changed: HasChanged::No,
-                    stalled_on: Some(stalled_on),
+                    stalled_on,
                 },
             ));
         }
@@ -476,16 +477,6 @@ where
                 HasChanged::No => {
                     let mut stalled_vars = orig_values;
 
-                    // Remove the canonicalized universal vars, since we only care about stalled existentials.
-                    stalled_vars.retain(|arg| match arg.kind() {
-                        ty::GenericArgKind::Type(ty) => matches!(ty.kind(), ty::Infer(_)),
-                        ty::GenericArgKind::Const(ct) => {
-                            matches!(ct.kind(), ty::ConstKind::Infer(_))
-                        }
-                        // Lifetimes can never stall goals.
-                        ty::GenericArgKind::Lifetime(_) => false,
-                    });
-
                     // Remove the unconstrained RHS arg, which is expected to have changed.
                     if let Some(normalizes_to) = goal.predicate.as_normalizes_to() {
                         let normalizes_to = normalizes_to.skip_binder();
@@ -497,6 +488,27 @@ where
                         stalled_vars.swap_remove(idx);
                     }
 
+                    // Remove the canonicalized universal vars, since we only care about stalled existentials.
+                    let mut sub_roots = Vec::new();
+                    stalled_vars.retain(|arg| match arg.kind() {
+                        // Lifetimes can never stall goals.
+                        ty::GenericArgKind::Lifetime(_) => false,
+                        ty::GenericArgKind::Type(ty) => match ty.kind() {
+                            ty::Infer(ty::TyVar(vid)) => {
+                                sub_roots.push(self.delegate.sub_unification_table_root_var(vid));
+                                true
+                            }
+                            ty::Infer(_) => true,
+                            ty::Param(_) | ty::Placeholder(_) => false,
+                            _ => unreachable!("unexpected orig_value: {ty:?}"),
+                        },
+                        ty::GenericArgKind::Const(ct) => match ct.kind() {
+                            ty::ConstKind::Infer(_) => true,
+                            ty::ConstKind::Param(_) | ty::ConstKind::Placeholder(_) => false,
+                            _ => unreachable!("unexpected orig_value: {ct:?}"),
+                        },
+                    });
+
                     Some(GoalStalledOn {
                         num_opaques: canonical_goal
                             .canonical
@@ -505,6 +517,7 @@ where
                             .opaque_types
                             .len(),
                         stalled_vars,
+                        sub_roots,
                         stalled_cause,
                     })
                 }
diff --git a/compiler/rustc_next_trait_solver/src/solve/mod.rs b/compiler/rustc_next_trait_solver/src/solve/mod.rs
index db3460c4ff6..cd27c9c26c1 100644
--- a/compiler/rustc_next_trait_solver/src/solve/mod.rs
+++ b/compiler/rustc_next_trait_solver/src/solve/mod.rs
@@ -24,7 +24,7 @@ mod trait_goals;
 use derive_where::derive_where;
 use rustc_type_ir::inherent::*;
 pub use rustc_type_ir::solve::*;
-use rustc_type_ir::{self as ty, Interner, TypingMode};
+use rustc_type_ir::{self as ty, Interner, TyVid, TypingMode};
 use tracing::instrument;
 
 pub use self::eval_ctxt::{
@@ -418,11 +418,6 @@ pub struct GoalEvaluation<I: Interner> {
     pub has_changed: HasChanged,
     /// If the [`Certainty`] was `Maybe`, then keep track of whether the goal has changed
     /// before rerunning it.
-    ///
-    /// We knowingly ignore the `sub_root` of our inference variables here. This means we
-    /// may not reevaluate a goal even though a change to the `sub_root` could cause a goal
-    /// to make progress. Tracking them adds additional complexity for an incredibly minor
-    /// type inference improvement. We could look into properly handling this in the future.
     pub stalled_on: Option<GoalStalledOn<I>>,
 }
 
@@ -431,6 +426,7 @@ pub struct GoalEvaluation<I: Interner> {
 pub struct GoalStalledOn<I: Interner> {
     pub num_opaques: usize,
     pub stalled_vars: Vec<I::GenericArg>,
+    pub sub_roots: Vec<TyVid>,
     /// The cause that will be returned on subsequent evaluations if this goal remains stalled.
     pub stalled_cause: MaybeCause,
 }