proc_macro: stop using a remote object handle for Ident

Doing this for all unicode identifiers would require a dependency on
`unicode-normalization` and `rustc_lexer`, which is currently not
possible for `proc_macro` due to it being built concurrently with `std`
and `core`. Instead, ASCII identifiers are validated locally, and an RPC
message is used to validate unicode identifiers when needed.

String values are interned on the both the server and client when
deserializing, to avoid unnecessary copies and keep Ident cheap to copy and
move. This appears to be important for performance.

The client-side interner is based roughly on the one from rustc_span, and uses
an arena inspired by rustc_arena.

RPC messages passing symbols always include the full value. This could
potentially be optimized in the future if it is revealed to be a
performance bottleneck.

Despite now having a relevant implementaion of Display for Ident, ToString is
still specialized, as it is a hot-path for this object.

The symbol infrastructure will also be used for literals in the next
part.

Author: mystor

proc_macro/src/bridge/arena.rs

@@ -0,0 +1,113 @@
+//! A minimal arena allocator inspired by `rustc_arena::DroplessArena`.
+//!
+//! This is unfortunately a minimal re-implementation rather than a dependency
+//! as it is difficult to depend on crates from within `proc_macro`, due to it
+//! being built at the same time as `std`.
+
+use std::cell::{Cell, RefCell};
+use std::cmp;
+use std::mem::MaybeUninit;
+use std::ops::Range;
+use std::ptr;
+use std::slice;
+use std::str;
+
+// The arenas start with PAGE-sized chunks, and then each new chunk is twice as
+// big as its predecessor, up until we reach HUGE_PAGE-sized chunks, whereupon
+// we stop growing. This scales well, from arenas that are barely used up to
+// arenas that are used for 100s of MiBs. Note also that the chosen sizes match
+// the usual sizes of pages and huge pages on Linux.
+const PAGE: usize = 4096;
+const HUGE_PAGE: usize = 2 * 1024 * 1024;
+
+/// A minimal arena allocator inspired by `rustc_arena::DroplessArena`.
+///
+/// This is unfortunately a complete re-implementation rather than a dependency
+/// as it is difficult to depend on crates from within `proc_macro`, due to it
+/// being built at the same time as `std`.
+///
+/// This arena doesn't have support for allocating anything other than byte
+/// slices, as that is all that is necessary.
+pub(crate) struct Arena {
+    start: Cell<*mut MaybeUninit<u8>>,
+    end: Cell<*mut MaybeUninit<u8>>,
+    chunks: RefCell<Vec<Box<[MaybeUninit<u8>]>>>,
+}
+
+impl Arena {
+    pub(crate) fn new() -> Self {
+        Arena {
+            start: Cell::new(ptr::null_mut()),
+            end: Cell::new(ptr::null_mut()),
+            chunks: RefCell::new(Vec::new()),
+        }
+    }
+
+    /// Add a new chunk with at least `additional` free bytes.
+    #[inline(never)]
+    #[cold]
+    fn grow(&self, additional: usize) {
+        let mut chunks = self.chunks.borrow_mut();
+        let mut new_cap;
+        if let Some(last_chunk) = chunks.last_mut() {
+            // If the previous chunk's len is less than HUGE_PAGE
+            // bytes, then this chunk will be least double the previous
+            // chunk's size.
+            new_cap = last_chunk.len().min(HUGE_PAGE / 2);
+            new_cap *= 2;
+        } else {
+            new_cap = PAGE;
+        }
+        // Also ensure that this chunk can fit `additional`.
+        new_cap = cmp::max(additional, new_cap);
+
+        let mut chunk = Box::new_uninit_slice(new_cap);
+        let Range { start, end } = chunk.as_mut_ptr_range();
+        self.start.set(start);
+        self.end.set(end);
+        chunks.push(chunk);
+    }
+
+    /// Allocates a byte slice with specified size from the current memory
+    /// chunk. Returns `None` if there is no free space left to satisfy the
+    /// request.
+    fn alloc_raw_without_grow(&self, bytes: usize) -> Option<&mut [MaybeUninit<u8>]> {
+        let start = self.start.get().addr();
+        let old_end = self.end.get();
+        let end = old_end.addr();
+
+        let new_end = end.checked_sub(bytes)?;
+        if start <= new_end {
+            let new_end = old_end.with_addr(new_end);
+            self.end.set(new_end);
+            // SAFETY: `bytes` bytes starting at `new_end` were just reserved.
+            Some(unsafe { slice::from_raw_parts_mut(new_end, bytes) })
+        } else {
+            None
+        }
+    }
+
+    fn alloc_raw(&self, bytes: usize) -> &mut [MaybeUninit<u8>] {
+        if bytes == 0 {
+            return &mut [];
+        }
+
+        loop {
+            if let Some(a) = self.alloc_raw_without_grow(bytes) {
+                break a;
+            }
+            // No free space left. Allocate a new chunk to satisfy the request.
+            // On failure the grow will panic or abort.
+            self.grow(bytes);
+        }
+    }
+
+    pub(crate) fn alloc_str<'a>(&'a self, string: &str) -> &'a mut str {
+        let alloc = self.alloc_raw(string.len());
+        let bytes = MaybeUninit::write_slice(alloc, string.as_bytes());
+
+        // SAFETY: we convert from `&str` to `&[u8]`, clone it into the arena,
+        // and immediately convert the clone back to `&str`.
+        unsafe { str::from_utf8_unchecked_mut(bytes) }
+    }
+}

