Rollup merge of #155547 - aerooneqq:better-disambiguators, r=oli-obk

Use per-parent disambiguators everywhere

This PR addressing the following concerns about per-parent disambiguators (#153955):

- DisambiguatorState is removed, PerParentDisambiguatorState is now used everywhere,
- Steals were removed from every per-parent disambiguator in resolver,
- It adds `parent` field in `PerParentDisambiguatorState` in `#[cfg(debug_assertions)]` for asserting that per-parent disambiguator corresponds to the same `LocalDefId` which is passed into `create_def`,
- ~Removes `Disambiguator` trait replacing it with `Disambiguator` enum, with this change we no longer expose `next` method (as trait should be public otherwise the warning will be emitted). It may affect perf in a negative way though.~

~Those changes should not fix perf issues that were [reported](#153955 (comment)), perf run that was attempted [before](#153955 (comment)) showed much better results. Performance can be probably fixed by removing per-parent disambiguators replacing them with a single one as it was before, then it will be passed to AST -> HIR lowering and modified. For delayed owners we can store ~followup disambiguators as it was in the beginning of the #153955~ per-parent disambiguators. This solution should save achievements from #153955 (removed `DefPathData` variants).
However, I would prefer to keep per-parent disambiguators as it seems a better architectural solution for me.~

r? @petrochenkov
cc @oli-obk

Author: JonathanBrouwer

compiler/rustc_ast_lowering/src/item.rs

@@ -1,13 +1,16 @@
 use std::mem;
+use std::sync::Arc;
 
 use rustc_abi::ExternAbi;
 use rustc_ast::visit::AssocCtxt;
 use rustc_ast::*;
 use rustc_data_structures::fx::FxIndexMap;
+use rustc_data_structures::steal::Steal;
 use rustc_errors::{E0570, ErrorGuaranteed, struct_span_code_err};
 use rustc_hir::attrs::{AttributeKind, EiiImplResolution};
 use rustc_hir::def::{DefKind, PerNS, Res};
-use rustc_hir::def_id::{CRATE_DEF_ID, LocalDefId};
+use rustc_hir::def_id::{CRATE_DEF_ID, LocalDefId, LocalDefIdMap};
+use rustc_hir::definitions::PerParentDisambiguatorState;
 use rustc_hir::{
     self as hir, HirId, ImplItemImplKind, LifetimeSource, PredicateOrigin, Target, find_attr,
 };
@@ -48,11 +51,21 @@ impl<'hir> Owners<'_, 'hir> {
     }
 }
 
+/// Default disambiguators are used during default lowering, when we lower
+/// AST owners in a loop we can use the whole map, in contrast delayed lowering
+/// lowers each AST owner separately, so we use readonly disambiguators map
+/// with `Steal`s to get disambiguators.
+pub(super) enum Disambiguators {
+    Default(LocalDefIdMap<PerParentDisambiguatorState>),
+    Delayed(Arc<LocalDefIdMap<Steal<PerParentDisambiguatorState>>>),
+}
+
 pub(super) struct ItemLowerer<'a, 'hir, R> {
     pub(super) tcx: TyCtxt<'hir>,
     pub(super) resolver: &'a mut R,
     pub(super) ast_index: &'a IndexSlice<LocalDefId, AstOwner<'a>>,
     pub(super) owners: Owners<'a, 'hir>,
+    pub(super) disambiguators: &'a mut Disambiguators,
 }
 
 /// When we have a ty alias we *may* have two where clauses. To give the best diagnostics, we set the span
