Skip to content

Commit afc36cc

Browse files
mystormystor
committed
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.
1 parent d5f2ff5 commit afc36cc

File tree

5 files changed

+164
-102
lines changed

5 files changed

+164
-102
lines changed

‎compiler/rustc_expand/src/proc_macro_server.rs

+18-12
Original file line numberDiff line numberDiff line change
@@ -388,7 +388,10 @@ impl<'a> Rustc<'a> {
388388
}
389389

390390
fn lit(&mut self, kind: token::LitKind, symbol: Symbol, suffix: Option<Symbol>) -> Literal {
391-
Literal { lit: token::Lit::new(kind, symbol, suffix), span: server::Span::call_site(self) }
391+
Literal {
392+
lit: token::Lit::new(kind, symbol, suffix),
393+
span: server::Context::call_site(self),
394+
}
392395
}
393396
}
394397

@@ -484,7 +487,7 @@ impl server::Group for Rustc<'_> {
484487
Group {
485488
delimiter,
486489
stream,
487-
span: DelimSpan::from_single(server::Span::call_site(self)),
490+
span: DelimSpan::from_single(server::Context::call_site(self)),
488491
flatten: false,
489492
}
490493
}
@@ -510,7 +513,7 @@ impl server::Group for Rustc<'_> {
510513

511514
impl server::Punct for Rustc<'_> {
512515
fn new(&mut self, ch: char, spacing: Spacing) -> Self::Punct {
513-
Punct::new(ch, spacing == Spacing::Joint, server::Span::call_site(self))
516+
Punct::new(ch, spacing == Spacing::Joint, server::Context::call_site(self))
514517
}
515518
fn as_char(&mut self, punct: Self::Punct) -> char {
516519
punct.ch
@@ -712,15 +715,6 @@ impl server::Span for Rustc<'_> {
712715
format!("{:?} bytes({}..{})", span.ctxt(), span.lo().0, span.hi().0)
713716
}
714717
}
715-
fn def_site(&mut self) -> Self::Span {
716-
self.def_site
717-
}
718-
fn call_site(&mut self) -> Self::Span {
719-
self.call_site
720-
}
721-
fn mixed_site(&mut self) -> Self::Span {
722-
self.mixed_site
723-
}
724718
fn source_file(&mut self, span: Self::Span) -> Self::SourceFile {
725719
self.sess.source_map().lookup_char_pos(span.lo()).file
726720
}
@@ -801,6 +795,18 @@ impl server::Span for Rustc<'_> {
801795
}
802796
}
803797

798+
impl server::Context for Rustc<'_> {
799+
fn def_site(&mut self) -> Self::Span {
800+
self.def_site
801+
}
802+
fn call_site(&mut self) -> Self::Span {
803+
self.call_site
804+
}
805+
fn mixed_site(&mut self) -> Self::Span {
806+
self.mixed_site
807+
}
808+
}
809+
804810
// See issue #74616 for details
805811
fn ident_name_compatibility_hack(
806812
nt: &Nonterminal,

‎library/proc_macro/src/bridge/client.rs

+87-57
Original file line numberDiff line numberDiff line change
@@ -222,6 +222,20 @@ impl Clone for SourceFile {
222222
}
223223
}
224224

