Auto merge of #108820 - cjgillot:ensure-on-disk, r=oli-obk

Ensure value is on the on-disk cache before returning from `ensure()`.

The current logic for `ensure()` a query just checks that the node is green in the dependency graph.
However, a lot of places use `ensure()` to prevent the query from being called later. This is the case before stealing a query result.

If the query is actually green but the value is not available in the on-disk cache, `ensure` would return, but a subsequent call to the full query would run the code, and attempt to read from a stolen value.

This PR conforms the query system to the usage by checking whether the queried value is loadable from disk before returning.

Sadly, I can't manage to craft a proper test...

Should fix all instances of "attempted to read from stolen value".

compiler/rustc_ast_lowering/src/lib.rs

@@ -435,8 +435,9 @@ fn compute_hir_hash(
 
 pub fn lower_to_hir(tcx: TyCtxt<'_>, (): ()) -> hir::Crate<'_> {
     let sess = tcx.sess;
-    tcx.ensure().output_filenames(());
+    // Queries that borrow `resolver_for_lowering`.
-    let _ = tcx.early_lint_checks(()); // Borrows `resolver_for_lowering`.
+    tcx.ensure_with_value().output_filenames(());
+    tcx.ensure_with_value().early_lint_checks(());
     let (mut resolver, krate) = tcx.resolver_for_lowering(()).steal();
 
     let ast_index = index_crate(&resolver.node_id_to_def_id, &krate);

compiler/rustc_metadata/src/rmeta/encoder.rs

@@ -2050,13 +2050,13 @@ fn prefetch_mir(tcx: TyCtxt<'_>) {
         let (encode_const, encode_opt) = should_encode_mir(tcx, def_id);
 
         if encode_const {
-            tcx.ensure().mir_for_ctfe(def_id);
+            tcx.ensure_with_value().mir_for_ctfe(def_id);
         }
         if encode_opt {
-            tcx.ensure().optimized_mir(def_id);
+            tcx.ensure_with_value().optimized_mir(def_id);
         }
         if encode_opt || encode_const {
-            tcx.ensure().promoted_mir(def_id);
+            tcx.ensure_with_value().promoted_mir(def_id);
         }
     })
 }

compiler/rustc_middle/src/ty/query.rs

@@ -99,6 +99,11 @@ pub struct TyCtxtEnsure<'tcx> {
     pub tcx: TyCtxt<'tcx>,
 }
 
+#[derive(Copy, Clone)]
+pub struct TyCtxtEnsureWithValue<'tcx> {
+    pub tcx: TyCtxt<'tcx>,
+}
+
 impl<'tcx> TyCtxt<'tcx> {
     /// Returns a transparent wrapper for `TyCtxt`, which ensures queries
     /// are executed instead of just returning their results.
@@ -107,6 +112,15 @@ impl<'tcx> TyCtxt<'tcx> {
         TyCtxtEnsure { tcx: self }
     }
 
+    /// Returns a transparent wrapper for `TyCtxt`, which ensures queries
+    /// are executed instead of just returning their results.
+    ///
+    /// This version verifies that the computed result exists in the cache before returning.
+    #[inline(always)]
+    pub fn ensure_with_value(self) -> TyCtxtEnsureWithValue<'tcx> {
+        TyCtxtEnsureWithValue { tcx: self }
+    }
+
     /// Returns a transparent wrapper for `TyCtxt` which uses
     /// `span` as the location of queries performed through it.
     #[inline(always)]
