track relevant sub_relations in canonical queries

This allows canonical queries to also rely on them in the future.
It also means it would now be sound to rely on `sub_relations` in
the generalizer.

Author: lcnr

compiler/rustc_infer/src/infer/canonical/canonicalizer.rs

@@ -6,6 +6,7 @@
 //! [c]: https://rust-lang.github.io/chalk/book/canonical_queries/canonicalization.html
 
 use rustc_data_structures::fx::FxHashMap;
+use rustc_data_structures::sso::SsoHashMap;
 use rustc_index::Idx;
 use rustc_middle::bug;
 use rustc_middle::ty::{
@@ -298,6 +299,7 @@ struct Canonicalizer<'cx, 'tcx> {
     // Note that indices is only used once `var_values` is big enough to be
     // heap-allocated.
     indices: FxHashMap<GenericArg<'tcx>, BoundVar>,
+    sub_root_lookup_table: SsoHashMap<ty::TyVid, usize>,
     canonicalize_mode: &'cx dyn CanonicalizeMode,
     needs_canonical_flags: TypeFlags,
 
@@ -366,8 +368,11 @@ impl<'cx, 'tcx> TypeFolder<TyCtxt<'tcx>> for Canonicalizer<'cx, 'tcx> {
                             // FIXME: perf problem described in #55921.
                             ui = ty::UniverseIndex::ROOT;
                         }
+                        let sub_root = self.get_or_insert_sub_root(vid);
                         self.canonicalize_ty_var(
-                            CanonicalVarInfo { kind: CanonicalVarKind::Ty(ui) },
+                            CanonicalVarInfo {
+                                kind: CanonicalVarKind::Ty { universe: ui, sub_root },
+                            },
                             t,
                         )
                     }
@@ -567,6 +572,7 @@ impl<'cx, 'tcx> Canonicalizer<'cx, 'tcx> {
             variables: SmallVec::from_slice(base.variables),
             query_state,
             indices: FxHashMap::default(),
+            sub_root_lookup_table: Default::default(),
             binder_index: ty::INNERMOST,
         };
         if canonicalizer.query_state.var_values.spilled() {
@@ -661,6 +667,13 @@ impl<'cx, 'tcx> Canonicalizer<'cx, 'tcx> {
         }
     }
 
+    fn get_or_insert_sub_root(&mut self, vid: ty::TyVid) -> ty::BoundVar {
+        let root_vid = self.infcx.unwrap().sub_root_var(vid);
+        let idx =
+            *self.sub_root_lookup_table.entry(root_vid).or_insert_with(|| self.variables.len());
+        ty::BoundVar::from(idx)
+    }
+
     /// Replaces the universe indexes used in `var_values` with their index in
     /// `query_state.universe_map`. This minimizes the maximum universe used in
     /// the canonicalized value.
