Skip to content

Commit

Permalink
proc_macro: cache static spans in client's thread-local state
Browse files Browse the repository at this point in the history 
This greatly improves the performance of the very frequently called
`call_site()` macro when running in a cross-thread configuration.
  • Loading branch information
@mystor
mystor committed Jul 2, 2021
1 parent d5f2ff5 commit afc36cc
Show file tree
Hide file tree
Showing 5 changed files with 164 additions and 102 deletions.
30 changes: 18 additions & 12 deletions compiler/rustc_expand/src/proc_macro_server.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -388,7 +388,10 @@ impl<'a> Rustc<'a> {
}

fn lit(&mut self, kind: token::LitKind, symbol: Symbol, suffix: Option<Symbol>) -> Literal {
Literal { lit: token::Lit::new(kind, symbol, suffix), span: server::Span::call_site(self) }
Literal {
lit: token::Lit::new(kind, symbol, suffix),
span: server::Context::call_site(self),
}
}
}

Expand Down Expand Up @@ -484,7 +487,7 @@ impl server::Group for Rustc<'_> {
Group {
delimiter,
stream,
span: DelimSpan::from_single(server::Span::call_site(self)),
span: DelimSpan::from_single(server::Context::call_site(self)),
flatten: false,
}
}
Expand All @@ -510,7 +513,7 @@ impl server::Group for Rustc<'_> {

impl server::Punct for Rustc<'_> {
fn new(&mut self, ch: char, spacing: Spacing) -> Self::Punct {
Punct::new(ch, spacing == Spacing::Joint, server::Span::call_site(self))
Punct::new(ch, spacing == Spacing::Joint, server::Context::call_site(self))
}
fn as_char(&mut self, punct: Self::Punct) -> char {
punct.ch
Expand Down Expand Up @@ -712,15 +715,6 @@ impl server::Span for Rustc<'_> {
format!("{:?} bytes({}..{})", span.ctxt(), span.lo().0, span.hi().0)
}
}
fn def_site(&mut self) -> Self::Span {
self.def_site
}
fn call_site(&mut self) -> Self::Span {
self.call_site
}
fn mixed_site(&mut self) -> Self::Span {
self.mixed_site
}
fn source_file(&mut self, span: Self::Span) -> Self::SourceFile {
self.sess.source_map().lookup_char_pos(span.lo()).file
}
Expand Down Expand Up @@ -801,6 +795,18 @@ impl server::Span for Rustc<'_> {
}
}

impl server::Context for Rustc<'_> {
fn def_site(&mut self) -> Self::Span {
self.def_site
}
fn call_site(&mut self) -> Self::Span {
self.call_site
}
fn mixed_site(&mut self) -> Self::Span {
self.mixed_site
}
}

// See issue #74616 for details
fn ident_name_compatibility_hack(
nt: &Nonterminal,
Expand Down
144 changes: 87 additions & 57 deletions library/proc_macro/src/bridge/client.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -222,6 +222,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())
Expand Down Expand Up @@ -255,6 +269,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<u8>,

/// Server-side function that the client uses to make requests.
dispatch: closure::Closure<'a, Buffer<u8>, Buffer<u8>>,

/// 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,
Expand Down Expand Up @@ -297,34 +326,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 => {
Expand All @@ -338,6 +339,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 "global object" (usually a function pointer),
/// which may be using a different `proc_macro` from the one
/// used by the server, but can be interacted with compatibly.
Expand All @@ -352,44 +360,66 @@ pub struct Client<F> {
// FIXME(eddyb) use a reference to the `static COUNTERS`, instead of
// 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<'_>, F) -> Buffer<u8>,
pub(super) run: extern "C" fn(BridgeConfig<'_>, F) -> Buffer<u8>,
pub(super) f: F,
}

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<u8> {
// The initial `cached_buffer` contains the input.
let mut b = bridge.cached_buffer.take();
let BridgeConfig { input: mut b, dispatch, force_show_panics } = config;

panic::catch_unwind(panic::AssertUnwindSafe(|| {
bridge.enter(|| {
let reader = &mut &b[..];
let input = A::decode(reader, &mut ());

// Put the `cached_buffer` back in the `Bridge`, for requests.
Bridge::with(|bridge| bridge.cached_buffer = b.take());

let output = f(input);

// Take the `cached_buffer` back out, for the output value.
b = 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 may not
// at the moment, so this is also potentially preventing UB).
b.clear();
Ok::<_, ()>(output).encode(&mut b, &mut ());
maybe_install_panic_hook(force_show_panics);

let reader = &mut &b[..];
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: b.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.
b = 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 may not
// at the moment, so this is also potentially preventing UB).
b.clear();
Ok::<_, ()>(output).encode(&mut b, &mut ());
})
})
}))
.map_err(PanicMessage::from)
Expand All @@ -404,7 +434,7 @@ impl Client<fn(crate::TokenStream) -> crate::TokenStream> {
#[rustc_allow_const_fn_unstable(const_fn)]
pub const fn expand1(f: fn(crate::TokenStream) -> crate::TokenStream) -> Self {
extern "C" fn run(
bridge: Bridge<'_>,
bridge: BridgeConfig<'_>,
f: impl FnOnce(crate::TokenStream) -> crate::TokenStream,
) -> Buffer<u8> {
run_client(bridge, |input| f(crate::TokenStream(input)).0)
Expand All @@ -419,7 +449,7 @@ impl Client<fn(crate::TokenStream, crate::TokenStream) -> crate::TokenStream> {
f: fn(crate::TokenStream, crate::TokenStream) -> crate::TokenStream,
) -> Self {
extern "C" fn run(
bridge: Bridge<'_>,
bridge: BridgeConfig<'_>,
f: impl FnOnce(crate::TokenStream, crate::TokenStream) -> crate::TokenStream,
) -> Buffer<u8> {
run_client(bridge, |(input, input2)| {
Expand Down
35 changes: 22 additions & 13 deletions library/proc_macro/src/bridge/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -152,9 +152,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;
Expand Down Expand Up @@ -210,16 +207,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<u8>,
pub struct BridgeConfig<'a> {
/// Buffer used to pass initial input to the client.
input: Buffer<u8>,

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

impl<'a> !Sync for Bridge<'a> {}
impl<'a> !Send for Bridge<'a> {}
impl<'a> !Sync for BridgeConfig<'a> {}
impl<'a> !Send for BridgeConfig<'a> {}

#[forbid(unsafe_code)]
#[allow(non_camel_case_types)]
Expand Down Expand Up @@ -469,3 +465,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 }
);
Loading

0 comments on commit afc36cc

Please sign in to comment.