mathieu/os-rust
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Auto merge of #122939 - joboet:proc_macro_bridge_state, r=petrochenkov

Browse source 
Simplify proc macro bridge state

Currently, `proc_macro` uses a `ScopedCell` to store the client-side proc-macro bridge. Unfortunately, this requires the `Bridge`, which has non-negligible size, to be copied out and back again on every access. In some cases, the optimizer might be able to elide these copies, but in general, this is suboptimal.

This PR removes `ScopedCell` and employs a similar trick as in [`scoped_tls`](https://crates.io/crates/scoped-tls), meaning that the only thing stored in TLS is a pointer to the state, which now is a `RefCell`. Access to the pointer is then scoped so that it is always within the lifetime of the reference to the state. Unfortunately, `scoped_tls` requires the referenced type to be `'static`, which `Bridge` is not, therefore we cannot simply copy that macro but have to reimplement its TLS trick.

This removes the `#[forbid(unsafe_code)]` on the `client` module so that we do not have to export `Bridge`, which currently is private, to the whole crate. I can change that, if necessary.
This commit is contained in:
bors 2024-03-26 13:05:59 +00:00
commit 536606bc5d
3 changed files with 58 additions and 132 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

122
library/proc_macro/src/bridge/client.rs
View file

@ -2,6 +2,7 @@
use super::*;
use std::cell::RefCell;
use std::marker::PhantomData;
use std::sync::atomic::AtomicU32;
@ -189,61 +190,61 @@ struct Bridge<'a> {
impl<'a> !Send for Bridge<'a> {}
impl<'a> !Sync for Bridge<'a> {}
enum BridgeState<'a> {
/// No server is currently connected to this client.
NotConnected,
#[allow(unsafe_code)]
mod state {
use super::Bridge;
use std::cell::{Cell, RefCell};
use std::ptr;
/// A server is connected and available for requests.
Connected(Bridge<'a>),
thread_local! {
static BRIDGE_STATE: Cell<*const ()> = const { Cell::new(ptr::null()) };
}
/// Access to the bridge is being exclusively acquired
/// (e.g., during `BridgeState::with`).
InUse,
}
pub(super) fn set<'bridge, R>(state: &RefCell<Bridge<'bridge>>, f: impl FnOnce() -> R) -> R {
struct RestoreOnDrop(*const ());
impl Drop for RestoreOnDrop {
fn drop(&mut self) {
BRIDGE_STATE.set(self.0);
}
}
enum BridgeStateL {}
let inner = ptr::from_ref(state).cast();
let outer = BRIDGE_STATE.replace(inner);
let _restore = RestoreOnDrop(outer);
impl<'a> scoped_cell::ApplyL<'a> for BridgeStateL {
type Out = BridgeState<'a>;
}
f()
}
thread_local! {
static BRIDGE_STATE: scoped_cell::ScopedCell<BridgeStateL> =
const { scoped_cell::ScopedCell::new(BridgeState::NotConnected) };
}
impl BridgeState<'_> {
/// Take exclusive control of the thread-local
/// `BridgeState`, and pass it to `f`, mutably.
/// The state will be restored after `f` exits, even
/// by panic, including modifications made to it by `f`.
///
/// N.B., while `f` is running, the thread-local state
/// is `BridgeState::InUse`.
fn with<R>(f: impl FnOnce(&mut BridgeState<'_>) -> R) -> R {
BRIDGE_STATE.with(|state| state.replace(BridgeState::InUse, f))
pub(super) fn with<R>(
f: impl for<'bridge> FnOnce(Option<&RefCell<Bridge<'bridge>>>) -> R,
) -> R {
let state = BRIDGE_STATE.get();
// SAFETY: the only place where the pointer is set is in `set`. It puts
// back the previous value after the inner call has returned, so we know
// that as long as the pointer is not null, it came from a reference to
// a `RefCell<Bridge>` that outlasts the call to this function. Since `f`
// works the same for any lifetime of the bridge, including the actual
// one, we can lie here and say that the lifetime is `'static` without
// anyone noticing.
let bridge = unsafe { state.cast::<RefCell<Bridge<'static>>>().as_ref() };
f(bridge)
}
}
impl Bridge<'_> {
fn with<R>(f: impl FnOnce(&mut Bridge<'_>) -> R) -> R {
BridgeState::with(|state| match state {
BridgeState::NotConnected => {
panic!("procedural macro API is used outside of a procedural macro");
}
BridgeState::InUse => {
panic!("procedural macro API is used while it's already in use");
}
BridgeState::Connected(bridge) => f(bridge),
state::with(|state| {
let bridge = state.expect("procedural macro API is used outside of a procedural macro");
let mut bridge = bridge
.try_borrow_mut()
.expect("procedural macro API is used while it's already in use");
f(&mut bridge)
})
}
}
pub(crate) fn is_available() -> bool {
BridgeState::with(|state| match state {
BridgeState::Connected(_) | BridgeState::InUse => true,
BridgeState::NotConnected => false,
})
state::with(|s| s.is_some())
}
/// A client-side RPC entry-point, which may be using a different `proc_macro`
@ -282,11 +283,7 @@ fn maybe_install_panic_hook(force_show_panics: bool) {
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 {
if force_show_panics || !is_available() {
prev(info)
}
}));
@ -312,29 +309,24 @@ fn run_client<A: for<'a, 's> DecodeMut<'a, 's, ()>, R: Encode<()>>(
let (globals, input) = <(ExpnGlobals<Span>, A)>::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, globals });
let state = RefCell::new(Bridge { cached_buffer: buf.take(), dispatch, globals });
BRIDGE_STATE.with(|state| {
state.set(new_state, || {
let output = f(input);
let output = state::set(&state, || f(input));
// Take the `cached_buffer` back out, for the output value.
buf = Bridge::with(|bridge| bridge.cached_buffer.take());
// Take the `cached_buffer` back out, for the output value.
buf = RefCell::into_inner(state).cached_buffer;
// 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 ());
})
})
// 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)
.unwrap_or_else(|e| {

4
library/proc_macro/src/bridge/mod.rs
View file

@ -154,7 +154,7 @@ macro_rules! reverse_decode {
mod arena;
#[allow(unsafe_code)]
mod buffer;
#[forbid(unsafe_code)]
#[deny(unsafe_code)]
pub mod client;
#[allow(unsafe_code)]
mod closure;
@ -166,8 +166,6 @@ mod handle;
#[forbid(unsafe_code)]
mod rpc;
#[allow(unsafe_code)]
mod scoped_cell;
#[allow(unsafe_code)]
mod selfless_reify;
#[forbid(unsafe_code)]
pub mod server;

64
library/proc_macro/src/bridge/scoped_cell.rs
View file

@ -1,64 +0,0 @@
//! `Cell` variant for (scoped) existential lifetimes.
use std::cell::Cell;
use std::mem;
/// Type lambda application, with a lifetime.
#[allow(unused_lifetimes)]
pub trait ApplyL<'a> {
type Out;
}
/// Type lambda taking a lifetime, i.e., `Lifetime -> Type`.
pub trait LambdaL: for<'a> ApplyL<'a> {}
impl<T: for<'a> ApplyL<'a>> LambdaL for T {}
pub struct ScopedCell<T: LambdaL>(Cell<<T as ApplyL<'static>>::Out>);
impl<T: LambdaL> ScopedCell<T> {
pub const fn new(value: <T as ApplyL<'static>>::Out) -> Self {
ScopedCell(Cell::new(value))
}
/// Sets the value in `self` to `replacement` while
/// running `f`, which gets the old value, mutably.
/// The old value will be restored after `f` exits, even
/// by panic, including modifications made to it by `f`.
#[rustc_confusables("swap")]
pub fn replace<'a, R>(
&self,
replacement: <T as ApplyL<'a>>::Out,
f: impl for<'b, 'c> FnOnce(&'b mut <T as ApplyL<'c>>::Out) -> R,
) -> R {
/// Wrapper that ensures that the cell always gets filled
/// (with the original state, optionally changed by `f`),
/// even if `f` had panicked.
struct PutBackOnDrop<'a, T: LambdaL> {
cell: &'a ScopedCell<T>,
value: Option<<T as ApplyL<'static>>::Out>,
}
impl<'a, T: LambdaL> Drop for PutBackOnDrop<'a, T> {
fn drop(&mut self) {
self.cell.0.set(self.value.take().unwrap());
}
}
let mut put_back_on_drop = PutBackOnDrop {
cell: self,
value: Some(self.0.replace(unsafe {
let erased = mem::transmute_copy(&replacement);
mem::forget(replacement);
erased
})),
};
f(put_back_on_drop.value.as_mut().unwrap())
}
/// Sets the value in `self` to `value` while running `f`.
pub fn set<R>(&self, value: <T as ApplyL<'_>>::Out, f: impl FnOnce() -> R) -> R {
self.replace(value, |_| f())
}
}