Auto merge of #129102 - futile:experimental/proc-macro-caching, r=<try>

Experimental: Add Derive Proc-Macro Caching

# On-Disk Caching For Derive Proc-Macro Invocations

This PR adds on-disk caching for derive proc-macro invocations using rustc's query system to speed up incremental compilation.

The implementation is (intentionally) a bit rough/incomplete, as I wanted to see whether this helps with performance before fully implementing it/RFCing etc.

I did some ad-hoc performance testing.

## Rough, Preliminary Eval Results:

Using a version built through `DEPLOY=1 src/ci/docker/run.sh dist-x86_64-linux` (which I got from [here](https://rustc-dev-guide.rust-lang.org/building/optimized-build.html#profile-guided-optimization)).

### [Some Small Personal Project](https://github.com/futile/ultra-game):

```console
# with -Zthreads=0 as well
$ touch src/main.rs && cargo +dist check
```

Caused a re-check of 1 crate (the only one).

Result:
| Configuration | Time (avg. ~5 runs) |
|--------|--------|
| Uncached | ~0.54s |
| Cached | ~0.54s |

No visible difference.

### [Bevy](https://github.com/bevyengine/bevy):

```console
$ touch crates/bevy_ecs/src/lib.rs && cargo +dist check
```

Caused a re-check of 29 crates.

Result:
| Configuration | Time (avg. ~5 runs) |
|--------|--------|
| Uncached | ~6.4s |
| Cached | ~5.3s |

Roughly 1s, or ~17% speedup.

### [Polkadot-Sdk](https://github.com/paritytech/polkadot-sdk):

Basically this script (not mine): https://github.com/coderemotedotdev/rustc-profiles/blob/d61ad38c496459d82e35d8bdb0a154fbb83de903/scripts/benchmark_incremental_builds_polkadot_sdk.sh

TL;DR: Two full `cargo check` runs to fill the incremental caches (for cached & uncached). Then 10 repetitions of `touch $some_file && cargo +uncached check && cargo +cached check`.

```console
$ cargo update # `time` didn't build because compiler too new/dep too old
$ ./benchmark_incremental_builds_polkadot_sdk.sh # see above
```

_Huge_ workspace with ~190 crates. Not sure how many were re-built/re-checkd on each invocation.

Result:
| Configuration | Time (avg. 10 runs) |
|--------|--------|
| Uncached | 99.4s |
| Cached | 67.5s |

Very visible speedup of 31.9s or ~32%.

---

**-> Based on these results I think it makes sense to do a rustc-perf run and see what that reports.**

---

## Current Limitations/TODOs

I left some `FIXME(pr-time)`s in the code for things I wanted to bring up/draw attention to in this PR. Usually when I wasn't sure if I found a (good) solution or when I knew that there might be a better way to do something; See the diff for these.

### High-Level Overview of What's Missing For "Real" Usage:

* [ ] Add caching for `Bang`- and `Attr`-proc macros (currently only `Derive`).
  * Not a big change, I just focused on `derive`-proc macros for now, since I felt like these should be most cacheable and are used very often in practice.
* [ ] Allow marking specific macros as "do not cache" (currently only all-or-nothing).
  * Extend the unstable option to support, e.g., `-Z cache-derive-macros=some_pm_crate::some_derive_macro_fn` for easy testing using the nightly compiler.
  * After Testing: Add a `#[proc_macro_cacheable]` annotation to allow proc-macro authors to "opt-in" to caching (or sth. similar). Would probably need an RFC?
  * Might make sense to try to combine this with #99515, so that external dependencies can be picked up and be taken into account as well.

---

So, just since you were in the loop on the attempt to cache declarative macro expansions:

r? `@petrochenkov`

Please feel free to re-/unassign!

Finally: I hope this isn't too big a PR, I'll also show up in Zulip since I read that that is usually appreciated. Thanks a lot for taking a look! :)

(Kind of related/very similar approach, old declarative macro caching PR: #128747)

Cargo.lock

@@ -3688,6 +3688,7 @@ dependencies = [
  "rustc_lexer",
  "rustc_lint_defs",
  "rustc_macros",
+ "rustc_middle",
  "rustc_parse",
  "rustc_serialize",
  "rustc_session",

compiler/rustc_ast/src/tokenstream.rs

@@ -14,12 +14,14 @@
 //! ownership of the original.
 
 use std::borrow::Cow;
+use std::hash::Hash;
 use std::{cmp, fmt, iter};
 
 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
 use rustc_data_structures::sync::{self, Lrc};
 use rustc_macros::{Decodable, Encodable, HashStable_Generic};
-use rustc_serialize::{Decodable, Encodable};
+use rustc_serialize::{Decodable, Encodable, Encoder};
+use rustc_span::def_id::{CrateNum, DefIndex};
 use rustc_span::{sym, Span, SpanDecoder, SpanEncoder, Symbol, DUMMY_SP};
 
 use crate::ast::{AttrStyle, StmtKind};
@@ -140,6 +142,11 @@ impl fmt::Debug for LazyAttrTokenStream {
 
 impl<S: SpanEncoder> Encodable<S> for LazyAttrTokenStream {
     fn encode(&self, _s: &mut S) {
+        // FIXME(pr-time): Just a reminder that this exists/was tried out,
+        // but probably not necessary anymore (see below).
+        // self.to_attr_token_stream().encode(s)
+        // We should not need to anymore, now that we `flatten`?
+        // Yep, that seems to be true! :)
         panic!("Attempted to encode LazyAttrTokenStream");
     }
 }
@@ -296,6 +303,96 @@ pub struct AttrsTarget {
 #[derive(Clone, Debug, Default, Encodable, Decodable)]
 pub struct TokenStream(pub(crate) Lrc<Vec<TokenTree>>);
 
+struct HashEncoder<H: std::hash::Hasher> {
+    hasher: H,
+}
+
+impl<H: std::hash::Hasher> Encoder for HashEncoder<H> {
+    fn emit_usize(&mut self, v: usize) {
+        self.hasher.write_usize(v)
+    }
+
+    fn emit_u128(&mut self, v: u128) {
+        self.hasher.write_u128(v)
+    }
+
+    fn emit_u64(&mut self, v: u64) {
+        self.hasher.write_u64(v)
+    }
+
+    fn emit_u32(&mut self, v: u32) {
+        self.hasher.write_u32(v)
+    }
+
+    fn emit_u16(&mut self, v: u16) {
+        self.hasher.write_u16(v)
+    }
+
+    fn emit_u8(&mut self, v: u8) {
+        self.hasher.write_u8(v)
+    }
+
+    fn emit_isize(&mut self, v: isize) {
+        self.hasher.write_isize(v)
+    }
+
+    fn emit_i128(&mut self, v: i128) {
+        self.hasher.write_i128(v)
+    }
+
+    fn emit_i64(&mut self, v: i64) {
+        self.hasher.write_i64(v)
+    }
+
+    fn emit_i32(&mut self, v: i32) {
+        self.hasher.write_i32(v)
+    }
+
+    fn emit_i16(&mut self, v: i16) {
+        self.hasher.write_i16(v)
+    }
+
+    fn emit_raw_bytes(&mut self, s: &[u8]) {
+        self.hasher.write(s)
+    }
+}
+
+impl<H: std::hash::Hasher> SpanEncoder for HashEncoder<H> {
+    fn encode_span(&mut self, span: Span) {
+        span.hash(&mut self.hasher)
+    }
+
+    fn encode_symbol(&mut self, symbol: Symbol) {
+        symbol.hash(&mut self.hasher)
+    }
+
+    fn encode_expn_id(&mut self, expn_id: rustc_span::ExpnId) {
+        expn_id.hash(&mut self.hasher)
+    }
+
+    fn encode_syntax_context(&mut self, syntax_context: rustc_span::SyntaxContext) {
+        syntax_context.hash(&mut self.hasher)
+    }
+
+    fn encode_crate_num(&mut self, crate_num: CrateNum) {
+        crate_num.hash(&mut self.hasher)
+    }
+
+    fn encode_def_index(&mut self, def_index: DefIndex) {
+        def_index.hash(&mut self.hasher)
+    }
+
+    fn encode_def_id(&mut self, def_id: rustc_span::def_id::DefId) {
+        def_id.hash(&mut self.hasher)
+    }
+}
+
+impl Hash for TokenStream {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        Encodable::encode(self, &mut HashEncoder { hasher: state });
+    }
+}
+
 /// Indicates whether a token can join with the following token to form a
 /// compound token. Used for conversions to `proc_macro::Spacing`. Also used to
 /// guide pretty-printing, which is where the `JointHidden` value (which isn't

compiler/rustc_expand/Cargo.toml

@@ -20,6 +20,7 @@ rustc_fluent_macro = { path = "../rustc_fluent_macro" }
 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" }
 rustc_serialize = { path = "../rustc_serialize" }
 rustc_session = { path = "../rustc_session" }

compiler/rustc_expand/src/derive_macro_expansion.rs

@@ -0,0 +1,125 @@
+use std::cell::Cell;
+use std::ptr;
+
+use rustc_ast::tokenstream::TokenStream;
+use rustc_data_structures::svh::Svh;
+use rustc_middle::ty::TyCtxt;
+use rustc_span::profiling::SpannedEventArgRecorder;
+use rustc_span::LocalExpnId;
+
+use crate::base::ExtCtxt;
+use crate::errors;
+
+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;
+
+    let res = with_context(|(ecx, client)| {
+        let span = invoc_id.expn_data().call_site;
+        let _timer = ecx.sess.prof.generic_activity_with_arg_recorder(
+            "expand_derive_proc_macro_inner",
+            |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 = crate::proc_macro::exec_strategy(ecx);
+        let server = crate::proc_macro_server::Rustc::new(ecx);
+        let res = match client.run(&strategy, server, input.clone(), proc_macro_backtrace) {
+            // FIXME(pr-time): without flattened some (weird) tests fail, but no idea if it's correct/enough
+            Ok(stream) => Ok(tcx.arena.alloc(stream.flattened()) as &TokenStream),
+            Err(e) => {
+                ecx.dcx().emit_err({
+                    errors::ProcMacroDerivePanicked {
+                        span,
+                        message: e.as_str().map(|message| errors::ProcMacroDerivePanickedHelp {
+                            message: message.into(),
+                        }),
+                    }
+                });
+                Err(())
+            }
+        };
+        res
+    });
+
+    res
+}
+
+type CLIENT = pm::bridge::client::Client<pm::TokenStream, pm::TokenStream>;
+
+// 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((ptr::null_mut(), None)) };
+}
+
+#[inline]
+fn erase(context: &mut ExtCtxt<'_>) -> *mut () {
+    context as *mut _ as *mut ()
+}
+
+#[inline]
+unsafe fn downcast<'a>(context: *mut ()) -> &'a mut ExtCtxt<'a> {
+    unsafe { &mut *(context as *mut ExtCtxt<'a>) }
+}
+
+#[inline]
+fn enter_context_erased<F, R>(erased: (*mut (), Option<CLIENT>), f: F) -> R
+where
+    F: FnOnce() -> R,
+{
+    TLV.with(|tlv| {
+        let old = tlv.replace(erased);
+        let _reset = rustc_data_structures::defer(move || tlv.set(old));
+        f()
+    })
+}
+
+/// Sets `context` as the new current `CONTEXT` for the duration of the function `f`.
+#[inline]
+pub fn enter_context<'a, F, R>(context: (&mut ExtCtxt<'a>, CLIENT), f: F) -> R
+where
+    F: FnOnce() -> R,
+{
+    let (ectx, client) = context;
+    let erased = (erase(ectx), Some(client));
+    enter_context_erased(erased, f)
+}
+
+/// Allows access to the current `CONTEXT` in a closure if one is available.
+#[inline]
+#[track_caller]
+pub fn with_context_opt<F, R>(f: F) -> R
+where
+    F: for<'a, 'b> FnOnce(Option<&'b mut (&mut ExtCtxt<'a>, CLIENT)>) -> R,
+{
+    let (ectx, client_opt) = TLV.get();
+    if ectx.is_null() {
+        f(None)
+    } else {
+        // 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)?
+        enter_context_erased((ptr::null_mut(), None), || unsafe {
+            let ectx = downcast(ectx);
+            f(Some(&mut (ectx, client_opt.unwrap())))
+        })
+    }
+}
+
+/// Allows access to the current `CONTEXT`.
+/// Panics if there is no `CONTEXT` available.
+#[inline]
+pub fn with_context<F, R>(f: F) -> R
+where
+    F: for<'a, 'b> FnOnce(&'b mut (&mut ExtCtxt<'a>, CLIENT)) -> R,
+{
+    with_context_opt(|opt_context| f(opt_context.expect("no CONTEXT stored in tls")))
+}

