Improve canonicalization performance by introducing TypingMode::ErasedNotCoherence

compiler/rustc_const_eval/src/const_eval/eval_queries.rs

@@ -8,7 +8,7 @@ use rustc_middle::mir::{self, ConstAlloc, ConstValue};
 use rustc_middle::query::TyCtxtAt;
 use rustc_middle::ty::layout::{HasTypingEnv, TyAndLayout};
 use rustc_middle::ty::print::with_no_trimmed_paths;
-use rustc_middle::ty::{self, Ty, TyCtxt};
+use rustc_middle::ty::{self, MayBeErased, Ty, TyCtxt};
 use rustc_middle::{bug, throw_inval};
 use rustc_span::Span;
 use rustc_span::def_id::LocalDefId;
@@ -384,6 +384,7 @@ pub fn eval_to_allocation_raw_provider<'tcx>(
                     "Const eval should always happens in PostAnalysis mode. See the comment in `InterpCx::new` for more details."
                 )
             }
+            ty::TypingMode::ErasedNotCoherence(MayBeErased) => todo!(),
         }
 
         // Make sure we format the instance even if we do not print it.

compiler/rustc_const_eval/src/const_eval/valtrees.rs

@@ -3,7 +3,7 @@ use rustc_data_structures::stack::ensure_sufficient_stack;
 use rustc_middle::mir::interpret::{EvalToValTreeResult, GlobalId, ValTreeCreationError};
 use rustc_middle::traits::ObligationCause;
 use rustc_middle::ty::layout::{LayoutCx, TyAndLayout};
-use rustc_middle::ty::{self, Ty, TyCtxt};
+use rustc_middle::ty::{self, MayBeErased, Ty, TyCtxt};
 use rustc_middle::{bug, mir};
 use rustc_span::DUMMY_SP;
 use tracing::{debug, instrument, trace};
@@ -247,6 +247,7 @@ pub(crate) fn eval_to_valtree<'tcx>(
                     "Const eval should always happens in PostAnalysis mode. See the comment in `InterpCx::new` for more details."
                 )
             }
+            ty::TypingMode::ErasedNotCoherence(MayBeErased) => unreachable!(),
         }
     }
     let const_alloc = tcx.eval_to_allocation_raw(typing_env.as_query_input(cid))?;

compiler/rustc_const_eval/src/interpret/eval_context.rs

@@ -9,8 +9,8 @@ use rustc_middle::ty::layout::{
     LayoutOfHelpers, TyAndLayout,
 };
 use rustc_middle::ty::{
-    self, GenericArgsRef, Ty, TyCtxt, TypeFoldable, TypeVisitableExt, TypingEnv, TypingMode,
-    Variance,
+    self, GenericArgsRef, MayBeErased, Ty, TyCtxt, TypeFoldable, TypeVisitableExt, TypingEnv,
+    TypingMode, Variance,
 };
 use rustc_middle::{bug, mir, span_bug};
 use rustc_span::Span;
@@ -251,6 +251,7 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
                 | TypingMode::PostBorrowckAnalysis { .. } => {
                     bug!("Const eval should always happens in PostAnalysis mode.");
                 }
+                TypingMode::ErasedNotCoherence(MayBeErased) => unreachable!(),
             }
         }
 

compiler/rustc_hir_typeck/src/callee.rs

@@ -62,6 +62,7 @@ enum CallStep<'tcx> {
 }
 
 impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
+    #[tracing::instrument(skip(self))]
     pub(crate) fn check_expr_call(
         &self,
         call_expr: &'tcx hir::Expr<'tcx>,

compiler/rustc_hir_typeck/src/fn_ctxt/mod.rs

@@ -17,7 +17,9 @@ use rustc_hir_analysis::hir_ty_lowering::{
 };
 use rustc_infer::infer::{self, RegionVariableOrigin};
 use rustc_infer::traits::{DynCompatibilityViolation, Obligation};
-use rustc_middle::ty::{self, Const, Flags, Ty, TyCtxt, TypeVisitableExt, Unnormalized};
+use rustc_middle::ty::{
+    self, CantBeErased, Const, Flags, Ty, TyCtxt, TypeVisitableExt, TypingMode, Unnormalized,
+};
 use rustc_session::Session;
 use rustc_span::{self, DUMMY_SP, ErrorGuaranteed, Ident, Span};
 use rustc_trait_selection::error_reporting::TypeErrCtxt;
@@ -161,6 +163,12 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
         }
     }
 
