Allow `ensure` queries to return `Result<(), ErrorGuaranteed>`

5 files changed, 75 insertions, 9 deletions

diff --git a/compiler/rustc_hir_analysis/src/check/wfcheck.rs b/compiler/rustc_hir_analysis/src/check/wfcheck.rs
index ae12e7a07fe..3f31ce7aa58 100644
--- a/compiler/rustc_hir_analysis/src/check/wfcheck.rs
+++ b/compiler/rustc_hir_analysis/src/check/wfcheck.rs
@@ -1899,10 +1899,10 @@ impl<'tcx> WfCheckingCtxt<'_, 'tcx> {
 
 fn check_mod_type_wf(tcx: TyCtxt<'_>, module: LocalModDefId) -> Result<(), ErrorGuaranteed> {
     let items = tcx.hir_module_items(module);
-    let mut res = items.par_items(|item| tcx.check_well_formed(item.owner_id));
-    res = res.and(items.par_impl_items(|item| tcx.check_well_formed(item.owner_id)));
-    res = res.and(items.par_trait_items(|item| tcx.check_well_formed(item.owner_id)));
-    res.and(items.par_foreign_items(|item| tcx.check_well_formed(item.owner_id)))
+    let mut res = items.par_items(|item| tcx.ensure().check_well_formed(item.owner_id));
+    res = res.and(items.par_impl_items(|item| tcx.ensure().check_well_formed(item.owner_id)));
+    res = res.and(items.par_trait_items(|item| tcx.ensure().check_well_formed(item.owner_id)));
+    res.and(items.par_foreign_items(|item| tcx.ensure().check_well_formed(item.owner_id)))
 }
 
 fn error_392(
diff --git a/compiler/rustc_hir_analysis/src/lib.rs b/compiler/rustc_hir_analysis/src/lib.rs
index b66508df610..88f3db03a4e 100644
--- a/compiler/rustc_hir_analysis/src/lib.rs
+++ b/compiler/rustc_hir_analysis/src/lib.rs
@@ -206,7 +206,7 @@ pub fn check_crate(tcx: TyCtxt<'_>) -> Result<(), ErrorGuaranteed> {
     }
 
     tcx.sess.time("wf_checking", || {
-        tcx.hir().try_par_for_each_module(|module| tcx.check_mod_type_wf(module))
+        tcx.hir().try_par_for_each_module(|module| tcx.ensure().check_mod_type_wf(module))
     })?;
 
     // NOTE: This is copy/pasted in librustdoc/core.rs and should be kept in sync.
diff --git a/compiler/rustc_macros/src/query.rs b/compiler/rustc_macros/src/query.rs
index d0d41c614d6..d8a695b131b 100644
--- a/compiler/rustc_macros/src/query.rs
+++ b/compiler/rustc_macros/src/query.rs
@@ -114,6 +114,11 @@ struct QueryModifiers {
 
     /// Generate a `feed` method to set the query's value from another query.
     feedable: Option<Ident>,
+
+    /// Forward the result on ensure if the query gets recomputed, and
+    /// return `Ok(())` otherwise. Only applicable to queries returning
+    /// `Result<(), ErrorGuaranteed>`
+    ensure_forwards_result_if_red: Option<Ident>,
 }
 
 fn parse_query_modifiers(input: ParseStream<'_>) -> Result<QueryModifiers> {
@@ -128,6 +133,7 @@ fn parse_query_modifiers(input: ParseStream<'_>) -> Result<QueryModifiers> {
     let mut depth_limit = None;
     let mut separate_provide_extern = None;
     let mut feedable = None;
+    let mut ensure_forwards_result_if_red = None;
 
     while !input.is_empty() {
         let modifier: Ident = input.parse()?;
@@ -187,6 +193,8 @@ fn parse_query_modifiers(input: ParseStream<'_>) -> Result<QueryModifiers> {
             try_insert!(separate_provide_extern = modifier);
         } else if modifier == "feedable" {
             try_insert!(feedable = modifier);
+        } else if modifier == "ensure_forwards_result_if_red" {
+            try_insert!(ensure_forwards_result_if_red = modifier);
         } else {
             return Err(Error::new(modifier.span(), "unknown query modifier"));
         }
@@ -206,6 +214,7 @@ fn parse_query_modifiers(input: ParseStream<'_>) -> Result<QueryModifiers> {
         depth_limit,
         separate_provide_extern,
         feedable,
+        ensure_forwards_result_if_red,
     })
 }
 
@@ -325,6 +334,7 @@ pub fn rustc_queries(input: TokenStream) -> TokenStream {
             eval_always,
             depth_limit,
             separate_provide_extern,
+            ensure_forwards_result_if_red,
         );
 
         if modifiers.cache.is_some() {
diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
index c78c4b74f3a..dffce8a6ec2 100644
--- a/compiler/rustc_middle/src/query/mod.rs
+++ b/compiler/rustc_middle/src/query/mod.rs
@@ -25,7 +25,9 @@ use crate::mir::interpret::{
 use crate::mir::interpret::{LitToConstError, LitToConstInput};
 use crate::mir::mono::CodegenUnit;
 use crate::query::erase::{erase, restore, Erase};
-use crate::query::plumbing::{query_ensure, query_get_at, CyclePlaceholder, DynamicQuery};
+use crate::query::plumbing::{
+    query_ensure, query_ensure_error_guaranteed, query_get_at, CyclePlaceholder, DynamicQuery,
+};
 use crate::thir;
 use crate::traits::query::{
     CanonicalPredicateGoal, CanonicalProjectionGoal, CanonicalTyGoal,
@@ -967,6 +969,7 @@ rustc_queries! {
 
     query check_mod_type_wf(key: LocalModDefId) -> Result<(), ErrorGuaranteed> {
         desc { |tcx| "checking that types are well-formed in {}", describe_as_module(key, tcx) }
+        ensure_forwards_result_if_red
     }
 
     query collect_mod_item_types(key: LocalModDefId) -> () {
@@ -1501,6 +1504,7 @@ rustc_queries! {
 
     query check_well_formed(key: hir::OwnerId) -> Result<(), ErrorGuaranteed> {
         desc { |tcx| "checking that `{}` is well-formed", tcx.def_path_str(key) }
+        ensure_forwards_result_if_red
     }
 
     // The `DefId`s of all non-generic functions and statics in the given crate
diff --git a/compiler/rustc_middle/src/query/plumbing.rs b/compiler/rustc_middle/src/query/plumbing.rs
index 34e5b02ba5b..f4a8ada8f68 100644
--- a/compiler/rustc_middle/src/query/plumbing.rs
+++ b/compiler/rustc_middle/src/query/plumbing.rs
@@ -173,6 +173,45 @@ pub fn query_ensure<'tcx, Cache>(
     }
 }
 
+#[inline]
+pub fn query_ensure_error_guaranteed<'tcx, Cache>(
+    tcx: TyCtxt<'tcx>,
+    execute_query: fn(TyCtxt<'tcx>, Span, Cache::Key, QueryMode) -> Option<Cache::Value>,
+    query_cache: &Cache,
+    key: Cache::Key,
+    check_cache: bool,
+) -> Result<(), ErrorGuaranteed>
+where
+    Cache: QueryCache<Value = super::erase::Erase<Result<(), ErrorGuaranteed>>>,
+{
+    let key = key.into_query_param();
+    if let Some(res) = try_get_cached(tcx, query_cache, &key) {
+        super::erase::restore(res)
+    } else {
+        execute_query(tcx, DUMMY_SP, key, QueryMode::Ensure { check_cache })
+            .map(super::erase::restore)
+            // Either we actually executed the query, which means we got a full `Result`,
+            // or we can just assume the query succeeded, because it was green in the
+            // incremental cache. If it is green, that means that the previous compilation
+            // that wrote to the incremental cache compiles successfully. That is only
+            // possible if the cache entry was `Ok(())`, so we emit that here, without
+            // actually encoding the `Result` in the cache or loading it from there.
+            .unwrap_or(Ok(()))
+    }
+}
+
+macro_rules! query_ensure {
+    ([]$($args:tt)*) => {
+        query_ensure($($args)*)
+    };
+    ([(ensure_forwards_result_if_red) $($rest:tt)*]$($args:tt)*) => {
+        query_ensure_error_guaranteed($($args)*)
+    };
+    ([$other:tt $($modifiers:tt)*]$($args:tt)*) => {
+        query_ensure!([$($modifiers)*]$($args)*)
+    };
+}
+
 macro_rules! query_helper_param_ty {
     (DefId) => { impl IntoQueryParam<DefId> };
     (LocalDefId) => { impl IntoQueryParam<LocalDefId> };
@@ -220,6 +259,18 @@ macro_rules! separate_provide_extern_decl {
     };
 }
 
+macro_rules! ensure_result {
+    ([][$ty:ty]) => {
+        ()
+    };
+    ([(ensure_forwards_result_if_red) $($rest:tt)*][$ty:ty]) => {
+        Result<(), ErrorGuaranteed>
+    };
+    ([$other:tt $($modifiers:tt)*][$($args:tt)*]) => {
+        ensure_result!([$($modifiers)*][$($args)*])
+    };
+}
+
 macro_rules! separate_provide_extern_default {
     ([][$name:ident]) => {
         ()
@@ -343,14 +394,15 @@ macro_rules! define_callbacks {
         impl<'tcx> TyCtxtEnsure<'tcx> {
             $($(#[$attr])*
             #[inline(always)]
-            pub fn $name(self, key: query_helper_param_ty!($($K)*)) {
-                query_ensure(
+            pub fn $name(self, key: query_helper_param_ty!($($K)*)) -> ensure_result!([$($modifiers)*][$V]) {
+                query_ensure!(
+                    [$($modifiers)*]
                     self.tcx,
                     self.tcx.query_system.fns.engine.$name,
                     &self.tcx.query_system.caches.$name,
                     key.into_query_param(),
                     false,
-                );
+                )
             })*
         }