proc_macro: simplify bridge state

3 files changed, 61 insertions, 135 deletions

diff --git a/library/proc_macro/src/bridge/client.rs b/library/proc_macro/src/bridge/client.rs
index 8a1ba436f72..f3cfc41bac7 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<Bridge<'bridge>>, 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<BridgeStateL> =
-        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<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 });
-
-        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 67c72cf98d4..8553e8d5e4f 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 a8b34143908..00000000000
--- 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<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())
-    }
-}