Auto merge of #140375 - lcnr:subrelations-infcx, r=BoxyUwU

eagerly compute `sub_unification_table` again

Previously called `sub_relations`. We still only using them for diagnostics right now. This mostly reverts rust-lang/rust#119989. Necessary for type inference guidance due to not-yet defined opaque types, cc https://github.com/rust-lang/trait-system-refactor-initiative/issues/182.

We could use them for cycle detection in generalization and it seems desirable to do so in the future. However, this is unsound with the old trait solver as its cache does not track these `sub_unification_table` in any way.

We now properly track the `sub_unification_table` when canonicalizing so using them in the new solver is totally sound and the performance impact is far more manageable than I thought back in rust-lang/rust#119989.

r? `@compiler-errors`

5 files changed, 91 insertions, 56 deletions

diff --git a/compiler/rustc_next_trait_solver/src/canonicalizer.rs b/compiler/rustc_next_trait_solver/src/canonicalizer.rs
index da05c49756f..a8f2b4e8db6 100644
--- a/compiler/rustc_next_trait_solver/src/canonicalizer.rs
+++ b/compiler/rustc_next_trait_solver/src/canonicalizer.rs
@@ -2,9 +2,8 @@ use rustc_type_ir::data_structures::{HashMap, ensure_sufficient_stack};
 use rustc_type_ir::inherent::*;
 use rustc_type_ir::solve::{Goal, QueryInput};
 use rustc_type_ir::{
-    self as ty, Canonical, CanonicalParamEnvCacheEntry, CanonicalTyVarKind, CanonicalVarKind,
-    Flags, InferCtxtLike, Interner, TypeFlags, TypeFoldable, TypeFolder, TypeSuperFoldable,
-    TypeVisitableExt,
+    self as ty, Canonical, CanonicalParamEnvCacheEntry, CanonicalVarKind, Flags, InferCtxtLike,
+    Interner, TypeFlags, TypeFoldable, TypeFolder, TypeSuperFoldable, TypeVisitableExt,
 };
 
 use crate::delegate::SolverDelegate;
@@ -68,6 +67,13 @@ pub struct Canonicalizer<'a, D: SolverDelegate<Interner = I>, I: Interner> {
     variables: &'a mut Vec<I::GenericArg>,
     var_kinds: Vec<CanonicalVarKind<I>>,
     variable_lookup_table: HashMap<I::GenericArg, usize>,
+    /// Maps each `sub_unification_table_root_var` to the index of the first
+    /// variable which used it.
+    ///
+    /// This means in case two type variables have the same sub relations root,
+    /// we set the `sub_root` of the second variable to the position of the first.
+    /// Otherwise the `sub_root` of each type variable is just its own position.
+    sub_root_lookup_table: HashMap<ty::TyVid, usize>,
     binder_index: ty::DebruijnIndex,
 
     /// We only use the debruijn index during lookup. We don't need to
@@ -89,6 +95,7 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
 
             variables,
             variable_lookup_table: Default::default(),
+            sub_root_lookup_table: Default::default(),
             var_kinds: Vec::new(),
             binder_index: ty::INNERMOST,
 
@@ -133,6 +140,7 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
 
                         variables: &mut variables,
                         variable_lookup_table: Default::default(),
+                        sub_root_lookup_table: Default::default(),
                         var_kinds: Vec::new(),
                         binder_index: ty::INNERMOST,
 
@@ -140,6 +148,7 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
                     };
                     let param_env = param_env.fold_with(&mut env_canonicalizer);
                     debug_assert_eq!(env_canonicalizer.binder_index, ty::INNERMOST);
+                    debug_assert!(env_canonicalizer.sub_root_lookup_table.is_empty());
                     CanonicalParamEnvCacheEntry {
                         param_env,
                         variable_lookup_table: env_canonicalizer.variable_lookup_table,
@@ -165,6 +174,7 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
 
                 variables,
                 variable_lookup_table: Default::default(),
+                sub_root_lookup_table: Default::default(),
                 var_kinds: Vec::new(),
                 binder_index: ty::INNERMOST,
 
@@ -172,6 +182,7 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
             };
             let param_env = param_env.fold_with(&mut env_canonicalizer);
             debug_assert_eq!(env_canonicalizer.binder_index, ty::INNERMOST);
+            debug_assert!(env_canonicalizer.sub_root_lookup_table.is_empty());
             (param_env, env_canonicalizer.variable_lookup_table, env_canonicalizer.var_kinds)
         }
     }
@@ -200,6 +211,7 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
 
             variables,
             variable_lookup_table,
+            sub_root_lookup_table: Default::default(),
             var_kinds,
             binder_index: ty::INNERMOST,
 
@@ -266,6 +278,13 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
         ty::BoundVar::from(idx)
     }
 
