Rollup of 3 perf-sensitive pull requests

compiler/rustc_ast/src/tokenstream.rs

@@ -7,7 +7,7 @@
 use std::borrow::Cow;
 use std::hash::Hash;
 use std::ops::Range;
-use std::sync::Arc;
+use std::sync::{Arc, LazyLock};
 use std::{cmp, fmt, iter, mem};
 
 use rustc_data_structures::stable_hash::{StableHash, StableHashCtxt, StableHasher};
@@ -419,6 +419,11 @@ fn make_attr_token_stream(
     AttrTokenStream::new(stack_top.inner)
 }
 
+/// A shared empty attr token stream, used to avoid an unnecessary `Arc` allocation for every empty
+/// token stream.
+static EMPTY_ATTR_TOKEN_STREAM: LazyLock<AttrTokenStream> =
+    LazyLock::new(|| AttrTokenStream(Arc::new(Vec::new())));
+
 /// Like `TokenTree`, but for `AttrTokenStream`.
 #[derive(Clone, Debug, Encodable, Decodable)]
 pub enum AttrTokenTree {
@@ -431,8 +436,12 @@ pub enum AttrTokenTree {
 }
 
 impl AttrTokenStream {
-    pub fn new(tokens: Vec<AttrTokenTree>) -> AttrTokenStream {
-        AttrTokenStream(Arc::new(tokens))
+    pub fn new(tts: Vec<AttrTokenTree>) -> AttrTokenStream {
+        if tts.is_empty() {
+            EMPTY_ATTR_TOKEN_STREAM.clone()
+        } else {
+            AttrTokenStream(Arc::new(tts))
+        }
     }
 
     /// Converts this `AttrTokenStream` to a plain `Vec<TokenTree>`. During
@@ -620,13 +629,18 @@ pub enum Spacing {
     JointHidden,
 }
 
+/// A shared empty token stream, used to avoid an unnecessary `Arc` allocation for every empty
+/// token stream.
+static EMPTY_TOKEN_STREAM: LazyLock<TokenStream> =
+    LazyLock::new(|| TokenStream(Arc::new(Vec::new())));
+
 /// A `TokenStream` is an abstract sequence of tokens, organized into [`TokenTree`]s.
-#[derive(Clone, Debug, Default, PartialEq, Eq, Hash, Encodable, Decodable)]
+#[derive(Clone, Debug, PartialEq, Eq, Hash, Encodable, Decodable)]
 pub struct TokenStream(Arc<Vec<TokenTree>>);
 
 impl TokenStream {
     pub fn new(tts: Vec<TokenTree>) -> TokenStream {
-        TokenStream(Arc::new(tts))
+        if tts.is_empty() { TokenStream::default() } else { TokenStream(Arc::new(tts)) }
     }
 
     pub fn is_empty(&self) -> bool {
@@ -840,6 +854,12 @@ impl TokenStream {
     }
 }
 
+impl Default for TokenStream {
+    fn default() -> TokenStream {
+        EMPTY_TOKEN_STREAM.clone()
+    }
+}
+
 impl FromIterator<TokenTree> for TokenStream {
     fn from_iter<I: IntoIterator<Item = TokenTree>>(iter: I) -> Self {
         TokenStream::new(iter.into_iter().collect::<Vec<TokenTree>>())

compiler/rustc_borrowck/src/nll.rs

@@ -144,6 +144,8 @@ pub(crate) fn compute_regions<'tcx>(
         &lowered_constraints,
     );
 
+    let num_points = location_map.num_points();
+
     // If requested for `-Zpolonius=next`, compute loan liveness information.
     // This is done prior to `RegionInferenceContext::new`, because we may add
     // additional liveness constraints.
@@ -155,6 +157,7 @@ pub(crate) fn compute_regions<'tcx>(
             &universal_region_relations.universal_regions,
             body,
             borrow_set,
+            num_points,
         );
     }
 

compiler/rustc_borrowck/src/polonius/mod.rs

@@ -41,7 +41,7 @@ mod liveness_constraints;
 use std::collections::BTreeMap;
 
 use rustc_data_structures::fx::FxHashSet;
-use rustc_index::bit_set::SparseBitMatrix;
+use rustc_index::bit_set::DenseBitSet;
 use rustc_middle::mir::{Body, Local};
 use rustc_middle::ty::RegionVid;
 use rustc_mir_dataflow::points::PointIndex;
@@ -55,7 +55,28 @@ use crate::dataflow::BorrowIndex;
 use crate::region_infer::values::LivenessValues;
 use crate::universal_regions::UniversalRegions;
 
-pub(crate) type LiveLoans = SparseBitMatrix<PointIndex, BorrowIndex>;
+#[derive(Clone)]
+pub(crate) struct LiveLoans {
+    num_points: usize,
+    // This matrix always has more rows (PointIndex) than columns (BorrowIndex),
+    // and the borrow dimension is usually very low (single digit in 90% of cases in our benchmark suite),
+    // so we store it packed in a single bitset. Rows are points, columns are borrows.
+    flat_matrix: DenseBitSet<usize>,
+}
+
+impl LiveLoans {
+    pub(crate) fn new(num_points: usize, num_borrows: usize) -> Self {
+        Self { num_points, flat_matrix: DenseBitSet::new_empty(num_points * num_borrows) }
+    }
+    pub(crate) fn insert(&mut self, row: PointIndex, col: BorrowIndex) {
+        let bit_index = row.index() + self.num_points * col.index();
+        self.flat_matrix.insert(bit_index);
+    }
+    pub(crate) fn contains(&self, row: PointIndex, col: BorrowIndex) -> bool {
+        let bit_index = row.index() + self.num_points * col.index();
+        self.flat_matrix.contains(bit_index)
+    }
+}
 
 /// This struct holds the necessary
 ///  - liveness data, created during MIR typeck, and which will be used to lazily compute the
@@ -109,6 +130,7 @@ impl PoloniusContext {
         universal_regions: &UniversalRegions<'tcx>,
         body: &Body<'tcx>,
         borrow_set: &BorrowSet<'tcx>,
+        num_points: usize,
     ) {
         // We don't need to prepare the graph (index NLL constraints, etc.) if we have no loans to
         // trace throughout localized constraints.
@@ -118,7 +140,7 @@ impl PoloniusContext {
             // step in the chain (the NLL loan scope and active loans computations).
             let graph = LocalizedConstraintGraph::new(liveness, outlives_constraints);
 
-            let mut live_loans = LiveLoans::new(borrow_set.len());
+            let mut live_loans = LiveLoans::new(num_points, borrow_set.len());
             let mut visitor = LoanLivenessVisitor { liveness, live_loans: &mut live_loans };
             graph.traverse(
                 body,

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,
@@ -698,6 +699,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>,
@@ -1831,7 +1831,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>) {
     (