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

compiler/rustc_expand/src/proc_macro_server.rs

@@ -11,13 +11,13 @@ use rustc_parse::lexer::nfc_normalize;
 use rustc_parse::parse_stream_from_source_str;
 use rustc_session::parse::ParseSess;
 use rustc_span::def_id::CrateNum;
-use rustc_span::symbol::{self, kw, sym, Symbol};
+use rustc_span::symbol::{self, sym, Symbol};
 use rustc_span::{BytePos, FileName, Pos, SourceFile, Span};
 
-use pm::bridge::{server, DelimSpan, ExpnGlobals, Group, Punct, TokenTree};
+use pm::bridge::{server, DelimSpan, ExpnGlobals, Group, Ident, Punct, TokenTree};
 use pm::{Delimiter, Level, LineColumn};
+use std::ascii;
 use std::ops::Bound;
-use std::{ascii, panic};
 
 trait FromInternal<T> {
     fn from_internal(x: T) -> Self;
@@ -50,7 +50,7 @@ impl ToInternal<token::Delimiter> for Delimiter {
 }
 
 impl FromInternal<(TokenStream, &mut Rustc<'_, '_>)>
-    for Vec<TokenTree<TokenStream, Span, Ident, Literal>>
+    for Vec<TokenTree<TokenStream, Span, Symbol, Literal>>
 {
     fn from_internal((stream, rustc): (TokenStream, &mut Rustc<'_, '_>)) -> Self {
         use rustc_ast::token::*;
@@ -135,13 +135,12 @@ impl FromInternal<(TokenStream, &mut Rustc<'_, '_>)>
                 Question => op("?"),
                 SingleQuote => op("'"),
 
-                Ident(name, false) if name == kw::DollarCrate => trees.push(TokenTree::Ident(Ident::dollar_crate(span))),
-                Ident(name, is_raw) => trees.push(TokenTree::Ident(Ident::new(rustc.sess(), name, is_raw, span))),
+                Ident(sym, is_raw) => trees.push(TokenTree::Ident(Ident { sym, is_raw, span })),
                 Lifetime(name) => {
                     let ident = symbol::Ident::new(name, span).without_first_quote();
                     trees.extend([
                         TokenTree::Punct(Punct { ch: b'\'', joint: true, span }),
-                        TokenTree::Ident(Ident::new(rustc.sess(), ident.name, false, span)),
+                        TokenTree::Ident(Ident { sym: ident.name, is_raw: false, span }),
                     ]);
                 }
                 Literal(lit) => trees.push(TokenTree::Literal(self::Literal { lit, span })),
@@ -170,7 +169,7 @@ impl FromInternal<(TokenStream, &mut Rustc<'_, '_>)>
                 }
 
                 Interpolated(nt) if let NtIdent(ident, is_raw) = *nt => {
-                    trees.push(TokenTree::Ident(Ident::new(rustc.sess(), ident.name, is_raw, ident.span)))
+                    trees.push(TokenTree::Ident(Ident { sym: ident.name, is_raw, span: ident.span }))
                 }
 
                 Interpolated(nt) => {
@@ -200,11 +199,14 @@ impl FromInternal<(TokenStream, &mut Rustc<'_, '_>)>
     }
 }
 
-impl ToInternal<TokenStream> for TokenTree<TokenStream, Span, Ident, Literal> {
+impl ToInternal<TokenStream>
+    for (TokenTree<TokenStream, Span, Symbol, Literal>, &mut Rustc<'_, '_>)
+{
     fn to_internal(self) -> TokenStream {
         use rustc_ast::token::*;
 
-        let (ch, joint, span) = match self {
+        let (tree, rustc) = self;
+        let (ch, joint, span) = match tree {
             TokenTree::Punct(Punct { ch, joint, span }) => (ch, joint, span),
             TokenTree::Group(Group { delimiter, stream, span: DelimSpan { open, close, .. } }) => {
                 return tokenstream::TokenTree::Delimited(
@@ -215,6 +217,7 @@ impl ToInternal<TokenStream> for TokenTree<TokenStream, Span, Ident, Literal> {
                 .into();
             }
             TokenTree::Ident(self::Ident { sym, is_raw, span }) => {
+                rustc.sess().symbol_gallery.insert(sym, span);
                 return tokenstream::TokenTree::token(Ident(sym, is_raw), span).into();
             }
             TokenTree::Literal(self::Literal {
@@ -289,33 +292,6 @@ impl ToInternal<rustc_errors::Level> for Level {
 
 pub struct FreeFunctions;
 
-#[derive(Copy, Clone, PartialEq, Eq, Hash)]
-pub struct Ident {
-    sym: Symbol,
-    is_raw: bool,
-    span: Span,
-}
-
-impl Ident {
-    fn new(sess: &ParseSess, sym: Symbol, is_raw: bool, span: Span) -> Ident {
-        let sym = nfc_normalize(sym.as_str());
-        let string = sym.as_str();
-        if !rustc_lexer::is_ident(string) {
-            panic!("`{:?}` is not a valid identifier", string)
-        }
-        if is_raw && !sym.can_be_raw() {
-            panic!("`{}` cannot be a raw identifier", string);
-        }
-        sess.symbol_gallery.insert(sym, span);
-        Ident { sym, is_raw, span }
-    }
-
-    fn dollar_crate(span: Span) -> Ident {
-        // `$crate` is accepted as an ident only if it comes from the compiler.
-        Ident { sym: kw::DollarCrate, is_raw: false, span }
-    }
-}
-
 // FIXME(eddyb) `Literal` should not expose internal `Debug` impls.
 #[derive(Clone, Debug)]
 pub struct Literal {
@@ -357,12 +333,12 @@ impl<'a, 'b> Rustc<'a, 'b> {
 impl server::Types for Rustc<'_, '_> {
     type FreeFunctions = FreeFunctions;
     type TokenStream = TokenStream;
-    type Ident = Ident;
     type Literal = Literal;
     type SourceFile = Lrc<SourceFile>;
     type MultiSpan = Vec<Span>;
     type Diagnostic = Diagnostic;
     type Span = Span;
+    type Symbol = Symbol;
 }
 
 impl server::FreeFunctions for Rustc<'_, '_> {
@@ -453,22 +429,22 @@ impl server::TokenStream for Rustc<'_, '_> {
 
     fn from_token_tree(
         &mut self,
-        tree: TokenTree<Self::TokenStream, Self::Span, Self::Ident, Self::Literal>,
+        tree: TokenTree<Self::TokenStream, Self::Span, Self::Symbol, Self::Literal>,
     ) -> Self::TokenStream {
-        tree.to_internal()
+        (tree, &mut *self).to_internal()
     }
 
     fn concat_trees(
         &mut self,
         base: Option<Self::TokenStream>,
-        trees: Vec<TokenTree<Self::TokenStream, Self::Span, Self::Ident, Self::Literal>>,
+        trees: Vec<TokenTree<Self::TokenStream, Self::Span, Self::Symbol, Self::Literal>>,
     ) -> Self::TokenStream {
         let mut builder = tokenstream::TokenStreamBuilder::new();
         if let Some(base) = base {
             builder.push(base);
         }
         for tree in trees {
-            builder.push(tree.to_internal());
+            builder.push((tree, &mut *self).to_internal());
         }
         builder.build()
     }
@@ -491,25 +467,11 @@ impl server::TokenStream for Rustc<'_, '_> {
     fn into_trees(
         &mut self,
         stream: Self::TokenStream,
-    ) -> Vec<TokenTree<Self::TokenStream, Self::Span, Self::Ident, Self::Literal>> {
+    ) -> Vec<TokenTree<Self::TokenStream, Self::Span, Self::Symbol, Self::Literal>> {
         FromInternal::from_internal((stream, self))
     }
 }
 
-impl server::Ident for Rustc<'_, '_> {
-    fn new(&mut self, string: &str, span: Self::Span, is_raw: bool) -> Self::Ident {
-        Ident::new(self.sess(), Symbol::intern(string), is_raw, span)
-    }
-
-    fn span(&mut self, ident: Self::Ident) -> Self::Span {
-        ident.span
-    }
-
-    fn with_span(&mut self, ident: Self::Ident, span: Self::Span) -> Self::Ident {
-        Ident { span, ..ident }
-    }
-}
-
 impl server::Literal for Rustc<'_, '_> {
     fn from_str(&mut self, s: &str) -> Result<Self::Literal, ()> {
         let name = FileName::proc_macro_source_code(s);
@@ -812,6 +774,13 @@ impl server::Span for Rustc<'_, '_> {
     }
 }
 
+impl server::Symbol for Rustc<'_, '_> {
+    fn normalize_and_validate_ident(&mut self, string: &str) -> Result<Self::Symbol, ()> {
+        let sym = nfc_normalize(string);
+        if rustc_lexer::is_ident(sym.as_str()) { Ok(sym) } else { Err(()) }
+    }
+}
+
 impl server::Server for Rustc<'_, '_> {
     fn globals(&mut self) -> ExpnGlobals<Self::Span> {
         ExpnGlobals {
@@ -820,4 +789,12 @@ impl server::Server for Rustc<'_, '_> {
             mixed_site: self.mixed_site,
         }
     }
+
+    fn intern_symbol(string: &str) -> Self::Symbol {
+        Symbol::intern(string)
+    }
+
+    fn with_symbol_string(symbol: &Self::Symbol, f: impl FnOnce(&str)) {
+        f(&symbol.as_str())
+    }
 }

library/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) }
+    }
+}

library/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
 }