@@ -80,7 +93,7 @@ impl<'hir, R: ResolverAstLoweringExt<'hir>> ItemLowerer<'_, 'hir, R> {
         owner: NodeId,
         f: impl FnOnce(&mut LoweringContext<'_, 'hir, R>) -> hir::OwnerNode<'hir>,
     ) {
-        let mut lctx = LoweringContext::new(self.tcx, self.resolver);
+        let mut lctx = LoweringContext::new(self.tcx, self.resolver, self.disambiguators);
         lctx.with_hir_id_owner(owner, |lctx| f(lctx));
 
         for (def_id, info) in lctx.children {

compiler/rustc_ast_lowering/src/lib.rs

@@ -70,7 +70,7 @@ use thin_vec::ThinVec;
 use tracing::{debug, instrument, trace};
 
 use crate::errors::{AssocTyParentheses, AssocTyParenthesesSub, MisplacedImplTrait};
-use crate::item::Owners;
+use crate::item::{Disambiguators, Owners};
 
 macro_rules! arena_vec {
     ($this:expr; $($x:expr),*) => (
@@ -94,7 +94,8 @@ pub mod stability;
 struct LoweringContext<'a, 'hir, R> {
     tcx: TyCtxt<'hir>,
     resolver: &'a mut R,
-    disambiguator: PerParentDisambiguatorState,
+    disambiguators: &'a mut Disambiguators,
+    current_disambiguator: PerParentDisambiguatorState,
 
     /// Used to allocate HIR nodes.
     arena: &'hir hir::Arena<'hir>,
@@ -154,12 +155,13 @@ struct LoweringContext<'a, 'hir, R> {
 }
 
 impl<'a, 'hir, R: ResolverAstLoweringExt<'hir>> LoweringContext<'a, 'hir, R> {
-    fn new(tcx: TyCtxt<'hir>, resolver: &'a mut R) -> Self {
+    fn new(tcx: TyCtxt<'hir>, resolver: &'a mut R, disambiguators: &'a mut Disambiguators) -> Self {
         let registered_tools = tcx.registered_tools(()).iter().map(|x| x.name).collect();
         Self {
             tcx,
             resolver,
-            disambiguator: Default::default(),
+            disambiguators,
+            current_disambiguator: Default::default(),
             arena: tcx.hir_arena,
 
             // HirId handling.
@@ -302,10 +304,6 @@ impl<'a, 'tcx> ResolverAstLoweringExt<'tcx> for ResolverDelayedAstLowering<'a, '
     fn next_node_id(&mut self) -> NodeId {
         next_node_id(&mut self.next_node_id)
     }
-
-    fn steal_or_create_disambiguator(&self, parent: LocalDefId) -> PerParentDisambiguatorState {
-        self.base.steal_or_create_disambiguator(parent)
-    }
 }
 
 fn next_node_id(current_id: &mut NodeId) -> NodeId {
@@ -408,10 +406,6 @@ impl<'tcx> ResolverAstLowering<'tcx> {
     fn next_node_id(&mut self) -> NodeId {
         next_node_id(&mut self.next_node_id)
     }
-
-    fn steal_or_create_disambiguator(&self, parent: LocalDefId) -> PerParentDisambiguatorState {
-        self.per_parent_disambiguators.get(&parent).map(|s| s.steal()).unwrap_or_default()
-    }
 }
 
 /// How relaxed bounds `?Trait` should be treated.
@@ -632,11 +626,13 @@ pub fn lower_to_hir(tcx: TyCtxt<'_>, (): ()) -> mid_hir::Crate<'_> {
         tcx.definitions_untracked().def_index_count(),
     );
 
+    let mut disambiguators = Disambiguators::Default(resolver.disambiguators.steal());
     let mut lowerer = item::ItemLowerer {
         tcx,
         resolver: &mut resolver,
         ast_index: &ast_index,
         owners: Owners::IndexVec(&mut owners),
+        disambiguators: &mut disambiguators,
     };
 
     let mut delayed_ids: FxIndexSet<LocalDefId> = Default::default();
@@ -655,15 +651,19 @@ pub fn lower_to_hir(tcx: TyCtxt<'_>, (): ()) -> mid_hir::Crate<'_> {
     let opt_hir_hash =
         if tcx.needs_crate_hash() { Some(compute_hir_hash(tcx, &owners)) } else { None };
 
-    let delayed_resolver = Steal::new((resolver, krate));
+    let Disambiguators::Default(disambiguators) = disambiguators else { unreachable!() };
+    let delayed_disambigs =
+        Arc::new(disambiguators.into_items().map(|(id, d)| (id, Steal::new(d))).collect());
+
+    let delayed_resolver = Steal::new((resolver, krate, delayed_disambigs));
     mid_hir::Crate::new(owners, delayed_ids, delayed_resolver, opt_hir_hash)
 }
 
 /// Lowers an AST owner corresponding to `def_id`, now only delegations are lowered this way.
 pub fn lower_delayed_owner(tcx: TyCtxt<'_>, def_id: LocalDefId) {
     let krate = tcx.hir_crate(());
 
-    let (resolver, krate) = &*krate.delayed_resolver.borrow();
+    let (resolver, krate, delayed_disambigs) = &*krate.delayed_resolver.borrow();
 
     // FIXME!!!(fn_delegation): make ast index lifetime same as resolver,
     // as it is too bad to reindex whole crate on each delegation lowering.
@@ -677,12 +677,12 @@ pub fn lower_delayed_owner(tcx: TyCtxt<'_>, def_id: LocalDefId) {
     };
 
     let mut map = Default::default();
-
     let mut lowerer = item::ItemLowerer {
         tcx,
         resolver: &mut resolver,
         ast_index: &ast_index,
         owners: Owners::Map(&mut map),
+        disambiguators: &mut Disambiguators::Delayed(Arc::clone(delayed_disambigs)),
     };
 
     lowerer.lower_node(def_id);
@@ -740,7 +740,7 @@ impl<'hir, R: ResolverAstLoweringExt<'hir>> LoweringContext<'_, 'hir, R> {
         let def_id = self
             .tcx
             .at(span)
-            .create_def(parent, name, def_kind, None, &mut self.disambiguator)
+            .create_def(parent, name, def_kind, None, &mut self.current_disambiguator)
             .def_id();
 
         debug!("create_def: def_id_to_node_id[{:?}] <-> {:?}", def_id, node_id);
@@ -780,9 +780,16 @@ impl<'hir, R: ResolverAstLoweringExt<'hir>> LoweringContext<'_, 'hir, R> {
         f: impl FnOnce(&mut Self) -> hir::OwnerNode<'hir>,
     ) {
         let owner_id = self.owner_id(owner);
+        let def_id = owner_id.def_id;
+
+        let new_disambig = match &mut self.disambiguators {
+            Disambiguators::Default(map) => map.remove(&def_id),
+            Disambiguators::Delayed(map) => map.get(&def_id).map(Steal::steal),
+        };
+
+        let new_disambig = new_disambig.unwrap_or_else(|| PerParentDisambiguatorState::new(def_id));
 
-        let new_disambig = self.resolver.steal_or_create_disambiguator(owner_id.def_id);
-        let disambiguator = std::mem::replace(&mut self.disambiguator, new_disambig);
+        let disambiguator = std::mem::replace(&mut self.current_disambiguator, new_disambig);
         let current_attrs = std::mem::take(&mut self.attrs);
         let current_bodies = std::mem::take(&mut self.bodies);
         let current_define_opaque = std::mem::take(&mut self.define_opaque);
@@ -817,7 +824,7 @@ impl<'hir, R: ResolverAstLoweringExt<'hir>> LoweringContext<'_, 'hir, R> {
         assert!(self.impl_trait_bounds.is_empty());
         let info = self.make_owner_info(item);
 
-        self.disambiguator = disambiguator;
+        self.current_disambiguator = disambiguator;
         self.attrs = current_attrs;
         self.bodies = current_bodies;
         self.define_opaque = current_define_opaque;

compiler/rustc_const_eval/src/interpret/intern.rs

@@ -16,8 +16,11 @@
 use hir::def::DefKind;
 use rustc_ast::Mutability;
 use rustc_data_structures::fx::{FxHashSet, FxIndexMap};
-use rustc_hir as hir;
-use rustc_hir::definitions::{DefPathData, DisambiguatorState};
+use rustc_hir::def_id::LocalDefIdMap;
+use rustc_hir::definitions::{
+    DefPathData, PerParentDisambiguatorState, PerParentDisambiguatorsMap,
+};
+use rustc_hir::{self as hir};
 use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrs;
 use rustc_middle::mir::interpret::{
     AllocBytes, ConstAllocation, CtfeProvenance, InterpResult, Provenance,
@@ -105,7 +108,7 @@ fn intern_shallow<'tcx, M: CompileTimeMachine<'tcx>>(
     ecx: &mut InterpCx<'tcx, M>,
     alloc_id: AllocId,
     mutability: Mutability,
-    disambiguator: Option<&mut DisambiguatorState>,
+    disambiguators: Option<&mut LocalDefIdMap<PerParentDisambiguatorState>>,
 ) -> Result<impl Iterator<Item = CtfeProvenance> + 'tcx, InternError> {
     trace!("intern_shallow {:?}", alloc_id);
     // remove allocation
@@ -129,7 +132,7 @@ fn intern_shallow<'tcx, M: CompileTimeMachine<'tcx>>(
             static_id,
             alloc_id,
             alloc,
-            disambiguator.expect("disambiguator needed"),
+            disambiguators.expect("disambiguators needed"),
         );
     } else {
         ecx.tcx.set_alloc_id_memory(alloc_id, alloc);
@@ -144,7 +147,7 @@ fn intern_as_new_static<'tcx>(
     static_id: LocalDefId,
     alloc_id: AllocId,
     alloc: ConstAllocation<'tcx>,
-    disambiguator: &mut DisambiguatorState,
+    disambiguators: &mut LocalDefIdMap<PerParentDisambiguatorState>,
 ) {
     // `intern_const_alloc_recursive` is called once per static and it contains the `DisambiguatorState`.
     //  The `<static_id>::{{nested}}` path is thus unique to `intern_const_alloc_recursive` and the
@@ -155,7 +158,8 @@ fn intern_as_new_static<'tcx>(
         None,
         DefKind::Static { safety: hir::Safety::Safe, mutability: alloc.0.mutability, nested: true },
         Some(DefPathData::NestedStatic),
-        disambiguator,
+        //FIXME(oli-obk): cleanup (https://github.com/rust-lang/rust/pull/155547#discussion_r3110792640)
+        disambiguators.get_or_create(static_id),
     );
     tcx.set_nested_alloc_id_static(alloc_id, feed.def_id());
 
@@ -205,7 +209,7 @@ pub fn intern_const_alloc_recursive<'tcx, M: CompileTimeMachine<'tcx>>(
     intern_kind: InternKind,
     ret: &MPlaceTy<'tcx>,
 ) -> Result<(), InternError> {
-    let mut disambiguator = DisambiguatorState::new();
+    let mut disambiguators = Default::default();
 
     // We are interning recursively, and for mutability we are distinguishing the "root" allocation
     // that we are starting in, and all other allocations that we are encountering recursively.
@@ -250,7 +254,7 @@ pub fn intern_const_alloc_recursive<'tcx, M: CompileTimeMachine<'tcx>>(
         prepare_alloc(*ecx.tcx, *kind, alloc, base_mutability)?;
         alloc.provenance().ptrs().iter().map(|&(_, prov)| prov).collect()
     } else {
-        intern_shallow(ecx, base_alloc_id, base_mutability, Some(&mut disambiguator))?.collect()
+        intern_shallow(ecx, base_alloc_id, base_mutability, Some(&mut disambiguators))?.collect()
     };
     // We need to distinguish "has just been interned" from "was already in `tcx`",
     // so we track this in a separate set.
@@ -332,7 +336,7 @@ pub fn intern_const_alloc_recursive<'tcx, M: CompileTimeMachine<'tcx>>(
         // okay with losing some potential for immutability here. This can anyway only affect
         // `static mut`.
         just_interned.insert(alloc_id);
-        let next = intern_shallow(ecx, alloc_id, inner_mutability, Some(&mut disambiguator))?;
+        let next = intern_shallow(ecx, alloc_id, inner_mutability, Some(&mut disambiguators))?;
         todo.extend(next);
     }
     if found_bad_mutable_ptr {

compiler/rustc_hir/src/definitions.rs

@@ -7,11 +7,12 @@
 use std::fmt::{self, Write};
 use std::hash::Hash;
 
+use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::stable_hasher::StableHasher;
-use rustc_data_structures::unord::UnordMap;
 use rustc_hashes::Hash64;
 use rustc_index::IndexVec;
-use rustc_macros::{BlobDecodable, Decodable, Encodable};
+use rustc_macros::{BlobDecodable, Decodable, Encodable, extension};
+use rustc_span::def_id::LocalDefIdMap;
 use rustc_span::{Symbol, kw, sym};
 use tracing::{debug, instrument};
 
@@ -97,45 +98,28 @@ impl DefPathTable {
     }
 }
 
-pub trait Disambiguator {
-    fn entry(&mut self, parent: LocalDefId, data: DefPathData) -> &mut u32;
-}
-
 #[derive(Debug, Default, Clone)]
 pub struct PerParentDisambiguatorState {
-    next: UnordMap<DefPathData, u32>,
-}
-
-impl Disambiguator for PerParentDisambiguatorState {
-    #[inline]
-    fn entry(&mut self, _: LocalDefId, data: DefPathData) -> &mut u32 {
-        self.next.entry(data).or_insert(0)
-    }
-}
-
-#[derive(Debug, Default, Clone)]
-pub struct DisambiguatorState {
-    next: UnordMap<(LocalDefId, DefPathData), u32>,
+    #[cfg(debug_assertions)]
+    parent: Option<LocalDefId>,
+    next: FxHashMap<DefPathData, u32>,
 }
 
-impl Disambiguator for DisambiguatorState {
-    #[inline]
-    fn entry(&mut self, parent: LocalDefId, data: DefPathData) -> &mut u32 {
-        self.next.entry((parent, data)).or_insert(0)
+impl PerParentDisambiguatorState {
+    #[inline(always)]
+    pub fn new(_parent: LocalDefId) -> PerParentDisambiguatorState {
+        PerParentDisambiguatorState {
+            #[cfg(debug_assertions)]
+            parent: Some(_parent),
+            next: Default::default(),
+        }
     }
 }
 
-impl DisambiguatorState {
-    pub const fn new() -> Self {
-        Self { next: Default::default() }
-    }
-
-    /// Creates a `DisambiguatorState` where the next allocated `(LocalDefId, DefPathData)` pair
-    /// will have `index` as the disambiguator.
-    pub fn with(def_id: LocalDefId, data: DefPathData, index: u32) -> Self {
-        let mut this = Self::new();
-        this.next.insert((def_id, data), index);
-        this
+#[extension(pub trait PerParentDisambiguatorsMap)]
+impl LocalDefIdMap<PerParentDisambiguatorState> {
+    fn get_or_create(&mut self, parent: LocalDefId) -> &mut PerParentDisambiguatorState {
+        self.entry(parent).or_insert_with(|| PerParentDisambiguatorState::new(parent))
     }
 }
 
@@ -408,7 +392,7 @@ impl Definitions {
         &mut self,
         parent: LocalDefId,
         data: DefPathData,
-        disambiguator: &mut impl Disambiguator,
+        disambiguator: &mut PerParentDisambiguatorState,
     ) -> LocalDefId {
         // We can't use `Debug` implementation for `LocalDefId` here, since it tries to acquire a
         // reference to `Definitions` and we're already holding a mutable reference.
@@ -422,7 +406,14 @@ impl Definitions {
 
         // Find the next free disambiguator for this key.
         let disambiguator = {
-            let next_disamb = disambiguator.entry(parent, data);
+            #[cfg(debug_assertions)]
+            debug_assert_eq!(
+                parent,
+                disambiguator.parent.expect("must be set"),
+                "provided parent and parent in disambiguator must be the same"
+            );
+
+            let next_disamb = disambiguator.next.entry(data).or_insert(0);
             let disambiguator = *next_disamb;
             *next_disamb = next_disamb.checked_add(1).expect("disambiguator overflow");
             disambiguator