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 (rust-lang/rust#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](https://github.com/rust-lang/rust/pull/153955#issuecomment-4269223191), perf run that was attempted [before](https://github.com/rust-lang/rust/pull/153955#issuecomment-4214516698) 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 rust-lang/rust#153955~ per-parent disambiguators. This solution should save achievements from rust-lang/rust#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

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 @@ fn get_or_insert_mut(&mut self, def_id: LocalDefId) -> &mut hir::MaybeOwner<'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 @@ fn with_lctx(
         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 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 @@ macro_rules! arena_vec {
 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 @@ fn trait_candidates(&self, node_id: NodeId) -> Option<&'tcx [hir::TraitCandidate
     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 @@ fn trait_candidates(&self, node_id: NodeId) -> Option<&'tcx [hir::TraitCandidate
     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,7 +651,11 @@ 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)
 }
 
@@ -663,7 +663,7 @@ pub fn lower_to_hir(tcx: TyCtxt<'_>, (): ()) -> mid_hir::Crate<'_> {
 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 @@ fn create_def(
         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 @@ fn with_hir_id_owner(
         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 = self.resolver.steal_or_create_disambiguator(owner_id.def_id);
-        let disambiguator = std::mem::replace(&mut self.disambiguator, new_disambig);
+        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 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 @@ fn with_hir_id_owner(
         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 @@ pub fn enumerated_keys_and_path_hashes(
     }
 }
 
-pub trait Disambiguator {
-    fn entry(&mut self, parent: LocalDefId, data: DefPathData) -> &mut u32;
-}
-
 #[derive(Debug, Default, Clone)]
 pub struct PerParentDisambiguatorState {
-    next: UnordMap<DefPathData, u32>,
+    #[cfg(debug_assertions)]
+    parent: Option<LocalDefId>,
+    next: FxHashMap<DefPathData, u32>,
 }
 
-impl Disambiguator for PerParentDisambiguatorState {
-    #[inline]
-    fn entry(&mut self, _: LocalDefId, data: DefPathData) -> &mut u32 {
-        self.next.entry(data).or_insert(0)
+impl PerParentDisambiguatorState {
+    #[inline(always)]
+    pub fn new(_parent: LocalDefId) -> PerParentDisambiguatorState {
+        PerParentDisambiguatorState {
+            #[cfg(debug_assertions)]
+            parent: Some(_parent),
+            next: Default::default(),
+        }
     }
 }
 
-#[derive(Debug, Default, Clone)]
-pub struct DisambiguatorState {
-    next: UnordMap<(LocalDefId, 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 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 @@ pub fn create_def(
         &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 @@ pub fn create_def(
 
         // 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

compiler/rustc_hir_analysis/src/collect/resolve_bound_vars.rs

@@ -14,7 +14,8 @@
 use rustc_data_structures::fx::{FxHashSet, FxIndexMap, FxIndexSet};
 use rustc_errors::ErrorGuaranteed;
 use rustc_hir::def::{DefKind, Res};
-use rustc_hir::definitions::{DefPathData, DisambiguatorState};
+use rustc_hir::def_id::LocalDefIdMap;
+use rustc_hir::definitions::{DefPathData, PerParentDisambiguatorsMap};
 use rustc_hir::intravisit::{self, InferKind, Visitor, VisitorExt};
 use rustc_hir::{
     self as hir, AmbigArg, GenericArg, GenericParam, GenericParamKind, HirId, LifetimeKind, Node,
@@ -30,6 +31,7 @@
 use tracing::{debug, debug_span, instrument};
 
 use crate::errors;
+use crate::hir::definitions::PerParentDisambiguatorState;
 
 #[extension(trait RegionExt)]
 impl ResolvedArg {
@@ -64,7 +66,7 @@ fn shifted(self, amount: u32) -> ResolvedArg {
 struct BoundVarContext<'a, 'tcx> {
     tcx: TyCtxt<'tcx>,
     rbv: &'a mut ResolveBoundVars<'tcx>,
-    disambiguator: &'a mut DisambiguatorState,
+    disambiguators: &'a mut LocalDefIdMap<PerParentDisambiguatorState>,
     scope: ScopeRef<'a, 'tcx>,
     opaque_capture_errors: RefCell<Option<OpaqueHigherRankedLifetimeCaptureErrors>>,
 }
@@ -259,7 +261,7 @@ fn resolve_bound_vars(tcx: TyCtxt<'_>, local_def_id: hir::OwnerId) -> ResolveBou
         tcx,
         rbv: &mut rbv,
         scope: &Scope::Root { opt_parent_item: None },
-        disambiguator: &mut DisambiguatorState::new(),
+        disambiguators: &mut Default::default(),
         opaque_capture_errors: RefCell::new(None),
     };
     match tcx.hir_owner_node(local_def_id) {
@@ -1104,12 +1106,12 @@ fn with<F>(&mut self, wrap_scope: Scope<'_, 'tcx>, f: F)
     where
         F: for<'b> FnOnce(&mut BoundVarContext<'b, 'tcx>),
     {
-        let BoundVarContext { tcx, rbv, disambiguator, .. } = self;
+        let BoundVarContext { tcx, rbv, disambiguators, .. } = self;
         let nested_errors = RefCell::new(self.opaque_capture_errors.borrow_mut().take());
         let mut this = BoundVarContext {
             tcx: *tcx,
             rbv,
-            disambiguator,
+            disambiguators,
             scope: &wrap_scope,
             opaque_capture_errors: nested_errors,
         };
@@ -1523,7 +1525,7 @@ fn remap_opaque_captures(
                     None,
                     DefKind::LifetimeParam,
                     Some(DefPathData::OpaqueLifetime(ident.name)),
-                    self.disambiguator,
+                    self.disambiguators.get_or_create(opaque_def_id),
                 );
                 feed.def_span(ident.span);
                 feed.def_ident_span(Some(ident.span));

compiler/rustc_middle/src/hir/map.rs

@@ -1251,7 +1251,7 @@ fn force_delayed_owners_lowering(tcx: TyCtxt<'_>) {
         tcx.ensure_done().lower_delayed_owner(id);
     }
 
-    let (_, krate) = krate.delayed_resolver.steal();
+    let (_, krate, _) = krate.delayed_resolver.steal();
     let prof = tcx.sess.prof.clone();
 
     // Drop AST to free memory. It can be expensive so try to drop it on a separate thread.

compiler/rustc_middle/src/hir/mod.rs

@@ -16,7 +16,8 @@
 use rustc_data_structures::steal::Steal;
 use rustc_data_structures::sync::{DynSend, DynSync, try_par_for_each_in};
 use rustc_hir::def::{DefKind, Res};
-use rustc_hir::def_id::{DefId, LocalDefId, LocalModDefId};
+use rustc_hir::def_id::{DefId, LocalDefId, LocalDefIdMap, LocalModDefId};
+use rustc_hir::definitions::PerParentDisambiguatorState;
 use rustc_hir::*;
 use rustc_index::IndexVec;
 use rustc_macros::{Decodable, Encodable, HashStable};
@@ -39,7 +40,11 @@ pub struct Crate<'hir> {
     pub delayed_ids: FxIndexSet<LocalDefId>,
     // The resolver and AST crate which are set in the end of the `hir_crate` query
     // and then stolen and dropped in `force_delayed_owners_lowering`.
-    pub delayed_resolver: Steal<(ResolverAstLowering<'hir>, Arc<ast::Crate>)>,
+    pub delayed_resolver: Steal<(
+        ResolverAstLowering<'hir>,
+        Arc<ast::Crate>,
+        Arc<LocalDefIdMap<Steal<PerParentDisambiguatorState>>>,
+    )>,
     // Only present when incr. comp. is enabled.
     pub opt_hir_hash: Option<Fingerprint>,
 }
@@ -48,7 +53,11 @@ impl<'hir> Crate<'hir> {
     pub fn new(
         owners: IndexVec<LocalDefId, MaybeOwner<'hir>>,
         delayed_ids: FxIndexSet<LocalDefId>,
-        delayed_resolver: Steal<(ResolverAstLowering<'hir>, Arc<ast::Crate>)>,
+        delayed_resolver: Steal<(
+            ResolverAstLowering<'hir>,
+            Arc<ast::Crate>,
+            Arc<LocalDefIdMap<Steal<PerParentDisambiguatorState>>>,
+        )>,
         opt_hir_hash: Option<Fingerprint>,
     ) -> Crate<'hir> {
         Crate { owners, delayed_ids, delayed_resolver, opt_hir_hash }

compiler/rustc_middle/src/ty/context.rs

@@ -32,7 +32,7 @@
 use rustc_errors::{Applicability, Diag, DiagCtxtHandle, Diagnostic, MultiSpan};
 use rustc_hir::def::DefKind;
 use rustc_hir::def_id::{CrateNum, DefId, LOCAL_CRATE, LocalDefId};
-use rustc_hir::definitions::{DefPathData, Definitions, Disambiguator};
+use rustc_hir::definitions::{DefPathData, Definitions, PerParentDisambiguatorState};
 use rustc_hir::intravisit::VisitorExt;
 use rustc_hir::lang_items::LangItem;
 use rustc_hir::limit::Limit;
@@ -1399,7 +1399,7 @@ pub fn create_def(
         name: Option<Symbol>,
         def_kind: DefKind,
         override_def_path_data: Option<DefPathData>,
-        disambiguator: &mut impl Disambiguator,
+        disambiguator: &mut PerParentDisambiguatorState,
     ) -> TyCtxtFeed<'tcx, LocalDefId> {
         let feed =
             self.tcx.create_def(parent, name, def_kind, override_def_path_data, disambiguator);
@@ -1417,7 +1417,7 @@ pub fn create_def(
         name: Option<Symbol>,
         def_kind: DefKind,
         override_def_path_data: Option<DefPathData>,
-        disambiguator: &mut impl Disambiguator,
+        disambiguator: &mut PerParentDisambiguatorState,
     ) -> TyCtxtFeed<'tcx, LocalDefId> {
         let data = override_def_path_data.unwrap_or_else(|| def_kind.def_path_data(name));
         // The following call has the side effect of modifying the tables inside `definitions`.

compiler/rustc_middle/src/ty/mod.rs

@@ -227,7 +227,7 @@ pub struct ResolverAstLowering<'tcx> {
     // Information about delegations which is used when handling recursive delegations
     pub delegation_infos: LocalDefIdMap<DelegationInfo>,
 
-    pub per_parent_disambiguators: LocalDefIdMap<Steal<PerParentDisambiguatorState>>,
+    pub disambiguators: Steal<LocalDefIdMap<PerParentDisambiguatorState>>,
 }
 
 #[derive(Debug)]

compiler/rustc_mir_transform/src/coroutine/by_move_body.rs

@@ -70,10 +70,10 @@
 use rustc_abi::{FieldIdx, VariantIdx};
 use rustc_data_structures::steal::Steal;
 use rustc_data_structures::unord::UnordMap;
-use rustc_hir as hir;
 use rustc_hir::def::DefKind;
 use rustc_hir::def_id::{DefId, LocalDefId};
-use rustc_hir::definitions::DisambiguatorState;
+use rustc_hir::definitions::PerParentDisambiguatorState;
+use rustc_hir::{self as hir};
 use rustc_middle::bug;
 use rustc_middle::hir::place::{Projection, ProjectionKind};
 use rustc_middle::mir::visit::MutVisitor;
@@ -221,7 +221,7 @@ pub(crate) fn coroutine_by_move_body_def_id<'tcx>(
         None,
         DefKind::SyntheticCoroutineBody,
         None,
-        &mut DisambiguatorState::new(),
+        &mut PerParentDisambiguatorState::new(parent_def_id),
     );
     by_move_body.source =
         mir::MirSource::from_instance(InstanceKind::Item(body_def.def_id().to_def_id()));

compiler/rustc_resolve/src/lib.rs

@@ -57,7 +57,7 @@
     PerNS,
 };
 use rustc_hir::def_id::{CRATE_DEF_ID, CrateNum, DefId, LOCAL_CRATE, LocalDefId, LocalDefIdMap};
-use rustc_hir::definitions::PerParentDisambiguatorState;
+use rustc_hir::definitions::{PerParentDisambiguatorState, PerParentDisambiguatorsMap};
 use rustc_hir::{PrimTy, TraitCandidate, find_attr};
 use rustc_index::bit_set::DenseBitSet;
 use rustc_metadata::creader::CStore;
@@ -1417,7 +1417,7 @@ pub struct Resolver<'ra, 'tcx> {
 
     node_id_to_def_id: NodeMap<Feed<'tcx, LocalDefId>>,
 
-    per_parent_disambiguators: LocalDefIdMap<PerParentDisambiguatorState>,
+    disambiguators: LocalDefIdMap<PerParentDisambiguatorState>,
 
     /// Indices of unnamed struct or variant fields with unresolved attributes.
     placeholder_field_indices: FxHashMap<NodeId, usize> = default::fx_hash_map(),
@@ -1620,14 +1620,10 @@ fn create_def(
             self.tcx.definitions_untracked().def_key(self.node_id_to_def_id[&node_id].key()),
         );
 
+        let disambiguator = self.disambiguators.get_or_create(parent);
+
         // FIXME: remove `def_span` body, pass in the right spans here and call `tcx.at().create_def()`
-        let feed = self.tcx.create_def(
-            parent,
-            name,
-            def_kind,
-            None,
-            self.per_parent_disambiguators.entry(parent).or_default(),
-        );
+        let feed = self.tcx.create_def(parent, name, def_kind, None, disambiguator);
         let def_id = feed.def_id();
 
         // Create the definition.
@@ -1811,7 +1807,7 @@ pub fn new(
             doc_link_resolutions: Default::default(),
             doc_link_traits_in_scope: Default::default(),
             current_crate_outer_attr_insert_span,
-            per_parent_disambiguators: Default::default(),
+            disambiguators: Default::default(),
             ..
         };
 
@@ -1951,11 +1947,7 @@ pub fn into_outputs(self) -> ResolverOutputs<'tcx> {
             lifetime_elision_allowed: self.lifetime_elision_allowed,
             lint_buffer: Steal::new(self.lint_buffer),
             delegation_infos: self.delegation_infos,
-            per_parent_disambiguators: self
-                .per_parent_disambiguators
-                .into_items()
-                .map(|(k, d)| (k, Steal::new(d)))
-                .collect(),
+            disambiguators: Steal::new(self.disambiguators),
         };
         ResolverOutputs { global_ctxt, ast_lowering }
     }

compiler/rustc_ty_utils/src/assoc.rs

@@ -1,6 +1,8 @@
 use rustc_hir::def::DefKind;
-use rustc_hir::def_id::{DefId, DefIdMap, LocalDefId};
-use rustc_hir::definitions::{DefPathData, DisambiguatorState};
+use rustc_hir::def_id::{DefId, DefIdMap, LocalDefId, LocalDefIdMap};
+use rustc_hir::definitions::{
+    DefPathData, PerParentDisambiguatorState, PerParentDisambiguatorsMap,
+};
 use rustc_hir::intravisit::{self, Visitor};
 use rustc_hir::{self as hir, ConstItemRhs, ImplItemImplKind, ItemKind};
 use rustc_middle::query::Providers;
@@ -129,7 +131,7 @@ struct RPITVisitor<'a, 'tcx> {
     tcx: TyCtxt<'tcx>,
     synthetics: Vec<LocalDefId>,
     data: DefPathData,
-    disambiguator: &'a mut DisambiguatorState,
+    disambiguators: &'a mut LocalDefIdMap<PerParentDisambiguatorState>,
 }
 
 impl<'tcx> Visitor<'tcx> for RPITVisitor<'_, 'tcx> {
@@ -138,7 +140,7 @@ fn visit_opaque_ty(&mut self, opaque: &'tcx hir::OpaqueTy<'tcx>) -> Self::Result
             self.tcx,
             opaque.def_id,
             self.data,
-            &mut self.disambiguator,
+            &mut self.disambiguators,
         ));
         intravisit::walk_opaque_ty(self, opaque)
     }
@@ -149,7 +151,7 @@ fn associated_types_for_impl_traits_in_trait_or_impl<'tcx>(
     def_id: LocalDefId,
 ) -> DefIdMap<Vec<DefId>> {
     let item = tcx.hir_expect_item(def_id);
-    let disambiguator = &mut DisambiguatorState::new();
+    let disambiguators = &mut Default::default();
     match item.kind {
         ItemKind::Trait(.., trait_item_refs) => trait_item_refs
             .iter()
@@ -163,7 +165,7 @@ fn associated_types_for_impl_traits_in_trait_or_impl<'tcx>(
                 };
                 let def_name = tcx.item_name(fn_def_id.to_def_id());
                 let data = DefPathData::AnonAssocTy(def_name);
-                let mut visitor = RPITVisitor { tcx, synthetics: vec![], data, disambiguator };
+                let mut visitor = RPITVisitor { tcx, synthetics: vec![], data, disambiguators };
                 visitor.visit_fn_ret_ty(output);
                 let defs = visitor
                     .synthetics
@@ -197,7 +199,7 @@ fn associated_types_for_impl_traits_in_trait_or_impl<'tcx>(
                         return Some((did, vec![]));
                     };
                     let iter = in_trait_def[&trait_item_def_id].iter().map(|&id| {
-                        associated_type_for_impl_trait_in_impl(tcx, id, item, disambiguator)
+                        associated_type_for_impl_trait_in_impl(tcx, id, item, disambiguators)
                             .to_def_id()
                     });
                     Some((did, iter.collect()))
@@ -221,7 +223,7 @@ fn associated_type_for_impl_trait_in_trait(
     tcx: TyCtxt<'_>,
     opaque_ty_def_id: LocalDefId,
     data: DefPathData,
-    disambiguator: &mut DisambiguatorState,
+    disambiguators: &mut LocalDefIdMap<PerParentDisambiguatorState>,
 ) -> LocalDefId {
     let (hir::OpaqueTyOrigin::FnReturn { parent: fn_def_id, .. }
     | hir::OpaqueTyOrigin::AsyncFn { parent: fn_def_id, .. }) =
@@ -240,7 +242,7 @@ fn associated_type_for_impl_trait_in_trait(
         None,
         DefKind::AssocTy,
         Some(data),
-        disambiguator,
+        disambiguators.get_or_create(trait_def_id),
     );
 
     let local_def_id = trait_assoc_ty.def_id();
@@ -283,7 +285,7 @@ fn associated_type_for_impl_trait_in_impl(
     tcx: TyCtxt<'_>,
     trait_assoc_def_id: DefId,
     impl_fn: &hir::ImplItem<'_>,
-    disambiguator: &mut DisambiguatorState,
+    disambiguators: &mut LocalDefIdMap<PerParentDisambiguatorState>,
 ) -> LocalDefId {
     let impl_local_def_id = tcx.local_parent(impl_fn.owner_id.def_id);
 
@@ -306,7 +308,7 @@ fn associated_type_for_impl_trait_in_impl(
         None,
         DefKind::AssocTy,
         Some(data),
-        disambiguator,
+        disambiguators.get_or_create(impl_local_def_id),
     );
 
     let local_def_id = impl_assoc_ty.def_id();