proc_macro: integrate RefCell into state pointer

Author: joboet

library/proc_macro/src/bridge/client.rs

@@ -2,7 +2,6 @@
 
 use super::*;
 
-use std::cell::RefCell;
 use std::marker::PhantomData;
 use std::sync::atomic::AtomicU32;
 
@@ -175,6 +174,13 @@ macro_rules! define_client_side {
 }
 with_api!(self, self, define_client_side);
 
+enum BridgeState<'a, 'bridge> {
+    NotConnected,
+    Connected(&'a mut Bridge<'bridge>),
+    InUse,
+}
+
+#[repr(align(2))]
 struct Bridge<'a> {
     /// Reusable buffer (only `clear`-ed, never shrunk), primarily
     /// used for making requests.
@@ -192,59 +198,70 @@ impl<'a> !Sync for Bridge<'a> {}
 
 #[allow(unsafe_code)]
 mod state {
-    use super::Bridge;
-    use std::cell::{Cell, RefCell};
+    use super::{Bridge, BridgeState};
+    use std::cell::Cell;
     use std::ptr;
 
-    thread_local! {
-        static BRIDGE_STATE: Cell<*const ()> = const { Cell::new(ptr::null()) };
-    }
+    struct RestoreOnDrop(*mut ());
 
-    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 Drop for RestoreOnDrop {
+        fn drop(&mut self) {
+            BRIDGE_STATE.set(self.0);
         }
+    }
 
-        let inner = ptr::from_ref(state).cast();
+    thread_local! {
+        static BRIDGE_STATE: Cell<*mut ()> = const { Cell::new(ptr::null_mut()) };
+    }
+
+    pub(super) fn connect<'bridge, R>(bridge: &mut Bridge<'bridge>, f: impl FnOnce() -> R) -> R {
+        let inner = ptr::from_mut(bridge).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() };
+    pub(super) fn with<R>(f: impl for<'bridge> FnOnce(BridgeState<'_, 'bridge>) -> R) -> R {
+        // As `Bridge` has an alignment of 2, this cannot be a valid pointer.
+        // Use it to indicate that the bridge is in use.
+        let state = BRIDGE_STATE.replace(ptr::without_provenance_mut(1));
+        let _restore = RestoreOnDrop(state);
+        let bridge = if state.is_null() {
+            BridgeState::NotConnected
+        } else if state.addr() == 1 {
+            BridgeState::InUse
+        } else {
+            // SAFETY: the only place where the pointer is set is in `connect`.
+            // It puts back the previous value after the inner call has returned,
+            // so we know that as long as the pointer is valid, it came from a
+            // mutable reference to a `Bridge` that outlasts the call to this
+            // function. Since `f` must work for any lifetime of the bridge,
+            // including the actual one, we can lie here and say that the
+            // lifetime is `'static` without anyone noticing.
+            BridgeState::Connected(unsafe { &mut *(state as *mut Bridge<'static>) })
+        };
+
         f(bridge)
     }
 }
 
 impl Bridge<'_> {
     fn with<R>(f: impl FnOnce(&mut Bridge<'_>) -> R) -> R {
-        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)
+        state::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),
         })
     }
 }
 
 pub(crate) fn is_available() -> bool {
-    state::with(|s| s.is_some())
+    state::with(|s| matches!(s, BridgeState::Connected(_) | BridgeState::InUse))
 }
 
 /// A client-side RPC entry-point, which may be using a different `proc_macro`
@@ -309,12 +326,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 state = RefCell::new(Bridge { cached_buffer: buf.take(), dispatch, globals });
+        let mut bridge = Bridge { cached_buffer: buf.take(), dispatch, globals };
 
-        let output = state::set(&state, || f(input));
+        let output = state::connect(&mut bridge, || f(input));
 
         // Take the `cached_buffer` back out, for the output value.
-        buf = RefCell::into_inner(state).cached_buffer;
+        buf = bridge.cached_buffer;
 
         // HACK(eddyb) Separate encoding a success value (`Ok(output)`)
         // from encoding a panic (`Err(e: PanicMessage)`) to avoid