@@ -684,7 +697,9 @@ impl<'cx, 'tcx> Canonicalizer<'cx, 'tcx> {
                     CanonicalVarKind::Int | CanonicalVarKind::Float => {
                         return *v;
                     }
-                    CanonicalVarKind::Ty(u) => CanonicalVarKind::Ty(reverse_universe_map[&u]),
+                    CanonicalVarKind::Ty { universe, sub_root } => {
+                        CanonicalVarKind::Ty { universe: reverse_universe_map[&universe], sub_root }
+                    }
                     CanonicalVarKind::Region(u) => {
                         CanonicalVarKind::Region(reverse_universe_map[&u])
                     }

compiler/rustc_infer/src/infer/canonical/mod.rs

@@ -84,13 +84,12 @@ impl<'tcx> InferCtxt<'tcx> {
         variables: &List<CanonicalVarInfo<'tcx>>,
         universe_map: impl Fn(ty::UniverseIndex) -> ty::UniverseIndex,
     ) -> CanonicalVarValues<'tcx> {
-        CanonicalVarValues {
-            var_values: self.tcx.mk_args_from_iter(
-                variables
-                    .iter()
-                    .map(|info| self.instantiate_canonical_var(span, info, &universe_map)),
-            ),
+        let mut var_values = Vec::new();
+        for info in variables.iter() {
+            let value = self.instantiate_canonical_var(span, info, &var_values, &universe_map);
+            var_values.push(value);
         }
+        CanonicalVarValues { var_values: self.tcx.mk_args(&var_values) }
     }
 
     /// Given the "info" about a canonical variable, creates a fresh
@@ -105,10 +104,22 @@ impl<'tcx> InferCtxt<'tcx> {
         &self,
         span: Span,
         cv_info: CanonicalVarInfo<'tcx>,
+        previous_var_values: &[GenericArg<'tcx>],
         universe_map: impl Fn(ty::UniverseIndex) -> ty::UniverseIndex,
     ) -> GenericArg<'tcx> {
         match cv_info.kind {
-            CanonicalVarKind::Ty(ui) => self.next_ty_var_in_universe(span, universe_map(ui)).into(),
+            CanonicalVarKind::Ty { universe, sub_root } => {
+                let vid = self.next_ty_var_id_in_universe(span, universe_map(universe));
+                if let Some(prev) = previous_var_values.get(sub_root.as_usize()) {
+                    // We cannot simply assume that previous `var_values` get instantiated
+                    // with inference variables as we may reuse the generic arguments from the
+                    // input which may have gotten constrained after we've canonicalized it.
+                    if let &ty::Infer(ty::TyVar(sub_root)) = prev.expect_ty().kind() {
+                        self.inner.borrow_mut().type_variables().sub(vid, sub_root);
+                    }
+                }
+                Ty::new_var(self.tcx, vid).into()
+            }
             CanonicalVarKind::Int => self.next_int_var().into(),
             CanonicalVarKind::Float => self.next_float_var().into(),
 

compiler/rustc_infer/src/infer/canonical/query_response.rs

@@ -13,7 +13,9 @@ use std::iter;
 use rustc_index::{Idx, IndexVec};
 use rustc_middle::arena::ArenaAllocatable;
 use rustc_middle::mir::ConstraintCategory;
-use rustc_middle::ty::{self, BoundVar, GenericArg, GenericArgKind, Ty, TyCtxt, TypeFoldable};
+use rustc_middle::ty::{
+    self, BoundVar, CanonicalVarKind, GenericArg, GenericArgKind, Ty, TyCtxt, TypeFoldable,
+};
 use rustc_middle::{bug, span_bug};
 use tracing::{debug, instrument};
 
@@ -455,32 +457,54 @@ impl<'tcx> InferCtxt<'tcx> {
         // Create result arguments: if we found a value for a
         // given variable in the loop above, use that. Otherwise, use
         // a fresh inference variable.
-        let result_args = CanonicalVarValues {
-            var_values: self.tcx.mk_args_from_iter(
-                query_response.variables.iter().enumerate().map(|(index, info)| {
-                    if info.universe() != ty::UniverseIndex::ROOT {
-                        // A variable from inside a binder of the query. While ideally these shouldn't
-                        // exist at all, we have to deal with them for now.
-                        self.instantiate_canonical_var(cause.span, info, |u| {
-                            universe_map[u.as_usize()]
-                        })
-                    } else if info.is_existential() {
-                        match opt_values[BoundVar::new(index)] {
-                            Some(k) => k,
-                            None => self.instantiate_canonical_var(cause.span, info, |u| {
-                                universe_map[u.as_usize()]
-                            }),
+        let mut var_values = Vec::new();
+        for (index, info) in query_response.variables.iter().enumerate() {
+            let value = if info.universe() != ty::UniverseIndex::ROOT {
+                // A variable from inside a binder of the query. While ideally these shouldn't
+                // exist at all, we have to deal with them for now.
+                self.instantiate_canonical_var(cause.span, info, &var_values, |u| {
+                    universe_map[u.as_usize()]
+                })
+            } else if info.is_existential() {
+                // As an optimization we sometimes avoid creating a new inference variable here.
+                // We need to still make sure to register any subtype relations returned by the
+                // query.
+                match opt_values[BoundVar::new(index)] {
+                    Some(v) => {
+                        if let CanonicalVarKind::Ty { universe: _, sub_root } = info.kind {
+                            if let Some(prev) = var_values.get(sub_root.as_usize()) {
+                                // We cannot simply assume that previous `var_values`
+                                // are inference variables, see the comment in
+                                // `instantiate_canonical_var`.
+                                let v = self.shallow_resolve(v.expect_ty());
+                                let prev = self.shallow_resolve(prev.expect_ty());
+                                match (v.kind(), prev.kind()) {
+                                    (
+                                        &ty::Infer(ty::TyVar(vid)),
+                                        &ty::Infer(ty::TyVar(sub_root)),
+                                    ) => {
+                                        self.inner.borrow_mut().type_variables().sub(vid, sub_root)
+                                    }
+                                    _ => {}
+                                }
+                            }
                         }
-                    } else {
-                        // For placeholders which were already part of the input, we simply map this
-                        // universal bound variable back the placeholder of the input.
-                        opt_values[BoundVar::new(index)].expect(
-                            "expected placeholder to be unified with itself during response",
-                        )
+                        v
                     }
-                }),
-            ),
-        };
+                    None => self.instantiate_canonical_var(cause.span, info, &var_values, |u| {
+                        universe_map[u.as_usize()]
+                    }),
+                }
+            } else {
+                // For placeholders which were already part of the input, we simply map this
+                // universal bound variable back the placeholder of the input.
+                opt_values[BoundVar::new(index)]
+                    .expect("expected placeholder to be unified with itself during response")
+            };
+            var_values.push(value)
+        }
+
+        let result_args = CanonicalVarValues { var_values: self.tcx.mk_args(&var_values) };
 
         let mut obligations = PredicateObligations::new();
 

compiler/rustc_infer/src/infer/context.rs

@@ -55,6 +55,9 @@ impl<'tcx> rustc_type_ir::InferCtxtLike for InferCtxt<'tcx> {
     fn root_ty_var(&self, var: ty::TyVid) -> ty::TyVid {
         self.root_var(var)
     }
+    fn sub_root_ty_var(&self, var: ty::TyVid) -> ty::TyVid {
+        self.sub_root_var(var)
+    }
 
     fn root_const_var(&self, var: ty::ConstVid) -> ty::ConstVid {
         self.root_const_var(var)

compiler/rustc_next_trait_solver/src/canonicalizer.rs

@@ -52,6 +52,7 @@ pub struct Canonicalizer<'a, D: SolverDelegate<Interner = I>, I: Interner> {
     variables: &'a mut Vec<I::GenericArg>,
     primitive_var_infos: Vec<CanonicalVarInfo<I>>,
     variable_lookup_table: HashMap<I::GenericArg, usize>,
+    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
@@ -73,6 +74,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(),
             primitive_var_infos: Vec::new(),
             binder_index: ty::INNERMOST,
 
@@ -106,6 +108,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(),
             primitive_var_infos: Vec::new(),
             binder_index: ty::INNERMOST,
 
@@ -123,6 +126,7 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
             // We're able to reuse the `variable_lookup_table` as whether or not
             // it already contains an entry for `'static` does not matter.
             variable_lookup_table: env_canonicalizer.variable_lookup_table,
+            sub_root_lookup_table: Default::default(),
             primitive_var_infos: env_canonicalizer.primitive_var_infos,
             binder_index: ty::INNERMOST,
 
@@ -177,6 +181,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_root_ty_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::CanonicalVars) {
         let mut var_infos = self.primitive_var_infos;
         // See the rustc-dev-guide section about how we deal with universes
@@ -323,11 +334,12 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
                         "ty vid should have been resolved fully before canonicalization"
                     );
 
-                    CanonicalVarKind::Ty(
-                        self.delegate
-                            .universe_of_ty(vid)
-                            .unwrap_or_else(|| panic!("ty var should have been resolved: {t:?}")),
-                    )
+                    let universe = 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);
+                    CanonicalVarKind::Ty { universe, sub_root }
                 }
                 ty::IntVar(vid) => {
                     assert_eq!(

compiler/rustc_next_trait_solver/src/delegate.rs

@@ -59,6 +59,7 @@ pub trait SolverDelegate: Deref<Target = Self::Infcx> + Sized {
         &self,
         cv_info: ty::CanonicalVarInfo<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;
 

compiler/rustc_next_trait_solver/src/solve/eval_ctxt/canonical.rs

@@ -15,7 +15,8 @@ use rustc_index::IndexVec;
 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};
 
@@ -354,37 +355,55 @@ where
             }
         }
 
-        let var_values = delegate.cx().mk_args_from_iter(
-            response.variables.iter().enumerate().map(|(index, info)| {
-                if info.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(info, span, |idx| {
-                        prev_universe + idx.index()
-                    })
-                } else if info.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(info, span, |_| prev_universe)
+        let mut var_values = Vec::new();
+        for (index, info) in response.variables.iter().enumerate() {
+            let value = if info.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(info, span, &var_values, |idx| {
+                    prev_universe + idx.index()
+                })
+            } else if info.is_existential() {
+                // As an optimization we sometimes avoid creating a new inference variable here.
+                // We need to still make sure to register any subtype relations returned by the
+                // query.
+                if let Some(v) = opt_values[ty::BoundVar::from_usize(index)] {
+                    if let CanonicalVarKind::Ty { universe: _, sub_root } = info.kind {
+                        if let Some(prev) = var_values.get(sub_root.as_usize()) {
+                            // We cannot simply assume that previous `var_values`
+                            // are inference variables, see the comment in
+                            // `instantiate_canonical_var`.
+                            let v = delegate.shallow_resolve(v.expect_ty());
+                            let prev = delegate.shallow_resolve(prev.expect_ty());
+                            match (v.kind(), prev.kind()) {
+                                (ty::Infer(ty::TyVar(vid)), ty::Infer(ty::TyVar(sub_root))) => {
+                                    delegate.sub_ty_vids_raw(vid, sub_root)
+                                }
+                                _ => {}
+                            }
+                        }
                     }
+                    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[info.expect_placeholder_index()]
+                    // 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.
+                    delegate.instantiate_canonical_var_with_infer(info, 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[info.expect_placeholder_index()]
+            };
+            var_values.push(value)
+        }
+        CanonicalVarValues { var_values: delegate.cx().mk_args(&var_values) }
     }
 
     /// Unify the `original_values` with the `var_values` returned by the canonical query..

compiler/rustc_trait_selection/src/solve/delegate.rs

@@ -148,9 +148,10 @@ impl<'tcx> rustc_next_trait_solver::delegate::SolverDelegate for SolverDelegate<
         &self,
         cv_info: CanonicalVarInfo<'tcx>,
         span: Span,
+        var_values: &[ty::GenericArg<'tcx>],
         universe_map: impl Fn(ty::UniverseIndex) -> ty::UniverseIndex,
     ) -> ty::GenericArg<'tcx> {
-        self.0.instantiate_canonical_var(span, cv_info, universe_map)
+        self.0.instantiate_canonical_var(span, cv_info, var_values, universe_map)
     }
 
     fn register_hidden_type_in_storage(