mikros/rust
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

proc_macro: simplify bridge state

Browse Source
This commit is contained in:
joboet 2024-03-23 12:15:11 +01:00
parent c3b05c6e5b
commit ac770f7bd5
No known key found for this signature in database
GPG Key ID: 704E0149B0194B3C
3 changed files with 58 additions and 132 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

100
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,
/// A server is connected and available for requests.
Connected(Bridge<'a>),
/// Access to the bridge is being exclusively acquired
/// (e.g., during `BridgeState::with`).
InUse,
}
enum BridgeStateL {}
impl<'a> scoped_cell::ApplyL<'a> for BridgeStateL {
type Out = BridgeState<'a>;
}
#[allow(unsafe_code)]
mod state {
use super::Bridge;
use std::cell::{Cell, RefCell};
use std::ptr;
thread_local! {
static BRIDGE_STATE: scoped_cell::ScopedCell<BridgeStateL> =
const { scoped_cell::ScopedCell::new(BridgeState::NotConnected) };
static BRIDGE_STATE: Cell<*const ()> = const { Cell::new(ptr::null()) };
}
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 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);
}
}
let inner = ptr::from_ref(state).cast();
let outer = BRIDGE_STATE.replace(inner);
let _restore = RestoreOnDrop(outer);
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,15 +309,12 @@ 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());
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
@ -333,8 +327,6 @@ fn run_client<A: for<'a, 's> DecodeMut<'a, 's, ()>, R: Encode<()>>(
// 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 @@ macro_rules! reverse_decode {
#[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())
}
}