Better disambiguators

Author: aerooneqq

compiler/rustc_ast_lowering/src/lib.rs

@@ -51,7 +51,7 @@ use rustc_data_structures::tagged_ptr::TaggedRef;
 use rustc_errors::{DiagArgFromDisplay, DiagCtxtHandle};
 use rustc_hir::def::{DefKind, LifetimeRes, Namespace, PartialRes, PerNS, Res};
 use rustc_hir::def_id::{CRATE_DEF_ID, LOCAL_CRATE, LocalDefId};
-use rustc_hir::definitions::PerParentDisambiguatorState;
+use rustc_hir::definitions::{Disambiguator, PerParentDisambiguatorState};
 use rustc_hir::lints::{AttributeLint, DelayedLint, DynAttribute};
 use rustc_hir::{
     self as hir, AngleBrackets, ConstArg, GenericArg, HirId, ItemLocalMap, LifetimeSource,
@@ -410,7 +410,12 @@ impl<'tcx> ResolverAstLowering<'tcx> {
     }
 
     fn steal_or_create_disambiguator(&self, parent: LocalDefId) -> PerParentDisambiguatorState {
-        self.per_parent_disambiguators.get(&parent).map(|s| s.steal()).unwrap_or_default()
+        self.per_parent_disambiguators.get(&parent).map(|s| s.steal()).unwrap_or_else(|| {
+            PerParentDisambiguatorState::new(
+                #[cfg(debug_assertions)]
+                parent,
+            )
+        })
     }
 }
 
@@ -740,7 +745,13 @@ 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,
+                Disambiguator::PerParent(&mut self.disambiguator),
+            )
             .def_id();
 
         debug!("create_def: def_id_to_node_id[{:?}] <-> {:?}", def_id, node_id);

compiler/rustc_const_eval/src/interpret/intern.rs

@@ -17,7 +17,7 @@ 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::definitions::{DefPathData, Disambiguator, DisambiguatorState};
 use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrs;
 use rustc_middle::mir::interpret::{
     AllocBytes, ConstAllocation, CtfeProvenance, InterpResult, Provenance,
@@ -155,7 +155,7 @@ fn intern_as_new_static<'tcx>(
         None,
         DefKind::Static { safety: hir::Safety::Safe, mutability: alloc.0.mutability, nested: true },
         Some(DefPathData::NestedStatic),
-        disambiguator,
+        Disambiguator::Full(disambiguator),
     );
     tcx.set_nested_alloc_id_static(alloc_id, feed.def_id());
 

compiler/rustc_hir/src/definitions.rs

@@ -7,8 +7,8 @@
 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};
@@ -97,36 +97,35 @@ impl DefPathTable {
     }
 }
 
