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

Author: Zalathar

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)
+    }
+}

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

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>)
                 ()
             ),)*
         }