proc_macro: cache static spans in client's thread-local state

This greatly improves the performance of the very frequently called
`call_site()` macro when running in a cross-thread configuration.

5 files changed, 170 insertions, 86 deletions

diff --git a/library/proc_macro/src/bridge/client.rs b/library/proc_macro/src/bridge/client.rs
index d7cc700495b..f1f8ae33898 100644
--- a/library/proc_macro/src/bridge/client.rs
+++ b/library/proc_macro/src/bridge/client.rs
@@ -230,6 +230,20 @@ impl Clone for SourceFile {
     }
 }
 
+impl Span {
+    pub(crate) fn def_site() -> Span {
+        Bridge::with(|bridge| bridge.context.def_site)
+    }
+
+    pub(crate) fn call_site() -> Span {
+        Bridge::with(|bridge| bridge.context.call_site)
+    }
+
+    pub(crate) fn mixed_site() -> Span {
+        Bridge::with(|bridge| bridge.context.mixed_site)
+    }
+}
+
 impl fmt::Debug for Span {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.write_str(&self.debug())
@@ -263,6 +277,21 @@ macro_rules! define_client_side {
 }
 with_api!(self, self, define_client_side);
 
+struct Bridge<'a> {
+    /// Reusable buffer (only `clear`-ed, never shrunk), primarily
+    /// used for making requests.
+    cached_buffer: Buffer,
+
+    /// Server-side function that the client uses to make requests.
+    dispatch: closure::Closure<'a, Buffer, Buffer>,
+
+    /// Provided context for this macro expansion.
+    context: ExpnContext<Span>,
+}
+
+impl<'a> !Send for Bridge<'a> {}
+impl<'a> !Sync for Bridge<'a> {}
+
 enum BridgeState<'a> {
     /// No server is currently connected to this client.
     NotConnected,
@@ -305,34 +334,6 @@ impl BridgeState<'_> {
 }
 
 impl Bridge<'_> {
