Skip to content

Clean up query macros for cache_on_disk_if #151943

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Zalathar wants to merge 1 commit into
base: main
Choose a base branch
from
+81 −62
Open

Clean up query macros for cache_on_disk_if #151943

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
65 changes: 41 additions & 24 deletions compiler/rustc_macros/src/query.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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>,
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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);
Expand All @@ -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" {
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Copy link
Contributor

@petrochenkov petrochenkov Feb 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there multiple functions/descriptions, or one function/description?
There's some naming inconsistency, query_description_stream below says that it's one, but "descs" says that there are multiple.

Copy link
Member Author

@Zalathar Zalathar Feb 3, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Each stream contains multiple descriptions or multiple functions (0-1 of each per query).

The existing name query_description_stream is misleading; I'll update this PR to change it.

}

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`
Expand All @@ -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 {
Expand All @@ -307,7 +316,7 @@ fn add_query_desc_cached_impl(
}
};

descs.extend(quote! {
streams.description_fns_stream.extend(quote! {
#desc
});
}
Expand All @@ -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! {};

Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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).
Expand Down Expand Up @@ -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?
Copy link
Contributor

@Zoxc Zoxc Feb 3, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We kind of want to keep rustc_query_impl small (for compile-time reasons), so having it point to functions (which won't be inlined) in rustc_middle isn't terrible.

Copy link
Member Author

@Zalathar Zalathar Feb 3, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was thinking that we also want to keep rustc_middle small (since it's a bottleneck for most other compiler crates), though I don't have a strong sense of how shuffling things between the two crates affects overall build times (especially with metadata pipelining).

Copy link
Contributor

@Zoxc Zoxc Feb 3, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, but there's other way to reduce rustc_middle (particularly if other crates can also define queries). rustc_query_impl will mostly keep growing as we add queries.

Copy link
Member Author

@Zalathar Zalathar Feb 3, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The primary purpose of rustc_query_impl is to allow rustc_middle to be smaller.

If we wanted to improve query-impl build times by enabling more parallelism, we wouldn't do so by moving code into middle; we'd move it into an intermediate crate between the two.

So I don't think keeping query-impl small is a reason to not move code from middle to query-impl.

pub mod _cache_on_disk_if_fns {
use super::*;
#query_cached_stream
#cache_on_disk_if_fns_stream
}

#errors
})
}
78 changes: 40 additions & 38 deletions compiler/rustc_query_impl/src/plumbing.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)* }
};
}

Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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;
}

Expand All @@ -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)
})
} {
Expand Down Expand Up @@ -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>,
Expand All @@ -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(
Expand Down Expand Up @@ -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>)
Expand Down
Loading