Rollup merge of #135911 - Zalathar:arena-cache-option, r=compiler-errors

Allow `arena_cache` queries to return `Option<&'tcx T>`

Currently, `arena_cache` queries always have to return `&'tcx T`[^deref]. This means that if an arena-cached query wants to return an optional value, it has to return `&'tcx Option<T>`, which has a few negative consequences:

- It goes against normal Rust style, where `Option<&T>` is preferred over `&Option<T>`.
- Callers that actually want an `Option<&T>` have to manually call `.as_ref()` on the query result.
- When the query result is `None`, a full-sized `Option<T>` still needs to be stored in the arena.

This PR solves that problem by introducing a helper trait `ArenaCached` that is implemented for both `&T` and `Option<&T>`, and takes care of bridging between the provided type, the arena-allocated type, and the declared query return type.

---

To demonstrate that this works, I have converted the two existing arena-cached queries that currently return `&Option<T>`: `mir_coroutine_witnesses` and `diagnostic_hir_wf_check`. Only the query declarations need to be modified; existing providers and callers continue to work with the new query return type.

(My real goal is to apply this to `coverage_ids_info`, which will return Option as of #135873, but that PR hasn't landed yet.)

[^deref]: Technically they could return other types that implement `Deref`, but it's hard to imagine this working well with anything other than `&T`.

3 files changed, 64 insertions, 9 deletions

diff --git a/compiler/rustc_middle/src/query/arena_cached.rs b/compiler/rustc_middle/src/query/arena_cached.rs
new file mode 100644
index 00000000000..ec6e466ff68
--- /dev/null
+++ b/compiler/rustc_middle/src/query/arena_cached.rs
@@ -0,0 +1,47 @@
+/// Helper trait that allows `arena_cache` queries to return `Option<&T>`
+/// instead of `&Option<T>`, and avoid allocating `None` in the arena.
+///
+/// An arena-cached query must be declared to return a type that implements
+/// this trait, i.e. either `&'tcx T` or `Option<&'tcx T>`. This trait then
+/// determines the types returned by the provider and stored in the arena,
+/// and provides a function to bridge between the three types.
+pub trait ArenaCached<'tcx>: Sized {
+    /// Type that is returned by the query provider.
+    type Provided;
+    /// Type that is stored in the arena.
+    type Allocated: 'tcx;
+
+    /// Takes a provided value, and allocates it in the arena (if appropriate)
+    /// with the help of the given `arena_alloc` closure.
+    fn alloc_in_arena(
+        arena_alloc: impl Fn(Self::Allocated) -> &'tcx Self::Allocated,
+        value: Self::Provided,
+    ) -> Self;
+}
+
+impl<'tcx, T> ArenaCached<'tcx> for &'tcx T {
+    type Provided = T;
+    type Allocated = T;
+
+    fn alloc_in_arena(
+        arena_alloc: impl Fn(Self::Allocated) -> &'tcx Self::Allocated,
+        value: Self::Provided,
+    ) -> Self {
+        // Just allocate in the arena normally.
+        arena_alloc(value)
+    }
+}
+
+impl<'tcx, T> ArenaCached<'tcx> for Option<&'tcx T> {
+    type Provided = Option<T>;
+    /// The provide value is `Option<T>`, but we only store `T` in the arena.
+    type Allocated = T;
+
+    fn alloc_in_arena(
+        arena_alloc: impl Fn(Self::Allocated) -> &'tcx Self::Allocated,
+        value: Self::Provided,
+    ) -> Self {
+        // Don't store None in the arena, and wrap the allocated reference in Some.
+        value.map(arena_alloc)
+    }
+}
diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
index 65e93c3a1cc..05ded71dbeb 100644
--- a/compiler/rustc_middle/src/query/mod.rs
+++ b/compiler/rustc_middle/src/query/mod.rs
@@ -7,7 +7,6 @@
 #![allow(unused_parens)]
 
 use std::mem;
-use std::ops::Deref;
 use std::path::PathBuf;
 use std::sync::Arc;
 
@@ -85,6 +84,7 @@ use crate::ty::{
 };
 use crate::{dep_graph, mir, thir};
 
+mod arena_cached;
 pub mod erase;
 mod keys;
 pub use keys::{AsLocalKey, Key, LocalCrate};
@@ -586,7 +586,7 @@ rustc_queries! {
         separate_provide_extern
     }
 
-    query mir_coroutine_witnesses(key: DefId) -> &'tcx Option<mir::CoroutineLayout<'tcx>> {
+    query mir_coroutine_witnesses(key: DefId) -> Option<&'tcx mir::CoroutineLayout<'tcx>> {
         arena_cache
         desc { |tcx| "coroutine witness types for `{}`", tcx.def_path_str(key) }
         cache_on_disk_if { key.is_local() }
@@ -2415,7 +2415,7 @@ rustc_queries! {
     /// because the `ty::Ty`-based wfcheck is always run.
     query diagnostic_hir_wf_check(
         key: (ty::Predicate<'tcx>, WellFormedLoc)
-    ) -> &'tcx Option<ObligationCause<'tcx>> {
+    ) -> Option<&'tcx ObligationCause<'tcx>> {
         arena_cache
         eval_always
         no_hash
diff --git a/compiler/rustc_middle/src/query/plumbing.rs b/compiler/rustc_middle/src/query/plumbing.rs
index 2cb6f6d8c6e..1c157f33a81 100644
--- a/compiler/rustc_middle/src/query/plumbing.rs
+++ b/compiler/rustc_middle/src/query/plumbing.rs
@@ -289,10 +289,10 @@ macro_rules! define_callbacks {
 
                 /// This type alias specifies the type returned from query providers and the type
                 /// used for decoding. For regular queries this is the declared returned type `V`,
-                /// but `arena_cache` will use `<V as Deref>::Target` instead.
+                /// but `arena_cache` will use `<V as ArenaCached>::Provided` instead.
                 pub type ProvidedValue<'tcx> = query_if_arena!(
                     [$($modifiers)*]
-                    (<$V as Deref>::Target)
+                    (<$V as $crate::query::arena_cached::ArenaCached<'tcx>>::Provided)
                     ($V)
                 );
 
@@ -307,10 +307,18 @@ macro_rules! define_callbacks {
                 ) -> Erase<Value<'tcx>> {
                     erase(query_if_arena!([$($modifiers)*]
                         {
-                            if mem::needs_drop::<ProvidedValue<'tcx>>() {
-                                &*_tcx.query_system.arenas.$name.alloc(value)
+                            use $crate::query::arena_cached::ArenaCached;
+
+                            if mem::needs_drop::<<$V as ArenaCached<'tcx>>::Allocated>() {
+                                <$V as ArenaCached>::alloc_in_arena(
+                                    |v| _tcx.query_system.arenas.$name.alloc(v),
+                                    value,
+                                )
                             } else {
-                                &*_tcx.arena.dropless.alloc(value)
+                                <$V as ArenaCached>::alloc_in_arena(
+                                    |v| _tcx.arena.dropless.alloc(v),
+                                    value,
+                                )
                             }
                         }
                         (value)
@@ -354,7 +362,7 @@ macro_rules! define_callbacks {
 
         pub struct QueryArenas<'tcx> {
             $($(#[$attr])* pub $name: query_if_arena!([$($modifiers)*]
-                (TypedArena<<$V as Deref>::Target>)
+                (TypedArena<<$V as $crate::query::arena_cached::ArenaCached<'tcx>>::Allocated>)
                 ()
             ),)*
         }