From 4c448d5163628b0aa5ff4c3bbdb8c6108c1119e9 Mon Sep 17 00:00:00 2001 From: Zalathar Date: Thu, 23 Jan 2025 00:03:47 +1100 Subject: [PATCH] Allow `arena_cache` queries to return `Option<&'tcx T>` --- .../rustc_middle/src/query/arena_cached.rs | 47 +++++++++++++++++++ compiler/rustc_middle/src/query/mod.rs | 6 +-- compiler/rustc_middle/src/query/plumbing.rs | 20 +++++--- 3 files changed, 64 insertions(+), 9 deletions(-) create mode 100644 compiler/rustc_middle/src/query/arena_cached.rs 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 0000000000000..ec6e466ff688c --- /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`, 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; + /// The provide value is `Option`, 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 65e93c3a1cc75..05ded71dbeb51 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> { + 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> { + ) -> 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 2cb6f6d8c6e87..1c157f33a813b 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 `::Target` instead. + /// but `arena_cache` will use `::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> { erase(query_if_arena!([$($modifiers)*] { - if mem::needs_drop::>() { - &*_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>) () ),)* }