-pub trait Disambiguator {
-    fn entry(&mut self, parent: LocalDefId, data: DefPathData) -> &mut u32;
+pub enum Disambiguator<'a> {
+    PerParent(&'a mut PerParentDisambiguatorState),
+    Full(&'a mut DisambiguatorState),
 }
 
 #[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 {
+    pub fn new(#[cfg(debug_assertions)] 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)
-    }
+    next: FxHashMap<(LocalDefId, DefPathData), u32>,
 }
 
 impl DisambiguatorState {
-    pub const fn new() -> Self {
+    pub fn new() -> Self {
         Self { next: Default::default() }
     }
 
@@ -408,7 +407,7 @@ impl Definitions {
         &mut self,
         parent: LocalDefId,
         data: DefPathData,
-        disambiguator: &mut impl Disambiguator,
+        disambiguator: Disambiguator<'_>,
     ) -> 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 +421,19 @@ impl Definitions {
 
         // Find the next free disambiguator for this key.
         let disambiguator = {
-            let next_disamb = disambiguator.entry(parent, data);
+            let next_disamb = match disambiguator {
+                Disambiguator::PerParent(d) => {
+                    debug_assert_eq!(
+                        parent,
+                        d.parent.expect("must be set"),
+                        "provided parent and parent in disambiguator must be the same"
+                    );
+
+                    d.next.entry(data).or_insert(0)
+                }
+                Disambiguator::Full(d) => d.next.entry((parent, 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,7 @@ use rustc_ast::visit::walk_list;
 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::definitions::{DefPathData, Disambiguator, DisambiguatorState};
 use rustc_hir::intravisit::{self, InferKind, Visitor, VisitorExt};
 use rustc_hir::{
     self as hir, AmbigArg, GenericArg, GenericParam, GenericParamKind, HirId, LifetimeKind, Node,
@@ -1523,7 +1523,7 @@ impl<'a, 'tcx> BoundVarContext<'a, 'tcx> {
                     None,
                     DefKind::LifetimeParam,
                     Some(DefPathData::OpaqueLifetime(ident.name)),
-                    self.disambiguator,
+                    Disambiguator::Full(self.disambiguator),
                 );
                 feed.def_span(ident.span);
                 feed.def_ident_span(Some(ident.span));

compiler/rustc_middle/src/ty/context.rs

@@ -1398,7 +1398,7 @@ impl<'tcx> TyCtxtAt<'tcx> {
         name: Option<Symbol>,
         def_kind: DefKind,
         override_def_path_data: Option<DefPathData>,
-        disambiguator: &mut impl Disambiguator,
+        disambiguator: Disambiguator<'_>,
     ) -> TyCtxtFeed<'tcx, LocalDefId> {
         let feed =
             self.tcx.create_def(parent, name, def_kind, override_def_path_data, disambiguator);
@@ -1416,7 +1416,7 @@ impl<'tcx> TyCtxt<'tcx> {
         name: Option<Symbol>,
         def_kind: DefKind,
         override_def_path_data: Option<DefPathData>,
-        disambiguator: &mut impl Disambiguator,
+        disambiguator: Disambiguator<'_>,
     ) -> 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_mir_transform/src/coroutine/by_move_body.rs

@@ -73,7 +73,7 @@ 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::{Disambiguator, DisambiguatorState};
 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(),
+        Disambiguator::Full(&mut DisambiguatorState::new()),
     );
     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 @@ use rustc_hir::def::{
     PerNS,
 };
 use rustc_hir::def_id::{CRATE_DEF_ID, CrateNum, DefId, LOCAL_CRATE, LocalDefId, LocalDefIdMap};
-use rustc_hir::definitions::PerParentDisambiguatorState;
+use rustc_hir::definitions::{Disambiguator, PerParentDisambiguatorState};
 use rustc_hir::{PrimTy, TraitCandidate, find_attr};
 use rustc_index::bit_set::DenseBitSet;
 use rustc_metadata::creader::CStore;
@@ -1620,13 +1620,20 @@ impl<'tcx> Resolver<'_, 'tcx> {
             self.tcx.definitions_untracked().def_key(self.node_id_to_def_id[&node_id].key()),
         );
 
+        let disambiguator = self.per_parent_disambiguators.entry(parent).or_insert_with(|| {
+            PerParentDisambiguatorState::new(
+                #[cfg(debug_assertions)]
+                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(),
+            Disambiguator::PerParent(disambiguator),
         );
         let def_id = feed.def_id();
 

compiler/rustc_ty_utils/src/assoc.rs

@@ -1,6 +1,6 @@
 use rustc_hir::def::DefKind;
 use rustc_hir::def_id::{DefId, DefIdMap, LocalDefId};
-use rustc_hir::definitions::{DefPathData, DisambiguatorState};
+use rustc_hir::definitions::{DefPathData, Disambiguator, DisambiguatorState};
 use rustc_hir::intravisit::{self, Visitor};
 use rustc_hir::{self as hir, ConstItemRhs, ImplItemImplKind, ItemKind};
 use rustc_middle::query::Providers;
@@ -240,7 +240,7 @@ fn associated_type_for_impl_trait_in_trait(
         None,
         DefKind::AssocTy,
         Some(data),
-        disambiguator,
+        Disambiguator::Full(disambiguator),
     );
 
     let local_def_id = trait_assoc_ty.def_id();
@@ -306,7 +306,7 @@ fn associated_type_for_impl_trait_in_impl(
         None,
         DefKind::AssocTy,
         Some(data),
-        disambiguator,
+        Disambiguator::Full(disambiguator),
     );
 
     let local_def_id = impl_assoc_ty.def_id();