Reduce next-solver memory usage by interning CanonicalQueryInput

compiler/rustc_middle/src/query/keys.rs

@@ -346,6 +346,24 @@ impl<'tcx, T: QueryKeyBounds> QueryKey for (CanonicalQueryInput<'tcx, T>, usize)
     }
 }
 
+impl<'tcx> QueryKey for crate::traits::solve::CanonicalInput<'tcx> {
+    fn default_span(&self, _tcx: TyCtxt<'_>) -> Span {
+        DUMMY_SP
+    }
+}
+
+impl<'tcx> QueryKey for (crate::traits::solve::CanonicalInput<'tcx>, bool) {
+    fn default_span(&self, _tcx: TyCtxt<'_>) -> Span {
+        DUMMY_SP
+    }
+}
+
+impl<'tcx> QueryKey for (crate::traits::solve::CanonicalInput<'tcx>, usize) {
+    fn default_span(&self, _tcx: TyCtxt<'_>) -> Span {
+        DUMMY_SP
+    }
+}
+
 impl<'tcx> QueryKey for (Ty<'tcx>, rustc_abi::VariantIdx) {
     fn default_span(&self, _tcx: TyCtxt<'_>) -> Span {
         DUMMY_SP

compiler/rustc_middle/src/traits/solve.rs

@@ -12,7 +12,6 @@ pub type Goal<'tcx, P> = ir::solve::Goal<TyCtxt<'tcx>, P>;
 pub type QueryInput<'tcx, P> = ir::solve::QueryInput<TyCtxt<'tcx>, P>;
 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>>;
@@ -23,6 +22,24 @@ pub type SucceededInErased<'tcx> = ir::solve::SucceededInErased<TyCtxt<'tcx>>;
 
 pub type PredefinedOpaques<'tcx> = &'tcx ty::List<(ty::OpaqueTypeKey<'tcx>, Ty<'tcx>)>;
 
+// Interning CanonicalInput drastically reduces max memory usage when compiling a crate that has
+// trait solver recursion depth overflows with next-solver deduplicating individual inputs.
+// This mostly fixes #161748 where it reduced the memory usage for compiling bevy_render from
+// ~14GiB to ~4GiB
+// Main improved types:
+//   - rustc_type_ir::search_graph::GlobalCache
+//   - rustc_type_ir::search_graph::NestedGoals
+#[derive(Debug, PartialEq, Eq, Copy, Clone, Hash, StableHash)]
+pub struct CanonicalInput<'tcx>(pub(crate) Interned<'tcx, CanonicalInputData<TyCtxt<'tcx>>>);
+
+impl<'tcx> std::ops::Deref for CanonicalInput<'tcx> {
+    type Target = CanonicalInputData<TyCtxt<'tcx>>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
 #[derive(Debug, PartialEq, Eq, Copy, Clone, Hash, StableHash)]
 pub struct ExternalConstraints<'tcx>(
     pub(crate) Interned<'tcx, ExternalConstraintsData<TyCtxt<'tcx>>>,

compiler/rustc_middle/src/ty/context.rs

@@ -63,7 +63,10 @@ use crate::mir::{Body, Local, Place, PlaceElem, ProjectionKind, Promoted};
 use crate::query::{IntoQueryKey, LocalCrate, Providers, QuerySystem, TyCtxtAt};
 use crate::thir::Thir;
 use crate::traits;
-use crate::traits::solve::{ExternalConstraints, ExternalConstraintsData, PredefinedOpaques};
+use crate::traits::solve::{
+    CanonicalInput, CanonicalInputData, ExternalConstraints, ExternalConstraintsData,
+    PredefinedOpaques,
+};
 use crate::ty::predicate::ExistentialPredicateStableCmpExt as _;
 use crate::ty::region::RegionExt;
 use crate::ty::{
@@ -162,6 +165,7 @@ pub struct CtxtInterners<'tcx> {
     valtree: InternedSet<'tcx, ty::ValTreeKind<TyCtxt<'tcx>>>,
     patterns: InternedSet<'tcx, List<ty::Pattern<'tcx>>>,
     outlives: InternedSet<'tcx, List<ty::ArgOutlivesClause<'tcx>>>,
+    canonical_inputs: InternedSet<'tcx, CanonicalInputData<TyCtxt<'tcx>>>,
 }
 
 impl<'tcx> CtxtInterners<'tcx> {
@@ -200,6 +204,7 @@ impl<'tcx> CtxtInterners<'tcx> {
             valtree: InternedSet::with_capacity(N),
             patterns: InternedSet::with_capacity(N),
             outlives: InternedSet::with_capacity(N),
+            canonical_inputs: InternedSet::with_capacity(N),
         }
     }
 
@@ -1990,6 +1995,7 @@ direct_interners! {
     adt_def: pub mk_adt_def_from_data(AdtDefData): AdtDef -> AdtDef<'tcx>,
     external_constraints: pub mk_external_constraints(ExternalConstraintsData<TyCtxt<'tcx>>):
         ExternalConstraints -> ExternalConstraints<'tcx>,
+    canonical_inputs: intern_canonical_input(CanonicalInputData<TyCtxt<'tcx>>): CanonicalInput -> CanonicalInput<'tcx>,
 }
 
 macro_rules! slice_interners {

compiler/rustc_middle/src/ty/context/impl_interner.rs

@@ -10,6 +10,7 @@ use rustc_hir::def::{CtorKind, DefKind};
 use rustc_hir::def_id::{DefId, LocalDefId};
 use rustc_span::{DUMMY_SP, Span, Symbol};
 use rustc_type_ir::lang_items::{SolverAdtLangItem, SolverProjectionLangItem, SolverTraitLangItem};
+use rustc_type_ir::solve::CanonicalInputData;
 use rustc_type_ir::{
     BoundVar, CollectAndApply, DebruijnIndex, Interner, TypeFoldable, Unnormalized, VisitorResult,
     search_graph, try_visit,
@@ -664,6 +665,10 @@ impl<'tcx> Interner for TyCtxt<'tcx> {
     fn mk_probe(self, probe: inspect::Probe<Self>) -> &'tcx inspect::Probe<TyCtxt<'tcx>> {
         self.arena.alloc(probe)
     }
+    type CanonicalInput = CanonicalInput<'tcx>;
+    fn mk_canonical_input(self, data: CanonicalInputData<Self>) -> CanonicalInput<'tcx> {
+        self.intern_canonical_input(data)
+    }
     fn evaluate_root_goal_for_proof_tree_raw(
         self,
         canonical_goal: CanonicalInput<'tcx>,

compiler/rustc_next_trait_solver/src/canonical/mod.rs

@@ -26,9 +26,8 @@ use tracing::instrument;
 
 use crate::delegate::SolverDelegate;
 use crate::solve::{
-    CanonicalInput, CanonicalResponse, Certainty, ExternalConstraintsData,
-    ExternalRegionConstraints, Goal, NestedNormalizationGoals, QueryInput, Response,
-    VisibleForLeakCheck, inspect,
+    CanonicalResponse, Certainty, ExternalConstraintsData, ExternalRegionConstraints, Goal,
+    NestedNormalizationGoals, QueryInput, Response, VisibleForLeakCheck, inspect,
 };
 
 pub mod canonicalizer;
@@ -58,7 +57,7 @@ pub(super) fn canonicalize_goal<D, I>(
     goal: Goal<I, I::Predicate>,
     opaque_types: &[(ty::OpaqueTypeKey<I>, I::Ty)],
     typing_mode: TypingMode<I>,
-) -> (ThinVec<I::GenericArg>, CanonicalInput<I, I::Predicate>)
+) -> (ThinVec<I::GenericArg>, I::CanonicalInput)
 where
     D: SolverDelegate<Interner = I>,
     I: Interner,
@@ -71,8 +70,10 @@ where
         },
     );
 
-    let query_input =
-        ty::CanonicalQueryInput { canonical, typing_mode: TypingModeEqWrapper(typing_mode) };
+    let query_input = delegate.cx().mk_canonical_input(ty::CanonicalQueryInput {
+        canonical,
+        typing_mode: TypingModeEqWrapper(typing_mode),
+    });
     (orig_values, query_input)
 }
 

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

@@ -41,8 +41,8 @@ use crate::solve::fast_path::compute_goal_fast_path_cold;
 use crate::solve::search_graph::SearchGraph;
 use crate::solve::ty::may_use_unstable_feature;
 use crate::solve::{
-    CanonicalInput, CanonicalResponse, Certainty, ExternalConstraintsData, FIXPOINT_STEP_LIMIT,
-    Goal, GoalEvaluation, GoalSource, GoalStalledOn, GoalStalledOnOpaques, HasChanged, MaybeCause,
+    CanonicalResponse, Certainty, ExternalConstraintsData, FIXPOINT_STEP_LIMIT, Goal,
+    GoalEvaluation, GoalSource, GoalStalledOn, GoalStalledOnOpaques, HasChanged, MaybeCause,
     NestedNormalizationGoals, NoSolution, QueryInput, QueryResult, Response, SucceededInErased,
     VisibleForLeakCheck, inspect,
 };
@@ -516,7 +516,7 @@ where
     pub(super) fn enter_canonical<T>(
         cx: I,
         search_graph: &'a mut SearchGraph<D>,
-        canonical_input: CanonicalInput<I>,
+        canonical_input: I::CanonicalInput,
         proof_tree_builder: &mut inspect::ProofTreeBuilder<D>,
         f: impl FnOnce(
             &mut EvalCtxt<'_, D>,
@@ -833,7 +833,7 @@ where
 
     fn build_stalled_on(
         &self,
-        canonical_goal: CanonicalInput<I>,
+        canonical_goal: I::CanonicalInput,
         maybe_info: MaybeInfo,
         stalled_vars: ThinVec<I::GenericArg>,
         previously_succeeded_in_erased: SucceededInErased<I>,
@@ -1827,7 +1827,7 @@ pub fn evaluate_root_goal_for_proof_tree_raw_provider<
     I: Interner,
 >(
     cx: I,
-    canonical_goal: CanonicalInput<I>,
+    canonical_goal: I::CanonicalInput,
     root_depth: usize,
 ) -> (QueryResult<I>, I::Probe, RequiredDepth) {
     let mut inspect = inspect::ProofTreeBuilder::new();

compiler/rustc_next_trait_solver/src/solve/search_graph.rs

@@ -2,9 +2,7 @@ use std::convert::Infallible;
 use std::marker::PhantomData;
 
 use rustc_type_ir::search_graph::{self, PathKind};
-use rustc_type_ir::solve::{
-    AccessedOpaques, CanonicalInput, Certainty, NoSolution, QueryResult, RerunResultExt,
-};
+use rustc_type_ir::solve::{AccessedOpaques, Certainty, NoSolution, QueryResult, RerunResultExt};
 use rustc_type_ir::{Interner, MayBeErased, TypingMode};
 
 use crate::canonical::response_no_constraints_raw;
@@ -30,7 +28,7 @@ where
     type ValidationScope = Infallible;
     fn enter_validation_scope(
         _cx: Self::Cx,
-        _input: CanonicalInput<I>,
+        _input: I::CanonicalInput,
     ) -> Option<Self::ValidationScope> {
         None
     }
@@ -47,7 +45,7 @@ where
     fn initial_provisional_result(
         cx: I,
         kind: PathKind,
-        input: CanonicalInput<I>,
+        input: I::CanonicalInput,
     ) -> (QueryResult<I>, AccessedOpaques<I>) {
         match kind {
             PathKind::Coinductive => response_no_constraints(cx, input, Certainty::Yes),
@@ -101,15 +99,15 @@ where
 
     fn stack_overflow_result(
         cx: I,
-        input: CanonicalInput<I>,
+        input: I::CanonicalInput,
     ) -> (QueryResult<I>, AccessedOpaques<I>) {
         response_no_constraints(cx, input, Certainty::overflow(true))
     }
 
     const FIXPOINT_OVERFLOW_AMBIGUITY_KIND: Certainty = Certainty::overflow(false);
     fn fixpoint_overflow_result(
         cx: I,
-        input: CanonicalInput<I>,
+        input: I::CanonicalInput,
     ) -> (QueryResult<I>, AccessedOpaques<I>) {
         response_no_constraints(cx, input, Certainty::overflow(false))
     }
@@ -129,7 +127,7 @@ where
     fn compute_goal(
         search_graph: &mut SearchGraph<D>,
         cx: I,
-        input: CanonicalInput<I>,
+        input: I::CanonicalInput,
         inspect: &mut Self::ProofTreeBuilder,
     ) -> (QueryResult<I>, AccessedOpaques<I>) {
         EvalCtxt::enter_canonical(cx, search_graph, input, inspect, |ecx, goal| {
@@ -144,7 +142,7 @@ where
 
 fn response_no_constraints<I: Interner>(
     cx: I,
-    input: CanonicalInput<I>,
+    input: I::CanonicalInput,
     certainty: Certainty,
 ) -> (QueryResult<I>, AccessedOpaques<I>) {
     (

compiler/rustc_trait_selection/src/solve.rs

@@ -19,7 +19,7 @@ pub use select::InferCtxtSelectExt;
 
 fn evaluate_root_goal_for_proof_tree_raw<'tcx>(
     tcx: TyCtxt<'tcx>,
-    key: (CanonicalInput<TyCtxt<'tcx>>, usize),
+    key: (rustc_middle::traits::solve::CanonicalInput<'tcx>, usize),
 ) -> (QueryResult<TyCtxt<'tcx>>, &'tcx inspect::Probe<TyCtxt<'tcx>>, RequiredDepth) {
     evaluate_root_goal_for_proof_tree_raw_provider::<SolverDelegate<'tcx>, TyCtxt<'tcx>>(
         tcx, key.0, key.1,

compiler/rustc_type_ir/src/interner.rs

@@ -17,7 +17,7 @@ use crate::lang_items::{SolverAdtLangItem, SolverProjectionLangItem, SolverTrait
 use crate::relate::Relate;
 use crate::search_graph::RequiredDepth;
 use crate::solve::{
-    AccessedOpaques, CanonicalInput, Certainty, ExternalConstraintsData, QueryResult, inspect,
+    AccessedOpaques, CanonicalInputData, Certainty, ExternalConstraintsData, QueryResult, inspect,
 };
 use crate::visit::{Flags, TypeVisitable};
 use crate::{
@@ -499,7 +499,7 @@ pub trait Interner:
     fn mk_probe(self, probe: inspect::Probe<Self>) -> Self::Probe;
     fn evaluate_root_goal_for_proof_tree_raw(
         self,
-        canonical_goal: CanonicalInput<Self>,
+        canonical_goal: Self::CanonicalInput,
         root_depth: usize,
     ) -> (QueryResult<Self>, Self::Probe, RequiredDepth);
 
@@ -520,6 +520,9 @@ pub trait Interner:
     ) -> Region<Self>;
 
     fn intern_canonical_bound(self, var: BoundVar) -> Region<Self>;
+
+    type CanonicalInput: Copy + Debug + Hash + Eq + Deref<Target = CanonicalInputData<Self>>;
+    fn mk_canonical_input(self, data: CanonicalInputData<Self>) -> Self::CanonicalInput;
 }
 
 macro_rules! declare_lift_into {
@@ -711,7 +714,7 @@ impl<T, R, E> CollectAndApply<T, R> for Result<T, E> {
 }
 
 impl<I: Interner> search_graph::Cx for I {
-    type Input = CanonicalInput<I>;
+    type Input = I::CanonicalInput;
     type Result = (QueryResult<I>, AccessedOpaques<I>);
     type AmbiguityKind = Certainty;
 

compiler/rustc_type_ir/src/solve/mod.rs

@@ -22,8 +22,8 @@ use crate::{
     InferConst, IntVid, Interner, TermKind, TyVid, TypingMode, Upcast,
 };
 
-pub type CanonicalInput<I, T = <I as Interner>::Predicate> =
-    ty::CanonicalQueryInput<I, QueryInput<I, T>>;
+pub type CanonicalInputData<I> =
+    ty::CanonicalQueryInput<I, QueryInput<I, <I as Interner>::Predicate>>;
 pub type CanonicalResponse<I> = Canonical<I, Response<I>>;
 /// The result of evaluating a canonical query.
 ///