@@ -314,7 +328,31 @@ macro_rules! define_callbacks {
 
                 match try_get_cached(self.tcx, &self.tcx.query_system.caches.$name, &key) {
                     Some(_) => return,
-                    None => self.tcx.queries.$name(self.tcx, DUMMY_SP, key, QueryMode::Ensure),
+                    None => self.tcx.queries.$name(
+                        self.tcx,
+                        DUMMY_SP,
+                        key,
+                        QueryMode::Ensure { check_cache: false },
+                    ),
+                };
+            })*
+        }
+
+        impl<'tcx> TyCtxtEnsureWithValue<'tcx> {
+            $($(#[$attr])*
+            #[inline(always)]
+            pub fn $name(self, key: query_helper_param_ty!($($K)*)) {
+                let key = key.into_query_param();
+                opt_remap_env_constness!([$($modifiers)*][key]);
+
+                match try_get_cached(self.tcx, &self.tcx.query_system.caches.$name, &key) {
+                    Some(_) => return,
+                    None => self.tcx.queries.$name(
+                        self.tcx,
+                        DUMMY_SP,
+                        key,
+                        QueryMode::Ensure { check_cache: true },
+                    ),
                 };
             })*
         }

compiler/rustc_mir_build/src/build/mod.rs

@@ -48,17 +48,14 @@ pub(crate) fn mir_built(
 /// Construct the MIR for a given `DefId`.
 fn mir_build(tcx: TyCtxt<'_>, def: ty::WithOptConstParam<LocalDefId>) -> Body<'_> {
     // Ensure unsafeck and abstract const building is ran before we steal the THIR.
-    // We can't use `ensure()` for `thir_abstract_const` as it doesn't compute the query
-    // if inputs are green. This can cause ICEs when calling `thir_abstract_const` after
-    // THIR has been stolen if we haven't computed this query yet.
     match def {
         ty::WithOptConstParam { did, const_param_did: Some(const_param_did) } => {
-            tcx.ensure().thir_check_unsafety_for_const_arg((did, const_param_did));
+            tcx.ensure_with_value().thir_check_unsafety_for_const_arg((did, const_param_did));
-            drop(tcx.thir_abstract_const_of_const_arg((did, const_param_did)));
+            tcx.ensure_with_value().thir_abstract_const_of_const_arg((did, const_param_did));
         }
         ty::WithOptConstParam { did, const_param_did: None } => {
-            tcx.ensure().thir_check_unsafety(did);
+            tcx.ensure_with_value().thir_check_unsafety(did);
-            drop(tcx.thir_abstract_const(did));
+            tcx.ensure_with_value().thir_abstract_const(did);
         }
     }
 

compiler/rustc_mir_transform/src/lib.rs

@@ -278,14 +278,14 @@ fn mir_const(tcx: TyCtxt<'_>, def: ty::WithOptConstParam<LocalDefId>) -> &Steal<
     // Unsafety check uses the raw mir, so make sure it is run.
     if !tcx.sess.opts.unstable_opts.thir_unsafeck {
         if let Some(param_did) = def.const_param_did {
-            tcx.ensure().unsafety_check_result_for_const_arg((def.did, param_did));
+            tcx.ensure_with_value().unsafety_check_result_for_const_arg((def.did, param_did));
         } else {
-            tcx.ensure().unsafety_check_result(def.did);
+            tcx.ensure_with_value().unsafety_check_result(def.did);
         }
     }
 
     // has_ffi_unwind_calls query uses the raw mir, so make sure it is run.
-    tcx.ensure().has_ffi_unwind_calls(def.did);
+    tcx.ensure_with_value().has_ffi_unwind_calls(def.did);
 
     let mut body = tcx.mir_built(def).steal();
 
@@ -433,7 +433,7 @@ fn mir_drops_elaborated_and_const_checked(
     if tcx.sess.opts.unstable_opts.drop_tracking_mir
         && let DefKind::Generator = tcx.def_kind(def.did)
     {
-        tcx.ensure().mir_generator_witnesses(def.did);
+        tcx.ensure_with_value().mir_generator_witnesses(def.did);
     }
     let mir_borrowck = tcx.mir_borrowck_opt_const_arg(def);
 
@@ -444,7 +444,7 @@ fn mir_drops_elaborated_and_const_checked(
 
         // Do not compute the mir call graph without said call graph actually being used.
         if inline::Inline.is_enabled(&tcx.sess) {
-            let _ = tcx.mir_inliner_callees(ty::InstanceDef::Item(def));
+            tcx.ensure_with_value().mir_inliner_callees(ty::InstanceDef::Item(def));
         }
     }
 
@@ -613,7 +613,7 @@ fn inner_optimized_mir(tcx: TyCtxt<'_>, did: LocalDefId) -> Body<'_> {
         // Run the `mir_for_ctfe` query, which depends on `mir_drops_elaborated_and_const_checked`
         // which we are going to steal below. Thus we need to run `mir_for_ctfe` first, so it
         // computes and caches its result.
-        Some(hir::ConstContext::ConstFn) => tcx.ensure().mir_for_ctfe(did),
+        Some(hir::ConstContext::ConstFn) => tcx.ensure_with_value().mir_for_ctfe(did),
         None => {}
         Some(other) => panic!("do not use `optimized_mir` for constants: {:?}", other),
     }

compiler/rustc_query_impl/src/lib.rs

@@ -31,6 +31,7 @@ use rustc_span::Span;
 #[macro_use]
 mod plumbing;
 pub use plumbing::QueryCtxt;
+use rustc_query_system::dep_graph::SerializedDepNodeIndex;
 use rustc_query_system::query::*;
 #[cfg(parallel_compiler)]
 pub use rustc_query_system::query::{deadlock, QueryContext};

compiler/rustc_query_impl/src/on_disk_cache.rs

@@ -388,6 +388,12 @@ impl<'sess> OnDiskCache<'sess> {
         debug_assert!(prev.is_none());
     }
 
+    /// Return whether the cached query result can be decoded.
+    pub fn loadable_from_disk(&self, dep_node_index: SerializedDepNodeIndex) -> bool {
+        self.query_result_index.contains_key(&dep_node_index)
+        // with_decoder is infallible, so we can stop here
+    }
+
     /// Returns the cached query result if there is something in the cache for
     /// the given `SerializedDepNodeIndex`; otherwise returns `None`.
     pub fn try_load_query_result<'tcx, T>(
@@ -398,7 +404,9 @@ impl<'sess> OnDiskCache<'sess> {
     where
         T: for<'a> Decodable<CacheDecoder<'a, 'tcx>>,
     {
-        self.load_indexed(tcx, dep_node_index, &self.query_result_index)
+        let opt_value = self.load_indexed(tcx, dep_node_index, &self.query_result_index);
+        debug_assert_eq!(opt_value.is_some(), self.loadable_from_disk(dep_node_index));
+        opt_value
     }
 
     /// Stores side effect emitted during computation of an anonymous query.
@@ -428,8 +436,8 @@ impl<'sess> OnDiskCache<'sess> {
         T: for<'a> Decodable<CacheDecoder<'a, 'tcx>>,
     {
         let pos = index.get(&dep_node_index).cloned()?;
-
+        let value = self.with_decoder(tcx, pos, |decoder| decode_tagged(decoder, dep_node_index));
-        self.with_decoder(tcx, pos, |decoder| Some(decode_tagged(decoder, dep_node_index)))
+        Some(value)
     }
 
     fn with_decoder<'a, 'tcx, T, F: for<'s> FnOnce(&mut CacheDecoder<'s, 'tcx>) -> T>(

compiler/rustc_query_impl/src/plumbing.rs

@@ -364,6 +364,14 @@ where
     }
 }
 
+pub(crate) fn loadable_from_disk<'tcx>(tcx: QueryCtxt<'tcx>, id: SerializedDepNodeIndex) -> bool {
+    if let Some(cache) = tcx.on_disk_cache().as_ref() {
+        cache.loadable_from_disk(id)
+    } else {
+        false
+    }
+}
+
 pub(crate) fn try_load_from_disk<'tcx, V>(
     tcx: QueryCtxt<'tcx>,
     id: SerializedDepNodeIndex,
@@ -535,6 +543,21 @@ macro_rules! define_queries {
                 })
             }
 
+            #[inline]
+            fn loadable_from_disk(
+                self,
+                _qcx: QueryCtxt<'tcx>,
+                _key: &Self::Key,
+                _index: SerializedDepNodeIndex,
+            ) -> bool {
+                should_ever_cache_on_disk!([$($modifiers)*] {
+                    self.cache_on_disk(_qcx.tcx, _key) &&
+                        $crate::plumbing::loadable_from_disk(_qcx, _index)
+                } {
+                    false
+                })
+            }
+
             #[inline(always)]
             fn anon(self) -> bool {
                 is_anon!([$($modifiers)*])

compiler/rustc_query_system/src/query/config.rs

@@ -43,6 +43,8 @@ pub trait QueryConfig<Qcx: QueryContext>: Copy {
 
     fn try_load_from_disk(self, qcx: Qcx, idx: &Self::Key) -> TryLoadFromDisk<Qcx, Self::Value>;
 
+    fn loadable_from_disk(self, qcx: Qcx, key: &Self::Key, idx: SerializedDepNodeIndex) -> bool;
+
     fn anon(self) -> bool;
     fn eval_always(self) -> bool;
     fn depth_limit(self) -> bool;

compiler/rustc_query_system/src/query/plumbing.rs

@@ -564,10 +564,17 @@ where
         // can be forced from `DepNode`.
         debug_assert!(
             !qcx.dep_context().fingerprint_style(dep_node.kind).reconstructible(),
-            "missing on-disk cache entry for {dep_node:?}"
+            "missing on-disk cache entry for reconstructible {dep_node:?}"
         );
     }
 
+    // Sanity check for the logic in `ensure`: if the node is green and the result loadable,
+    // we should actually be able to load it.
+    debug_assert!(
+        !query.loadable_from_disk(qcx, &key, prev_dep_node_index),
+        "missing on-disk cache entry for loadable {dep_node:?}"
+    );
+
     // We could not load a result from the on-disk cache, so
     // recompute.
     let prof_timer = qcx.dep_context().profiler().query_provider();
@@ -718,6 +725,7 @@ fn ensure_must_run<Q, Qcx>(
     query: Q,
     qcx: Qcx,
     key: &Q::Key,
+    check_cache: bool,
 ) -> (bool, Option<DepNode<Qcx::DepKind>>)
 where
     Q: QueryConfig<Qcx>,
@@ -733,7 +741,7 @@ where
     let dep_node = query.construct_dep_node(*qcx.dep_context(), key);
 
     let dep_graph = qcx.dep_context().dep_graph();
-    match dep_graph.try_mark_green(qcx, &dep_node) {
+    let serialized_dep_node_index = match dep_graph.try_mark_green(qcx, &dep_node) {
         None => {
             // A None return from `try_mark_green` means that this is either
             // a new dep node or that the dep node has already been marked red.
@@ -741,20 +749,28 @@ where
             // DepNodeIndex. We must invoke the query itself. The performance cost
             // this introduces should be negligible as we'll immediately hit the
             // in-memory cache, or another query down the line will.
-            (true, Some(dep_node))
+            return (true, Some(dep_node));
         }
-        Some((_, dep_node_index)) => {
+        Some((serialized_dep_node_index, dep_node_index)) => {
             dep_graph.read_index(dep_node_index);
             qcx.dep_context().profiler().query_cache_hit(dep_node_index.into());
-            (false, None)
+            serialized_dep_node_index
         }
+    };
+
+    // We do not need the value at all, so do not check the cache.
+    if !check_cache {
+        return (false, None);
     }
+
+    let loadable = query.loadable_from_disk(qcx, key, serialized_dep_node_index);
+    (!loadable, Some(dep_node))
 }
 
 #[derive(Debug)]
 pub enum QueryMode {
     Get,
-    Ensure,
+    Ensure { check_cache: bool },
 }
 
 #[inline(always)]
@@ -769,8 +785,8 @@ where
     Q: QueryConfig<Qcx>,
     Qcx: QueryContext,
 {
-    let dep_node = if let QueryMode::Ensure = mode {
+    let dep_node = if let QueryMode::Ensure { check_cache } = mode {
-        let (must_run, dep_node) = ensure_must_run(query, qcx, &key);
+        let (must_run, dep_node) = ensure_must_run(query, qcx, &key, check_cache);
         if !must_run {
             return None;
         }