Make create_def invoke a helper query if used from a recorded query

oli-obk · oli-obk · commit 17bd715e55f3 · 2025-04-25T13:34:59.000Z
diff --git a/compiler/rustc_hir/src/definitions.rs b/compiler/rustc_hir/src/definitions.rs
@@ -10,8 +10,8 @@ use std::hash::Hash;
 use rustc_data_structures::stable_hasher::StableHasher;
 use rustc_data_structures::unord::UnordMap;
 use rustc_hashes::Hash64;
-use rustc_index::IndexVec;
-use rustc_macros::{Decodable, Encodable};
+use rustc_index::{IndexVec, static_assert_size};
+use rustc_macros::{Decodable, Encodable, HashStable_Generic};
 use rustc_span::{Symbol, kw, sym};
 use tracing::{debug, instrument};
 
@@ -252,7 +252,7 @@ impl DefPath {
 }
 
 /// New variants should only be added in synchronization with `enum DefKind`.
-#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, Encodable, Decodable)]
+#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, Encodable, Decodable, HashStable_Generic)]
 pub enum DefPathData {
     // Root: these should only be used for the root nodes, because
     // they are treated specially by the `def_path` function.
@@ -293,6 +293,8 @@ pub enum DefPathData {
     SyntheticCoroutineBody,
 }
 
+static_assert_size!(DefPathData, 8);
+
 impl Definitions {
     pub fn def_path_table(&self) -> &DefPathTable {
         &self.table
@@ -346,7 +348,12 @@ impl Definitions {
     }
 
     /// Adds a definition with a parent definition.
-    pub fn create_def(&mut self, parent: LocalDefId, data: DefPathData) -> LocalDefId {
+    pub fn create_def(
+        &mut self,
+        parent: LocalDefId,
+        data: DefPathData,
+        query_local_disambiguator: Option<u32>,
+    ) -> 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.
         debug!(
@@ -361,8 +368,20 @@ impl Definitions {
         let disambiguator = {
             let next_disamb = self.next_disambiguator.entry((parent, data)).or_insert(0);
             let disambiguator = *next_disamb;
-            *next_disamb = next_disamb.checked_add(1).expect("disambiguator overflow");
-            disambiguator
+            // The disambiguator number space is split into the ones counting from 0 upwards,
+            match query_local_disambiguator {
+                None => {
+                    assert!(disambiguator < u32::MAX / 2);
+                    *next_disamb = next_disamb.checked_add(1).expect("disambiguator overflow");
+                    disambiguator
+                }
+                // and the ones counting from MAX downwards.
+                Some(local) => {
+                    // Ensure these two number spaces do not collide. 2^31 disambiguators should be enough for everyone.
+                    assert!(local < u32::MAX / 2);
+                    u32::MAX - local
+                }
+            }
         };
         let key = DefKey {
             parent: Some(parent.local_def_index),
diff --git a/compiler/rustc_interface/src/callbacks.rs b/compiler/rustc_interface/src/callbacks.rs
@@ -25,7 +25,10 @@ fn track_span_parent(def_id: rustc_span::def_id::LocalDefId) {
             // Skip doing anything if we aren't tracking dependencies.
             let tracks_deps = match icx.task_deps {
                 TaskDepsRef::Allow(..) => true,
-                TaskDepsRef::EvalAlways | TaskDepsRef::Ignore | TaskDepsRef::Forbid => false,
+                TaskDepsRef::Replay { .. }
+                | TaskDepsRef::EvalAlways
+                | TaskDepsRef::Ignore
+                | TaskDepsRef::Forbid => false,
             };
             if tracks_deps {
                 let _span = icx.tcx.source_span(def_id);
diff --git a/compiler/rustc_middle/src/hir/mod.rs b/compiler/rustc_middle/src/hir/mod.rs
@@ -239,4 +239,11 @@ pub fn provide(providers: &mut Providers) {
     providers.in_scope_traits_map = |tcx, id| {
         tcx.hir_crate(()).owners[id.def_id].as_owner().map(|owner_info| &owner_info.trait_map)
     };
+    providers.create_def_raw = |tcx, (parent, data, query_local_disambiguator)| {
+        tcx.untracked().definitions.write().create_def(
+            parent,
+            data,
+            Some(query_local_disambiguator),
+        )
+    }
 }
diff --git a/compiler/rustc_middle/src/query/keys.rs b/compiler/rustc_middle/src/query/keys.rs
@@ -3,6 +3,7 @@
 use std::ffi::OsStr;
 
 use rustc_hir::def_id::{CrateNum, DefId, LOCAL_CRATE, LocalDefId, LocalModDefId, ModDefId};
+use rustc_hir::definitions::DefPathData;
 use rustc_hir::hir_id::{HirId, OwnerId};
 use rustc_query_system::dep_graph::DepNodeIndex;
 use rustc_query_system::query::{DefIdCache, DefaultCache, SingleCache, VecCache};
@@ -265,6 +266,14 @@ impl Key for (LocalDefId, LocalDefId) {
     }
 }
 
+impl Key for (LocalDefId, DefPathData, u32) {
+    type Cache<V> = DefaultCache<Self, V>;
+
+    fn default_span(&self, tcx: TyCtxt<'_>) -> Span {
+        self.0.default_span(tcx)
+    }
+}
+
 impl Key for (DefId, Ident) {
     type Cache<V> = DefaultCache<Self, V>;
 
diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
@@ -159,6 +159,13 @@ rustc_queries! {
         desc { "getting the source span" }
     }
 
+    /// Used to handle incremental replays of [`TyCtxt::create_def``] invocations from tracked queries.
+    query create_def_raw(key: (LocalDefId, rustc_hir::definitions::DefPathData, u32)) -> LocalDefId {
+        // Accesses untracked data
+        eval_always
+        desc { "generating a new def id" }
+    }
+
     /// Represents crate as a whole (as distinct from the top-level crate module).
     ///
     /// If you call `tcx.hir_crate(())` we will have to assume that any change
diff --git a/compiler/rustc_middle/src/ty/context.rs b/compiler/rustc_middle/src/ty/context.rs
@@ -41,7 +41,7 @@ use rustc_hir::{self as hir, Attribute, HirId, Node, TraitCandidate};
 use rustc_index::IndexVec;
 use rustc_macros::{HashStable, TyDecodable, TyEncodable};
 use rustc_query_system::cache::WithDepNode;
-use rustc_query_system::dep_graph::DepNodeIndex;
+use rustc_query_system::dep_graph::{DepNodeIndex, TaskDepsRef};
 use rustc_query_system::ich::StableHashingContext;
 use rustc_serialize::opaque::{FileEncodeResult, FileEncoder};
 use rustc_session::config::CrateType;
@@ -1955,27 +1955,42 @@ impl<'tcx> TyCtxt<'tcx> {
         def_kind: DefKind,
     ) -> TyCtxtFeed<'tcx, LocalDefId> {
         let data = def_kind.def_path_data(name);
-        // The following call has the side effect of modifying the tables inside `definitions`.
-        // These very tables are relied on by the incr. comp. engine to decode DepNodes and to
-        // decode the on-disk cache.
-        //
-        // Any LocalDefId which is used within queries, either as key or result, either:
-        // - has been created before the construction of the TyCtxt;
-        // - has been created by this call to `create_def`.
-        // As a consequence, this LocalDefId is always re-created before it is needed by the incr.
-        // comp. engine itself.
-        //
-        // This call also writes to the value of the `source_span` query.
-        // This is fine because:
-        // - that query is `eval_always` so we won't miss its result changing;
-        // - this write will have happened before that query is called.
-        let def_id = self.untracked.definitions.write().create_def(parent, data);
-
-        // This function modifies `self.definitions` using a side-effect.
-        // We need to ensure that these side effects are re-run by the incr. comp. engine.
-        // Depending on the forever-red node will tell the graph that the calling query
-        // needs to be re-evaluated.
-        self.dep_graph.read_index(DepNodeIndex::FOREVER_RED_NODE);
+
+        let def_id = tls::with_context(|icx| {
+            match icx.task_deps {
+                // Always gets rerun anyway, so nothing to replay
+                TaskDepsRef::EvalAlways |
+                // Top-level queries like the resolver get rerun every time anyway
+                TaskDepsRef::Ignore => {
+                    // The following call has the side effect of modifying the tables inside `definitions`.
+                    // These very tables are relied on by the incr. comp. engine to decode DepNodes and to
+                    // decode the on-disk cache.
+                    //
+                    // Any LocalDefId which is used within queries, either as key or result, either:
+                    // - has been created before the construction of the TyCtxt;
+                    // - has been created by this call to `create_def`.
+                    // As a consequence, this LocalDefId is always re-created before it is needed by the incr.
+                    // comp. engine itself.
+                    //
+                    // This call also writes to the value of the `source_span` query.
+                    // This is fine because:
+                    // - that query is `eval_always` so we won't miss its result changing;
+                    // - this write will have happened before that query is called.
+                    self.untracked.definitions.write().create_def(parent, data, None)
+                }
+                TaskDepsRef::Forbid => bug!(
+                    "cannot create definition {parent:?} {data:?} without being able to register task dependencies"
+                ),
+                TaskDepsRef::Allow(deps) => {
+                    let idx = deps.lock().next_def_id_idx();
+                    self.create_def_raw((parent, data, idx))
+                }
+                TaskDepsRef::Replay { created_def_ids } => {
+                    let idx = created_def_ids.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
+                    self.create_def_raw((parent, data, idx))
+                }
+            }
+        });
 
         let feed = TyCtxtFeed { tcx: self, key: def_id };
         feed.def_kind(def_kind);
diff --git a/compiler/rustc_query_system/src/dep_graph/graph.rs b/compiler/rustc_query_system/src/dep_graph/graph.rs
@@ -207,6 +207,10 @@ impl<D: Deps> DepGraph<D> {
         }
     }
 
+    pub(crate) fn with_replay<R>(&self, created_def_ids: &AtomicU32, op: impl FnOnce() -> R) -> R {
+        D::with_deps(TaskDepsRef::Replay { created_def_ids }, op)
+    }
+
     pub fn with_ignore<OP, R>(&self, op: OP) -> R
     where
         OP: FnOnce() -> R,
@@ -361,6 +365,7 @@ impl<D: Deps> DepGraphData<D> {
                 node: Some(key),
                 reads: EdgesVec::new(),
                 read_set: Default::default(),
+                created_def_ids: 0,
             });
             (with_deps(TaskDepsRef::Allow(&task_deps)), task_deps.into_inner().reads)
         };
@@ -476,6 +481,12 @@ impl<D: Deps> DepGraph<D> {
                         return;
                     }
                     TaskDepsRef::Ignore => return,
+                    // We don't need to record dependencies when rerunning a query
+                    // because we have no disk cache entry to load. The dependencies
+                    // are preserved.
+                    // FIXME: assert that the dependencies don't change instead of
+                    // recording them.
+                    TaskDepsRef::Replay { .. } => return,
                     TaskDepsRef::Forbid => {
                         // Reading is forbidden in this context. ICE with a useful error message.
                         panic_on_forbidden_read(data, dep_node_index)
@@ -526,7 +537,9 @@ impl<D: Deps> DepGraph<D> {
     pub fn record_diagnostic<Qcx: QueryContext>(&self, qcx: Qcx, diagnostic: &DiagInner) {
         if let Some(ref data) = self.data {
             D::read_deps(|task_deps| match task_deps {
-                TaskDepsRef::EvalAlways | TaskDepsRef::Ignore => return,
+                TaskDepsRef::EvalAlways | TaskDepsRef::Ignore | TaskDepsRef::Replay { .. } => {
+                    return;
+                }
                 TaskDepsRef::Forbid | TaskDepsRef::Allow(..) => {
                     self.read_index(data.encode_diagnostic(qcx, diagnostic));
                 }
@@ -607,7 +620,7 @@ impl<D: Deps> DepGraph<D> {
                     edges.push(DepNodeIndex::FOREVER_RED_NODE);
                 }
                 TaskDepsRef::Ignore => {}
-                TaskDepsRef::Forbid => {
+                TaskDepsRef::Replay { .. } | TaskDepsRef::Forbid => {
                     panic!("Cannot summarize when dependencies are not recorded.")
                 }
             });
@@ -1317,6 +1330,17 @@ pub enum TaskDepsRef<'a> {
     /// to ensure that the decoding process doesn't itself
     /// require the execution of any queries.
     Forbid,
+    /// Side effects from the previous run made available to
+    /// queries when they are reexecuted because their result was not
+    /// available in the cache. Whenever the query creates a new `DefId`,
+    /// it is checked against the entries in `QuerySideEffects::definitions`
+    /// to ensure that the new `DefId`s are the same as the ones that were
+    /// created the last time the query was executed.
+    Replay {
+        /// Every new `DefId` created increases this counter so that we produce the same ones
+        /// as the original execution created.
+        created_def_ids: &'a AtomicU32,
+    },
 }
 
 #[derive(Debug)]
@@ -1325,6 +1349,8 @@ pub struct TaskDeps {
     node: Option<DepNode>,
     reads: EdgesVec,
     read_set: FxHashSet<DepNodeIndex>,
+    /// Every new `DefId` created increases this counter so that they can be replayed.
+    created_def_ids: u32,
 }
 
 impl Default for TaskDeps {
@@ -1334,9 +1360,19 @@ impl Default for TaskDeps {
             node: None,
             reads: EdgesVec::new(),
             read_set: FxHashSet::with_capacity_and_hasher(128, Default::default()),
+            created_def_ids: 0,
         }
     }
 }
+
+impl TaskDeps {
+    pub fn next_def_id_idx(&mut self) -> u32 {
+        let idx = self.created_def_ids;
+        self.created_def_ids += 1;
+        idx
+    }
+}
+
 // A data structure that stores Option<DepNodeColor> values as a contiguous
 // array, using one u32 per entry.
 pub(super) struct DepNodeColorMap {
diff --git a/compiler/rustc_query_system/src/query/plumbing.rs b/compiler/rustc_query_system/src/query/plumbing.rs
@@ -6,6 +6,7 @@ use std::cell::Cell;
 use std::fmt::Debug;
 use std::hash::Hash;
 use std::mem;
+use std::sync::atomic::AtomicU32;
 
 use hashbrown::hash_table::Entry;
 use rustc_data_structures::fingerprint::Fingerprint;
@@ -633,8 +634,10 @@ where
     // recompute.
     let prof_timer = qcx.dep_context().profiler().query_provider();
 
+    let created_def_ids = AtomicU32::new(0);
     // The dep-graph for this computation is already in-place.
-    let result = qcx.dep_context().dep_graph().with_ignore(|| query.compute(qcx, *key));
+    let result =
+        qcx.dep_context().dep_graph().with_replay(&created_def_ids, || query.compute(qcx, *key));
 
     prof_timer.finish_with_query_invocation_id(dep_node_index.into());
 
