From 55ba7a7c622fb1ac95ef05e4a652cedf923a8164 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 | 50 +++++++++++++++++++++++
 compiler/rustc_session/src/options.rs     |  4 +-
 4 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/compiler/rustc_interface/src/queries.rs b/compiler/rustc_interface/src/queries.rs
index da11d090b74..c2218822696 100644
--- a/compiler/rustc_interface/src/queries.rs
+++ b/compiler/rustc_interface/src/queries.rs
@@ -321,6 +321,8 @@ pub fn enter<F, T>(&self, f: F) -> T
             }
 
             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 33116737a42..52767155532 100644
--- a/compiler/rustc_query_impl/src/lib.rs
+++ b/compiler/rustc_query_impl/src/lib.rs
@@ -41,7 +41,7 @@
 
 #[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 b06d75be390..aa965779731 100644
--- a/compiler/rustc_query_impl/src/plumbing.rs
+++ b/compiler/rustc_query_impl/src/plumbing.rs
@@ -7,6 +7,7 @@
 use crate::QueryConfigRestored;
 use rustc_data_structures::stable_hasher::{Hash64, HashStable, StableHasher};
 use rustc_data_structures::sync::Lock;
+use rustc_data_structures::unord::UnordMap;
 use rustc_errors::DiagInner;
 
 use rustc_index::Idx;
@@ -189,6 +190,16 @@ pub(super) fn encode_all_query_results<'tcx>(
     }
 }
 
+pub fn query_key_hash_verify_all<'tcx>(tcx: TyCtxt<'tcx>) {
+    if tcx.sess().opts.unstable_opts.incremental_verify_ich || cfg!(debug_assertions) {
+        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 +381,34 @@ 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 = UnordMap::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:\n\
+                `{:?}`\n\
+                and key:\n\
+                `{:?}`\n\
+                mapped to the same dep node:\n\
+                {:?}",
+                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 +730,13 @@ pub fn encode_query_results<'tcx>(
                     )
                 }
             }}
+
+            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 +776,10 @@ pub fn dynamic_queries<'tcx>() -> DynamicQueries<'tcx> {
             >
         ] = &[$(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::*;
diff --git a/compiler/rustc_session/src/options.rs b/compiler/rustc_session/src/options.rs
index 93bef82e4ba..0122b9ec33e 100644
--- a/compiler/rustc_session/src/options.rs
+++ b/compiler/rustc_session/src/options.rs
@@ -1674,7 +1674,9 @@ pub(crate) fn parse_function_return(slot: &mut FunctionReturn, v: Option<&str>)
         "print high-level information about incremental reuse (or the lack thereof) \
         (default: no)"),
     incremental_verify_ich: bool = (false, parse_bool, [UNTRACKED],
-        "verify incr. comp. hashes of green query instances (default: no)"),
+        "verify extended properties for incr. comp. (default: no):
+        - hashes of green query instances
+        - hash collisions of query keys"),
     inline_in_all_cgus: Option<bool> = (None, parse_opt_bool, [TRACKED],
         "control whether `#[inline]` functions are in all CGUs"),
     inline_llvm: bool = (true, parse_bool, [TRACKED],