Implement incremental caching for derive macro expansions

futile · Kobzol · commit cc292c843688 · 2025-09-02T19:05:27.000+02:00
diff --git a/Cargo.lock b/Cargo.lock
@@ -3748,6 +3748,7 @@ dependencies = [
  "rustc_lexer",
  "rustc_lint_defs",
  "rustc_macros",
+ "rustc_middle",
  "rustc_parse",
  "rustc_proc_macro",
  "rustc_serialize",
diff --git a/compiler/rustc_ast/src/tokenstream.rs b/compiler/rustc_ast/src/tokenstream.rs
@@ -5,15 +5,17 @@
 //! which are themselves a single [`Token`] or a `Delimited` subsequence of tokens.
 
 use std::borrow::Cow;
+use std::hash::Hash;
 use std::ops::Range;
 use std::sync::Arc;
 use std::{cmp, fmt, iter, mem};
 
 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
 use rustc_data_structures::sync;
 use rustc_macros::{Decodable, Encodable, HashStable_Generic, Walkable};
-use rustc_serialize::{Decodable, Encodable};
-use rustc_span::{DUMMY_SP, Span, SpanDecoder, SpanEncoder, Symbol, sym};
+use rustc_serialize::{Decodable, Encodable, Encoder};
+use rustc_span::def_id::{CrateNum, DefIndex};
+use rustc_span::{ByteSymbol, DUMMY_SP, Span, SpanDecoder, SpanEncoder, Symbol, sym};
 use thin_vec::ThinVec;
 
 use crate::ast::AttrStyle;
diff --git a/compiler/rustc_expand/Cargo.toml b/compiler/rustc_expand/Cargo.toml
@@ -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.
diff --git a/compiler/rustc_expand/src/lib.rs b/compiler/rustc_expand/src/lib.rs
@@ -31,4 +31,8 @@ pub mod module;
 #[allow(rustc::untranslatable_diagnostic)]
 pub mod proc_macro;
 
+pub fn provide(providers: &mut rustc_middle::util::Providers) {
+    providers.derive_macro_expansion = proc_macro::provide_derive_macro_expansion;
+}
+
 rustc_fluent_macro::fluent_messages! { "../messages.ftl" }
diff --git a/compiler/rustc_expand/src/proc_macro.rs b/compiler/rustc_expand/src/proc_macro.rs
@@ -1,9 +1,15 @@
+use std::cell::Cell;
+use std::ptr::NonNull;
+
 use rustc_ast::tokenstream::TokenStream;
+use rustc_data_structures::svh::Svh;
 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 +36,9 @@ impl<T> pm::bridge::server::MessagePipe<T> for MessagePipe<T> {
     }
 }
 
