Clean up query-forcing functions

compiler/rustc_query_impl/src/dep_kind_vtables.rs

@@ -4,7 +4,7 @@ use rustc_middle::dep_graph::{DepKindVTable, DepNodeKey, KeyFingerprintStyle};
 use rustc_middle::query::QueryCache;
 
 use crate::GetQueryVTable;
-use crate::plumbing::{force_from_dep_node_inner, promote_from_disk_inner};
+use crate::plumbing::promote_from_disk_inner;
 
 /// [`DepKindVTable`] constructors for special dep kinds that aren't queries.
 #[expect(non_snake_case, reason = "use non-snake case to avoid collision with query names")]
@@ -111,10 +111,16 @@ where
     DepKindVTable {
         is_eval_always,
         key_fingerprint_style,
-        force_from_dep_node_fn: (can_recover && !is_no_force)
-            .then_some(force_from_dep_node_inner::<Q>),
-        promote_from_disk_fn: (can_recover && is_cache_on_disk)
-            .then_some(promote_from_disk_inner::<Q>),
+        force_from_dep_node_fn: (can_recover && !is_no_force).then_some(
+            |tcx, dep_node, _prev_index| {
+                let query = Q::query_vtable(tcx);
+                crate::execution::force_query_dep_node(tcx, query, dep_node)
+            },
+        ),
+        promote_from_disk_fn: (can_recover && is_cache_on_disk).then_some(|tcx, dep_node| {
+            let query = Q::query_vtable(tcx);
+            promote_from_disk_inner(tcx, query, dep_node)
+        }),
     }
 }
 

compiler/rustc_query_impl/src/execution.rs

@@ -680,20 +680,26 @@ pub(super) fn execute_query_incr_inner<'tcx, C: QueryCache>(
     Some(result)
 }
 
-pub(crate) fn force_query<'tcx, C: QueryCache>(
-    query: &'tcx QueryVTable<'tcx, C>,
+/// Inner implementation of [`DepKindVTable::force_from_dep_node_fn`][force_fn]
+/// for query nodes.
+///
+/// [force_fn]: rustc_middle::dep_graph::DepKindVTable::force_from_dep_node_fn
+pub(crate) fn force_query_dep_node<'tcx, C: QueryCache>(
     tcx: TyCtxt<'tcx>,
-    key: C::Key,
+    query: &'tcx QueryVTable<'tcx, C>,
     dep_node: DepNode,
-) {
-    // We may be concurrently trying both execute and force a query.
-    // Ensure that only one of them runs the query.
-    if let Some((_, index)) = query.cache.lookup(&key) {
-        tcx.prof.query_cache_hit(index.into());
-        return;
-    }
+) -> bool {
+    let Some(key) = C::Key::try_recover_key(tcx, &dep_node) else {
+        // We couldn't recover a key from the node's key fingerprint.
+        // Tell the caller that we couldn't force the node.
+        return false;
+    };
 
     ensure_sufficient_stack(|| {
         try_execute_query::<C, true>(query, tcx, DUMMY_SP, key, Some(dep_node))
     });
+
+    // We did manage to recover a key and force the node, though it's up to
+    // the caller to check whether the node ended up marked red or green.
+    true
 }

compiler/rustc_query_impl/src/plumbing.rs

@@ -6,7 +6,7 @@ use rustc_index::Idx;
 use rustc_middle::bug;
 #[expect(unused_imports, reason = "used by doc comments")]
 use rustc_middle::dep_graph::DepKindVTable;
-use rustc_middle::dep_graph::{DepKind, DepNode, DepNodeIndex, DepNodeKey, SerializedDepNodeIndex};
+use rustc_middle::dep_graph::{DepNode, DepNodeIndex, DepNodeKey, SerializedDepNodeIndex};
 use rustc_middle::query::erase::{Erasable, Erased};
 use rustc_middle::query::on_disk_cache::{
     AbsoluteBytePos, CacheDecoder, CacheEncoder, EncodedDepNodeIndex,
@@ -20,10 +20,10 @@ use rustc_span::DUMMY_SP;
 use rustc_span::def_id::LOCAL_CRATE;
 
 use crate::error::{QueryOverflow, QueryOverflowNote};
-use crate::execution::{all_inactive, force_query};
+use crate::execution::all_inactive;
 use crate::job::find_dep_kind_root;
 use crate::query_impl::for_each_query_vtable;
-use crate::{CollectActiveJobsKind, GetQueryVTable, collect_active_query_jobs};
+use crate::{CollectActiveJobsKind, collect_active_query_jobs};
 
 fn depth_limit_error<'tcx>(tcx: TyCtxt<'tcx>, job: QueryJobId) {
     let job_map = collect_active_query_jobs(tcx, CollectActiveJobsKind::Full);
@@ -151,15 +151,15 @@ fn verify_query_key_hashes_inner<'tcx, C: QueryCache>(
     });
 }
 
-/// Implementation of [`DepKindVTable::promote_from_disk_fn`] for queries.
-pub(crate) fn promote_from_disk_inner<'tcx, Q: GetQueryVTable<'tcx>>(
+/// Inner implementation of [`DepKindVTable::promote_from_disk_fn`] for queries.
+pub(crate) fn promote_from_disk_inner<'tcx, C: QueryCache>(
     tcx: TyCtxt<'tcx>,
+    query: &'tcx QueryVTable<'tcx, C>,
     dep_node: DepNode,
 ) {
-    let query = Q::query_vtable(tcx);
     debug_assert!(tcx.dep_graph.is_green(&dep_node));
 
-    let key = <Q::Cache as QueryCache>::Key::try_recover_key(tcx, &dep_node).unwrap_or_else(|| {
+    let key = C::Key::try_recover_key(tcx, &dep_node).unwrap_or_else(|| {
         panic!(
             "Failed to recover key for {dep_node:?} with key fingerprint {}",
             dep_node.key_fingerprint
@@ -220,38 +220,3 @@ where
 
     value
 }
-
-/// Implementation of [`DepKindVTable::force_from_dep_node_fn`] for queries.
-pub(crate) fn force_from_dep_node_inner<'tcx, Q: GetQueryVTable<'tcx>>(
-    tcx: TyCtxt<'tcx>,
-    dep_node: DepNode,
-    // Needed by the vtable function signature, but not used when forcing queries.
-    _prev_index: SerializedDepNodeIndex,
-) -> bool {
-    let query = Q::query_vtable(tcx);
-
-    // We must avoid ever having to call `force_from_dep_node()` for a
-    // `DepNode::codegen_unit`:
-    // Since we cannot reconstruct the query key of a `DepNode::codegen_unit`, we
-    // would always end up having to evaluate the first caller of the
-    // `codegen_unit` query that *is* reconstructible. This might very well be
-    // the `compile_codegen_unit` query, thus re-codegenning the whole CGU just
-    // to re-trigger calling the `codegen_unit` query with the right key. At
-    // that point we would already have re-done all the work we are trying to
-    // avoid doing in the first place.
-    // The solution is simple: Just explicitly call the `codegen_unit` query for
-    // each CGU, right after partitioning. This way `try_mark_green` will always
-    // hit the cache instead of having to go through `force_from_dep_node`.
-    // This assertion makes sure, we actually keep applying the solution above.
-    debug_assert!(
-        dep_node.kind != DepKind::codegen_unit,
-        "calling force_from_dep_node() on dep_kinds::codegen_unit"
-    );
-
-    if let Some(key) = <Q::Cache as QueryCache>::Key::try_recover_key(tcx, &dep_node) {
-        force_query(query, tcx, key, dep_node);
-        true
-    } else {
-        false
-    }
-}