From ac770f7bd57614e728980073a2a89702e78da2d8 Mon Sep 17 00:00:00 2001 From: joboet Date: Sat, 23 Mar 2024 12:15:11 +0100 Subject: [PATCH] proc_macro: simplify bridge state --- library/proc_macro/src/bridge/client.rs | 128 +++++++++---------- library/proc_macro/src/bridge/mod.rs | 4 +- library/proc_macro/src/bridge/scoped_cell.rs | 64 ---------- 3 files changed, 61 insertions(+), 135 deletions(-) delete mode 100644 library/proc_macro/src/bridge/scoped_cell.rs diff --git a/library/proc_macro/src/bridge/client.rs b/library/proc_macro/src/bridge/client.rs index 8a1ba436f7261..f3cfc41bac7c6 100644 --- a/library/proc_macro/src/bridge/client.rs +++ b/library/proc_macro/src/bridge/client.rs @@ -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>), - - /// Access to the bridge is being exclusively acquired - /// (e.g., during `BridgeState::with`). - InUse, -} + thread_local! { + static BRIDGE_STATE: Cell<*const ()> = const { Cell::new(ptr::null()) }; + } -enum BridgeStateL {} + pub(super) fn set<'bridge, R>(state: &RefCell>, f: impl FnOnce() -> R) -> R { + struct RestoreOnDrop(*const ()); + impl Drop for RestoreOnDrop { + fn drop(&mut self) { + BRIDGE_STATE.set(self.0); + } + } -impl<'a> scoped_cell::ApplyL<'a> for BridgeStateL { - type Out = BridgeState<'a>; -} + let inner = ptr::from_ref(state).cast(); + let outer = BRIDGE_STATE.replace(inner); + let _restore = RestoreOnDrop(outer); -thread_local! { - static BRIDGE_STATE: scoped_cell::ScopedCell = - const { scoped_cell::ScopedCell::new(BridgeState::NotConnected) }; -} + f() + } -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(f: impl FnOnce(&mut BridgeState<'_>) -> R) -> R { - BRIDGE_STATE.with(|state| state.replace(BridgeState::InUse, f)) + pub(super) fn with( + f: impl for<'bridge> FnOnce(Option<&RefCell>>) -> 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` 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::>>().as_ref() }; + f(bridge) } } impl Bridge<'_> { fn with(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 DecodeMut<'a, 's, ()>, R: Encode<()>>( let (globals, input) = <(ExpnGlobals, 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 }); - - BRIDGE_STATE.with(|state| { - state.set(new_state, || { - let output = f(input); - - // Take the `cached_buffer` back out, for the output value. - buf = 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 might not - // at the moment, so this is also potentially preventing UB). - buf.clear(); - Ok::<_, ()>(output).encode(&mut buf, &mut ()); - }) - }) + let state = RefCell::new(Bridge { cached_buffer: buf.take(), dispatch, globals }); + + let output = state::set(&state, || f(input)); + + // 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 ()); })) .map_err(PanicMessage::from) .unwrap_or_else(|e| { diff --git a/library/proc_macro/src/bridge/mod.rs b/library/proc_macro/src/bridge/mod.rs index 67c72cf98d405..8553e8d5e4fd6 100644 --- a/library/proc_macro/src/bridge/mod.rs +++ b/library/proc_macro/src/bridge/mod.rs @@ -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; diff --git a/library/proc_macro/src/bridge/scoped_cell.rs b/library/proc_macro/src/bridge/scoped_cell.rs deleted file mode 100644 index a8b341439084c..0000000000000 --- a/library/proc_macro/src/bridge/scoped_cell.rs +++ /dev/null @@ -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 ApplyL<'a>> LambdaL for T {} - -pub struct ScopedCell(Cell<>::Out>); - -impl ScopedCell { - pub const fn new(value: >::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: >::Out, - f: impl for<'b, 'c> FnOnce(&'b mut >::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, - value: Option<>::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(&self, value: >::Out, f: impl FnOnce() -> R) -> R { - self.replace(value, |_| f()) - } -}