-fn exec_strategy(ecx: &ExtCtxt<'_>) -> impl pm::bridge::server::ExecutionStrategy + 'static {
+pub 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 +60,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 +91,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| {
@@ -113,6 +119,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 +136,39 @@ 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 res = ty::tls::with(|tcx| {
+            let input = tcx.arena.alloc(input) as &TokenStream;
+            let invoc_id = ecx.current_expansion.id;
+            let invoc_expn_data = invoc_id.expn_data();
+
+            assert_eq!(invoc_expn_data.call_site, span);
+
+            // FIXME(pr-time): Is this the correct way to check for incremental compilation (as
+            // well as for `cache_proc_macros`)?
+            if tcx.sess.opts.incremental.is_some() && tcx.sess.opts.unstable_opts.cache_proc_macros
+            {
+                // FIXME(pr-time): Just using the crate hash to notice when the proc-macro code has
+                // changed. How to *correctly* depend on exactly the macro definition?
+                // I.e., depending on the crate hash is just a HACK, and ideally the dependency would be
+                // more narrow.
+                let macro_def_id = invoc_expn_data.macro_def_id.unwrap();
+                let proc_macro_crate_hash = tcx.crate_hash(macro_def_id.krate);
+
+                let key = (invoc_id, proc_macro_crate_hash, input);
+
+                enter_context((ecx, self.client), move || tcx.derive_macro_expansion(key).cloned())
+            } else {
+                expand_derive_macro(tcx, invoc_id, input, ecx, self.client).cloned()
             }
+        });
+
+        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 +196,102 @@ impl MultiItemModifier for DeriveProcMacro {
         ExpandResult::Ready(items)
     }
 }
+
+pub(super) fn provide_derive_macro_expansion<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    key: (LocalExpnId, Svh, &'tcx TokenStream),
+) -> Result<&'tcx TokenStream, ()> {
+    let (invoc_id, _macro_crate_hash, input) = key;
+
+    with_context(|(ecx, client)| expand_derive_macro(tcx, invoc_id, input, ecx, *client))
+}
+
+type CLIENT = pm::bridge::client::Client<pm::TokenStream, pm::TokenStream>;
+
+fn expand_derive_macro<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    invoc_id: LocalExpnId,
+    input: &'tcx TokenStream,
+    ecx: &mut ExtCtxt<'_>,
+    client: CLIENT,
+) -> Result<&'tcx TokenStream, ()> {
+    let invoc_expn_data = invoc_id.expn_data();
+    let span = invoc_expn_data.call_site;
+    let event_arg = invoc_expn_data.kind.descr();
+    let _timer = tcx.sess.prof.generic_activity_with_arg_recorder(
+        "expand_derive_proc_macro_inner",
+        |recorder| {
+            recorder.record_arg_with_span(tcx.sess.source_map(), event_arg.clone(), span);
+        },
+    );
+
+    let proc_macro_backtrace = ecx.ecfg.proc_macro_backtrace;
+    let strategy = crate::proc_macro::exec_strategy(tcx.sess);
+    let server = crate::proc_macro_server::Rustc::new(ecx);
+
+    match client.run(&strategy, server, input.clone(), proc_macro_backtrace) {
+        Ok(stream) => Ok(tcx.arena.alloc(stream) as &TokenStream),
+        Err(e) => {
+            tcx.dcx().emit_err({
+                errors::ProcMacroDerivePanicked {
+                    span,
+                    message: e.as_str().map(|message| errors::ProcMacroDerivePanickedHelp {
+                        message: message.into(),
+                    }),
+                }
+            });
+            Err(())
+        }
+    }
+}
+
+// based on rust/compiler/rustc_middle/src/ty/context/tls.rs
+thread_local! {
+    /// A thread local variable that stores a pointer to the current `CONTEXT`.
+    static TLV: Cell<(*mut (), Option<CLIENT>)> = const { Cell::new((std::ptr::null_mut(), None)) };
+}
+
+/// Sets `context` as the new current `CONTEXT` for the duration of the function `f`.
+#[inline]
+pub(crate) fn enter_context<'a, F, R>(context: (&mut ExtCtxt<'a>, CLIENT), f: F) -> R
+where
+    F: FnOnce() -> R,
+{
+    let (ectx, client) = context;
+    let erased = (ectx as *mut _ as *mut (), Some(client));
+    TLV.with(|tlv| {
+        let old = tlv.replace(erased);
+        let _reset = rustc_data_structures::defer(move || tlv.set(old));
+        f()
+    })
+}
+
+/// Allows access to the current `CONTEXT`.
+/// Panics if there is no `CONTEXT` available.
+#[inline]
+#[track_caller]
+fn with_context<F, R>(f: F) -> R
+where
+    F: for<'a, 'b> FnOnce(&'b mut (&mut ExtCtxt<'a>, CLIENT)) -> R,
+{
+    let (ectx, client_opt) = TLV.get();
+    let ectx = NonNull::new(ectx).expect("no CONTEXT stored in tls");
+
+    // We could get an `CONTEXT` pointer from another thread.
+    // Ensure that `CONTEXT` is `DynSync`.
+    // FIXME(pr-time): we should not be able to?
+    // sync::assert_dyn_sync::<CONTEXT<'_>>();
+
+    // prevent double entering, as that would allow creating two `&mut ExtCtxt`s
+    // FIXME(pr-time): probably use a RefCell instead (which checks this properly)?
+    TLV.with(|tlv| {
+        let old = tlv.replace((std::ptr::null_mut(), None));
+        let _reset = rustc_data_structures::defer(move || tlv.set(old));
+        let ectx = {
+            let mut casted = ectx.cast::<ExtCtxt<'_>>();
+            unsafe { casted.as_mut() }
+        };
+
+        f(&mut (ectx, client_opt.unwrap()))
+    })
+}
diff --git a/compiler/rustc_interface/src/passes.rs b/compiler/rustc_interface/src/passes.rs
@@ -881,6 +881,7 @@ pub static DEFAULT_QUERY_PROVIDERS: LazyLock<Providers> = LazyLock::new(|| {
     providers.env_var_os = env_var_os;
     limits::provide(providers);
     proc_macro_decls::provide(providers);
+    rustc_expand::provide(providers);
     rustc_const_eval::provide(providers);
     rustc_middle::hir::provide(providers);
     rustc_borrowck::provide(providers);
diff --git a/compiler/rustc_middle/src/arena.rs b/compiler/rustc_middle/src/arena.rs
@@ -116,6 +116,7 @@ macro_rules! arena_types {
             [decode] specialization_graph: rustc_middle::traits::specialization_graph::Graph,
             [] crate_inherent_impls: rustc_middle::ty::CrateInherentImpls,
             [] hir_owner_nodes: rustc_hir::OwnerNodes<'tcx>,
+            [decode] token_stream: rustc_ast::tokenstream::TokenStream,
         ]);
     )
 }