+    fn get_or_insert_sub_root(&mut self, vid: ty::TyVid) -> ty::BoundVar {
+        let root_vid = self.delegate.sub_unification_table_root_var(vid);
+        let idx =
+            *self.sub_root_lookup_table.entry(root_vid).or_insert_with(|| self.variables.len());
+        ty::BoundVar::from(idx)
+    }
+
     fn finalize(self) -> (ty::UniverseIndex, I::CanonicalVarKinds) {
         let mut var_kinds = self.var_kinds;
         // See the rustc-dev-guide section about how we deal with universes
@@ -313,18 +332,15 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
                         "ty vid should have been resolved fully before canonicalization"
                     );
 
-                    match self.canonicalize_mode {
-                        CanonicalizeMode::Input { .. } => CanonicalVarKind::Ty(
-                            CanonicalTyVarKind::General(ty::UniverseIndex::ROOT),
-                        ),
-                        CanonicalizeMode::Response { .. } => {
-                            CanonicalVarKind::Ty(CanonicalTyVarKind::General(
-                                self.delegate.universe_of_ty(vid).unwrap_or_else(|| {
-                                    panic!("ty var should have been resolved: {t:?}")
-                                }),
-                            ))
-                        }
-                    }
+                    let sub_root = self.get_or_insert_sub_root(vid);
+                    let ui = match self.canonicalize_mode {
+                        CanonicalizeMode::Input { .. } => ty::UniverseIndex::ROOT,
+                        CanonicalizeMode::Response { .. } => self
+                            .delegate
+                            .universe_of_ty(vid)
+                            .unwrap_or_else(|| panic!("ty var should have been resolved: {t:?}")),
+                    };
+                    CanonicalVarKind::Ty { ui, sub_root }
                 }
                 ty::IntVar(vid) => {
                     debug_assert_eq!(
@@ -332,7 +348,7 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
                         t,
                         "ty vid should have been resolved fully before canonicalization"
                     );
-                    CanonicalVarKind::Ty(CanonicalTyVarKind::Int)
+                    CanonicalVarKind::Int
                 }
                 ty::FloatVar(vid) => {
                     debug_assert_eq!(
@@ -340,7 +356,7 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
                         t,
                         "ty vid should have been resolved fully before canonicalization"
                     );
-                    CanonicalVarKind::Ty(CanonicalTyVarKind::Float)
+                    CanonicalVarKind::Float
                 }
                 ty::FreshTy(_) | ty::FreshIntTy(_) | ty::FreshFloatTy(_) => {
                     panic!("fresh vars not expected in canonicalization")
diff --git a/compiler/rustc_next_trait_solver/src/delegate.rs b/compiler/rustc_next_trait_solver/src/delegate.rs
index a7ecbfcb7c4..41c2e5f60ec 100644
--- a/compiler/rustc_next_trait_solver/src/delegate.rs
+++ b/compiler/rustc_next_trait_solver/src/delegate.rs
@@ -57,12 +57,14 @@ pub trait SolverDelegate: Deref<Target = Self::Infcx> + Sized {
     where
         V: TypeFoldable<Self::Interner>;
 
-    fn instantiate_canonical_var_with_infer(
+    fn instantiate_canonical_var(
         &self,
         kind: ty::CanonicalVarKind<Self::Interner>,
         span: <Self::Interner as Interner>::Span,
+        var_values: &[<Self::Interner as Interner>::GenericArg],
         universe_map: impl Fn(ty::UniverseIndex) -> ty::UniverseIndex,
     ) -> <Self::Interner as Interner>::GenericArg;
+
     fn add_item_bounds_for_hidden_type(
         &self,
         def_id: <Self::Interner as Interner>::DefId,
diff --git a/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/canonical.rs b/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/canonical.rs
index 6f9f4067384..169832ca5fb 100644
--- a/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/canonical.rs
+++ b/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/canonical.rs
@@ -16,7 +16,8 @@ use rustc_type_ir::data_structures::HashSet;
 use rustc_type_ir::inherent::*;
 use rustc_type_ir::relate::solver_relating::RelateExt;
 use rustc_type_ir::{
-    self as ty, Canonical, CanonicalVarValues, InferCtxtLike, Interner, TypeFoldable,
+    self as ty, Canonical, CanonicalVarKind, CanonicalVarValues, InferCtxtLike, Interner,
+    TypeFoldable,
 };
 use tracing::{debug, instrument, trace};
 
@@ -336,7 +337,16 @@ where
         {
             match result_value.kind() {
                 ty::GenericArgKind::Type(t) => {
-                    if let ty::Bound(debruijn, b) = t.kind() {
+                    // We disable the instantiation guess for inference variables
+                    // and only use it for placeholders. We need to handle the
+                    // `sub_root` of type inference variables which would make this
+                    // more involved. They are also a lot rarer than region variables.
+                    if let ty::Bound(debruijn, b) = t.kind()
+                        && !matches!(
+                            response.variables.get(b.var().as_usize()).unwrap(),
+                            CanonicalVarKind::Ty { .. }
+                        )
+                    {
                         assert_eq!(debruijn, ty::INNERMOST);
                         opt_values[b.var()] = Some(*original_value);
                     }
@@ -355,39 +365,33 @@ where
                 }
             }
         }
-
-        let var_values = delegate.cx().mk_args_from_iter(
-            response.variables.iter().enumerate().map(|(index, var_kind)| {
-                if var_kind.universe() != ty::UniverseIndex::ROOT {
-                    // A variable from inside a binder of the query. While ideally these shouldn't
-                    // exist at all (see the FIXME at the start of this method), we have to deal with
-                    // them for now.
-                    delegate.instantiate_canonical_var_with_infer(var_kind, span, |idx| {
-                        prev_universe + idx.index()
-                    })
-                } else if var_kind.is_existential() {
-                    // As an optimization we sometimes avoid creating a new inference variable here.
-                    //
-                    // All new inference variables we create start out in the current universe of the caller.
-                    // This is conceptually wrong as these inference variables would be able to name
-                    // more placeholders then they should be able to. However the inference variables have
-                    // to "come from somewhere", so by equating them with the original values of the caller
-                    // later on, we pull them down into their correct universe again.
-                    if let Some(v) = opt_values[ty::BoundVar::from_usize(index)] {
-                        v
-                    } else {
-                        delegate
-                            .instantiate_canonical_var_with_infer(var_kind, span, |_| prev_universe)
-                    }
+        CanonicalVarValues::instantiate(delegate.cx(), response.variables, |var_values, kind| {
+            if kind.universe() != ty::UniverseIndex::ROOT {
+                // A variable from inside a binder of the query. While ideally these shouldn't
+                // exist at all (see the FIXME at the start of this method), we have to deal with
+                // them for now.
+                delegate.instantiate_canonical_var(kind, span, &var_values, |idx| {
+                    prev_universe + idx.index()
+                })
+            } else if kind.is_existential() {
+                // As an optimization we sometimes avoid creating a new inference variable here.
+                //
+                // All new inference variables we create start out in the current universe of the caller.
+                // This is conceptually wrong as these inference variables would be able to name
+                // more placeholders then they should be able to. However the inference variables have
+                // to "come from somewhere", so by equating them with the original values of the caller
+                // later on, we pull them down into their correct universe again.
+                if let Some(v) = opt_values[ty::BoundVar::from_usize(var_values.len())] {
+                    v
                 } else {
-                    // For placeholders which were already part of the input, we simply map this
-                    // universal bound variable back the placeholder of the input.
-                    original_values[var_kind.expect_placeholder_index()]
+                    delegate.instantiate_canonical_var(kind, span, &var_values, |_| prev_universe)
                 }
-            }),
-        );
-
-        CanonicalVarValues { var_values }
+            } else {
+                // For placeholders which were already part of the input, we simply map this
+                // universal bound variable back the placeholder of the input.
+                original_values[kind.expect_placeholder_index()]
+            }
+        })
     }
 
     /// Unify the `original_values` with the `var_values` returned by the canonical query..
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 99aa9bef5ac..5cd597d4fe6 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
@@ -900,6 +900,10 @@ where
             && goal.param_env.visit_with(&mut visitor).is_continue()
     }
 
+    pub(super) fn sub_unify_ty_vids_raw(&self, a: ty::TyVid, b: ty::TyVid) {
+        self.delegate.sub_unify_ty_vids_raw(a, b)
+    }
+
     #[instrument(level = "trace", skip(self, param_env), ret)]
     pub(super) fn eq<T: Relate<I>>(
         &mut self,
diff --git a/compiler/rustc_next_trait_solver/src/solve/mod.rs b/compiler/rustc_next_trait_solver/src/solve/mod.rs
index 85f9d852d95..db3460c4ff6 100644
--- a/compiler/rustc_next_trait_solver/src/solve/mod.rs
+++ b/compiler/rustc_next_trait_solver/src/solve/mod.rs
@@ -119,11 +119,15 @@ where
 
     #[instrument(level = "trace", skip(self))]
     fn compute_subtype_goal(&mut self, goal: Goal<I, ty::SubtypePredicate<I>>) -> QueryResult<I> {
-        if goal.predicate.a.is_ty_var() && goal.predicate.b.is_ty_var() {
-            self.evaluate_added_goals_and_make_canonical_response(Certainty::AMBIGUOUS)
-        } else {
-            self.sub(goal.param_env, goal.predicate.a, goal.predicate.b)?;
-            self.evaluate_added_goals_and_make_canonical_response(Certainty::Yes)
+        match (goal.predicate.a.kind(), goal.predicate.b.kind()) {
+            (ty::Infer(ty::TyVar(a_vid)), ty::Infer(ty::TyVar(b_vid))) => {
+                self.sub_unify_ty_vids_raw(a_vid, b_vid);
+                self.evaluate_added_goals_and_make_canonical_response(Certainty::AMBIGUOUS)
+            }
+            _ => {
+                self.sub(goal.param_env, goal.predicate.a, goal.predicate.b)?;
+                self.evaluate_added_goals_and_make_canonical_response(Certainty::Yes)
+            }
         }
     }
 
@@ -414,6 +418,11 @@ 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>>,
 }