225+
impl Span {
226+
pub(crate) fn def_site() -> Span {
227+
Bridge::with(|bridge| bridge.context.def_site)
228+
}
229+
230+
pub(crate) fn call_site() -> Span {
231+
Bridge::with(|bridge| bridge.context.call_site)
232+
}
233+
234+
pub(crate) fn mixed_site() -> Span {
235+
Bridge::with(|bridge| bridge.context.mixed_site)
236+
}
237+
}
238+
225239
impl fmt::Debug for Span {
226240
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
227241
f.write_str(&self.debug())
@@ -255,6 +269,21 @@ macro_rules! define_client_side {
255269
}
256270
with_api!(self, self, define_client_side);
257271

272+
struct Bridge<'a> {
273+
/// Reusable buffer (only `clear`-ed, never shrunk), primarily
274+
/// used for making requests.
275+
cached_buffer: Buffer<u8>,
276+
277+
/// Server-side function that the client uses to make requests.
278+
dispatch: closure::Closure<'a, Buffer<u8>, Buffer<u8>>,
279+
280+
/// Provided context for this macro expansion.
281+
context: ExpnContext<Span>,
282+
}
283+
284+
impl<'a> !Send for Bridge<'a> {}
285+
impl<'a> !Sync for Bridge<'a> {}
286+
258287
enum BridgeState<'a> {
259288
/// No server is currently connected to this client.
260289
NotConnected,
@@ -297,34 +326,6 @@ impl BridgeState<'_> {
297326
}
298327

