Skip to content

Commit b805b59

Browse files
borsbors
committed
Auto merge of #98187 - mystor:fast_span_call_site, r=eddyb
proc_macro/bridge: cache static spans in proc_macro's client thread-local state This is the second part of rust-lang/rust#86822, split off as requested in rust-lang/rust#86822 (review). This patch removes the RPC calls required for the very common operations of `Span::call_site()`, `Span::def_site()` and `Span::mixed_site()`. Some notes: This part is one of the ones I don't love as a final solution from a design standpoint, because I don't like how the spans are serialized immediately at macro invocation. I think a more elegant solution might've been to reserve special IDs for `call_site`, `def_site`, and `mixed_site` at compile time (either starting at 1 or from `u32::MAX`) and making reading a Span handle automatically map these IDs to the relevant values, rather than doing extra serialization. This would also have an advantage for potential future work to allow `proc_macro` to operate more independently from the compiler (e.g. to reduce the necessity of `proc-macro2`), as methods like `Span::call_site()` could be made to function without access to the compiler backend. That was unfortunately tricky to do at the time, as this was the first part I wrote of the patches. After the later part (#98188, #98189), the other uses of `InternedStore` are removed meaning that a custom serialization strategy for `Span` is easier to implement. If we want to go that path, we'll still need the majority of the work to split the bridge object and introduce the `Context` trait for free methods, and it will be easier to do after `Span` is the only user of `InternedStore` (after #98189).
2 parents 9b00fb8 + cc347df commit b805b59

File tree

5 files changed

+154
-86
lines changed

5 files changed

+154
-86
lines changed

‎proc_macro/src/bridge/client.rs

+85-55
Original file line numberDiff line numberDiff line change
@@ -230,6 +230,20 @@ impl Clone for SourceFile {
230230
}
231231
}
232232

233+
impl Span {
234+
pub(crate) fn def_site() -> Span {
235+
Bridge::with(|bridge| bridge.globals.def_site)
236+
}
237+
238+
pub(crate) fn call_site() -> Span {
239+
Bridge::with(|bridge| bridge.globals.call_site)
240+
}
241+
242+
pub(crate) fn mixed_site() -> Span {
243+
Bridge::with(|bridge| bridge.globals.mixed_site)
244+
}
245+
}
246+
233247
impl fmt::Debug for Span {
234248
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
235249
f.write_str(&self.debug())
@@ -263,6 +277,21 @@ macro_rules! define_client_side {
263277
}
264278
with_api!(self, self, define_client_side);
265279

280+
struct Bridge<'a> {
281+
/// Reusable buffer (only `clear`-ed, never shrunk), primarily
282+
/// used for making requests.
283+
cached_buffer: Buffer,
284+
285+
/// Server-side function that the client uses to make requests.
286+
dispatch: closure::Closure<'a, Buffer, Buffer>,
287+
288+
/// Provided globals for this macro expansion.
289+
globals: ExpnGlobals<Span>,
290+
}
291+
292+
impl<'a> !Send for Bridge<'a> {}
293+
impl<'a> !Sync for Bridge<'a> {}
294+
266295
enum BridgeState<'a> {
267296
/// No server is currently connected to this client.
268297
NotConnected,
@@ -305,34 +334,6 @@ impl BridgeState<'_> {
305334
}
306335