-    pub(crate) fn is_available() -> bool {
-        BridgeState::with(|state| match state {
-            BridgeState::Connected(_) | BridgeState::InUse => true,
-            BridgeState::NotConnected => false,
-        })
-    }
-
-    fn enter<R>(self, f: impl FnOnce() -> R) -> R {
-        let force_show_panics = self.force_show_panics;
-        // Hide the default panic output within `proc_macro` expansions.
-        // NB. the server can't do this because it may use a different libstd.
-        static HIDE_PANICS_DURING_EXPANSION: Once = Once::new();
-        HIDE_PANICS_DURING_EXPANSION.call_once(|| {
-            let prev = panic::take_hook();
-            panic::set_hook(Box::new(move |info| {
-                let show = BridgeState::with(|state| match state {
-                    BridgeState::NotConnected => true,
-                    BridgeState::Connected(_) | BridgeState::InUse => force_show_panics,
-                });
-                if show {
-                    prev(info)
-                }
-            }));
-        });
-
-        BRIDGE_STATE.with(|state| state.set(BridgeState::Connected(self), f))
-    }
-
     fn with<R>(f: impl FnOnce(&mut Bridge<'_>) -> R) -> R {
         BridgeState::with(|state| match state {
             BridgeState::NotConnected => {
@@ -346,6 +347,13 @@ impl Bridge<'_> {
     }
 }
 
+pub(crate) fn is_available() -> bool {
+    BridgeState::with(|state| match state {
+        BridgeState::Connected(_) | BridgeState::InUse => true,
+        BridgeState::NotConnected => false,
+    })
+}
+
 /// A client-side RPC entry-point, which may be using a different `proc_macro`
 /// from the one used by the server, but can be invoked compatibly.
 ///
@@ -363,7 +371,7 @@ pub struct Client<I, O> {
     // a wrapper `fn` pointer, once `const fn` can reference `static`s.
     pub(super) get_handle_counters: extern "C" fn() -> &'static HandleCounters,
 
-    pub(super) run: extern "C" fn(Bridge<'_>) -> Buffer,
+    pub(super) run: extern "C" fn(BridgeConfig<'_>) -> Buffer,
 
     pub(super) _marker: PhantomData<fn(I) -> O>,
 }
@@ -375,40 +383,62 @@ impl<I, O> Clone for Client<I, O> {
     }
 }
 
+fn maybe_install_panic_hook(force_show_panics: bool) {
+    // Hide the default panic output within `proc_macro` expansions.
+    // NB. the server can't do this because it may use a different libstd.
+    static HIDE_PANICS_DURING_EXPANSION: Once = Once::new();
+    HIDE_PANICS_DURING_EXPANSION.call_once(|| {
+        let prev = panic::take_hook();
+        panic::set_hook(Box::new(move |info| {
+            let show = BridgeState::with(|state| match state {
+                BridgeState::NotConnected => true,
+                BridgeState::Connected(_) | BridgeState::InUse => force_show_panics,
+            });
+            if show {
+                prev(info)
+            }
+        }));
+    });
+}
+
 /// Client-side helper for handling client panics, entering the bridge,
 /// deserializing input and serializing output.
 // FIXME(eddyb) maybe replace `Bridge::enter` with this?
 fn run_client<A: for<'a, 's> DecodeMut<'a, 's, ()>, R: Encode<()>>(
-    mut bridge: Bridge<'_>,
+    config: BridgeConfig<'_>,
     f: impl FnOnce(A) -> R,
 ) -> Buffer {
-    // The initial `cached_buffer` contains the input.
-    let mut buf = bridge.cached_buffer.take();
+    let BridgeConfig { input: mut buf, dispatch, force_show_panics, .. } = config;
 
     panic::catch_unwind(panic::AssertUnwindSafe(|| {
-        bridge.enter(|| {
-            let reader = &mut &buf[..];
-            let input = A::decode(reader, &mut ());
-
-            // Put the `cached_buffer` back in the `Bridge`, for requests.
-            Bridge::with(|bridge| bridge.cached_buffer = buf.take());
-
-            let output = f(input);
-
-            // Take the `cached_buffer` back out, for the output value.
-            buf = Bridge::with(|bridge| bridge.cached_buffer.take());
-
-            // HACK(eddyb) Separate encoding a success value (`Ok(output)`)
-            // from encoding a panic (`Err(e: PanicMessage)`) to avoid
-            // having handles outside the `bridge.enter(|| ...)` scope, and
-            // to catch panics that could happen while encoding the success.
-            //
-            // Note that panics should be impossible beyond this point, but
-            // this is defensively trying to avoid any accidental panicking
-            // reaching the `extern "C"` (which should `abort` but might not
-            // at the moment, so this is also potentially preventing UB).
-            buf.clear();
-            Ok::<_, ()>(output).encode(&mut buf, &mut ());
+        maybe_install_panic_hook(force_show_panics);
+
+        let reader = &mut &buf[..];
+        let (input, context) = <(A, ExpnContext<Span>)>::decode(reader, &mut ());
+
+        // Put the buffer we used for input back in the `Bridge` for requests.
+        let new_state =
+            BridgeState::Connected(Bridge { cached_buffer: buf.take(), dispatch, context });
+
+        BRIDGE_STATE.with(|state| {
+            state.set(new_state, || {
+                let output = f(input);
+
+                // Take the `cached_buffer` back out, for the output value.
+                buf = Bridge::with(|bridge| bridge.cached_buffer.take());
+
+                // HACK(eddyb) Separate encoding a success value (`Ok(output)`)
+                // from encoding a panic (`Err(e: PanicMessage)`) to avoid
+                // having handles outside the `bridge.enter(|| ...)` scope, and
+                // to catch panics that could happen while encoding the success.
+                //
+                // Note that panics should be impossible beyond this point, but
+                // this is defensively trying to avoid any accidental panicking
+                // reaching the `extern "C"` (which should `abort` but might not
+                // at the moment, so this is also potentially preventing UB).
+                buf.clear();
+                Ok::<_, ()>(output).encode(&mut buf, &mut ());
+            })
         })
     }))
     .map_err(PanicMessage::from)
diff --git a/library/proc_macro/src/bridge/mod.rs b/library/proc_macro/src/bridge/mod.rs
index 4e931569ef6..91d2d0f3ccc 100644
--- a/library/proc_macro/src/bridge/mod.rs
+++ b/library/proc_macro/src/bridge/mod.rs
@@ -151,9 +151,6 @@ macro_rules! with_api {
             },
             Span {
                 fn debug($self: $S::Span) -> String;
-                fn def_site() -> $S::Span;
-                fn call_site() -> $S::Span;
-                fn mixed_site() -> $S::Span;
                 fn source_file($self: $S::Span) -> $S::SourceFile;
                 fn parent($self: $S::Span) -> Option<$S::Span>;
                 fn source($self: $S::Span) -> $S::Span;
@@ -213,16 +210,15 @@ use buffer::Buffer;
 pub use rpc::PanicMessage;
 use rpc::{Decode, DecodeMut, Encode, Reader, Writer};
 
-/// An active connection between a server and a client.
-/// The server creates the bridge (`Bridge::run_server` in `server.rs`),
-/// then passes it to the client through the function pointer in the `run`
-/// field of `client::Client`. The client holds its copy of the `Bridge`
+/// Configuration for establishing an active connection between a server and a
+/// client.  The server creates the bridge config (`run_server` in `server.rs`),
+/// then passes it to the client through the function pointer in the `run` field
+/// of `client::Client`. The client constructs a local `Bridge` from the config
 /// in TLS during its execution (`Bridge::{enter, with}` in `client.rs`).
 #[repr(C)]
-pub struct Bridge<'a> {
-    /// Reusable buffer (only `clear`-ed, never shrunk), primarily
-    /// used for making requests, but also for passing input to client.
-    cached_buffer: Buffer,
+pub struct BridgeConfig<'a> {
+    /// Buffer used to pass initial input to the client.
+    input: Buffer,
 
     /// Server-side function that the client uses to make requests.
     dispatch: closure::Closure<'a, Buffer, Buffer>,
@@ -379,6 +375,25 @@ rpc_encode_decode!(
 );
 
 macro_rules! mark_compound {
+    (struct $name:ident <$($T:ident),+> { $($field:ident),* $(,)? }) => {
+        impl<$($T: Mark),+> Mark for $name <$($T),+> {
+            type Unmarked = $name <$($T::Unmarked),+>;
+            fn mark(unmarked: Self::Unmarked) -> Self {
+                $name {
+                    $($field: Mark::mark(unmarked.$field)),*
+                }
+            }
+        }
+
+        impl<$($T: Unmark),+> Unmark for $name <$($T),+> {
+            type Unmarked = $name <$($T::Unmarked),+>;
+            fn unmark(self) -> Self::Unmarked {
+                $name {
+                    $($field: Unmark::unmark(self.$field)),*
+                }
+            }
+        }
+    };
     (enum $name:ident <$($T:ident),+> { $($variant:ident $(($field:ident))?),* $(,)? }) => {
         impl<$($T: Mark),+> Mark for $name <$($T),+> {
             type Unmarked = $name <$($T::Unmarked),+>;
@@ -449,3 +464,16 @@ compound_traits!(
         Literal(tt),
     }
 );
+
+/// Context provided alongside the initial inputs for a macro expansion.
+/// Provides values such as spans which are used frequently to avoid RPC.
+#[derive(Clone)]
+struct ExpnContext<S> {
+    def_site: S,
+    call_site: S,
+    mixed_site: S,
+}
+
+compound_traits!(
+    struct ExpnContext<Sp> { def_site, call_site, mixed_site }
+);
diff --git a/library/proc_macro/src/bridge/selfless_reify.rs b/library/proc_macro/src/bridge/selfless_reify.rs
index 4ee4bb87c2b..907ad256e4b 100644
--- a/library/proc_macro/src/bridge/selfless_reify.rs
+++ b/library/proc_macro/src/bridge/selfless_reify.rs
@@ -75,9 +75,10 @@ macro_rules! define_reify_functions {
 define_reify_functions! {
     fn _reify_to_extern_c_fn_unary<A, R> for extern "C" fn(arg: A) -> R;
 
-    // HACK(eddyb) this abstraction is used with `for<'a> fn(Bridge<'a>) -> T`
-    // but that doesn't work with just `reify_to_extern_c_fn_unary` because of
-    // the `fn` pointer type being "higher-ranked" (i.e. the `for<'a>` binder).
-    // FIXME(eddyb) try to remove the lifetime from `Bridge`, that'd help.
-    fn reify_to_extern_c_fn_hrt_bridge<R> for extern "C" fn(bridge: super::Bridge<'_>) -> R;
+    // HACK(eddyb) this abstraction is used with `for<'a> fn(BridgeConfig<'a>)
+    // -> T` but that doesn't work with just `reify_to_extern_c_fn_unary`
+    // because of the `fn` pointer type being "higher-ranked" (i.e. the
+    // `for<'a>` binder).
+    // FIXME(eddyb) try to remove the lifetime from `BridgeConfig`, that'd help.
+    fn reify_to_extern_c_fn_hrt_bridge<R> for extern "C" fn(bridge: super::BridgeConfig<'_>) -> R;
 }
diff --git a/library/proc_macro/src/bridge/server.rs b/library/proc_macro/src/bridge/server.rs
index 3672299f18f..d9a6ce81e4e 100644
--- a/library/proc_macro/src/bridge/server.rs
+++ b/library/proc_macro/src/bridge/server.rs
@@ -30,6 +30,13 @@ macro_rules! associated_fn {
     ($($item:tt)*) => ($($item)*;)
 }
 
+/// Helper methods defined by `Server` types not invoked over RPC.
+pub trait Context: Types {
+    fn def_site(&mut self) -> Self::Span;
+    fn call_site(&mut self) -> Self::Span;
+    fn mixed_site(&mut self) -> Self::Span;
+}
+
 macro_rules! declare_server_traits {
     ($($name:ident {
         $(fn $method:ident($($arg:ident: $arg_ty:ty),* $(,)?) $(-> $ret_ty:ty)?;)*
@@ -38,14 +45,26 @@ macro_rules! declare_server_traits {
             $(associated_fn!(fn $method(&mut self, $($arg: $arg_ty),*) $(-> $ret_ty)?);)*
         })*
 
-        pub trait Server: Types $(+ $name)* {}
-        impl<S: Types $(+ $name)*> Server for S {}
+        pub trait Server: Types + Context $(+ $name)* {}
+        impl<S: Types + Context $(+ $name)*> Server for S {}
     }
 }
 with_api!(Self, self_, declare_server_traits);
 
 pub(super) struct MarkedTypes<S: Types>(S);
 
+impl<S: Context> Context for MarkedTypes<S> {
+    fn def_site(&mut self) -> Self::Span {
+        <_>::mark(Context::def_site(&mut self.0))
+    }
+    fn call_site(&mut self) -> Self::Span {
+        <_>::mark(Context::call_site(&mut self.0))
+    }
+    fn mixed_site(&mut self) -> Self::Span {
+        <_>::mark(Context::mixed_site(&mut self.0))
+    }
+}
+
 macro_rules! define_mark_types_impls {
     ($($name:ident {
         $(fn $method:ident($($arg:ident: $arg_ty:ty),* $(,)?) $(-> $ret_ty:ty)?;)*
@@ -120,7 +139,7 @@ pub trait ExecutionStrategy {
         &self,
         dispatcher: &mut impl DispatcherTrait,
         input: Buffer,
-        run_client: extern "C" fn(Bridge<'_>) -> Buffer,
+        run_client: extern "C" fn(BridgeConfig<'_>) -> Buffer,
         force_show_panics: bool,
     ) -> Buffer;
 }
@@ -132,13 +151,13 @@ impl ExecutionStrategy for SameThread {
         &self,
         dispatcher: &mut impl DispatcherTrait,
         input: Buffer,
-        run_client: extern "C" fn(Bridge<'_>) -> Buffer,
+        run_client: extern "C" fn(BridgeConfig<'_>) -> Buffer,
         force_show_panics: bool,
     ) -> Buffer {
         let mut dispatch = |buf| dispatcher.dispatch(buf);
 
-        run_client(Bridge {
-            cached_buffer: input,
+        run_client(BridgeConfig {
+            input,
             dispatch: (&mut dispatch).into(),
             force_show_panics,
             _marker: marker::PhantomData,
@@ -156,7 +175,7 @@ impl ExecutionStrategy for CrossThread1 {
         &self,
         dispatcher: &mut impl DispatcherTrait,
         input: Buffer,
-        run_client: extern "C" fn(Bridge<'_>) -> Buffer,
+        run_client: extern "C" fn(BridgeConfig<'_>) -> Buffer,
         force_show_panics: bool,
     ) -> Buffer {
         use std::sync::mpsc::channel;
@@ -170,8 +189,8 @@ impl ExecutionStrategy for CrossThread1 {
                 res_rx.recv().unwrap()
             };
 
-            run_client(Bridge {
-                cached_buffer: input,
+            run_client(BridgeConfig {
+                input,
                 dispatch: (&mut dispatch).into(),
                 force_show_panics,
                 _marker: marker::PhantomData,
@@ -193,7 +212,7 @@ impl ExecutionStrategy for CrossThread2 {
         &self,
         dispatcher: &mut impl DispatcherTrait,
         input: Buffer,
-        run_client: extern "C" fn(Bridge<'_>) -> Buffer,
+        run_client: extern "C" fn(BridgeConfig<'_>) -> Buffer,
         force_show_panics: bool,
     ) -> Buffer {
         use std::sync::{Arc, Mutex};
@@ -219,8 +238,8 @@ impl ExecutionStrategy for CrossThread2 {
                 }
             };
 
-            let r = run_client(Bridge {
-                cached_buffer: input,
+            let r = run_client(BridgeConfig {
+                input,
                 dispatch: (&mut dispatch).into(),
                 force_show_panics,
                 _marker: marker::PhantomData,
@@ -258,14 +277,20 @@ fn run_server<
     handle_counters: &'static client::HandleCounters,
     server: S,
     input: I,
-    run_client: extern "C" fn(Bridge<'_>) -> Buffer,
+    run_client: extern "C" fn(BridgeConfig<'_>) -> Buffer,
     force_show_panics: bool,
 ) -> Result<O, PanicMessage> {
     let mut dispatcher =
         Dispatcher { handle_store: HandleStore::new(handle_counters), server: MarkedTypes(server) };
 
+    let expn_context = ExpnContext {
+        def_site: dispatcher.server.def_site(),
+        call_site: dispatcher.server.call_site(),
+        mixed_site: dispatcher.server.mixed_site(),
+    };
+
     let mut buf = Buffer::new();
-    input.encode(&mut buf, &mut dispatcher.handle_store);
+    (input, expn_context).encode(&mut buf, &mut dispatcher.handle_store);
 
     buf = strategy.run_bridge_and_client(&mut dispatcher, buf, run_client, force_show_panics);
 
diff --git a/library/proc_macro/src/lib.rs b/library/proc_macro/src/lib.rs
index 5e1289ec79d..771ee50e138 100644
--- a/library/proc_macro/src/lib.rs
+++ b/library/proc_macro/src/lib.rs
@@ -60,7 +60,7 @@ use std::{error, fmt, iter};
 /// inside of a procedural macro, false if invoked from any other binary.
 #[stable(feature = "proc_macro_is_available", since = "1.57.0")]
 pub fn is_available() -> bool {
-    bridge::Bridge::is_available()
+    bridge::client::is_available()
 }
 
 /// The main type provided by this crate, representing an abstract stream of