Implement incremental caching for derive macro expansions

This commit is contained in:
Felix Rath 2025-08-13 13:05:07 +02:00 committed by Jakub Beránek
parent a6acf0f07f
commit 8fa2f693bb
No known key found for this signature in database
GPG key ID: 909CD0D26483516B
21 changed files with 284 additions and 72 deletions

View file

@ -21,6 +21,7 @@ rustc_hir = { path = "../rustc_hir" }
rustc_lexer = { path = "../rustc_lexer" }
rustc_lint_defs = { path = "../rustc_lint_defs" }
rustc_macros = { path = "../rustc_macros" }
rustc_middle = { path = "../rustc_middle" }
rustc_parse = { path = "../rustc_parse" }
# We must use the proc_macro version that we will compile proc-macros against,
# not the one from our own sysroot.
@ -28,6 +29,7 @@ rustc_proc_macro = { path = "../rustc_proc_macro" }
rustc_serialize = { path = "../rustc_serialize" }
rustc_session = { path = "../rustc_session" }
rustc_span = { path = "../rustc_span" }
scoped-tls = "1.0"
smallvec = { version = "1.8.1", features = ["union", "may_dangle"] }
thin-vec = "0.2.12"
tracing = "0.1"

View file

@ -29,4 +29,8 @@ pub mod module;
#[allow(rustc::untranslatable_diagnostic)]
pub mod proc_macro;
pub fn provide(providers: &mut rustc_middle::query::Providers) {
providers.derive_macro_expansion = proc_macro::provide_derive_macro_expansion;
}
rustc_fluent_macro::fluent_messages! { "../messages.ftl" }

View file