+    pub(crate) fn typing_mode(&self) -> TypingMode<'tcx, CantBeErased> {
+        // `FnCtxt` is never constructed in the trait solver, so we can safely use
+        // `assert_not_erased`.
+        self.infcx.typing_mode_raw().assert_not_erased()
+    }
+
     pub(crate) fn dcx(&self) -> DiagCtxtHandle<'a> {
         self.root_ctxt.infcx.dcx()
     }

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

@@ -72,7 +72,7 @@ impl<'tcx> InferCtxt<'tcx> {
             query_state,
         )
         .unchecked_map(|(param_env, value)| param_env.and(value));
-        CanonicalQueryInput { canonical, typing_mode: TypingModeEqWrapper(self.typing_mode()) }
+        CanonicalQueryInput { canonical, typing_mode: TypingModeEqWrapper(self.typing_mode_raw()) }
     }
 
     /// Canonicalizes a query *response* `V`. When we canonicalize a

compiler/rustc_infer/src/infer/context.rs

@@ -22,8 +22,8 @@ impl<'tcx> rustc_type_ir::InferCtxtLike for InferCtxt<'tcx> {
         self.next_trait_solver
     }
 
-    fn typing_mode(&self) -> ty::TypingMode<'tcx> {
-        self.typing_mode()
+    fn typing_mode_raw(&self) -> ty::TypingMode<'tcx> {
+        self.typing_mode_raw()
     }
 
     fn universe(&self) -> ty::UniverseIndex {

compiler/rustc_infer/src/infer/mod.rs

@@ -34,6 +34,7 @@ use rustc_middle::ty::{
     TypeSuperFoldable, TypeVisitable, TypeVisitableExt, TypingEnv, TypingMode, fold_regions,
 };
 use rustc_span::{DUMMY_SP, Span, Symbol};
+use rustc_type_ir::MayBeErased;
 use snapshot::undo_log::InferCtxtUndoLogs;
 use tracing::{debug, instrument};
 use type_variable::TypeVariableOrigin;
@@ -640,8 +641,28 @@ impl<'tcx> InferCtxt<'tcx> {
         self.next_trait_solver
     }
 
+    /// This method is deliberately called `..._raw`,
+    /// since the output may possibly include [`TypingMode::ErasedNotCoherence`](TypingMode::ErasedNotCoherence).
+    /// `ErasedNotCoherence` is an implementation detail of the next trait solver, see its docs for
+    /// more information.
+    ///
+    /// `InferCtxt` has two uses: the trait solver calls some methods on it, because the `InferCtxt`
+    /// works as a kind of store for for example type unification information.
+    /// `InferCtxt` is also often used outside the trait solver during typeck.
+    /// There, we don't care about the `ErasedNotCoherence` case and should never encounter it.
+    /// To make sure these two uses are never confused, we want to statically encode this information.
+    ///
+    /// The `FnCtxt`, for example, is only used in the outside-trait-solver case. It has a non-raw
+    /// version of the `typing_mode` method available that asserts `ErasedNotCoherence` is
+    /// impossible, and returns a `TypingMode` where `ErasedNotCoherence` is made uninhabited using
+    /// the [`CantBeErased`](rustc_type_ir::CantBeErased) enum. That way you don't even have to
+    /// match on the variant and can safely ignore it.
+    ///
+    /// Prefer non-raw apis if available. e.g.,
+    /// - On the `FnCtxt`
+    /// - on the `SelectionCtxt`
     #[inline(always)]
-    pub fn typing_mode(&self) -> TypingMode<'tcx> {
+    pub fn typing_mode_raw(&self) -> TypingMode<'tcx> {
         self.typing_mode
     }
 
