Make the check for cache opt-in.

4 files changed, 57 insertions, 13 deletions

diff --git a/compiler/rustc_middle/src/ty/query.rs b/compiler/rustc_middle/src/ty/query.rs
index 963cb29ffb3..05219efe5f5 100644
--- a/compiler/rustc_middle/src/ty/query.rs
+++ b/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 },
+                    ),
                 };
             })*
         }
diff --git a/compiler/rustc_mir_build/src/build/mod.rs b/compiler/rustc_mir_build/src/build/mod.rs
index d51b031d366..70d5fc2d958 100644
--- a/compiler/rustc_mir_build/src/build/mod.rs
+++ b/compiler/rustc_mir_build/src/build/mod.rs
@@ -50,12 +50,12 @@ fn mir_build(tcx: TyCtxt<'_>, def: ty::WithOptConstParam<LocalDefId>) -> Body<'_
     // Ensure unsafeck and abstract const building is ran before we steal the THIR.
     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().thir_abstract_const_of_const_arg((did, const_param_did));
+            tcx.ensure_with_value().thir_check_unsafety_for_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().thir_abstract_const(did);
+            tcx.ensure_with_value().thir_check_unsafety(did);
+            tcx.ensure_with_value().thir_abstract_const(did);
         }
     }
 
diff --git a/compiler/rustc_mir_transform/src/lib.rs b/compiler/rustc_mir_transform/src/lib.rs
index 5fd923190ef..18776926328 100644
--- a/compiler/rustc_mir_transform/src/lib.rs
+++ b/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);
 
@@ -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),
     }
diff --git a/compiler/rustc_query_system/src/query/plumbing.rs b/compiler/rustc_query_system/src/query/plumbing.rs
index 1229a5fe5e8..04ab2b767e2 100644
--- a/compiler/rustc_query_system/src/query/plumbing.rs
+++ b/compiler/rustc_query_system/src/query/plumbing.rs
@@ -711,6 +711,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>,
@@ -743,6 +744,11 @@ where
         }
     };
 
+    // 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))
 }
@@ -750,7 +756,7 @@ where
 #[derive(Debug)]
 pub enum QueryMode {
     Get,
-    Ensure,
+    Ensure { check_cache: bool },
 }
 
 #[inline(always)]
@@ -765,8 +771,8 @@ where
     Q: QueryConfig<Qcx>,
     Qcx: QueryContext,
 {
-    let dep_node = if let QueryMode::Ensure = mode {
-        let (must_run, dep_node) = ensure_must_run(query, qcx, &key);
+    let dep_node = if let QueryMode::Ensure { check_cache } = mode {
+        let (must_run, dep_node) = ensure_must_run(query, qcx, &key, check_cache);
         if !must_run {
             return None;
         }