@ -1,9 +1,11 @@
use rustc_ast::tokenstream::TokenStream;
use rustc_errors::ErrorGuaranteed;
use rustc_middle::ty::{self, TyCtxt};
use rustc_parse::parser::{ForceCollect, Parser};
use rustc_session::Session;
use rustc_session::config::ProcMacroExecutionStrategy;
use rustc_span::Span;
use rustc_span::profiling::SpannedEventArgRecorder;
use rustc_span::{LocalExpnId, Span};
use {rustc_ast as ast, rustc_proc_macro as pm};
use crate::base::{self, *};
@ -30,9 +32,9 @@ impl<T> pm::bridge::server::MessagePipe<T> for MessagePipe<T> {
}
}
fn exec_strategy(ecx: &ExtCtxt<'_>) -> impl pm::bridge::server::ExecutionStrategy + 'static {
fn exec_strategy(sess: &Session) -> impl pm::bridge::server::ExecutionStrategy + 'static {
pm::bridge::server::MaybeCrossThread::<MessagePipe<_>>::new(
ecx.sess.opts.unstable_opts.proc_macro_execution_strategy
sess.opts.unstable_opts.proc_macro_execution_strategy
== ProcMacroExecutionStrategy::CrossThread,
)
}
@ -54,7 +56,7 @@ impl base::BangProcMacro for BangProcMacro {
});
let proc_macro_backtrace = ecx.ecfg.proc_macro_backtrace;
let strategy = exec_strategy(ecx);
let strategy = exec_strategy(ecx.sess);
let server = proc_macro_server::Rustc::new(ecx);
self.client.run(&strategy, server, input, proc_macro_backtrace).map_err(|e| {
ecx.dcx().emit_err(errors::ProcMacroPanicked {
@ -85,7 +87,7 @@ impl base::AttrProcMacro for AttrProcMacro {
});
let proc_macro_backtrace = ecx.ecfg.proc_macro_backtrace;
let strategy = exec_strategy(ecx);
let strategy = exec_strategy(ecx.sess);
let server = proc_macro_server::Rustc::new(ecx);
self.client.run(&strategy, server, annotation, annotated, proc_macro_backtrace).map_err(
|e| {
@ -101,7 +103,7 @@ impl base::AttrProcMacro for AttrProcMacro {
}
pub struct DeriveProcMacro {
pub client: pm::bridge::client::Client<pm::TokenStream, pm::TokenStream>,
pub client: DeriveClient,
}
impl MultiItemModifier for DeriveProcMacro {
@ -113,6 +115,13 @@ impl MultiItemModifier for DeriveProcMacro {
item: Annotatable,
_is_derive_const: bool,
) -> ExpandResult<Vec<Annotatable>, Annotatable> {
let _timer = ecx.sess.prof.generic_activity_with_arg_recorder(
"expand_derive_proc_macro_outer",
|recorder| {
recorder.record_arg_with_span(ecx.sess.source_map(), ecx.expansion_descr(), span);
},
);
// We need special handling for statement items
// (e.g. `fn foo() { #[derive(Debug)] struct Bar; }`)
let is_stmt = matches!(item, Annotatable::Stmt(..));
@ -123,36 +132,31 @@ impl MultiItemModifier for DeriveProcMacro {
// altogether. See #73345.
crate::base::ann_pretty_printing_compatibility_hack(&item, &ecx.sess.psess);
let input = item.to_tokens();
let stream = {
let _timer =
ecx.sess.prof.generic_activity_with_arg_recorder("expand_proc_macro", |recorder| {
recorder.record_arg_with_span(
ecx.sess.source_map(),
ecx.expansion_descr(),
span,
);
});
let proc_macro_backtrace = ecx.ecfg.proc_macro_backtrace;
let strategy = exec_strategy(ecx);
let server = proc_macro_server::Rustc::new(ecx);
match self.client.run(&strategy, server, input, proc_macro_backtrace) {
Ok(stream) => stream,
Err(e) => {
ecx.dcx().emit_err({
errors::ProcMacroDerivePanicked {
span,
message: e.as_str().map(|message| {
errors::ProcMacroDerivePanickedHelp { message: message.into() }
}),
}
});
return ExpandResult::Ready(vec![]);
}
}
let invoc_id = ecx.current_expansion.id;
let res = if ecx.sess.opts.incremental.is_some()
&& ecx.sess.opts.unstable_opts.cache_proc_macros
{
ty::tls::with(|tcx| {
let input = &*tcx.arena.alloc(input);
let key: (LocalExpnId, &TokenStream) = (invoc_id, input);
QueryDeriveExpandCtx::enter(ecx, self.client, move || {
tcx.derive_macro_expansion(key).cloned()
})
})
} else {
expand_derive_macro(invoc_id, input, ecx, self.client)
};
let Ok(output) = res else {
// error will already have been emitted
return ExpandResult::Ready(vec![]);
};
let error_count_before = ecx.dcx().err_count();
let mut parser = Parser::new(&ecx.sess.psess, stream, Some("proc-macro derive"));
let mut parser = Parser::new(&ecx.sess.psess, output, Some("proc-macro derive"));
let mut items = vec![];
loop {
@ -180,3 +184,101 @@ impl MultiItemModifier for DeriveProcMacro {
ExpandResult::Ready(items)
}
}
/// Provide a query for computing the output of a derive macro.
pub(super) fn provide_derive_macro_expansion<'tcx>(
tcx: TyCtxt<'tcx>,
key: (LocalExpnId, &'tcx TokenStream),
) -> Result<&'tcx TokenStream, ()> {
let (invoc_id, input) = key;
// Make sure that we invalidate the query when the crate defining the proc macro changes
let _ = tcx.crate_hash(invoc_id.expn_data().macro_def_id.unwrap().krate);
QueryDeriveExpandCtx::with(|ecx, client| {
expand_derive_macro(invoc_id, input.clone(), ecx, client).map(|ts| &*tcx.arena.alloc(ts))
})
}
type DeriveClient = pm::bridge::client::Client<pm::TokenStream, pm::TokenStream>;
fn expand_derive_macro(
invoc_id: LocalExpnId,
input: TokenStream,
ecx: &mut ExtCtxt<'_>,
client: DeriveClient,
) -> Result<TokenStream, ()> {
let _timer =
ecx.sess.prof.generic_activity_with_arg_recorder("expand_proc_macro", |recorder| {
let invoc_expn_data = invoc_id.expn_data();
let span = invoc_expn_data.call_site;
let event_arg = invoc_expn_data.kind.descr();
recorder.record_arg_with_span(ecx.sess.source_map(), event_arg.clone(), span);
});
let proc_macro_backtrace = ecx.ecfg.proc_macro_backtrace;
let strategy = exec_strategy(ecx.sess);
let server = proc_macro_server::Rustc::new(ecx);
match client.run(&strategy, server, input, proc_macro_backtrace) {
Ok(stream) => Ok(stream),
Err(e) => {
let invoc_expn_data = invoc_id.expn_data();
let span = invoc_expn_data.call_site;
ecx.dcx().emit_err({
errors::ProcMacroDerivePanicked {
span,
message: e.as_str().map(|message| errors::ProcMacroDerivePanickedHelp {
message: message.into(),
}),
}
});
Err(())
}
}
}
/// Stores the context necessary to expand a derive proc macro via a query.
struct QueryDeriveExpandCtx {
/// Type-erased version of `&mut ExtCtxt`
expansion_ctx: *mut (),
client: DeriveClient,
}
impl QueryDeriveExpandCtx {
/// Store the extension context and the client into the thread local value.
/// It will be accessible via the `with` method while `f` is active.
fn enter<F, R>(ecx: &mut ExtCtxt<'_>, client: DeriveClient, f: F) -> R
where
F: FnOnce() -> R,
{
// We need erasure to get rid of the lifetime
let ctx = Self { expansion_ctx: ecx as *mut _ as *mut (), client };
DERIVE_EXPAND_CTX.set(&ctx, || f())
}
/// Accesses the thread local value of the derive expansion context.
/// Must be called while the `enter` function is active.
fn with<F, R>(f: F) -> R
where
F: for<'a, 'b> FnOnce(&'b mut ExtCtxt<'a>, DeriveClient) -> R,
{
DERIVE_EXPAND_CTX.with(|ctx| {
let ectx = {
let casted = ctx.expansion_ctx.cast::<ExtCtxt<'_>>();
// SAFETY: We can only get the value from `with` while the `enter` function
// is active (on the callstack), and that function's signature ensures that the
// lifetime is valid.
// If `with` is called at some other time, it will panic due to usage of
// `scoped_tls::with`.
unsafe { casted.as_mut().unwrap() }
};
f(ectx, ctx.client)
})
}
}
// When we invoke a query to expand a derive proc macro, we need to provide it with the expansion
// context and derive Client. We do that using a thread-local.
scoped_tls::scoped_thread_local!(static DERIVE_EXPAND_CTX: QueryDeriveExpandCtx);