compiler/rustc_expand/src/expand.rs

@@ -488,7 +488,7 @@ impl<'a, 'b> MacroExpander<'a, 'b> {
             self.cx.force_mode = force;
 
             let fragment_kind = invoc.fragment_kind;
-            match self.expand_invoc(invoc, &ext.kind) {
+            match self.expand_invoc(invoc, &ext) {
                 ExpandResult::Ready(fragment) => {
                     let mut derive_invocations = Vec::new();
                     let derive_placeholders = self
@@ -650,7 +650,7 @@ impl<'a, 'b> MacroExpander<'a, 'b> {
     fn expand_invoc(
         &mut self,
         invoc: Invocation,
-        ext: &SyntaxExtensionKind,
+        ext: &Lrc<SyntaxExtension>,
     ) -> ExpandResult<AstFragment, Invocation> {
         let recursion_limit = match self.cx.reduced_recursion_limit {
             Some((limit, _)) => limit,
@@ -671,7 +671,7 @@ impl<'a, 'b> MacroExpander<'a, 'b> {
 
         let (fragment_kind, span) = (invoc.fragment_kind, invoc.span());
         ExpandResult::Ready(match invoc.kind {
-            InvocationKind::Bang { mac, span } => match ext {
+            InvocationKind::Bang { mac, span } => match &ext.kind {
                 SyntaxExtensionKind::Bang(expander) => {
                     match expander.expand(self.cx, span, mac.args.tokens.clone()) {
                         Ok(tok_result) => {
@@ -701,7 +701,7 @@ impl<'a, 'b> MacroExpander<'a, 'b> {
                 }
                 _ => unreachable!(),
             },
-            InvocationKind::Attr { attr, pos, mut item, derives } => match ext {
+            InvocationKind::Attr { attr, pos, mut item, derives } => match &ext.kind {
                 SyntaxExtensionKind::Attr(expander) => {
                     self.gate_proc_macro_input(&item);
                     self.gate_proc_macro_attr_item(span, &item);
@@ -780,10 +780,10 @@ impl<'a, 'b> MacroExpander<'a, 'b> {
                 }
                 _ => unreachable!(),
             },
-            InvocationKind::Derive { path, item, is_const } => match ext {
+            InvocationKind::Derive { path, item, is_const } => match &ext.kind {
                 SyntaxExtensionKind::Derive(expander)
                 | SyntaxExtensionKind::LegacyDerive(expander) => {
-                    if let SyntaxExtensionKind::Derive(..) = ext {
+                    if let SyntaxExtensionKind::Derive(..) = ext.kind {
                         self.gate_proc_macro_input(&item);
                     }
                     // The `MetaItem` representing the trait to derive can't
@@ -794,6 +794,7 @@ impl<'a, 'b> MacroExpander<'a, 'b> {
                         span,
                         path,
                     };
+                    invoc.expansion_data.id.expn_data();
                     let items = match expander.expand(self.cx, span, &meta, item, is_const) {
                         ExpandResult::Ready(items) => items,
                         ExpandResult::Retry(item) => {
@@ -810,18 +811,19 @@ impl<'a, 'b> MacroExpander<'a, 'b> {
             },
             InvocationKind::GlobDelegation { item } => {
                 let AssocItemKind::DelegationMac(deleg) = &item.kind else { unreachable!() };
-                let suffixes = match ext {
-                    SyntaxExtensionKind::GlobDelegation(expander) => match expander.expand(self.cx)
-                    {
-                        ExpandResult::Ready(suffixes) => suffixes,
-                        ExpandResult::Retry(()) => {
-                            // Reassemble the original invocation for retrying.
-                            return ExpandResult::Retry(Invocation {
-                                kind: InvocationKind::GlobDelegation { item },
-                                ..invoc
-                            });
+                let suffixes = match &ext.kind {
+                    SyntaxExtensionKind::GlobDelegation(expander) => {
+                        match expander.expand(self.cx) {
+                            ExpandResult::Ready(suffixes) => suffixes,
+                            ExpandResult::Retry(()) => {
+                                // Reassemble the original invocation for retrying.
+                                return ExpandResult::Retry(Invocation {
+                                    kind: InvocationKind::GlobDelegation { item },
+                                    ..invoc
+                                });
+                            }
                         }
-                    },
+                    }
                     SyntaxExtensionKind::LegacyBang(..) => {
                         let msg = "expanded a dummy glob delegation";
                         let guar = self.cx.dcx().span_delayed_bug(span, msg);

compiler/rustc_expand/src/lib.rs

@@ -28,10 +28,15 @@ mod proc_macro_server;
 pub use mbe::macro_rules::compile_declarative_macro;
 pub mod base;
 pub mod config;
+pub(crate) mod derive_macro_expansion;
 pub mod expand;
 pub mod module;
 // FIXME(Nilstrieb) Translate proc_macro diagnostics
 #[allow(rustc::untranslatable_diagnostic)]
 pub mod proc_macro;
 
+pub fn provide(providers: &mut rustc_middle::util::Providers) {
+    providers.derive_macro_expansion = derive_macro_expansion::provide_derive_macro_expansion;
+}
+
 rustc_fluent_macro::fluent_messages! { "../messages.ftl" }