refactor: Move derive_macro_expansion.rs to end of proc_macro.rs

Author: futile

Diff for: compiler/rustc_expand/src/derive_macro_expansion.rs

@@ -18,7 +18,6 @@
 extern crate proc_macro as pm;
 
 mod build;
-mod derive_macro_expansion;
 mod errors;
 // FIXME(Nilstrieb) Translate macro_rules diagnostics
 #[allow(rustc::untranslatable_diagnostic)]
@@ -36,7 +35,7 @@ pub mod module;
 pub mod proc_macro;
 
 pub fn provide(providers: &mut rustc_middle::util::Providers) {
-    providers.derive_macro_expansion = derive_macro_expansion::provide_derive_macro_expansion;
+    providers.derive_macro_expansion = proc_macro::provide_derive_macro_expansion;
 }
 
 rustc_fluent_macro::fluent_messages! { "../messages.ftl" }

Diff for: compiler/rustc_expand/src/lib.rs

@@ -1,13 +1,17 @@
+use std::cell::Cell;
+use std::ptr::{self, NonNull};
+
 use rustc_ast as ast;
 use rustc_ast::ptr::P;
 use rustc_ast::tokenstream::TokenStream;
+use rustc_data_structures::svh::Svh;
 use rustc_errors::ErrorGuaranteed;
-use rustc_middle::ty;
+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 crate::base::{self, *};
 use crate::{errors, proc_macro_server};
@@ -154,11 +158,9 @@ impl MultiItemModifier for DeriveProcMacro {
             // well)?
             if tcx.sess.opts.incremental.is_some() && tcx.sess.opts.unstable_opts.cache_proc_macros
             {
-                crate::derive_macro_expansion::enter_context((ecx, self.client), move || {
-                    tcx.derive_macro_expansion(key).cloned()
-                })
+                enter_context((ecx, self.client), move || tcx.derive_macro_expansion(key).cloned())
             } else {
-                crate::derive_macro_expansion::provide_derive_macro_expansion(tcx, key).cloned()
+                provide_derive_macro_expansion(tcx, key).cloned()
             }
         });
         let Ok(output) = res else {
@@ -195,3 +197,94 @@ 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)| {
+        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(())
+            }
+        }
+    })
+}
+
+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)) };
+}
+
+/// 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((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()))
+    })
+}