From 0200bbfa1d2fbf175d952b22155f38d22e813ad4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?John=20K=C3=A5re=20Alsaker?= <john.kare.alsaker@gmail.com>
Date: Sat, 9 Mar 2024 09:29:16 +0100
Subject: [PATCH] Verify that query keys result in unique dep nodes

---
 compiler/rustc_interface/src/queries.rs   |  2 ++
 compiler/rustc_query_impl/src/lib.rs      |  2 +-
 compiler/rustc_query_impl/src/plumbing.rs | 43 +++++++++++++++++++++++
 3 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/compiler/rustc_interface/src/queries.rs b/compiler/rustc_interface/src/queries.rs
index da11d090b7497..c22188226961e 100644
--- a/compiler/rustc_interface/src/queries.rs
+++ b/compiler/rustc_interface/src/queries.rs
@@ -321,6 +321,8 @@ impl Compiler {
             }
 
             self.sess.time("serialize_dep_graph", || gcx.enter(rustc_incremental::save_dep_graph));
+
+            gcx.enter(rustc_query_impl::query_key_hash_verify_all);
         }
 
         // The timer's lifetime spans the dropping of `queries`, which contains
diff --git a/compiler/rustc_query_impl/src/lib.rs b/compiler/rustc_query_impl/src/lib.rs
index 33116737a4203..5276715553272 100644
--- a/compiler/rustc_query_impl/src/lib.rs
+++ b/compiler/rustc_query_impl/src/lib.rs
@@ -41,7 +41,7 @@ use rustc_span::{ErrorGuaranteed, Span};
 
 #[macro_use]
 mod plumbing;
-pub use crate::plumbing::QueryCtxt;
+pub use crate::plumbing::{query_key_hash_verify_all, QueryCtxt};
 
 mod profiling_support;
 pub use self::profiling_support::alloc_self_profile_query_strings;
diff --git a/compiler/rustc_query_impl/src/plumbing.rs b/compiler/rustc_query_impl/src/plumbing.rs
index b06d75be3902d..492971922ca45 100644
--- a/compiler/rustc_query_impl/src/plumbing.rs
+++ b/compiler/rustc_query_impl/src/plumbing.rs
@@ -5,6 +5,7 @@
 use crate::rustc_middle::dep_graph::DepContext;
 use crate::rustc_middle::ty::TyEncoder;
 use crate::QueryConfigRestored;
+use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::stable_hasher::{Hash64, HashStable, StableHasher};
 use rustc_data_structures::sync::Lock;
 use rustc_errors::DiagInner;
@@ -189,6 +190,14 @@ pub(super) fn encode_all_query_results<'tcx>(
     }
 }
 
+pub fn query_key_hash_verify_all<'tcx>(tcx: TyCtxt<'tcx>) {
+    tcx.sess.time("query_key_hash_verify_all", || {
+        for verify in super::QUERY_KEY_HASH_VERIFY.iter() {
+            verify(tcx);
+        }
+    })
+}
+
 macro_rules! handle_cycle_error {
     ([]) => {{
         rustc_query_system::HandleCycleError::Error
@@ -370,6 +379,29 @@ pub(crate) fn encode_query_results<'a, 'tcx, Q>(
     });
 }
 
+pub(crate) fn query_key_hash_verify<'tcx>(
+    query: impl QueryConfig<QueryCtxt<'tcx>>,
+    qcx: QueryCtxt<'tcx>,
+) {
+    let _timer =
+        qcx.profiler().generic_activity_with_arg("query_key_hash_verify_for", query.name());
+
+    let mut map = FxHashMap::default();
+
+    let cache = query.query_cache(qcx);
+    cache.iter(&mut |key, _, _| {
+        let node = DepNode::construct(qcx.tcx, query.dep_kind(), key);
+        if let Some(other_key) = map.insert(node, *key) {
+            bug!(
+                "query key `{:?}` and `{:?}` mapped to the same dep node {:?}",
+                key,
+                other_key,
+                node
+            );
+        }
+    });
+}
+
 fn try_load_from_on_disk_cache<'tcx, Q>(query: Q, tcx: TyCtxt<'tcx>, dep_node: DepNode)
 where
     Q: QueryConfig<QueryCtxt<'tcx>>,
@@ -691,6 +723,13 @@ macro_rules! define_queries {
                     )
                 }
             }}
+
+            pub fn query_key_hash_verify<'tcx>(tcx: TyCtxt<'tcx>) {
+                $crate::plumbing::query_key_hash_verify(
+                    query_impl::$name::QueryType::config(tcx),
+                    QueryCtxt::new(tcx),
+                )
+            }
         })*}
 
         pub(crate) fn engine(incremental: bool) -> QueryEngine {
@@ -730,6 +769,10 @@ macro_rules! define_queries {
             >
         ] = &[$(expand_if_cached!([$($modifiers)*], query_impl::$name::encode_query_results)),*];
 
+        const QUERY_KEY_HASH_VERIFY: &[
+            for<'tcx> fn(TyCtxt<'tcx>)
+        ] = &[$(query_impl::$name::query_key_hash_verify),*];
+
         #[allow(nonstandard_style)]
         mod query_callbacks {
             use super::*;