299328
impl Bridge<'_> {
300-
pub(crate) fn is_available() -> bool {
301-
BridgeState::with(|state| match state {
302-
BridgeState::Connected(_) | BridgeState::InUse => true,
303-
BridgeState::NotConnected => false,
304-
})
305-
}
306-
307-
fn enter<R>(self, f: impl FnOnce() -> R) -> R {
308-
let force_show_panics = self.force_show_panics;
309-
// Hide the default panic output within `proc_macro` expansions.
310-
// NB. the server can't do this because it may use a different libstd.
311-
static HIDE_PANICS_DURING_EXPANSION: Once = Once::new();
312-
HIDE_PANICS_DURING_EXPANSION.call_once(|| {
313-
let prev = panic::take_hook();
314-
panic::set_hook(Box::new(move |info| {
315-
let show = BridgeState::with(|state| match state {
316-
BridgeState::NotConnected => true,
317-
BridgeState::Connected(_) | BridgeState::InUse => force_show_panics,
318-
});
319-
if show {
320-
prev(info)
321-
}
322-
}));
323-
});
324-
325-
BRIDGE_STATE.with(|state| state.set(BridgeState::Connected(self), f))
326-
}
327-
328329
fn with<R>(f: impl FnOnce(&mut Bridge<'_>) -> R) -> R {
329330
BridgeState::with(|state| match state {
330331
BridgeState::NotConnected => {
@@ -338,6 +339,13 @@ impl Bridge<'_> {
338339
}
339340
}
340341

342+
pub(crate) fn is_available() -> bool {
343+
BridgeState::with(|state| match state {
344+
BridgeState::Connected(_) | BridgeState::InUse => true,
345+
BridgeState::NotConnected => false,
346+
})
347+
}
348+
341349
/// A client-side "global object" (usually a function pointer),
342350
/// which may be using a different `proc_macro` from the one
343351
/// used by the server, but can be interacted with compatibly.
@@ -352,44 +360,66 @@ pub struct Client<F> {
352360
// FIXME(eddyb) use a reference to the `static COUNTERS`, instead of
353361
// a wrapper `fn` pointer, once `const fn` can reference `static`s.
354362
pub(super) get_handle_counters: extern "C" fn() -> &'static HandleCounters,
355-
pub(super) run: extern "C" fn(Bridge<'_>, F) -> Buffer<u8>,
363+
pub(super) run: extern "C" fn(BridgeConfig<'_>, F) -> Buffer<u8>,
356364
pub(super) f: F,
357365
}
358366

367+
fn maybe_install_panic_hook(force_show_panics: bool) {
368+
// Hide the default panic output within `proc_macro` expansions.
369+
// NB. the server can't do this because it may use a different libstd.
370+
static HIDE_PANICS_DURING_EXPANSION: Once = Once::new();
371+
HIDE_PANICS_DURING_EXPANSION.call_once(|| {
372+
let prev = panic::take_hook();
373+
panic::set_hook(Box::new(move |info| {
374+
let show = BridgeState::with(|state| match state {
375+
BridgeState::NotConnected => true,
376+
BridgeState::Connected(_) | BridgeState::InUse => force_show_panics,
377+
});
378+
if show {
379+
prev(info)
380+
}
381+
}));
382+
});
383+
}
384+
359385
/// Client-side helper for handling client panics, entering the bridge,
360386
/// deserializing input and serializing output.
361387
// FIXME(eddyb) maybe replace `Bridge::enter` with this?
362388
fn run_client<A: for<'a, 's> DecodeMut<'a, 's, ()>, R: Encode<()>>(
363-
mut bridge: Bridge<'_>,
389+
config: BridgeConfig<'_>,
364390
f: impl FnOnce(A) -> R,
365391
) -> Buffer<u8> {
366-
// The initial `cached_buffer` contains the input.
367-
let mut b = bridge.cached_buffer.take();
392+
let BridgeConfig { input: mut b, dispatch, force_show_panics } = config;
368393

369394
panic::catch_unwind(panic::AssertUnwindSafe(|| {
370-
bridge.enter(|| {
371-
let reader = &mut &b[..];
372-
let input = A::decode(reader, &mut ());
373-
374-
// Put the `cached_buffer` back in the `Bridge`, for requests.
375-
Bridge::with(|bridge| bridge.cached_buffer = b.take());
376-
377-
let output = f(input);
378-
379-
// Take the `cached_buffer` back out, for the output value.
380-
b = Bridge::with(|bridge| bridge.cached_buffer.take());
381-
382-
// HACK(eddyb) Separate encoding a success value (`Ok(output)`)
383-
// from encoding a panic (`Err(e: PanicMessage)`) to avoid
384-
// having handles outside the `bridge.enter(|| ...)` scope, and
385-
// to catch panics that could happen while encoding the success.
386-
//
387-
// Note that panics should be impossible beyond this point, but
388-
// this is defensively trying to avoid any accidental panicking
389-
// reaching the `extern "C"` (which should `abort` but may not
390-
// at the moment, so this is also potentially preventing UB).
391-
b.clear();
392-
Ok::<_, ()>(output).encode(&mut b, &mut ());
395+
maybe_install_panic_hook(force_show_panics);
396+
397+
let reader = &mut &b[..];
398+
let (input, context) = <(A, ExpnContext<Span>)>::decode(reader, &mut ());
399+
400+
// Put the buffer we used for input back in the `Bridge` for requests.
401+
let new_state =
402+
BridgeState::Connected(Bridge { cached_buffer: b.take(), dispatch, context });
403+
404+
BRIDGE_STATE.with(|state| {
405+
state.set(new_state, || {
406+
let output = f(input);
407+
408+
// Take the `cached_buffer` back out, for the output value.
409+
b = Bridge::with(|bridge| bridge.cached_buffer.take());
410+
411+
// HACK(eddyb) Separate encoding a success value (`Ok(output)`)
412+
// from encoding a panic (`Err(e: PanicMessage)`) to avoid
413+
// having handles outside the `bridge.enter(|| ...)` scope, and
414+
// to catch panics that could happen while encoding the success.
415+
//
416+
// Note that panics should be impossible beyond this point, but
417+
// this is defensively trying to avoid any accidental panicking
418+
// reaching the `extern "C"` (which should `abort` but may not
419+
// at the moment, so this is also potentially preventing UB).
420+
b.clear();
421+
Ok::<_, ()>(output).encode(&mut b, &mut ());
422+
})
393423
})
394424
}))
395425
.map_err(PanicMessage::from)
@@ -404,7 +434,7 @@ impl Client<fn(crate::TokenStream) -> crate::TokenStream> {
404434
#[rustc_allow_const_fn_unstable(const_fn)]
405435
pub const fn expand1(f: fn(crate::TokenStream) -> crate::TokenStream) -> Self {
406436
extern "C" fn run(
407-
bridge: Bridge<'_>,
437+
bridge: BridgeConfig<'_>,
408438
f: impl FnOnce(crate::TokenStream) -> crate::TokenStream,
409439
) -> Buffer<u8> {
410440
run_client(bridge, |input| f(crate::TokenStream(input)).0)
@@ -419,7 +449,7 @@ impl Client<fn(crate::TokenStream, crate::TokenStream) -> crate::TokenStream> {
419449
f: fn(crate::TokenStream, crate::TokenStream) -> crate::TokenStream,
420450
) -> Self {
421451
extern "C" fn run(
422-
bridge: Bridge<'_>,
452+
bridge: BridgeConfig<'_>,
423453
f: impl FnOnce(crate::TokenStream, crate::TokenStream) -> crate::TokenStream,
424454
) -> Buffer<u8> {
425455
run_client(bridge, |(input, input2)| {

‎library/proc_macro/src/bridge/mod.rs

+22-13
Original file line numberDiff line numberDiff line change
@@ -152,9 +152,6 @@ macro_rules! with_api {
152152
},
153153
Span {
154154
fn debug($self: $S::Span) -> String;
155-
fn def_site() -> $S::Span;
156-
fn call_site() -> $S::Span;
157-
fn mixed_site() -> $S::Span;
158155
fn source_file($self: $S::Span) -> $S::SourceFile;
159156
fn parent($self: $S::Span) -> Option<$S::Span>;
160157
fn source($self: $S::Span) -> $S::Span;
@@ -210,16 +207,15 @@ use buffer::Buffer;
210207
pub use rpc::PanicMessage;
211208
use rpc::{Decode, DecodeMut, Encode, Reader, Writer};
212209

213-
/// An active connection between a server and a client.
214-
/// The server creates the bridge (`Bridge::run_server` in `server.rs`),
215-
/// then passes it to the client through the function pointer in the `run`
216-
/// field of `client::Client`. The client holds its copy of the `Bridge`
210+
/// Configuration for establishing an active connection between a server and a
211+
/// client. The server creates the bridge config (`run_server` in `server.rs`),
212+
/// then passes it to the client through the function pointer in the `run` field
213+
/// of `client::Client`. The client constructs a local `Bridge` from the config
217214
/// in TLS during its execution (`Bridge::{enter, with}` in `client.rs`).
218215
#[repr(C)]
219-
pub struct Bridge<'a> {
220-
/// Reusable buffer (only `clear`-ed, never shrunk), primarily
221-
/// used for making requests, but also for passing input to client.
222-
cached_buffer: Buffer<u8>,
216+
pub struct BridgeConfig<'a> {
217+
/// Buffer used to pass initial input to the client.
218+
input: Buffer<u8>,
223219

224220
/// Server-side function that the client uses to make requests.
225221
dispatch: closure::Closure<'a, Buffer<u8>, Buffer<u8>>,
@@ -228,8 +224,8 @@ pub struct Bridge<'a> {
228224
force_show_panics: bool,
229225
}
230226

231-
impl<'a> !Sync for Bridge<'a> {}
232-
impl<'a> !Send for Bridge<'a> {}
227+
impl<'a> !Sync for BridgeConfig<'a> {}
228+
impl<'a> !Send for BridgeConfig<'a> {}
233229

234230
#[forbid(unsafe_code)]
235231
#[allow(non_camel_case_types)]
@@ -469,3 +465,16 @@ compound_traits!(
469465
Literal(tt),
470466
}
471467
);
468+
469+
/// Context provided alongside the initial inputs for a macro expansion.
470+
/// Provides values such as spans which are used frequently to avoid RPC.
471+
#[derive(Clone)]
472+
struct ExpnContext<S> {
473+
def_site: S,
474+
call_site: S,
475+
mixed_site: S,
476+
}
477+
478+
compound_traits!(
479+
struct ExpnContext<Sp> { def_site, call_site, mixed_site }
480+
);

0 commit comments

Comments
 (0)