Store disambiguators for each DepNode in CurrentDepNode.

Author: cjgillot

compiler/rustc_middle/src/ty/context.rs

@@ -5,7 +5,7 @@
 pub mod tls;
 
 use crate::arena::Arena;
-use crate::dep_graph::{CurrentDepNode, DepGraph, DepKindStruct, DepNode};
+use crate::dep_graph::{CurrentDepNode, DepGraph, DepKindStruct};
 use crate::infer::canonical::CanonicalVarInfo;
 use crate::lint::struct_lint_level;
 use crate::metadata::ModChild;
@@ -902,40 +902,28 @@ impl<'tcx> TyCtxt<'tcx> {
     }
 }
 
-const ENSURE_FAKE_DEP_NODE_SYNC: () = if std::mem::size_of::<rustc_span::hygiene::FakeDepNode>()
-    != std::mem::size_of::<DepNode>()
-{
-    panic!(
-        "rustc_span::hygiene::FakeDepNode has come out of sync with DepNode and should be updated."
-    )
-};
-
 impl<'tcx> TyCtxt<'tcx> {
     #[instrument(level = "trace", skip(self), ret)]
     pub fn create_expansion(self, expn_data: ExpnData) -> LocalExpnId {
-        let current_node = tls::with_related_context(self, |icx| icx.current_node);
-        let CurrentDepNode::Regular(DepNode { kind, hash }) = current_node else {
+        let current_node = tls::with_related_context(self, |icx| icx.current_node.clone());
+        let CurrentDepNode::Regular { dep_node, expn_disambiguators } = current_node else {
             bug!("creating an expansion outside of a query")
         };
-        let _ = ENSURE_FAKE_DEP_NODE_SYNC;
-        let fake_dep_node = rustc_span::hygiene::FakeDepNode { kind: kind as _, hash };
         self.with_stable_hashing_context(|ctx| {
-            LocalExpnId::create_untracked_expansion(expn_data, fake_dep_node, ctx)
+            LocalExpnId::create_untracked_expansion(expn_data, dep_node, ctx, &expn_disambiguators)
         })
     }
 
     /// Fill an empty expansion. This method must not be used outside of the resolver.
     #[inline]
     #[instrument(level = "trace", skip(self))]
     pub fn finalize_expansion(self, expn_id: LocalExpnId, expn_data: ExpnData) {
-        let current_node = tls::with_related_context(self, |icx| icx.current_node);
-        let CurrentDepNode::Regular(DepNode { kind, hash }) = current_node else {
+        let current_node = tls::with_related_context(self, |icx| icx.current_node.clone());
+        let CurrentDepNode::Regular { dep_node, expn_disambiguators } = current_node else {
             bug!("creating an expansion outside of a query")
         };
-        let _ = ENSURE_FAKE_DEP_NODE_SYNC;
-        let fake_dep_node = rustc_span::hygiene::FakeDepNode { kind: kind as _, hash };
         self.with_stable_hashing_context(|ctx| {
-            expn_id.set_untracked_expn_data(expn_data, fake_dep_node, ctx)
+            expn_id.set_untracked_expn_data(expn_data, dep_node, ctx, &expn_disambiguators)
         });
     }
 

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

@@ -31,7 +31,7 @@ pub struct ImplicitCtxt<'a, 'tcx> {
     /// Used to prevent queries from calling too deeply.
     pub query_depth: usize,
 
-    /// The DepNode of the query being executed. This is updated by the dep-graph. Thisis used to
+    /// The DepNode of the query being executed. This is updated by the dep-graph. This is used to
     /// know which query created an expansion.
     pub current_node: CurrentDepNode,
 
@@ -46,7 +46,7 @@ impl<'a, 'tcx> ImplicitCtxt<'a, 'tcx> {
         let current_node = if tcx.dep_graph.is_fully_enabled() {
             CurrentDepNode::Untracked
         } else {
-            CurrentDepNode::Regular(DepNode::NULL)
+            CurrentDepNode::regular(DepNode::NULL)
         };
         ImplicitCtxt {
             tcx,

compiler/rustc_query_impl/src/plumbing.rs

@@ -142,7 +142,7 @@ impl QueryContext for QueryCtxt<'_> {
                 query: Some(token),
                 diagnostics,
                 query_depth: current_icx.query_depth + depth_limit as usize,
-                current_node: current_icx.current_node,
+                current_node: current_icx.current_node.clone(),
                 task_deps: current_icx.task_deps,
             };
 

compiler/rustc_query_system/src/dep_graph/graph.rs

@@ -355,7 +355,7 @@ impl<K: DepKind> DepGraphData<K> {
         );
 
         let with_deps =
-            |task_deps| K::with_deps(CurrentDepNode::Regular(key), task_deps, || task(cx, arg));
+            |task_deps| K::with_deps(CurrentDepNode::regular(key), task_deps, || task(cx, arg));
         let (result, edges) = if cx.dep_context().is_eval_always(key.kind) {
             (with_deps(TaskDepsRef::EvalAlways), smallvec![])
         } else {

compiler/rustc_query_system/src/dep_graph/mod.rs

@@ -14,6 +14,9 @@ pub use serialized::{SerializedDepGraph, SerializedDepNodeIndex};
 
 use crate::ich::StableHashingContext;
 use rustc_data_structures::profiling::SelfProfilerRef;
+use rustc_data_structures::stable_hasher::Hash64;
+use rustc_data_structures::sync::{Lock, Lrc};
+use rustc_data_structures::unhash::UnhashMap;
 use rustc_serialize::{opaque::FileEncoder, Encodable};
 use rustc_session::Session;
 
@@ -137,13 +140,27 @@ impl FingerprintStyle {
     }
 }
 
-#[derive(Copy, Clone, Hash)]
+#[derive(Clone)]
 pub enum CurrentDepNode<K> {
-    Regular(DepNode<K>),
+    Regular {
+        dep_node: DepNode<K>,
+        /// Disambiguation map for expansions created while executing
+        /// the query with this `DepNode`.
+        expn_disambiguators: Lrc<Lock<UnhashMap<Hash64, u32>>>,
+    },
     Anonymous,
     Untracked,
 }
 
+impl<K> CurrentDepNode<K> {
+    pub fn regular(dep_node: DepNode<K>) -> CurrentDepNode<K> {
+        CurrentDepNode::Regular {
+            dep_node,
+            expn_disambiguators: Lrc::new(Lock::new(UnhashMap::default())),
+        }
+    }
+}
+
 /// Describe the different families of dependency nodes.
 pub trait DepKind: Copy + fmt::Debug + Eq + Hash + Send + Encodable<FileEncoder> + 'static {
     /// DepKind to use when incr. comp. is turned off.

compiler/rustc_query_system/src/query/plumbing.rs

@@ -632,7 +632,7 @@ where
     // The dep-graph for this computation is already in-place, but any side effect due to this
     // query needs to know our DepNode.
     let result =
-        Qcx::DepKind::with_deps(CurrentDepNode::Regular(*dep_node), TaskDepsRef::Ignore, || {
+        Qcx::DepKind::with_deps(CurrentDepNode::regular(*dep_node), TaskDepsRef::Ignore, || {
             query.compute(qcx, *key)
         });
 

compiler/rustc_span/src/hygiene.rs

@@ -30,7 +30,7 @@ use crate::with_session_globals;
 use crate::{HashStableContext, Span, DUMMY_SP};
 
 use crate::def_id::{CrateNum, DefId, StableCrateId, CRATE_DEF_ID, LOCAL_CRATE};
-use rustc_data_structures::fingerprint::{Fingerprint, PackedFingerprint};
+use rustc_data_structures::fingerprint::Fingerprint;
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
 use rustc_data_structures::stable_hasher::HashingControls;
 use rustc_data_structures::stable_hasher::{Hash64, HashStable, StableHasher};
@@ -195,11 +195,12 @@ impl LocalExpnId {
     #[instrument(level = "trace", skip(ctx), ret)]
     pub fn create_untracked_expansion(
         mut expn_data: ExpnData,
-        dep_node: FakeDepNode,
+        hash_extra: impl Hash + Copy + fmt::Debug,
         ctx: impl HashStableContext,
+        disambiguation_map: &Lock<UnhashMap<Hash64, u32>>,
     ) -> LocalExpnId {
         debug_assert_eq!(expn_data.parent.krate, LOCAL_CRATE);
-        let expn_hash = update_disambiguator(&mut expn_data, dep_node, ctx);
+        let expn_hash = update_disambiguator(&mut expn_data, hash_extra, ctx, disambiguation_map);
         HygieneData::with(|data| {
             let expn_id = data.local_expn_data.push(Some(expn_data));
             let _eid = data.local_expn_hashes.push(expn_hash);
@@ -218,11 +219,12 @@ impl LocalExpnId {
     pub fn set_untracked_expn_data(
         self,
         mut expn_data: ExpnData,
-        dep_node: FakeDepNode,
+        hash_extra: impl Hash + Copy + fmt::Debug,
         ctx: impl HashStableContext,
+        disambiguation_map: &Lock<UnhashMap<Hash64, u32>>,
     ) {
         debug_assert_eq!(expn_data.parent.krate, LOCAL_CRATE);
-        let expn_hash = update_disambiguator(&mut expn_data, dep_node, ctx);
+        let expn_hash = update_disambiguator(&mut expn_data, hash_extra, ctx, disambiguation_map);
         HygieneData::with(|data| {
             let old_expn_data = &mut data.local_expn_data[self];
             assert!(old_expn_data.is_none(), "expansion data is reset for an expansion ID");
@@ -349,14 +351,6 @@ impl ExpnId {
     }
 }
 
-/// This struct is meant to be a surrogate for the actual `DepNode` in rustc_middle.
-/// Both types should be kept in sync.
-#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
-pub struct FakeDepNode {
-    pub kind: u16,
-    pub hash: PackedFingerprint,
-}
-
 #[derive(Debug)]
 pub struct HygieneData {
     /// Each expansion should have an associated expansion data, but sometimes there's a delay
@@ -371,12 +365,6 @@ pub struct HygieneData {
     expn_hash_to_expn_id: UnhashMap<ExpnHash, ExpnId>,
     syntax_context_data: Vec<SyntaxContextData>,
     syntax_context_map: FxHashMap<(SyntaxContext, ExpnId, Transparency), SyntaxContext>,
-    /// Maps the `local_hash` of an `ExpnData` to the next disambiguator value.
-    /// This is used by `update_disambiguator` to keep track of which `ExpnData`s
-    /// would have collisions without a disambiguator.
-    /// The keys of this map are always computed with `ExpnData.disambiguator`
-    /// set to 0.
-    expn_data_disambiguators: FxHashMap<FakeDepNode, FxHashMap<Hash64, u32>>,
 }
 
 impl HygieneData {
@@ -405,7 +393,6 @@ impl HygieneData {
                 dollar_crate_name: kw::DollarCrate,
             }],
             syntax_context_map: FxHashMap::default(),
-            expn_data_disambiguators: FxHashMap::default(),
         }
     }
 
@@ -1049,10 +1036,10 @@ impl ExpnData {
     }
 
     #[inline]
-    fn hash_expn(&self, dep_node: FakeDepNode, ctx: &mut impl HashStableContext) -> Hash64 {
+    fn hash_expn(&self, hash_extra: impl Hash, ctx: &mut impl HashStableContext) -> Hash64 {
         let mut hasher = StableHasher::new();
         self.hash_stable(ctx, &mut hasher);
-        dep_node.hash(&mut hasher);
+        hash_extra.hash(&mut hasher);
         hasher.finish()
     }
 }
@@ -1478,44 +1465,38 @@ impl<D: Decoder> Decodable<D> for SyntaxContext {
 #[instrument(level = "trace", skip(ctx), ret)]
 fn update_disambiguator(
     expn_data: &mut ExpnData,
-    dep_node: FakeDepNode,
+    hash_extra: impl Hash + Copy + fmt::Debug,
     mut ctx: impl HashStableContext,
+    disambiguation_map: &Lock<UnhashMap<Hash64, u32>>,
 ) -> ExpnHash {
     // This disambiguator should not have been set yet.
     assert_eq!(expn_data.disambiguator, 0, "Already set disambiguator for ExpnData: {expn_data:?}");
     assert_default_hashing_controls(&ctx, "ExpnData (disambiguator)");
-    let mut expn_hash = expn_data.hash_expn(dep_node, &mut ctx);
+    let mut expn_hash = expn_data.hash_expn(hash_extra, &mut ctx);
     debug!(?expn_hash);
 
-    let disambiguator = HygieneData::with(|data| {
+    let disambiguator = {
         // If this is the first ExpnData with a given hash, then keep our
         // disambiguator at 0 (the default u32 value)
-        let disambig = data
-            .expn_data_disambiguators
-            .entry(dep_node)
-            .or_default()
-            .entry(expn_hash)
-            .or_default();
+        let mut disambiguation_map = disambiguation_map.lock();
+        let disambig = disambiguation_map.entry(expn_hash).or_default();
         let disambiguator = *disambig;
         *disambig += 1;
         disambiguator
-    });
+    };
 
     if disambiguator != 0 {
         debug!("Set disambiguator for expn_data={:?} expn_hash={:?}", expn_data, expn_hash);
 
         expn_data.disambiguator = disambiguator;
-        expn_hash = expn_data.hash_expn(dep_node, &mut ctx);
+        expn_hash = expn_data.hash_expn(hash_extra, &mut ctx);
 
         // Verify that the new disambiguator makes the hash unique
-        #[cfg(debug_assertions)]
-        HygieneData::with(|data| {
-            assert_eq!(
-                data.expn_data_disambiguators[&dep_node].get(&expn_hash),
-                None,
-                "Hash collision after disambiguator update!",
-            );
-        });
+        debug_assert_eq!(
+            disambiguation_map.lock().get(&expn_hash),
+            None,
+            "Hash collision after disambiguator update!",
+        );
     }
 
     ExpnHash::new(ctx.def_path_hash(LOCAL_CRATE.as_def_id()).stable_crate_id(), expn_hash)