Rollup merge of rust-lang#157076 - bjorn3:proc_macro_refactors4, r=petrochenkov

Various proc-macro related code cleanups

Some are just misc cleanups. Others are to make the proc-macro ABI and RPC interface a bit less target dependent. I've got some local changes that change the ABI to what is effectively a single `&[extern "C" fn(BridgeConfig<'_>) -> Buffer]` export.

Author: JonathanBrouwer

compiler/rustc_metadata/src/rmeta/encoder.rs

@@ -2036,7 +2036,6 @@ impl<'a, 'tcx> EncodeContext<'a, 'tcx> {
 
                 let def_id = id.to_def_id();
                 self.tables.def_kind.set_some(def_id.index, DefKind::Macro(macro_kind.into()));
-                self.tables.proc_macro.set_some(def_id.index, macro_kind);
                 self.encode_attrs(id);
                 record!(self.tables.def_keys[def_id] <- def_key);
                 record!(self.tables.def_ident_span[def_id] <- span);

compiler/rustc_metadata/src/rmeta/mod.rs

@@ -464,7 +464,6 @@ define_tables! {
     variant_data: Table<DefIndex, LazyValue<VariantData>>,
     assoc_container: Table<DefIndex, LazyValue<ty::AssocContainer>>,
     macro_definition: Table<DefIndex, LazyValue<ast::DelimArgs>>,
-    proc_macro: Table<DefIndex, MacroKind>,
     deduced_param_attrs: Table<DefIndex, LazyArray<DeducedParamAttrs>>,
     collect_return_position_impl_trait_in_trait_tys: Table<DefIndex, LazyValue<DefIdMap<ty::EarlyBinder<'static, Ty<'static>>>>>,
     doc_link_resolutions: Table<DefIndex, LazyValue<DocLinkResMap>>,

library/proc_macro/src/bridge/client.rs

@@ -2,19 +2,9 @@
 
 use std::cell::RefCell;
 use std::marker::PhantomData;
-use std::sync::atomic::AtomicU32;
 
 use super::*;
 
-#[repr(C)]
-pub(super) struct HandleCounters {
-    pub(super) token_stream: AtomicU32,
-    pub(super) span: AtomicU32,
-}
-
-static COUNTERS: HandleCounters =
-    HandleCounters { token_stream: AtomicU32::new(1), span: AtomicU32::new(1) };
-
 pub(crate) struct TokenStream {
     handle: handle::Handle,
 }
@@ -47,6 +37,18 @@ impl<S> Decode<'_, '_, S> for TokenStream {
     }
 }
 
+impl Encode<()> for crate::TokenStream {
+    fn encode(self, w: &mut Buffer, s: &mut ()) {
+        self.0.encode(w, s)
+    }
+}
+
+impl Decode<'_, '_, ()> for crate::TokenStream {
+    fn decode(r: &mut &[u8], s: &mut ()) -> Self {
+        crate::TokenStream(Some(Decode::decode(r, s)))
+    }
+}
+
 #[derive(Copy, Clone, PartialEq, Eq, Hash)]
 pub(crate) struct Span {
     handle: handle::Handle,
@@ -209,8 +211,6 @@ pub(crate) fn is_available() -> bool {
 /// and forcing the use of APIs that take/return `S::TokenStream`, server-side.
 #[repr(C)]
 pub struct Client<I, O> {
-    pub(super) handle_counters: &'static HandleCounters,
-
     pub(super) run: extern "C" fn(BridgeConfig<'_>) -> Buffer,
 
     pub(super) _marker: PhantomData<fn(I) -> O>,
@@ -243,14 +243,13 @@ fn maybe_install_panic_hook(force_show_panics: bool) {
 
 /// Client-side helper for handling client panics, entering the bridge,
 /// deserializing input and serializing output.
-// FIXME(eddyb) maybe replace `Bridge::enter` with this?
-fn run_client<A: for<'a, 's> Decode<'a, 's, ()>, R: Encode<()>>(
+fn run_client<A: for<'a, 's> Decode<'a, 's, ()>>(
     config: BridgeConfig<'_>,
-    f: impl FnOnce(A) -> R,
+    f: impl FnOnce(A) -> crate::TokenStream,
 ) -> Buffer {
     let BridgeConfig { input: mut buf, dispatch, force_show_panics, .. } = config;
 
-    panic::catch_unwind(panic::AssertUnwindSafe(|| {
+    let res = panic::catch_unwind(panic::AssertUnwindSafe(|| {
         maybe_install_panic_hook(force_show_panics);
 
         // Make sure the symbol store is empty before decoding inputs.
@@ -267,23 +266,12 @@ fn run_client<A: for<'a, 's> Decode<'a, 's, ()>, R: Encode<()>>(
         // 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| {
-        buf.clear();
-        Err::<(), _>(e).encode(&mut buf, &mut ());
-    });
+        output
+    }));
+
+    // Serialize response of type `Result<R, PanicMessage>`.
+    buf.clear();
+    res.map_err(PanicMessage::from).encode(&mut buf, &mut ());
 
     // Now that a response has been serialized, invalidate all symbols
     // registered with the interner.
@@ -294,9 +282,8 @@ fn run_client<A: for<'a, 's> Decode<'a, 's, ()>, R: Encode<()>>(
 impl Client<crate::TokenStream, crate::TokenStream> {
     pub const fn expand1(f: impl Fn(crate::TokenStream) -> crate::TokenStream + Copy) -> Self {
         Client {
-            handle_counters: &COUNTERS,
             run: super::selfless_reify::reify_to_extern_c_fn_hrt_bridge(move |bridge| {
-                run_client(bridge, |input| f(crate::TokenStream(Some(input))).0)
+                run_client(bridge, |input| f(input))
             }),
             _marker: PhantomData,
         }
@@ -308,11 +295,8 @@ impl Client<(crate::TokenStream, crate::TokenStream), crate::TokenStream> {
         f: impl Fn(crate::TokenStream, crate::TokenStream) -> crate::TokenStream + Copy,
     ) -> Self {
         Client {
-            handle_counters: &COUNTERS,
             run: super::selfless_reify::reify_to_extern_c_fn_hrt_bridge(move |bridge| {
-                run_client(bridge, |(input, input2)| {
-                    f(crate::TokenStream(Some(input)), crate::TokenStream(Some(input2))).0
-                })
+                run_client(bridge, |(input, input2)| f(input, input2))
             }),
             _marker: PhantomData,
         }

library/proc_macro/src/bridge/rpc.rs

@@ -15,20 +15,30 @@ pub(super) trait Decode<'a, 's, S>: Sized {
 }
 
 macro_rules! rpc_encode_decode {
-    (le $ty:ty) => {
+    (le $ty:ident $size:literal) => {
         impl<S> Encode<S> for $ty {
             fn encode(self, w: &mut Buffer, _: &mut S) {
-                w.extend_from_array(&self.to_le_bytes());
+                const N: usize = size_of::<$ty>();
+
+                // We can pad with zeros without changing the value because of
+                // little endian encoding.
+                let mut bytes = [0; $size];
+                bytes[..N].copy_from_slice(&self.to_le_bytes());
+
+                w.extend_from_array(&bytes);
             }
         }
 
         impl<S> Decode<'_, '_, S> for $ty {
             fn decode(r: &mut &[u8], _: &mut S) -> Self {
                 const N: usize = size_of::<$ty>();
+                const {
+                    assert!(N <= $size);
+                }
 
                 let mut bytes = [0; N];
                 bytes.copy_from_slice(&r[..N]);
-                *r = &r[N..];
+                *r = &r[$size..];
 
                 Self::from_le_bytes(bytes)
             }
@@ -108,8 +118,8 @@ impl<S> Decode<'_, '_, S> for u8 {
     }
 }
 
-rpc_encode_decode!(le u32);
-rpc_encode_decode!(le usize);
+rpc_encode_decode!(le u32 4);
+rpc_encode_decode!(le usize 8);
 
 impl<S> Encode<S> for bool {
     fn encode(self, w: &mut Buffer, s: &mut S) {

library/proc_macro/src/bridge/server.rs

@@ -1,6 +1,7 @@
 //! Server-side traits.
 
 use std::cell::Cell;
+use std::sync::atomic::AtomicU32;
 use std::sync::mpsc;
 
 use super::*;
@@ -11,10 +12,13 @@ pub(super) struct HandleStore<S: Server> {
 }
 
 impl<S: Server> HandleStore<S> {
-    fn new(handle_counters: &'static client::HandleCounters) -> Self {
+    fn new() -> Self {
+        static TOKEN_STREAM: AtomicU32 = AtomicU32::new(1);
+        static SPAN: AtomicU32 = AtomicU32::new(1);
+
         HandleStore {
-            token_stream: handle::OwnedStore::new(&handle_counters.token_stream),
-            span: handle::InternedStore::new(&handle_counters.span),
+            token_stream: handle::OwnedStore::new(&TOKEN_STREAM),
+            span: handle::InternedStore::new(&SPAN),
         }
     }
 }
@@ -246,13 +250,12 @@ fn run_server<
     O: for<'a, 's> Decode<'a, 's, HandleStore<S>>,
 >(
     strategy: &impl ExecutionStrategy,
-    handle_counters: &'static client::HandleCounters,
     server: S,
     input: I,
     run_client: extern "C" fn(BridgeConfig<'_>) -> Buffer,
     force_show_panics: bool,
 ) -> Result<O, PanicMessage> {
-    let mut dispatcher = Dispatcher { handle_store: HandleStore::new(handle_counters), server };
+    let mut dispatcher = Dispatcher { handle_store: HandleStore::new(), server };
 
     let globals = dispatcher.server.globals();
 
@@ -276,16 +279,9 @@ impl client::Client<crate::TokenStream, crate::TokenStream> {
     where
         S: Server,
     {
-        let client::Client { handle_counters, run, _marker } = *self;
-        run_server(
-            strategy,
-            handle_counters,
-            server,
-            <MarkedTokenStream<S>>::mark(input),
-            run,
-            force_show_panics,
-        )
-        .map(|s| <Option<MarkedTokenStream<S>>>::unmark(s).unwrap_or_default())
+        let client::Client { run, _marker } = *self;
+        run_server(strategy, server, <MarkedTokenStream<S>>::mark(input), run, force_show_panics)
+            .map(|s| <Option<MarkedTokenStream<S>>>::unmark(s).unwrap_or_default())
     }
 }
 
@@ -301,10 +297,9 @@ impl client::Client<(crate::TokenStream, crate::TokenStream), crate::TokenStream
     where
         S: Server,
     {
-        let client::Client { handle_counters, run, _marker } = *self;
+        let client::Client { run, _marker } = *self;
         run_server(
             strategy,
-            handle_counters,
             server,
             (<MarkedTokenStream<S>>::mark(input), <MarkedTokenStream<S>>::mark(input2)),
             run,