Auto merge of rust-lang#122939 - joboet:proc_macro_bridge_state, r=<try>

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.

Author: bors

library/proc_macro/src/bridge/client.rs

@@ -174,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.
@@ -189,45 +196,59 @@ 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, BridgeState};
+    use std::cell::Cell;
+    use std::ptr;
 
-    /// A server is connected and available for requests.
-    Connected(Bridge<'a>),
+    struct RestoreOnDrop(*mut ());
 
-    /// Access to the bridge is being exclusively acquired
-    /// (e.g., during `BridgeState::with`).
-    InUse,
-}
+    impl Drop for RestoreOnDrop {
+        fn drop(&mut self) {
+            BRIDGE_STATE.set(self.0);
+        }
+    }
 
-enum BridgeStateL {}
+    thread_local! {
+        static BRIDGE_STATE: Cell<*mut ()> = const { Cell::new(ptr::null_mut()) };
+    }
 
-impl<'a> scoped_cell::ApplyL<'a> for BridgeStateL {
-    type Out = BridgeState<'a>;
-}
+    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);
 
-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(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 {
-        BridgeState::with(|state| match state {
+        state::with(|state| match state {
             BridgeState::NotConnected => {
                 panic!("procedural macro API is used outside of a procedural macro");
             }
@@ -240,10 +261,7 @@ impl Bridge<'_> {
 }
 
 pub(crate) fn is_available() -> bool {
-    BridgeState::with(|state| match state {
-        BridgeState::Connected(_) | BridgeState::InUse => true,
-        BridgeState::NotConnected => false,
-    })
+    state::with(|s| matches!(s, BridgeState::Connected(_) | BridgeState::InUse))
 }
 
 /// A client-side RPC entry-point, which may be using a different `proc_macro`
@@ -282,11 +300,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 +326,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 mut bridge = Bridge { cached_buffer: buf.take(), dispatch, globals };
+
+        let output = state::connect(&mut bridge, || f(input));
+
+        // Take the `cached_buffer` back out, for the output value.
+        buf = bridge.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| {

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;