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 #119989. Necessary for type inference guidance due to not-yet defined opaque types, cc rust-lang/trait-system-refactor-initiative#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 #119989.

r? `@compiler-errors`

Author: bors

Cargo.lock

@@ -4681,6 +4681,7 @@ dependencies = [
 name = "rustc_type_ir"
 version = "0.0.0"
 dependencies = [
+ "arrayvec",
  "bitflags",
  "derive-where",
  "ena",

compiler/rustc_hir_typeck/src/fn_ctxt/mod.rs

@@ -21,7 +21,6 @@ use rustc_middle::ty::{self, Const, Ty, TyCtxt, TypeVisitableExt};
 use rustc_session::Session;
 use rustc_span::{self, DUMMY_SP, ErrorGuaranteed, Ident, Span, sym};
 use rustc_trait_selection::error_reporting::TypeErrCtxt;
-use rustc_trait_selection::error_reporting::infer::sub_relations::SubRelations;
 use rustc_trait_selection::traits::{
     self, FulfillmentError, ObligationCause, ObligationCauseCode, ObligationCtxt,
 };
@@ -188,14 +187,8 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
     ///
     /// [`InferCtxtErrorExt::err_ctxt`]: rustc_trait_selection::error_reporting::InferCtxtErrorExt::err_ctxt
     pub(crate) fn err_ctxt(&'a self) -> TypeErrCtxt<'a, 'tcx> {
-        let mut sub_relations = SubRelations::default();
-        sub_relations.add_constraints(
-            self,
-            self.fulfillment_cx.borrow_mut().pending_obligations().iter().map(|o| o.predicate),
-        );
         TypeErrCtxt {
             infcx: &self.infcx,
-            sub_relations: RefCell::new(sub_relations),
             typeck_results: Some(self.typeck_results.borrow()),
             fallback_has_occurred: self.fallback_has_occurred.get(),
             normalize_fn_sig: Box::new(|fn_sig| {

compiler/rustc_hir_typeck/src/method/probe.rs

@@ -403,15 +403,13 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                 // special handling for this "trivial case" is a good idea.
 
                 let infcx = &self.infcx;
-                let (ParamEnvAnd { param_env: _, value: self_ty }, canonical_inference_vars) =
+                let (ParamEnvAnd { param_env: _, value: self_ty }, var_values) =
                     infcx.instantiate_canonical(span, &query_input.canonical);
                 debug!(?self_ty, ?query_input, "probe_op: Mode::Path");
                 MethodAutoderefStepsResult {
                     steps: infcx.tcx.arena.alloc_from_iter([CandidateStep {
-                        self_ty: self.make_query_response_ignoring_pending_obligations(
-                            canonical_inference_vars,
-                            self_ty,
-                        ),
+                        self_ty: self
+                            .make_query_response_ignoring_pending_obligations(var_values, self_ty),
                         autoderefs: 0,
                         from_unsafe_deref: false,
                         unsize: false,

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::{
@@ -17,7 +18,7 @@ use tracing::debug;
 
 use crate::infer::InferCtxt;
 use crate::infer::canonical::{
-    Canonical, CanonicalQueryInput, CanonicalTyVarKind, CanonicalVarKind, OriginalQueryValues,
+    Canonical, CanonicalQueryInput, CanonicalVarKind, OriginalQueryValues,
 };
 
 impl<'tcx> InferCtxt<'tcx> {
@@ -293,6 +294,13 @@ 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>,
+    /// 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: SsoHashMap<ty::TyVid, usize>,
     canonicalize_mode: &'cx dyn CanonicalizeMode,
     needs_canonical_flags: TypeFlags,
 
@@ -361,10 +369,8 @@ impl<'cx, 'tcx> TypeFolder<TyCtxt<'tcx>> for Canonicalizer<'cx, 'tcx> {
                             // FIXME: perf problem described in #55921.
                             ui = ty::UniverseIndex::ROOT;
                         }
-                        self.canonicalize_ty_var(
-                            CanonicalVarKind::Ty(CanonicalTyVarKind::General(ui)),
-                            t,
-                        )
+                        let sub_root = self.get_or_insert_sub_root(vid);
+                        self.canonicalize_ty_var(CanonicalVarKind::Ty { ui, sub_root }, t)
                     }
                 }
             }
@@ -374,15 +380,15 @@ impl<'cx, 'tcx> TypeFolder<TyCtxt<'tcx>> for Canonicalizer<'cx, 'tcx> {
                 if nt != t {
                     return self.fold_ty(nt);
                 } else {
-                    self.canonicalize_ty_var(CanonicalVarKind::Ty(CanonicalTyVarKind::Int), t)
+                    self.canonicalize_ty_var(CanonicalVarKind::Int, t)
                 }
             }
             ty::Infer(ty::FloatVar(vid)) => {
                 let nt = self.infcx.unwrap().opportunistic_resolve_float_var(vid);
                 if nt != t {
                     return self.fold_ty(nt);
                 } else {
-                    self.canonicalize_ty_var(CanonicalVarKind::Ty(CanonicalTyVarKind::Float), t)
+                    self.canonicalize_ty_var(CanonicalVarKind::Float, t)
                 }
             }
 
@@ -562,6 +568,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() {
@@ -660,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_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)
+    }
+
     /// 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.
@@ -679,11 +693,11 @@ impl<'cx, 'tcx> Canonicalizer<'cx, 'tcx> {
         self.variables
             .iter()
             .map(|&kind| match kind {
-                CanonicalVarKind::Ty(CanonicalTyVarKind::Int | CanonicalTyVarKind::Float) => {
+                CanonicalVarKind::Int | CanonicalVarKind::Float => {
                     return kind;
                 }
-                CanonicalVarKind::Ty(CanonicalTyVarKind::General(u)) => {
-                    CanonicalVarKind::Ty(CanonicalTyVarKind::General(reverse_universe_map[&u]))
+                CanonicalVarKind::Ty { ui, sub_root } => {
+                    CanonicalVarKind::Ty { ui: reverse_universe_map[&ui], sub_root }
                 }
                 CanonicalVarKind::Region(u) => CanonicalVarKind::Region(reverse_universe_map[&u]),
                 CanonicalVarKind::Const(u) => CanonicalVarKind::Const(reverse_universe_map[&u]),

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

@@ -24,7 +24,7 @@
 pub use instantiate::CanonicalExt;
 use rustc_index::IndexVec;
 pub use rustc_middle::infer::canonical::*;
-use rustc_middle::ty::{self, GenericArg, List, Ty, TyCtxt, TypeFoldable};
+use rustc_middle::ty::{self, GenericArg, Ty, TyCtxt, TypeFoldable};
 use rustc_span::Span;
 
 use crate::infer::{InferCtxt, RegionVariableOrigin};
@@ -67,30 +67,12 @@ impl<'tcx> InferCtxt<'tcx> {
             .chain((1..=canonical.max_universe.as_u32()).map(|_| self.create_next_universe()))
             .collect();
 
-        let canonical_inference_vars =
-            self.instantiate_canonical_vars(span, canonical.variables, |ui| universes[ui]);
-        let result = canonical.instantiate(self.tcx, &canonical_inference_vars);
-        (result, canonical_inference_vars)
-    }
-
-    /// Given the "infos" about the canonical variables from some
-    /// canonical, creates fresh variables with the same
-    /// characteristics (see `instantiate_canonical_var` for
-    /// details). You can then use `instantiate` to instantiate the
-    /// canonical variable with these inference variables.
-    fn instantiate_canonical_vars(
-        &self,
-        span: Span,
-        variables: &List<CanonicalVarKind<'tcx>>,
-        universe_map: impl Fn(ty::UniverseIndex) -> ty::UniverseIndex,
-    ) -> CanonicalVarValues<'tcx> {
-        CanonicalVarValues {
-            var_values: self.tcx.mk_args_from_iter(
-                variables
-                    .iter()
-                    .map(|kind| self.instantiate_canonical_var(span, kind, &universe_map)),
-            ),
-        }
+        let var_values =
+            CanonicalVarValues::instantiate(self.tcx, &canonical.variables, |var_values, info| {
+                self.instantiate_canonical_var(span, info, &var_values, |ui| universes[ui])
+            });
+        let result = canonical.instantiate(self.tcx, &var_values);
+        (result, var_values)
     }
 
     /// Given the "info" about a canonical variable, creates a fresh
@@ -105,21 +87,27 @@ impl<'tcx> InferCtxt<'tcx> {
         &self,
         span: Span,
         kind: CanonicalVarKind<'tcx>,
+        previous_var_values: &[GenericArg<'tcx>],
         universe_map: impl Fn(ty::UniverseIndex) -> ty::UniverseIndex,
     ) -> GenericArg<'tcx> {
         match kind {
-            CanonicalVarKind::Ty(ty_kind) => {
-                let ty = match ty_kind {
-                    CanonicalTyVarKind::General(ui) => {
-                        self.next_ty_var_in_universe(span, universe_map(ui))
+            CanonicalVarKind::Ty { ui, sub_root } => {
+                let vid = self.next_ty_vid_in_universe(span, universe_map(ui));
+                // If this inference variable is related to an earlier variable
+                // via subtyping, we need to add that info to the inference context.
+                if let Some(prev) = previous_var_values.get(sub_root.as_usize()) {
+                    if let &ty::Infer(ty::TyVar(sub_root)) = prev.expect_ty().kind() {
+                        self.sub_unify_ty_vids_raw(vid, sub_root);
+                    } else {
+                        unreachable!()
                     }
+                }
+                Ty::new_var(self.tcx, vid).into()
+            }
 
-                    CanonicalTyVarKind::Int => self.next_int_var(),
+            CanonicalVarKind::Int => self.next_int_var().into(),
 
-                    CanonicalTyVarKind::Float => self.next_float_var(),
-                };
-                ty.into()
-            }
+            CanonicalVarKind::Float => self.next_float_var().into(),
 
             CanonicalVarKind::PlaceholderTy(ty::PlaceholderType { universe, bound }) => {
                 let universe_mapped = universe_map(universe);

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

@@ -13,6 +13,7 @@ use std::iter;
 use rustc_index::{Idx, IndexVec};
 use rustc_middle::arena::ArenaAllocatable;
 use rustc_middle::bug;
+use rustc_middle::infer::canonical::CanonicalVarKind;
 use rustc_middle::ty::{self, BoundVar, GenericArg, GenericArgKind, Ty, TyCtxt, TypeFoldable};
 use tracing::{debug, instrument};
 
@@ -413,35 +414,34 @@ impl<'tcx> InferCtxt<'tcx> {
         let mut opt_values: IndexVec<BoundVar, Option<GenericArg<'tcx>>> =
             IndexVec::from_elem_n(None, query_response.variables.len());
 
-        // In terms of our example above, we are iterating over pairs like:
-        // [(?A, Vec<?0>), ('static, '?1), (?B, ?0)]
         for (original_value, result_value) in iter::zip(&original_values.var_values, result_values)
         {
             match result_value.kind() {
                 GenericArgKind::Type(result_value) => {
-                    // e.g., here `result_value` might be `?0` in the example above...
-                    if let ty::Bound(debruijn, b) = *result_value.kind() {
-                        // ...in which case we would set `canonical_vars[0]` to `Some(?U)`.
-
+                    // 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) = *result_value.kind()
+                        && !matches!(
+                            query_response.variables[b.var.as_usize()],
+                            CanonicalVarKind::Ty { .. }
+                        )
+                    {
                         // We only allow a `ty::INNERMOST` index in generic parameters.
                         assert_eq!(debruijn, ty::INNERMOST);
                         opt_values[b.var] = Some(*original_value);
                     }
                 }
                 GenericArgKind::Lifetime(result_value) => {
-                    // e.g., here `result_value` might be `'?1` in the example above...
                     if let ty::ReBound(debruijn, b) = result_value.kind() {
-                        // ... in which case we would set `canonical_vars[0]` to `Some('static)`.
-
                         // We only allow a `ty::INNERMOST` index in generic parameters.
                         assert_eq!(debruijn, ty::INNERMOST);
                         opt_values[b.var] = Some(*original_value);
                     }
                 }
                 GenericArgKind::Const(result_value) => {
                     if let ty::ConstKind::Bound(debruijn, b) = result_value.kind() {
-                        // ...in which case we would set `canonical_vars[0]` to `Some(const X)`.
-
                         // We only allow a `ty::INNERMOST` index in generic parameters.
                         assert_eq!(debruijn, ty::INNERMOST);
                         opt_values[b.var] = Some(*original_value);
@@ -453,39 +453,36 @@ 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, 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, we have to deal with them for now.
-                        self.instantiate_canonical_var(cause.span, var_kind, |u| {
-                            universe_map[u.as_usize()]
-                        })
-                    } else if var_kind.is_existential() {
-                        match opt_values[BoundVar::new(index)] {
-                            Some(k) => k,
-                            None => self.instantiate_canonical_var(cause.span, var_kind, |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",
-                        )
-                    }
-                }),
-            ),
-        };
+        let tcx = self.tcx;
+        let variables = query_response.variables;
+        let var_values = CanonicalVarValues::instantiate(tcx, 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, we have to deal with them for now.
+                self.instantiate_canonical_var(cause.span, kind, &var_values, |u| {
+                    universe_map[u.as_usize()]
+                })
+            } else if kind.is_existential() {
+                match opt_values[BoundVar::new(var_values.len())] {
+                    Some(k) => k,
+                    None => self.instantiate_canonical_var(cause.span, kind, &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(var_values.len())]
+                    .expect("expected placeholder to be unified with itself during response")
+            }
+        });
 
         let mut obligations = PredicateObligations::new();
 
         // Carry all newly resolved opaque types to the caller's scope
         for &(a, b) in &query_response.value.opaque_types {
-            let a = instantiate_value(self.tcx, &result_args, a);
-            let b = instantiate_value(self.tcx, &result_args, b);
+            let a = instantiate_value(self.tcx, &var_values, a);
+            let b = instantiate_value(self.tcx, &var_values, b);
             debug!(?a, ?b, "constrain opaque type");
             // We use equate here instead of, for example, just registering the
             // opaque type's hidden value directly, because the hidden type may have been an inference
@@ -502,7 +499,7 @@ impl<'tcx> InferCtxt<'tcx> {
             );
         }
 
-        Ok(InferOk { value: result_args, obligations })
+        Ok(InferOk { value: var_values, obligations })
     }
 
     /// Given a "guess" at the values for the canonical variables in

compiler/rustc_infer/src/infer/context.rs

@@ -59,6 +59,10 @@ impl<'tcx> rustc_type_ir::InferCtxtLike for InferCtxt<'tcx> {
         self.root_var(var)
     }
 
+    fn sub_unification_table_root_var(&self, var: ty::TyVid) -> ty::TyVid {
+        self.sub_unification_table_root_var(var)
+    }
+
     fn root_const_var(&self, var: ty::ConstVid) -> ty::ConstVid {
         self.root_const_var(var)
     }
@@ -179,6 +183,10 @@ impl<'tcx> rustc_type_ir::InferCtxtLike for InferCtxt<'tcx> {
         self.inner.borrow_mut().type_variables().equate(a, b);
     }
 
+    fn sub_unify_ty_vids_raw(&self, a: ty::TyVid, b: ty::TyVid) {
+        self.sub_unify_ty_vids_raw(a, b);
+    }
+
     fn equate_int_vids_raw(&self, a: ty::IntVid, b: ty::IntVid) {
         self.inner.borrow_mut().int_unification_table().union(a, b);
     }