Clean up query macros for cache_on_disk_if

This commit is contained in:
Zalathar 2026-02-01 17:37:22 +11:00
parent 9f4b56a5ae
commit 892665ba11
2 changed files with 81 additions and 62 deletions

View file

@ -90,7 +90,7 @@ struct QueryModifiers {
arena_cache: Option<Ident>,
/// Cache the query to disk if the `Block` returns true.
cache: Option<(Option<Pat>, Block)>,
cache_on_disk_if: Option<(Option<Pat>, Block)>,
/// A cycle error for this query aborting the compilation with a fatal error.
cycle_fatal: Option<Ident>,
@ -134,7 +134,7 @@ struct QueryModifiers {
fn parse_query_modifiers(input: ParseStream<'_>) -> Result<QueryModifiers> {
let mut arena_cache = None;
let mut cache = None;
let mut cache_on_disk_if = None;
let mut desc = None;
let mut cycle_fatal = None;
let mut cycle_delay_bug = None;
@ -175,8 +175,11 @@ fn parse_query_modifiers(input: ParseStream<'_>) -> Result<QueryModifiers> {
let list = attr_content.parse_terminated(Expr::parse, Token![,])?;
try_insert!(desc = (tcx, list));
} else if modifier == "cache_on_disk_if" {
// Parse a cache modifier like:
// `cache(tcx) { |tcx| key.is_local() }`
// Parse a cache-on-disk modifier like:
//
// `cache_on_disk_if { true }`
// `cache_on_disk_if { key.is_local() }`
// `cache_on_disk_if(tcx) { tcx.is_typeck_child(key.to_def_id()) }`
let args = if input.peek(token::Paren) {
let args;
parenthesized!(args in input);
@ -186,7 +189,7 @@ fn parse_query_modifiers(input: ParseStream<'_>) -> Result<QueryModifiers> {
None
};
let block = input.parse()?;
try_insert!(cache = (args, block));
try_insert!(cache_on_disk_if = (args, block));
} else if modifier == "arena_cache" {
try_insert!(arena_cache = modifier);
} else if modifier == "cycle_fatal" {
@ -218,7 +221,7 @@ fn parse_query_modifiers(input: ParseStream<'_>) -> Result<QueryModifiers> {
};
Ok(QueryModifiers {
arena_cache,
cache,
cache_on_disk_if,
desc,
cycle_fatal,
cycle_delay_bug,
@ -260,12 +263,18 @@ fn doc_comment_from_desc(list: &Punctuated<Expr, token::Comma>) -> Result<Attrib
Ok(parse_quote! { #[doc = #doc_string] })
}
/// Add the impl of QueryDescription for the query to `impls` if one is requested
fn add_query_desc_cached_impl(
query: &Query,
descs: &mut proc_macro2::TokenStream,
cached: &mut proc_macro2::TokenStream,
) {
/// Contains token streams that are used to accumulate per-query helper
/// functions, to be used by the final output of `rustc_queries!`.
///
/// Helper items typically have the same name as the query they relate to,
/// and expect to be interpolated into a dedicated module.
#[derive(Default)]
struct HelperTokenStreams {
description_fns_stream: proc_macro2::TokenStream,
cache_on_disk_if_fns_stream: proc_macro2::TokenStream,
}
fn make_helpers_for_query(query: &Query, streams: &mut HelperTokenStreams) {
let Query { name, key, modifiers, .. } = &query;
// This dead code exists to instruct rust-analyzer about the link between the `rustc_queries`
@ -281,12 +290,12 @@ fn add_query_desc_cached_impl(
};
// Generate a function to check whether we should cache the query to disk, for some key.
if let Some((args, expr)) = modifiers.cache.as_ref() {
if let Some((args, expr)) = modifiers.cache_on_disk_if.as_ref() {
let tcx = args.as_ref().map(|t| quote! { #t }).unwrap_or_else(|| quote! { _ });
// expr is a `Block`, meaning that `{ #expr }` gets expanded
// to `{ { stmts... } }`, which triggers the `unused_braces` lint.
// we're taking `key` by reference, but some rustc types usually prefer being passed by value
cached.extend(quote! {
streams.cache_on_disk_if_fns_stream.extend(quote! {
#[allow(unused_variables, unused_braces, rustc::pass_by_value)]
#[inline]
pub fn #name<'tcx>(#tcx: TyCtxt<'tcx>, #key: &crate::queries::#name::Key<'tcx>) -> bool {
@ -307,7 +316,7 @@ fn add_query_desc_cached_impl(
}
};
descs.extend(quote! {
streams.description_fns_stream.extend(quote! {
#desc
});
}
@ -316,8 +325,7 @@ pub(super) fn rustc_queries(input: TokenStream) -> TokenStream {
let queries = parse_macro_input!(input as List<Query>);
let mut query_stream = quote! {};
let mut query_description_stream = quote! {};
let mut query_cached_stream = quote! {};
let mut helpers = HelperTokenStreams::default();
let mut feedable_queries = quote! {};
let mut errors = quote! {};
@ -363,9 +371,11 @@ pub(super) fn rustc_queries(input: TokenStream) -> TokenStream {
return_result_from_ensure_ok,
);
// Pass on the cache modifier
if modifiers.cache.is_some() {
attributes.push(quote! { (cache) });
// If there was a `cache_on_disk_if` modifier in the real input, pass
// on a synthetic `(cache_on_disk)` modifier that can be inspected by
// macro-rules macros.
if modifiers.cache_on_disk_if.is_some() {
attributes.push(quote! { (cache_on_disk) });
}
// This uses the span of the query definition for the commas,
@ -399,9 +409,11 @@ pub(super) fn rustc_queries(input: TokenStream) -> TokenStream {
});
}
add_query_desc_cached_impl(&query, &mut query_description_stream, &mut query_cached_stream);
make_helpers_for_query(&query, &mut helpers);
}
let HelperTokenStreams { description_fns_stream, cache_on_disk_if_fns_stream } = helpers;
TokenStream::from(quote! {
/// Higher-order macro that invokes the specified macro with a prepared
/// list of all query signatures (including modifiers).
@ -431,12 +443,17 @@ pub(super) fn rustc_queries(input: TokenStream) -> TokenStream {
}
pub mod descs {
use super::*;
#query_description_stream
#description_fns_stream
}
pub mod cached {
// FIXME(Zalathar): Instead of declaring these functions directly, can
// we put them in a macro and then expand that macro downstream in
// `rustc_query_impl`, where the functions are actually used?
pub mod _cache_on_disk_if_fns {
use super::*;
#query_cached_stream
#cache_on_disk_if_fns_stream
}
#errors
})
}

View file

@ -300,15 +300,31 @@ macro_rules! call_provider {
};
}
macro_rules! should_ever_cache_on_disk {
([]$yes:tt $no:tt) => {{
/// Expands to one of two token trees, depending on whether the current query
/// has the `cache_on_disk_if` modifier.
macro_rules! if_cache_on_disk {
([] $yes:tt $no:tt) => {
$no
}};
([(cache) $($rest:tt)*]$yes:tt $no:tt) => {{
};
// The `cache_on_disk_if` modifier generates a synthetic `(cache_on_disk)`,
// modifier, for use by this macro and similar macros.
([(cache_on_disk) $($rest:tt)*] $yes:tt $no:tt) => {
$yes
}};
([$other:tt $($modifiers:tt)*]$yes:tt $no:tt) => {
should_ever_cache_on_disk!([$($modifiers)*]$yes $no)
};
([$other:tt $($modifiers:tt)*] $yes:tt $no:tt) => {
if_cache_on_disk!([$($modifiers)*] $yes $no)
};
}
/// Conditionally expands to some token trees, if the current query has the
/// `cache_on_disk_if` modifier.
macro_rules! item_if_cache_on_disk {
([] $($item:tt)*) => {};
([(cache_on_disk) $($rest:tt)*] $($item:tt)*) => {
$($item)*
};
([$other:tt $($modifiers:tt)*] $($item:tt)*) => {
item_if_cache_on_disk! { [$($modifiers)*] $($item)* }
};
}
@ -544,28 +560,6 @@ where
}
}
macro_rules! item_if_cached {
([] $tokens:tt) => {};
([(cache) $($rest:tt)*] { $($tokens:tt)* }) => {
$($tokens)*
};
([$other:tt $($modifiers:tt)*] $tokens:tt) => {
item_if_cached! { [$($modifiers)*] $tokens }
};
}
macro_rules! expand_if_cached {
([], $tokens:expr) => {{
None
}};
([(cache) $($rest:tt)*], $tokens:expr) => {{
Some($tokens)
}};
([$other:tt $($modifiers:tt)*], $tokens:expr) => {
expand_if_cached!([$($modifiers)*], $tokens)
};
}
// NOTE: `$V` isn't used here, but we still need to match on it so it can be passed to other macros
// invoked by `rustc_with_all_queries`.
macro_rules! define_queries {
@ -660,17 +654,17 @@ macro_rules! define_queries {
cycle_error_handling: cycle_error_handling!([$($modifiers)*]),
query_state: std::mem::offset_of!(QueryStates<'tcx>, $name),
query_cache: std::mem::offset_of!(QueryCaches<'tcx>, $name),
will_cache_on_disk_for_key_fn: should_ever_cache_on_disk!([$($modifiers)*] {
Some(queries::cached::$name)
will_cache_on_disk_for_key_fn: if_cache_on_disk!([$($modifiers)*] {
Some(::rustc_middle::queries::_cache_on_disk_if_fns::$name)
} {
None
}),
execute_query: |tcx, key| erase::erase_val(tcx.$name(key)),
compute_fn: self::compute_fn::__rust_begin_short_backtrace,
try_load_from_disk_fn: should_ever_cache_on_disk!([$($modifiers)*] {
try_load_from_disk_fn: if_cache_on_disk!([$($modifiers)*] {
Some(|tcx, key, prev_index, index| {
// Check the `cache_on_disk_if` condition for this key.
if !queries::cached::$name(tcx, key) {
if !::rustc_middle::queries::_cache_on_disk_if_fns::$name(tcx, key) {
return None;
}
@ -683,9 +677,9 @@ macro_rules! define_queries {
} {
None
}),
is_loadable_from_disk_fn: should_ever_cache_on_disk!([$($modifiers)*] {
is_loadable_from_disk_fn: if_cache_on_disk!([$($modifiers)*] {
Some(|tcx, key, index| -> bool {
::rustc_middle::queries::cached::$name(tcx, key) &&
::rustc_middle::queries::_cache_on_disk_if_fns::$name(tcx, key) &&
$crate::plumbing::loadable_from_disk(tcx, index)
})
} {
@ -780,7 +774,7 @@ macro_rules! define_queries {
)
}
item_if_cached! { [$($modifiers)*] {
item_if_cache_on_disk! { [$($modifiers)*]
pub(crate) fn encode_query_results<'tcx>(
tcx: TyCtxt<'tcx>,
encoder: &mut CacheEncoder<'_, 'tcx>,
@ -793,7 +787,7 @@ macro_rules! define_queries {
query_result_index,
)
}
}}
}
pub(crate) fn query_key_hash_verify<'tcx>(tcx: TyCtxt<'tcx>) {
$crate::plumbing::query_key_hash_verify(
@ -847,7 +841,15 @@ macro_rules! define_queries {
&mut CacheEncoder<'_, 'tcx>,
&mut EncodedDepNodeIndex)
>
] = &[$(expand_if_cached!([$($modifiers)*], query_impl::$name::encode_query_results)),*];
] = &[
$(
if_cache_on_disk!([$($modifiers)*] {
Some(query_impl::$name::encode_query_results)
} {
None
})
),*
];
const QUERY_KEY_HASH_VERIFY: &[
for<'tcx> fn(TyCtxt<'tcx>)