Auto merge of #101414 - mystor:pm_nested_cross_thread, r=eddyb

proc_macro/bridge: use the cross-thread executor for nested proc-macros

While working on some other changes in the bridge, I noticed that when
running a nested proc-macro (which is currently only possible using
the unstable `TokenStream::expand_expr`), any symbols held by the
proc-macro client would be invalidated, as the same thread would be used
for the nested macro by default, and the interner doesn't handle nested
use.

After discussing with `@eddyb,` we decided the best approach might be to
force the use of the cross-thread executor for nested invocations, as it
will never re-use thread-local storage, avoiding the issue. This
shouldn't impact performance, as expand_expr is still unstable, and
infrequently used.

This was chosen rather than making the client symbol interner handle
nested invocations, as that would require replacing the internal
interner `Vec` with a `BTreeMap` (as valid symbol id ranges could now be
disjoint), and the symbol interner is known to be fairly perf-sensitive.

This patch adds checks to the execution strategy to use the cross-thread
executor when doing nested invocations. An alternative implementation
strategy could be to track this information in the `ExtCtxt`, however a
thread-local in the `proc_macro` crate was chosen to add an assertion so
that `rust-analyzer` is aware of the issue if it implements
`expand_expr` in the future.

r? `@eddyb`

library/proc_macro/src/bridge/server.rs

@@ -2,6 +2,7 @@
 
 use super::*;
 
+use std::cell::Cell;
 use std::marker::PhantomData;
 
 // FIXME(eddyb) generate the definition of `HandleStore` in `server.rs`.
@@ -143,6 +144,38 @@ pub trait ExecutionStrategy {
     ) -> Buffer;
 }
 
+thread_local! {
+    /// While running a proc-macro with the same-thread executor, this flag will
+    /// be set, forcing nested proc-macro invocations (e.g. due to
+    /// `TokenStream::expand_expr`) to be run using a cross-thread executor.
+    ///
+    /// This is required as the thread-local state in the proc_macro client does
+    /// not handle being re-entered, and will invalidate all `Symbol`s when
+    /// entering a nested macro.
+    static ALREADY_RUNNING_SAME_THREAD: Cell<bool> = Cell::new(false);
+}
+
+/// Keep `ALREADY_RUNNING_SAME_THREAD` (see also its documentation)
+/// set to `true`, preventing same-thread reentrance.
+struct RunningSameThreadGuard(());
+
+impl RunningSameThreadGuard {
+    fn new() -> Self {
+        let already_running = ALREADY_RUNNING_SAME_THREAD.replace(true);
+        assert!(
+            !already_running,
+            "same-thread nesting (\"reentrance\") of proc macro executions is not supported"
+        );
+        RunningSameThreadGuard(())
+    }
+}
+
+impl Drop for RunningSameThreadGuard {
+    fn drop(&mut self) {
+        ALREADY_RUNNING_SAME_THREAD.set(false);
+    }
+}
+
 pub struct MaybeCrossThread<P> {
     cross_thread: bool,
     marker: PhantomData<P>,
@@ -165,7 +198,7 @@ where
         run_client: extern "C" fn(BridgeConfig<'_>) -> Buffer,
         force_show_panics: bool,
     ) -> Buffer {
-        if self.cross_thread {
+        if self.cross_thread || ALREADY_RUNNING_SAME_THREAD.get() {
             <CrossThread<P>>::new().run_bridge_and_client(
                 dispatcher,
                 input,
@@ -188,6 +221,8 @@ impl ExecutionStrategy for SameThread {
         run_client: extern "C" fn(BridgeConfig<'_>) -> Buffer,
         force_show_panics: bool,
     ) -> Buffer {
+        let _guard = RunningSameThreadGuard::new();
+
         let mut dispatch = |buf| dispatcher.dispatch(buf);
 
         run_client(BridgeConfig {

src/test/ui/proc-macro/auxiliary/expand-expr.rs

@@ -80,13 +80,21 @@ fn assert_ts_eq(lhs: &TokenStream, rhs: &TokenStream) {
 pub fn expand_expr_is(input: TokenStream) -> TokenStream {
     let mut iter = input.into_iter();
     let mut expected_tts = Vec::new();
-    loop {
+    let comma = loop {
         match iter.next() {
-            Some(TokenTree::Punct(ref p)) if p.as_char() == ',' => break,
+            Some(TokenTree::Punct(p)) if p.as_char() == ',' => break p,
             Some(tt) => expected_tts.push(tt),
             None => panic!("expected comma"),
         }
-    }
+    };
+
+    // Make sure that `Ident` and `Literal` objects from this proc-macro's
+    // environment are not invalidated when `expand_expr` recursively invokes
+    // another macro by taking a local copy, and checking it after the fact.
+    let pre_expand_span = comma.span();
+    let pre_expand_ident = Ident::new("ident", comma.span());
+    let pre_expand_literal = Literal::string("literal");
+    let pre_expand_call_site = Span::call_site();
 
     let expected = expected_tts.into_iter().collect::<TokenStream>();
     let expanded = iter.collect::<TokenStream>().expand_expr().expect("expand_expr failed");
@@ -100,6 +108,15 @@ pub fn expand_expr_is(input: TokenStream) -> TokenStream {
     // Also compare the raw tts to make sure they line up.
     assert_ts_eq(&expected, &expanded);
 
+    assert!(comma.span().eq(&pre_expand_span), "pre-expansion span is still equal");
+    assert_eq!(pre_expand_ident.to_string(), "ident", "pre-expansion identifier is still valid");
+    assert_eq!(
+        pre_expand_literal.to_string(),
+        "\"literal\"",
+        "pre-expansion literal is still valid"
+    );
+    assert!(Span::call_site().eq(&pre_expand_call_site), "pre-expansion call-site is still equal");
+
     TokenStream::new()
 }