diff --git a/compiler/rustc_middle/src/query/erase.rs b/compiler/rustc_middle/src/query/erase.rs
@@ -2,6 +2,7 @@ use std::ffi::OsStr;
 use std::intrinsics::transmute_unchecked;
 use std::mem::MaybeUninit;
 
+use rustc_ast::tokenstream::TokenStream;
 use rustc_span::ErrorGuaranteed;
 
 use crate::mir::interpret::EvalToValTreeResult;
@@ -170,6 +171,10 @@ impl EraseType for Result<ty::EarlyBinder<'_, Ty<'_>>, CyclePlaceholder> {
     type Result = [u8; size_of::<Result<ty::EarlyBinder<'static, Ty<'_>>, CyclePlaceholder>>()];
 }
 
+impl EraseType for Result<&'_ TokenStream, ()> {
+    type Result = [u8; size_of::<Result<&'static TokenStream, ()>>()];
+}
+
 impl<T> EraseType for Option<&'_ T> {
     type Result = [u8; size_of::<Option<&'static ()>>()];
 }
diff --git a/compiler/rustc_middle/src/query/keys.rs b/compiler/rustc_middle/src/query/keys.rs
@@ -2,11 +2,13 @@
 
 use std::ffi::OsStr;
 
+use rustc_ast::tokenstream::TokenStream;
+use rustc_data_structures::svh::Svh;
 use rustc_hir::def_id::{CrateNum, DefId, LOCAL_CRATE, LocalDefId, LocalModDefId, ModDefId};
 use rustc_hir::hir_id::{HirId, OwnerId};
 use rustc_query_system::dep_graph::DepNodeIndex;
 use rustc_query_system::query::{DefIdCache, DefaultCache, SingleCache, VecCache};
-use rustc_span::{DUMMY_SP, Ident, Span, Symbol};
+use rustc_span::{DUMMY_SP, Ident, LocalExpnId, Span, Symbol};
 
 use crate::infer::canonical::CanonicalQueryInput;
 use crate::mir::mono::CollectionMode;
@@ -616,6 +618,19 @@ impl Key for (LocalDefId, HirId) {
     }
 }
 
+impl<'tcx> Key for (LocalExpnId, Svh, &'tcx TokenStream) {
+    type Cache<V> = DefaultCache<Self, V>;
+
+    fn default_span(&self, _tcx: TyCtxt<'_>) -> Span {
+        self.0.expn_data().call_site
+    }
+
+    #[inline(always)]
+    fn key_as_def_id(&self) -> Option<DefId> {
+        None
+    }
+}
+
 impl<'tcx> Key for (ValidityRequirement, ty::PseudoCanonicalInput<'tcx, Ty<'tcx>>) {
     type Cache<V> = DefaultCache<Self, V>;
 
diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
@@ -70,6 +70,7 @@ use std::sync::Arc;
 use rustc_abi::Align;
 use rustc_arena::TypedArena;
 use rustc_ast::expand::allocator::AllocatorKind;