@@ -1064,7 +1085,7 @@ impl<'tcx> InferCtxt<'tcx> {
     #[inline(always)]
     pub fn can_define_opaque_ty(&self, id: impl Into<DefId>) -> bool {
         debug_assert!(!self.next_trait_solver());
-        match self.typing_mode() {
+        match self.typing_mode_raw() {
             TypingMode::Analysis {
                 defining_opaque_types_and_generators: defining_opaque_types,
             }
@@ -1076,7 +1097,8 @@ impl<'tcx> InferCtxt<'tcx> {
             // to support PostBorrowckAnalysis in the old solver as well.
             TypingMode::Coherence
             | TypingMode::PostBorrowckAnalysis { .. }
-            | TypingMode::PostAnalysis => false,
+            | TypingMode::PostAnalysis
+            | TypingMode::ErasedNotCoherence(MayBeErased) => false,
         }
     }
 
@@ -1395,7 +1417,7 @@ impl<'tcx> InferCtxt<'tcx> {
     /// which contains the necessary information to use the trait system without
     /// using canonicalization or carrying this inference context around.
     pub fn typing_env(&self, param_env: ty::ParamEnv<'tcx>) -> ty::TypingEnv<'tcx> {
-        let typing_mode = match self.typing_mode() {
+        let typing_mode = match self.typing_mode_raw() {
             // FIXME(#132279): This erases the `defining_opaque_types` as it isn't possible
             // to handle them without proper canonicalization. This means we may cause cycle
             // errors and fail to reveal opaques while inside of bodies. We should rename this
@@ -1407,6 +1429,7 @@ impl<'tcx> InferCtxt<'tcx> {
             mode @ (ty::TypingMode::Coherence
             | ty::TypingMode::PostBorrowckAnalysis { .. }
             | ty::TypingMode::PostAnalysis) => mode,
+            ty::TypingMode::ErasedNotCoherence(MayBeErased) => unreachable!(),
         };
         ty::TypingEnv::new(param_env, typing_mode)
     }

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

@@ -89,7 +89,7 @@ impl<'tcx> InferCtxt<'tcx> {
                 if def_id.is_local() =>
             {
                 let def_id = def_id.expect_local();
-                if self.typing_mode().is_coherence() {
+                if self.typing_mode_raw().is_coherence() {
                     // See comment on `insert_hidden_type` for why this is sufficient in coherence
                     return Some(self.register_hidden_type(
                         OpaqueTypeKey { def_id, args },
@@ -229,7 +229,9 @@ impl<'tcx> InferCtxt<'tcx> {
         // value being folded. In simple cases like `-> impl Foo`,
         // these are the same span, but not in cases like `-> (impl
         // Foo, impl Bar)`.
-        match self.typing_mode() {
+        //
+        // Note: we don't use this function in the next solver so we can safely call `assert_not_erased`
+        match self.typing_mode_raw().assert_not_erased() {
             ty::TypingMode::Coherence => {
                 // During intercrate we do not define opaque types but instead always
                 // force ambiguity unless the hidden type is known to not implement

compiler/rustc_infer/src/infer/relate/generalize.rs

@@ -604,7 +604,7 @@ impl<'tcx> TypeRelation<TyCtxt<'tcx>> for Generalizer<'_, 'tcx> {
                             //
                             // cc trait-system-refactor-initiative#108
                             if self.infcx.next_trait_solver()
-                                && !self.infcx.typing_mode().is_coherence()
+                                && !self.infcx.typing_mode_raw().is_coherence()
                                 && self.in_alias
                             {
                                 inner.type_variables().equate(vid, new_var_id);
@@ -736,7 +736,7 @@ impl<'tcx> TypeRelation<TyCtxt<'tcx>> for Generalizer<'_, 'tcx> {
                             // See the comment for type inference variables
                             // for more details.
                             if self.infcx.next_trait_solver()
-                                && !self.infcx.typing_mode().is_coherence()
+                                && !self.infcx.typing_mode_raw().is_coherence()
                                 && self.in_alias
                             {
                                 variable_table.union(vid, new_var_id);

compiler/rustc_middle/src/traits/solve.rs

@@ -14,6 +14,8 @@ pub type QueryResult<'tcx> = ir::solve::QueryResult<TyCtxt<'tcx>>;
 pub type CandidateSource<'tcx> = ir::solve::CandidateSource<TyCtxt<'tcx>>;
 pub type CanonicalInput<'tcx, P = ty::Predicate<'tcx>> = ir::solve::CanonicalInput<TyCtxt<'tcx>, P>;
 pub type CanonicalResponse<'tcx> = ir::solve::CanonicalResponse<TyCtxt<'tcx>>;
+pub type FetchEligibleAssocItemResponse<'tcx> =
+    ir::solve::FetchEligibleAssocItemResponse<TyCtxt<'tcx>>;
 
 pub type PredefinedOpaques<'tcx> = &'tcx ty::List<(ty::OpaqueTypeKey<'tcx>, Ty<'tcx>)>;
 

compiler/rustc_middle/src/ty/mod.rs

@@ -1063,7 +1063,7 @@ impl<'tcx> TypingEnv<'tcx> {
         def_id: impl IntoQueryKey<DefId>,
     ) -> TypingEnv<'tcx> {
         let def_id = def_id.into_query_key();
-        Self::new(tcx.param_env(def_id), TypingMode::non_body_analysis())
+        Self::new(tcx.param_env(def_id), TypingMode::non_body_analysis().into())
     }
 
     pub fn post_analysis(tcx: TyCtxt<'tcx>, def_id: impl IntoQueryKey<DefId>) -> TypingEnv<'tcx> {
@@ -1080,7 +1080,12 @@ impl<'tcx> TypingEnv<'tcx> {
             | TypingMode::Analysis { .. }
             | TypingMode::Borrowck { .. }
             | TypingMode::PostBorrowckAnalysis { .. } => {}
-            TypingMode::PostAnalysis => return self,
+            TypingMode::PostAnalysis
+            // If we get here in erased mode, we don't know the original typing mode.
+            // It could be post analysis, or something else.
+            // Regardless of what it was, we want to continue in erased mode,
+            // if we end up accessing any opaque types we bail out anyway.
+            | TypingMode::ErasedNotCoherence(MayBeErased) => return self,
         }
 
         // No need to reveal opaques with the new solver enabled,

compiler/rustc_middle/src/ty/sty.rs

@@ -17,7 +17,9 @@ use rustc_span::{DUMMY_SP, Span, Symbol, kw, sym};
 use rustc_type_ir::TyKind::*;
 use rustc_type_ir::solve::SizedTraitKind;
 use rustc_type_ir::walk::TypeWalker;
-use rustc_type_ir::{self as ir, BoundVar, CollectAndApply, TypeVisitableExt, elaborate};
+use rustc_type_ir::{
+    self as ir, BoundVar, CollectAndApply, MayBeErased, TypeVisitableExt, elaborate,
+};
 use tracing::instrument;
 use ty::util::IntTypeExt;
 
@@ -41,7 +43,7 @@ pub type FnSigKind<'tcx> = ir::FnSigKind<TyCtxt<'tcx>>;
 pub type Binder<'tcx, T> = ir::Binder<TyCtxt<'tcx>, T>;
 pub type EarlyBinder<'tcx, T> = ir::EarlyBinder<TyCtxt<'tcx>, T>;
 pub type Unnormalized<'tcx, T> = ir::Unnormalized<TyCtxt<'tcx>, T>;
-pub type TypingMode<'tcx> = ir::TypingMode<TyCtxt<'tcx>>;
+pub type TypingMode<'tcx, S = MayBeErased> = ir::TypingMode<TyCtxt<'tcx>, S>;
 pub type TypingModeEqWrapper<'tcx> = ir::TypingModeEqWrapper<TyCtxt<'tcx>>;
 pub type Placeholder<'tcx, T> = ir::Placeholder<TyCtxt<'tcx>, T>;
 pub type PlaceholderRegion<'tcx> = ir::PlaceholderRegion<TyCtxt<'tcx>>;

compiler/rustc_mir_transform/src/elaborate_drop.rs

@@ -7,7 +7,7 @@ use rustc_index::Idx;
 use rustc_middle::mir::*;
 use rustc_middle::ty::adjustment::PointerCoercion;
 use rustc_middle::ty::util::IntTypeExt;
-use rustc_middle::ty::{self, GenericArg, GenericArgsRef, Ty, TyCtxt, Unnormalized};
+use rustc_middle::ty::{self, GenericArg, GenericArgsRef, MayBeErased, Ty, TyCtxt, Unnormalized};
 use rustc_middle::{bug, span_bug, traits};
 use rustc_span::{DUMMY_SP, Spanned, dummy_spanned};
 use tracing::{debug, instrument};
@@ -556,6 +556,7 @@ where
                     | ty::TypingMode::PostBorrowckAnalysis { .. } => {
                         bug!()
                     }
+                    ty::TypingMode::ErasedNotCoherence(MayBeErased) => unreachable!(),
                 }
 
                 let field_ty = field.ty(tcx, args);

compiler/rustc_next_trait_solver/src/canonical/mod.rs

@@ -17,7 +17,7 @@ use rustc_type_ir::inherent::*;
 use rustc_type_ir::relate::solver_relating::RelateExt;
 use rustc_type_ir::{
     self as ty, Canonical, CanonicalVarKind, CanonicalVarValues, InferCtxtLike, Interner,
-    TypeFoldable, TypingModeEqWrapper,
+    MayBeErased, TypeFoldable, TypingMode, TypingModeEqWrapper,
 };
 use tracing::instrument;
 
@@ -46,6 +46,13 @@ impl<I: Interner, T> ResponseT<I> for inspect::State<I, T> {
     }
 }
 
+pub(super) enum EraseOpaqueTypes {
+    /// This setting erases opaque types, unless we're in coherence.
+    /// In `TypingMode::Coherence` we never erase opaque types
+    IfNotCoherence,
+    No,
+}
+
 /// Canonicalizes the goal remembering the original values
 /// for each bound variable.
 ///
@@ -54,22 +61,45 @@ pub(super) fn canonicalize_goal<D, I>(
     delegate: &D,
     goal: Goal<I, I::Predicate>,
     opaque_types: &[(ty::OpaqueTypeKey<I>, I::Ty)],
+    erase_opaque_types: EraseOpaqueTypes,
 ) -> (Vec<I::GenericArg>, CanonicalInput<I, I::Predicate>)
 where
     D: SolverDelegate<Interner = I>,
     I: Interner,
 {
+    let (opaque_types, typing_mode) = match (erase_opaque_types, delegate.typing_mode_raw()) {
+        // In `TypingMode::Coherence` there should not be any opaques, and we also don't change typing mode.
+        (_, TypingMode::Coherence) => {
+            assert!(opaque_types.is_empty());
+            (&[][..], TypingMode::Coherence)
+        }
+        // Make sure we're not recursively in `ErasedNotCoherence`.
+        (_, TypingMode::ErasedNotCoherence(MayBeErased)) => {
+            assert!(opaque_types.is_empty());
+            (&[][..], TypingMode::ErasedNotCoherence(MayBeErased))
+        }
+        // If we're supposed to erase opaque types, and we're in any typing mode other than coherence,
+        // do the erasing and change typing mode.
+        (
+            EraseOpaqueTypes::IfNotCoherence,
+            TypingMode::Analysis { .. }
+            | TypingMode::Borrowck { .. }
+            | TypingMode::PostBorrowckAnalysis { .. }
+            | TypingMode::PostAnalysis,
+        ) => (&[][..], TypingMode::ErasedNotCoherence(MayBeErased)),
+        (EraseOpaqueTypes::No, typing_mode) => (opaque_types, typing_mode),
+    };
+
     let (orig_values, canonical) = Canonicalizer::canonicalize_input(
         delegate,
         QueryInput {
             goal,
             predefined_opaques_in_body: delegate.cx().mk_predefined_opaques_in_body(opaque_types),
         },
     );
-    let query_input = ty::CanonicalQueryInput {
-        canonical,
-        typing_mode: TypingModeEqWrapper(delegate.typing_mode()),
-    };
+
+    let query_input =
+        ty::CanonicalQueryInput { canonical, typing_mode: TypingModeEqWrapper(typing_mode) };
     (orig_values, query_input)
 }
 

compiler/rustc_next_trait_solver/src/delegate.rs

@@ -1,6 +1,6 @@
 use std::ops::Deref;
 
-use rustc_type_ir::solve::{Certainty, Goal, NoSolution};
+use rustc_type_ir::solve::{Certainty, FetchEligibleAssocItemResponse, Goal, NoSolution};
 use rustc_type_ir::{self as ty, InferCtxtLike, Interner, TypeFoldable};
 
 pub trait SolverDelegate: Deref<Target = Self::Infcx> + Sized {
@@ -77,10 +77,7 @@ pub trait SolverDelegate: Deref<Target = Self::Infcx> + Sized {
         goal_trait_ref: ty::TraitRef<Self::Interner>,
         trait_assoc_def_id: <Self::Interner as Interner>::DefId,
         impl_def_id: <Self::Interner as Interner>::ImplId,
-    ) -> Result<
-        Option<<Self::Interner as Interner>::DefId>,
-        <Self::Interner as Interner>::ErrorGuaranteed,
-    >;
+    ) -> FetchEligibleAssocItemResponse<Self::Interner>;
 
     fn is_transmutable(
         &self,