307336
impl Bridge<'_> {
308-
pub(crate) fn is_available() -> bool {
309-
BridgeState::with(|state| match state {
310-
BridgeState::Connected(_) | BridgeState::InUse => true,
311-
BridgeState::NotConnected => false,
312-
})
313-
}
314-
315-
fn enter<R>(self, f: impl FnOnce() -> R) -> R {
316-
let force_show_panics = self.force_show_panics;
317-
// Hide the default panic output within `proc_macro` expansions.
318-
// NB. the server can't do this because it may use a different libstd.
319-
static HIDE_PANICS_DURING_EXPANSION: Once = Once::new();
320-
HIDE_PANICS_DURING_EXPANSION.call_once(|| {
321-
let prev = panic::take_hook();
322-
panic::set_hook(Box::new(move |info| {
323-
let show = BridgeState::with(|state| match state {
324-
BridgeState::NotConnected => true,
325-
BridgeState::Connected(_) | BridgeState::InUse => force_show_panics,
326-
});
327-
if show {
328-
prev(info)
329-
}
330-
}));
331-
});
332-
333-
BRIDGE_STATE.with(|state| state.set(BridgeState::Connected(self), f))
334-
}
335-
336337
fn with<R>(f: impl FnOnce(&mut Bridge<'_>) -> R) -> R {
337338
BridgeState::with(|state| match state {
338339
BridgeState::NotConnected => {
@@ -346,6 +347,13 @@ impl Bridge<'_> {
346347
}
347348
}
348349

350+
pub(crate) fn is_available() -> bool {
351+
BridgeState::with(|state| match state {
352+
BridgeState::Connected(_) | BridgeState::InUse => true,
353+
BridgeState::NotConnected => false,
354+
})
355+
}
356+
349357
/// A client-side RPC entry-point, which may be using a different `proc_macro`
350358
/// from the one used by the server, but can be invoked compatibly.
351359
///
@@ -363,7 +371,7 @@ pub struct Client<I, O> {
363371
// a wrapper `fn` pointer, once `const fn` can reference `static`s.
364372
pub(super) get_handle_counters: extern "C" fn() -> &'static HandleCounters,
365373

366-
pub(super) run: extern "C" fn(Bridge<'_>) -> Buffer,
374+
pub(super) run: extern "C" fn(BridgeConfig<'_>) -> Buffer,
367375

368376
pub(super) _marker: PhantomData<fn(I) -> O>,
369377
}
@@ -375,40 +383,62 @@ impl<I, O> Clone for Client<I, O> {
375383
}
376384
}
377385

386+
fn maybe_install_panic_hook(force_show_panics: bool) {
387+
// Hide the default panic output within `proc_macro` expansions.
388+
// NB. the server can't do this because it may use a different libstd.
389+
static HIDE_PANICS_DURING_EXPANSION: Once = Once::new();
390+
HIDE_PANICS_DURING_EXPANSION.call_once(|| {
391+
let prev = panic::take_hook();
392+
panic::set_hook(Box::new(move |info| {
393+
let show = BridgeState::with(|state| match state {
394+
BridgeState::NotConnected => true,
395+
BridgeState::Connected(_) | BridgeState::InUse => force_show_panics,
396+
});
397+
if show {
398+
prev(info)
399+
}
400+
}));
401+
});
402+
}
403+
378404
/// Client-side helper for handling client panics, entering the bridge,
379405
/// deserializing input and serializing output.
380406
// FIXME(eddyb) maybe replace `Bridge::enter` with this?
381407
fn run_client<A: for<'a, 's> DecodeMut<'a, 's, ()>, R: Encode<()>>(
382-
mut bridge: Bridge<'_>,
408+
config: BridgeConfig<'_>,
383409
f: impl FnOnce(A) -> R,
384410
) -> Buffer {
385-
// The initial `cached_buffer` contains the input.
386-
let mut buf = bridge.cached_buffer.take();
411+
let BridgeConfig { input: mut buf, dispatch, force_show_panics, .. } = config;
387412

388413
panic::catch_unwind(panic::AssertUnwindSafe(|| {
389-
bridge.enter(|| {
390-
let reader = &mut &buf[..];
391-
let input = A::decode(reader, &mut ());
392-
393-
// Put the `cached_buffer` back in the `Bridge`, for requests.
394-
Bridge::with(|bridge| bridge.cached_buffer = buf.take());
395-
396-
let output = f(input);
397-
398-
// Take the `cached_buffer` back out, for the output value.
399-
buf = Bridge::with(|bridge| bridge.cached_buffer.take());
400-
401-
// HACK(eddyb) Separate encoding a success value (`Ok(output)`)
402-
// from encoding a panic (`Err(e: PanicMessage)`) to avoid
403-
// having handles outside the `bridge.enter(|| ...)` scope, and
404-
// to catch panics that could happen while encoding the success.
405-
//
406-
// Note that panics should be impossible beyond this point, but
407-
// this is defensively trying to avoid any accidental panicking
408-
// reaching the `extern "C"` (which should `abort` but might not
409-
// at the moment, so this is also potentially preventing UB).
410-
buf.clear();
411-
Ok::<_, ()>(output).encode(&mut buf, &mut ());
414+
maybe_install_panic_hook(force_show_panics);
415+
416+
let reader = &mut &buf[..];
417+
let (globals, input) = <(ExpnGlobals<Span>, A)>::decode(reader, &mut ());
418+
419+
// Put the buffer we used for input back in the `Bridge` for requests.
420+
let new_state =
421+
BridgeState::Connected(Bridge { cached_buffer: buf.take(), dispatch, globals });
422+
423+
BRIDGE_STATE.with(|state| {
424+
state.set(new_state, || {
425+
let output = f(input);
426+
427+
// Take the `cached_buffer` back out, for the output value.
428+
buf = Bridge::with(|bridge| bridge.cached_buffer.take());
429+
430+
// HACK(eddyb) Separate encoding a success value (`Ok(output)`)
431+
// from encoding a panic (`Err(e: PanicMessage)`) to avoid
432+
// having handles outside the `bridge.enter(|| ...)` scope, and
433+
// to catch panics that could happen while encoding the success.
434+
//
435+
// Note that panics should be impossible beyond this point, but
436+
// this is defensively trying to avoid any accidental panicking
437+
// reaching the `extern "C"` (which should `abort` but might not
438+
// at the moment, so this is also potentially preventing UB).
439+
buf.clear();
440+
Ok::<_, ()>(output).encode(&mut buf, &mut ());
441+
})
412442
})
413443
}))
414444
.map_err(PanicMessage::from)

‎proc_macro/src/bridge/mod.rs

+39-11
Original file line numberDiff line numberDiff line change
@@ -151,9 +151,6 @@ macro_rules! with_api {
151151
},
152152
Span {
153153
fn debug($self: $S::Span) -> String;
154-
fn def_site() -> $S::Span;
155-
fn call_site() -> $S::Span;
156-
fn mixed_site() -> $S::Span;
157154
fn source_file($self: $S::Span) -> $S::SourceFile;
158155
fn parent($self: $S::Span) -> Option<$S::Span>;
159156
fn source($self: $S::Span) -> $S::Span;
@@ -213,16 +210,15 @@ use buffer::Buffer;
213210
pub use rpc::PanicMessage;
214211
use rpc::{Decode, DecodeMut, Encode, Reader, Writer};
215212

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

227223
/// Server-side function that the client uses to make requests.
228224
dispatch: closure::Closure<'a, Buffer, Buffer>,
@@ -379,6 +375,25 @@ rpc_encode_decode!(
379375
);
380376

381377
macro_rules! mark_compound {
378+
(struct $name:ident <$($T:ident),+> { $($field:ident),* $(,)? }) => {
379+
impl<$($T: Mark),+> Mark for $name <$($T),+> {
380+
type Unmarked = $name <$($T::Unmarked),+>;
381+
fn mark(unmarked: Self::Unmarked) -> Self {
382+
$name {
383+
$($field: Mark::mark(unmarked.$field)),*
384+
}
385+
}
386+
}
387+
388+
impl<$($T: Unmark),+> Unmark for $name <$($T),+> {
389+
type Unmarked = $name <$($T::Unmarked),+>;
390+
fn unmark(self) -> Self::Unmarked {
391+
$name {
392+
$($field: Unmark::unmark(self.$field)),*
393+
}
394+
}
395+
}
396+
};
382397
(enum $name:ident <$($T:ident),+> { $($variant:ident $(($field:ident))?),* $(,)? }) => {
383398
impl<$($T: Mark),+> Mark for $name <$($T),+> {
384399
type Unmarked = $name <$($T::Unmarked),+>;
@@ -449,3 +464,16 @@ compound_traits!(
449464
Literal(tt),
450465
}
451466
);
467+
468+
/// Globals provided alongside the initial inputs for a macro expansion.
469+
/// Provides values such as spans which are used frequently to avoid RPC.
470+
#[derive(Clone)]
471+
pub struct ExpnGlobals<S> {
472+
pub def_site: S,
473+
pub call_site: S,
474+
pub mixed_site: S,
475+
}
476+
477+
compound_traits!(
478+
struct ExpnGlobals<Sp> { def_site, call_site, mixed_site }
479+
);

‎proc_macro/src/bridge/selfless_reify.rs

+6-5
Original file line numberDiff line numberDiff line change
@@ -75,9 +75,10 @@ macro_rules! define_reify_functions {
7575
define_reify_functions! {
7676
fn _reify_to_extern_c_fn_unary<A, R> for extern "C" fn(arg: A) -> R;
7777

78-
// HACK(eddyb) this abstraction is used with `for<'a> fn(Bridge<'a>) -> T`
79-
// but that doesn't work with just `reify_to_extern_c_fn_unary` because of
80-
// the `fn` pointer type being "higher-ranked" (i.e. the `for<'a>` binder).
81-
// FIXME(eddyb) try to remove the lifetime from `Bridge`, that'd help.
82-
fn reify_to_extern_c_fn_hrt_bridge<R> for extern "C" fn(bridge: super::Bridge<'_>) -> R;
78+
// HACK(eddyb) this abstraction is used with `for<'a> fn(BridgeConfig<'a>)
79+
// -> T` but that doesn't work with just `reify_to_extern_c_fn_unary`
80+
// because of the `fn` pointer type being "higher-ranked" (i.e. the
81+
// `for<'a>` binder).
82+
// FIXME(eddyb) try to remove the lifetime from `BridgeConfig`, that'd help.
83+
fn reify_to_extern_c_fn_hrt_bridge<R> for extern "C" fn(bridge: super::BridgeConfig<'_>) -> R;
8384
}

‎proc_macro/src/bridge/server.rs

+23-14
Original file line numberDiff line numberDiff line change
@@ -38,14 +38,21 @@ macro_rules! declare_server_traits {
3838
$(associated_fn!(fn $method(&mut self, $($arg: $arg_ty),*) $(-> $ret_ty)?);)*
3939
})*
4040

41-
pub trait Server: Types $(+ $name)* {}
42-
impl<S: Types $(+ $name)*> Server for S {}
41+
pub trait Server: Types $(+ $name)* {
42+
fn globals(&mut self) -> ExpnGlobals<Self::Span>;
43+
}
4344
}
4445
}
4546
with_api!(Self, self_, declare_server_traits);
4647

4748
pub(super) struct MarkedTypes<S: Types>(S);
4849

50+
impl<S: Server> Server for MarkedTypes<S> {
51+
fn globals(&mut self) -> ExpnGlobals<Self::Span> {
52+
<_>::mark(Server::globals(&mut self.0))
53+
}
54+
}
55+
4956
macro_rules! define_mark_types_impls {
5057
($($name:ident {
5158
$(fn $method:ident($($arg:ident: $arg_ty:ty),* $(,)?) $(-> $ret_ty:ty)?;)*
@@ -120,7 +127,7 @@ pub trait ExecutionStrategy {
120127
&self,
121128
dispatcher: &mut impl DispatcherTrait,
122129
input: Buffer,
123-
run_client: extern "C" fn(Bridge<'_>) -> Buffer,
130+
run_client: extern "C" fn(BridgeConfig<'_>) -> Buffer,
124131
force_show_panics: bool,
125132
) -> Buffer;
126133
}
@@ -132,13 +139,13 @@ impl ExecutionStrategy for SameThread {
132139
&self,
133140
dispatcher: &mut impl DispatcherTrait,
134141
input: Buffer,
135-
run_client: extern "C" fn(Bridge<'_>) -> Buffer,
142+
run_client: extern "C" fn(BridgeConfig<'_>) -> Buffer,
136143
force_show_panics: bool,
137144
) -> Buffer {
138145
let mut dispatch = |buf| dispatcher.dispatch(buf);
139146

140-
run_client(Bridge {
141-
cached_buffer: input,
147+
run_client(BridgeConfig {
148+
input,
142149
dispatch: (&mut dispatch).into(),
143150
force_show_panics,
144151
_marker: marker::PhantomData,
@@ -156,7 +163,7 @@ impl ExecutionStrategy for CrossThread1 {
156163
&self,
157164
dispatcher: &mut impl DispatcherTrait,
158165
input: Buffer,
159-
run_client: extern "C" fn(Bridge<'_>) -> Buffer,
166+
run_client: extern "C" fn(BridgeConfig<'_>) -> Buffer,
160167
force_show_panics: bool,
161168
) -> Buffer {
162169
use std::sync::mpsc::channel;
@@ -170,8 +177,8 @@ impl ExecutionStrategy for CrossThread1 {
170177
res_rx.recv().unwrap()
171178
};
172179

173-
run_client(Bridge {
174-
cached_buffer: input,
180+
run_client(BridgeConfig {
181+
input,
175182
dispatch: (&mut dispatch).into(),
176183
force_show_panics,
177184
_marker: marker::PhantomData,
@@ -193,7 +200,7 @@ impl ExecutionStrategy for CrossThread2 {
193200
&self,
194201
dispatcher: &mut impl DispatcherTrait,
195202
input: Buffer,
196-
run_client: extern "C" fn(Bridge<'_>) -> Buffer,
203+
run_client: extern "C" fn(BridgeConfig<'_>) -> Buffer,
197204
force_show_panics: bool,
198205
) -> Buffer {
199206
use std::sync::{Arc, Mutex};
@@ -219,8 +226,8 @@ impl ExecutionStrategy for CrossThread2 {
219226
}
220227
};
221228

222-
let r = run_client(Bridge {
223-
cached_buffer: input,
229+
let r = run_client(BridgeConfig {
230+
input,
224231
dispatch: (&mut dispatch).into(),
225232
force_show_panics,
226233
_marker: marker::PhantomData,
@@ -258,14 +265,16 @@ fn run_server<
258265
handle_counters: &'static client::HandleCounters,
259266
server: S,
260267
input: I,
261-
run_client: extern "C" fn(Bridge<'_>) -> Buffer,
268+
run_client: extern "C" fn(BridgeConfig<'_>) -> Buffer,
262269
force_show_panics: bool,
263270
) -> Result<O, PanicMessage> {
264271
let mut dispatcher =
265272
Dispatcher { handle_store: HandleStore::new(handle_counters), server: MarkedTypes(server) };
266273

274+
let globals = dispatcher.server.globals();
275+
267276
let mut buf = Buffer::new();
268-
input.encode(&mut buf, &mut dispatcher.handle_store);
277+
(globals, input).encode(&mut buf, &mut dispatcher.handle_store);
269278

270279
buf = strategy.run_bridge_and_client(&mut dispatcher, buf, run_client, force_show_panics);
271280

0 commit comments

Comments
 (0)