+use rustc_ast::tokenstream::TokenStream;
 use rustc_data_structures::fingerprint::Fingerprint;
 use rustc_data_structures::fx::{FxIndexMap, FxIndexSet};
 use rustc_data_structures::sorted_map::SortedMap;
@@ -99,7 +100,7 @@ use rustc_session::cstore::{
 use rustc_session::lint::LintExpectationId;
 use rustc_span::def_id::LOCAL_CRATE;
 use rustc_span::source_map::Spanned;
-use rustc_span::{DUMMY_SP, Span, Symbol};
+use rustc_span::{DUMMY_SP, LocalExpnId, Span, Symbol};
 use rustc_target::spec::{PanicStrategy, SanitizerSet};
 use {rustc_abi as abi, rustc_ast as ast, rustc_hir as hir};
 
@@ -164,6 +165,13 @@ pub use plumbing::{IntoQueryParam, TyCtxtAt, TyCtxtEnsureDone, TyCtxtEnsureOk};
 // Queries marked with `fatal_cycle` do not need the latter implementation,
 // as they will raise an fatal error on query cycles instead.
 rustc_queries! {
+    query derive_macro_expansion(key: (LocalExpnId, Svh, &'tcx TokenStream)) -> Result<&'tcx TokenStream, ()> {
+        // eval_always
+        // no_hash
+        desc { "expanding a derive (proc) macro" }
+        cache_on_disk_if { true }
+    }
+
     /// This exists purely for testing the interactions between delayed bugs and incremental.
     query trigger_delayed_bug(key: DefId) {
         desc { "triggering a delayed bug for testing incremental" }
diff --git a/compiler/rustc_middle/src/query/on_disk_cache.rs b/compiler/rustc_middle/src/query/on_disk_cache.rs
@@ -780,6 +780,13 @@ impl<'a, 'tcx> Decodable<CacheDecoder<'a, 'tcx>>
     }
 }
 
+impl<'a, 'tcx> Decodable<CacheDecoder<'a, 'tcx>> for &'tcx rustc_ast::tokenstream::TokenStream {
+    #[inline]
+    fn decode(d: &mut CacheDecoder<'a, 'tcx>) -> Self {
+        RefDecodable::decode(d)
+    }
+}
+
 macro_rules! impl_ref_decoder {
     (<$tcx:tt> $($ty:ty,)*) => {
         $(impl<'a, $tcx> Decodable<CacheDecoder<'a, $tcx>> for &$tcx [$ty] {
diff --git a/compiler/rustc_session/src/options.rs b/compiler/rustc_session/src/options.rs
@@ -2157,6 +2157,8 @@ options! {
         "set options for branch target identification and pointer authentication on AArch64"),
     build_sdylib_interface: bool = (false, parse_bool, [UNTRACKED],
         "whether the stable interface is being built"),
+    cache_derive_macros: bool = (false, parse_bool, [TRACKED],
+        "cache the results of derive proc macro invocations (potentially unsound!) (default: no"),
     cf_protection: CFProtection = (CFProtection::None, parse_cfprotection, [TRACKED],
         "instrument control-flow architecture protection"),
     check_cfg_all_expected: bool = (false, parse_bool, [UNTRACKED],
diff --git a/compiler/rustc_span/src/hygiene.rs b/compiler/rustc_span/src/hygiene.rs
diff --git a/tests/incremental/derive_macro_expansion/auxiliary/derive_nothing.rs b/tests/incremental/derive_macro_expansion/auxiliary/derive_nothing.rs
diff --git a/tests/incremental/derive_macro_expansion/proc_macro_unchanged.rs b/tests/incremental/derive_macro_expansion/proc_macro_unchanged.rs

Original file line number	Diff line number	Diff line change
`@@ -116,6 +116,7 @@ macro_rules! arena_types {`
`116`	`116`	`[decode] specialization_graph: rustc_middle::traits::specialization_graph::Graph,`
`117`	`117`	`[] crate_inherent_impls: rustc_middle::ty::CrateInherentImpls,`
`118`	`118`	`[] hir_owner_nodes: rustc_hir::OwnerNodes<'tcx>,`
	`119`	`+ [decode] token_stream: rustc_ast::tokenstream::TokenStream,`
`119`	`120`	`]);`
`120`	`121`	`)`
`121`	`122`	`}`