proc_macro/src/bridge/client.rs

@@ -181,7 +181,6 @@ define_handles! {
     Diagnostic,
 
     'interned:
-    Ident,
     Span,
 }
 
@@ -242,6 +241,8 @@ impl fmt::Debug for Span {
     }
 }
 
+pub(crate) use super::symbol::Symbol;
+
 macro_rules! define_client_side {
     ($($name:ident {
         $(fn $method:ident($($arg:ident: $arg_ty:ty),* $(,)?) $(-> $ret_ty:ty)*;)*
@@ -405,6 +406,9 @@ fn run_client<A: for<'a, 's> DecodeMut<'a, 's, ()>, R: Encode<()>>(
     panic::catch_unwind(panic::AssertUnwindSafe(|| {
         maybe_install_panic_hook(force_show_panics);
 
+        // Make sure the symbol store is empty before decoding inputs.
+        Symbol::invalidate_all();
+
         let reader = &mut &buf[..];
         let (globals, input) = <(ExpnGlobals<Span>, A)>::decode(reader, &mut ());
 
@@ -438,6 +442,10 @@ fn run_client<A: for<'a, 's> DecodeMut<'a, 's, ()>, R: Encode<()>>(
         buf.clear();
         Err::<(), _>(e).encode(&mut buf, &mut ());
     });
+
+    // Now that a response has been serialized, invalidate all symbols
+    // registered with the interner.
+    Symbol::invalidate_all();
     buf
 }
 

proc_macro/src/bridge/mod.rs

@@ -65,24 +65,19 @@ macro_rules! with_api {
                 fn from_str(src: &str) -> $S::TokenStream;
                 fn to_string($self: &$S::TokenStream) -> String;
                 fn from_token_tree(
-                    tree: TokenTree<$S::TokenStream, $S::Span, $S::Ident, $S::Literal>,
+                    tree: TokenTree<$S::TokenStream, $S::Span, $S::Symbol, $S::Literal>,
                 ) -> $S::TokenStream;
                 fn concat_trees(
                     base: Option<$S::TokenStream>,
-                    trees: Vec<TokenTree<$S::TokenStream, $S::Span, $S::Ident, $S::Literal>>,
+                    trees: Vec<TokenTree<$S::TokenStream, $S::Span, $S::Symbol, $S::Literal>>,
                 ) -> $S::TokenStream;
                 fn concat_streams(
                     base: Option<$S::TokenStream>,
                     streams: Vec<$S::TokenStream>,
                 ) -> $S::TokenStream;
                 fn into_trees(
                     $self: $S::TokenStream
-                ) -> Vec<TokenTree<$S::TokenStream, $S::Span, $S::Ident, $S::Literal>>;
-            },
-            Ident {
-                fn new(string: &str, span: $S::Span, is_raw: bool) -> $S::Ident;
-                fn span($self: $S::Ident) -> $S::Span;
-                fn with_span($self: $S::Ident, span: $S::Span) -> $S::Ident;
+                ) -> Vec<TokenTree<$S::TokenStream, $S::Span, $S::Symbol, $S::Literal>>;
             },
             Literal {
                 fn drop($self: $S::Literal);
@@ -146,6 +141,9 @@ macro_rules! with_api {
                 fn save_span($self: $S::Span) -> usize;
                 fn recover_proc_macro_span(id: usize) -> $S::Span;
             },
+            Symbol {
+                fn normalize_and_validate_ident(string: &str) -> Result<$S::Symbol, ()>;
+            },
         }
     };
 }
@@ -170,6 +168,8 @@ macro_rules! reverse_decode {
     }
 }
 
+#[allow(unsafe_code)]
+mod arena;
 #[allow(unsafe_code)]
 mod buffer;
 #[forbid(unsafe_code)]
@@ -189,6 +189,8 @@ mod scoped_cell;
 mod selfless_reify;
 #[forbid(unsafe_code)]
 pub mod server;
+#[allow(unsafe_code)]
+mod symbol;
 
 use buffer::Buffer;
 pub use rpc::PanicMessage;
@@ -466,16 +468,25 @@ pub struct Punct<Span> {
 
 compound_traits!(struct Punct<Span> { ch, joint, span });
 
+#[derive(Copy, Clone, Eq, PartialEq)]
+pub struct Ident<Span, Symbol> {
+    pub sym: Symbol,
+    pub is_raw: bool,
+    pub span: Span,
+}
+
+compound_traits!(struct Ident<Span, Symbol> { sym, is_raw, span });
+
 #[derive(Clone)]
-pub enum TokenTree<TokenStream, Span, Ident, Literal> {
+pub enum TokenTree<TokenStream, Span, Symbol, Literal> {
     Group(Group<TokenStream, Span>),
     Punct(Punct<Span>),
-    Ident(Ident),
+    Ident(Ident<Span, Symbol>),
     Literal(Literal),
 }
 
 compound_traits!(
-    enum TokenTree<TokenStream, Span, Ident, Literal> {
+    enum TokenTree<TokenStream, Span, Symbol, Literal> {
         Group(tt),
         Punct(tt),
         Ident(tt),

proc_macro/src/bridge/server.rs

@@ -8,12 +8,12 @@ use super::client::HandleStore;
 pub trait Types {
     type FreeFunctions: 'static;
     type TokenStream: 'static + Clone;
-    type Ident: 'static + Copy + Eq + Hash;
     type Literal: 'static + Clone;
     type SourceFile: 'static + Clone;
     type MultiSpan: 'static;
     type Diagnostic: 'static;
     type Span: 'static + Copy + Eq + Hash;
+    type Symbol: 'static;
 }
 
 /// Declare an associated fn of one of the traits below, adding necessary
@@ -38,6 +38,12 @@ macro_rules! declare_server_traits {
 
         pub trait Server: Types $(+ $name)* {
             fn globals(&mut self) -> ExpnGlobals<Self::Span>;
+
+            /// Intern a symbol received from RPC
+            fn intern_symbol(ident: &str) -> Self::Symbol;
+
+            /// Recover the string value of a symbol, and invoke a callback with it.
+            fn with_symbol_string(symbol: &Self::Symbol, f: impl FnOnce(&str));
         }
     }
 }
@@ -49,6 +55,12 @@ impl<S: Server> Server for MarkedTypes<S> {
     fn globals(&mut self) -> ExpnGlobals<Self::Span> {
         <_>::mark(Server::globals(&mut self.0))
     }
+    fn intern_symbol(ident: &str) -> Self::Symbol {
+        <_>::mark(S::intern_symbol(ident))
+    }
+    fn with_symbol_string(symbol: &Self::Symbol, f: impl FnOnce(&str)) {
+        S::with_symbol_string(symbol.unmark(), f)
+    }
 }
 
 macro_rules! define_mark_types_impls {
@@ -81,11 +93,13 @@ macro_rules! define_dispatcher_impl {
         pub trait DispatcherTrait {
             // HACK(eddyb) these are here to allow `Self::$name` to work below.
             $(type $name;)*
+
             fn dispatch(&mut self, buf: Buffer) -> Buffer;
         }
 
         impl<S: Server> DispatcherTrait for Dispatcher<MarkedTypes<S>> {
             $(type $name = <MarkedTypes<S> as Types>::$name;)*
+
             fn dispatch(&mut self, mut buf: Buffer) -> Buffer {
                 let Dispatcher { handle_store, server } = self;