Auto merge of #98189 - mystor:fast_ident_literal, r=eddyb

proc_macro/bridge: stop using a remote object handle for proc_macro Ident and Literal

This is the fourth part of https://github.com/rust-lang/rust/pull/86822, split off as requested in https://github.com/rust-lang/rust/pull/86822#pullrequestreview-1008655452. This patch transforms the `Ident` and `Group` types into structs serialized over IPC rather than handles.

Symbol values are interned on both the client and server when deserializing, to avoid unnecessary string copies and keep the size of `TokenTree` down. To do the interning efficiently on the client, the proc-macro crate is given a vendored version of the fxhash hasher, as `SipHash` appeared to cause performance issues. This was done rather than depending on `rustc_hash` as it is unfortunately difficult to depend on crates from within `proc_macro` due to it being built at the same time as `std`.

In addition, a custom arena allocator and symbol store was also added, inspired by those in `rustc_arena` and `rustc_span`. To prevent symbol re-use across multiple invocations of a macro on the same thread, a new range of `Symbol` names are used for each invocation of the macro, and symbols from previous invocations are cleaned-up.

In order to keep `Ident` creation efficient, a special ASCII-only case was added to perform ident validation without using RPC for simple identifiers. Full identifier validation couldn't be easily added, as it would require depending on the `rustc_lexer` and `unicode-normalization` crates from within `proc_macro`. Unicode identifiers are validated and normalized using RPC.

See the individual commit messages for more details on trade-offs and design decisions behind these patches.

1 files changed, 173 insertions, 229 deletions

diff --git a/compiler/rustc_expand/src/proc_macro_server.rs b/compiler/rustc_expand/src/proc_macro_server.rs
index 411d6be9c44..176c77ca6ed 100644
--- a/compiler/rustc_expand/src/proc_macro_server.rs
+++ b/compiler/rustc_expand/src/proc_macro_server.rs
@@ -11,13 +11,14 @@ 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, LitKind, Literal, Punct, TokenTree,
+};
 use pm::{Delimiter, Level, LineColumn};
 use std::ops::Bound;
-use std::{ascii, panic};
 
 trait FromInternal<T> {
     fn from_internal(x: T) -> Self;
@@ -49,9 +50,40 @@ impl ToInternal<token::Delimiter> for Delimiter {
     }
 }
 
-impl FromInternal<(TokenStream, &mut Rustc<'_, '_>)>
-    for Vec<TokenTree<TokenStream, Span, Ident, Literal>>
-{
+impl FromInternal<token::LitKind> for LitKind {
+    fn from_internal(kind: token::LitKind) -> Self {
+        match kind {
+            token::Byte => LitKind::Byte,
+            token::Char => LitKind::Char,
+            token::Integer => LitKind::Integer,
+            token::Float => LitKind::Float,
+            token::Str => LitKind::Str,
+            token::StrRaw(n) => LitKind::StrRaw(n),
+            token::ByteStr => LitKind::ByteStr,
+            token::ByteStrRaw(n) => LitKind::ByteStrRaw(n),
+            token::Err => LitKind::Err,
+            token::Bool => unreachable!(),
+        }
+    }
+}
+
+impl ToInternal<token::LitKind> for LitKind {
+    fn to_internal(self) -> token::LitKind {
+        match self {
+            LitKind::Byte => token::Byte,
+            LitKind::Char => token::Char,
+            LitKind::Integer => token::Integer,
+            LitKind::Float => token::Float,
+            LitKind::Str => token::Str,
+            LitKind::StrRaw(n) => token::StrRaw(n),
+            LitKind::ByteStr => token::ByteStr,
+            LitKind::ByteStrRaw(n) => token::ByteStrRaw(n),
+            LitKind::Err => token::Err,
+        }
+    }
+}
+
+impl FromInternal<(TokenStream, &mut Rustc<'_, '_>)> for Vec<TokenTree<TokenStream, Span, Symbol>> {
     fn from_internal((stream, rustc): (TokenStream, &mut Rustc<'_, '_>)) -> Self {
         use rustc_ast::token::*;
 
@@ -135,16 +167,22 @@ 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 })),
+                Literal(token::Lit { kind, symbol, suffix }) => {
+                    trees.push(TokenTree::Literal(self::Literal {
+                        kind: FromInternal::from_internal(kind),
+                        symbol,
+                        suffix,
+                        span,
+                    }));
+                }
                 DocComment(_, attr_style, data) => {
                     let mut escaped = String::new();
                     for ch in data.as_str().chars() {
@@ -170,7 +208,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 +238,12 @@ impl FromInternal<(TokenStream, &mut Rustc<'_, '_>)>
     }
 }
 
-impl ToInternal<TokenStream> for TokenTree<TokenStream, Span, Ident, Literal> {
+impl ToInternal<TokenStream> for (TokenTree<TokenStream, Span, Symbol>, &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,10 +254,13 @@ 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 {
-                lit: token::Lit { kind: token::Integer, symbol, suffix },
+                kind: self::LitKind::Integer,
+                symbol,
+                suffix,
                 span,
             }) if symbol.as_str().starts_with('-') => {
                 let minus = BinOp(BinOpToken::Minus);
@@ -229,7 +271,9 @@ impl ToInternal<TokenStream> for TokenTree<TokenStream, Span, Ident, Literal> {
                 return [a, b].into_iter().collect();
             }
             TokenTree::Literal(self::Literal {
-                lit: token::Lit { kind: token::Float, symbol, suffix },
+                kind: self::LitKind::Float,
+                symbol,
+                suffix,
                 span,
             }) if symbol.as_str().starts_with('-') => {
                 let minus = BinOp(BinOpToken::Minus);
@@ -239,8 +283,12 @@ impl ToInternal<TokenStream> for TokenTree<TokenStream, Span, Ident, Literal> {
                 let b = tokenstream::TokenTree::token(float, span);
                 return [a, b].into_iter().collect();
             }
-            TokenTree::Literal(self::Literal { lit, span }) => {
-                return tokenstream::TokenTree::token(Literal(lit), span).into();
+            TokenTree::Literal(self::Literal { kind, symbol, suffix, span }) => {
+                return tokenstream::TokenTree::token(
+                    TokenKind::lit(kind.to_internal(), symbol, suffix),
+                    span,
+                )
+                .into();
             }
         };
 
@@ -289,40 +337,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 {
-    lit: token::Lit,
-    span: Span,
-}
-
 pub(crate) struct Rustc<'a, 'b> {
     ecx: &'a mut ExtCtxt<'b>,
     def_site: Span,
@@ -348,21 +362,16 @@ impl<'a, 'b> Rustc<'a, 'b> {
     fn sess(&self) -> &ParseSess {
         self.ecx.parse_sess()
     }
-
-    fn lit(&mut self, kind: token::LitKind, symbol: Symbol, suffix: Option<Symbol>) -> Literal {
-        Literal { lit: token::Lit::new(kind, symbol, suffix), span: self.call_site }
-    }
 }
 
 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<'_, '_> {
@@ -376,6 +385,57 @@ impl server::FreeFunctions for Rustc<'_, '_> {
     fn track_path(&mut self, path: &str) {
         self.sess().file_depinfo.borrow_mut().insert(Symbol::intern(path));
     }
+
+    fn literal_from_str(&mut self, s: &str) -> Result<Literal<Self::Span, Self::Symbol>, ()> {
+        let name = FileName::proc_macro_source_code(s);
+        let mut parser = rustc_parse::new_parser_from_source_str(self.sess(), name, s.to_owned());
+
+        let first_span = parser.token.span.data();
+        let minus_present = parser.eat(&token::BinOp(token::Minus));
+
+        let lit_span = parser.token.span.data();
+        let token::Literal(mut lit) = parser.token.kind else {
+            return Err(());
+        };
+
+        // Check no comment or whitespace surrounding the (possibly negative)
+        // literal, or more tokens after it.
+        if (lit_span.hi.0 - first_span.lo.0) as usize != s.len() {
+            return Err(());
+        }
+
+        if minus_present {
+            // If minus is present, check no comment or whitespace in between it
+            // and the literal token.
+            if first_span.hi.0 != lit_span.lo.0 {
+                return Err(());
+            }
+
+            // Check literal is a kind we allow to be negated in a proc macro token.
+            match lit.kind {
+                token::LitKind::Bool
+                | token::LitKind::Byte
+                | token::LitKind::Char
+                | token::LitKind::Str
+                | token::LitKind::StrRaw(_)
+                | token::LitKind::ByteStr
+                | token::LitKind::ByteStrRaw(_)
+                | token::LitKind::Err => return Err(()),
+                token::LitKind::Integer | token::LitKind::Float => {}
+            }
+
+            // Synthesize a new symbol that includes the minus sign.
+            let symbol = Symbol::intern(&s[..1 + lit.symbol.as_str().len()]);
+            lit = token::Lit::new(lit.kind, symbol, lit.suffix);
+        }
+        let token::Lit { kind, symbol, suffix } = lit;
+        Ok(Literal {
+            kind: FromInternal::from_internal(kind),
+            symbol,
+            suffix,
+            span: self.call_site,
+        })
+    }
 }
 
 impl server::TokenStream for Rustc<'_, '_> {
@@ -453,22 +513,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::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::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,178 +551,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>> {
         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);
-        let mut parser = rustc_parse::new_parser_from_source_str(self.sess(), name, s.to_owned());
-
-        let first_span = parser.token.span.data();
-        let minus_present = parser.eat(&token::BinOp(token::Minus));
-
-        let lit_span = parser.token.span.data();
-        let token::Literal(mut lit) = parser.token.kind else {
-            return Err(());
-        };
-
-        // Check no comment or whitespace surrounding the (possibly negative)
-        // literal, or more tokens after it.
-        if (lit_span.hi.0 - first_span.lo.0) as usize != s.len() {
-            return Err(());
-        }
-
-        if minus_present {
-            // If minus is present, check no comment or whitespace in between it
-            // and the literal token.
-            if first_span.hi.0 != lit_span.lo.0 {
-                return Err(());
-            }
-
-            // Check literal is a kind we allow to be negated in a proc macro token.
-            match lit.kind {
-                token::LitKind::Bool
-                | token::LitKind::Byte
-                | token::LitKind::Char
-                | token::LitKind::Str
-                | token::LitKind::StrRaw(_)
-                | token::LitKind::ByteStr
-                | token::LitKind::ByteStrRaw(_)
-                | token::LitKind::Err => return Err(()),
-                token::LitKind::Integer | token::LitKind::Float => {}
-            }
-
-            // Synthesize a new symbol that includes the minus sign.
-            let symbol = Symbol::intern(&s[..1 + lit.symbol.as_str().len()]);
-            lit = token::Lit::new(lit.kind, symbol, lit.suffix);
-        }
-
-        Ok(Literal { lit, span: self.call_site })
-    }
-
-    fn to_string(&mut self, literal: &Self::Literal) -> String {
-        literal.lit.to_string()
-    }
-
-    fn debug_kind(&mut self, literal: &Self::Literal) -> String {
-        format!("{:?}", literal.lit.kind)
-    }
-
-    fn symbol(&mut self, literal: &Self::Literal) -> String {
-        literal.lit.symbol.to_string()
-    }
-
-    fn suffix(&mut self, literal: &Self::Literal) -> Option<String> {
-        literal.lit.suffix.as_ref().map(Symbol::to_string)
-    }
-
-    fn integer(&mut self, n: &str) -> Self::Literal {
-        self.lit(token::Integer, Symbol::intern(n), None)
-    }
-
-    fn typed_integer(&mut self, n: &str, kind: &str) -> Self::Literal {
-        self.lit(token::Integer, Symbol::intern(n), Some(Symbol::intern(kind)))
-    }
-
-    fn float(&mut self, n: &str) -> Self::Literal {
-        self.lit(token::Float, Symbol::intern(n), None)
-    }
-
-    fn f32(&mut self, n: &str) -> Self::Literal {
-        self.lit(token::Float, Symbol::intern(n), Some(sym::f32))
-    }
-
-    fn f64(&mut self, n: &str) -> Self::Literal {
-        self.lit(token::Float, Symbol::intern(n), Some(sym::f64))
-    }
-
-    fn string(&mut self, string: &str) -> Self::Literal {
-        let quoted = format!("{:?}", string);
-        assert!(quoted.starts_with('"') && quoted.ends_with('"'));
-        let symbol = &quoted[1..quoted.len() - 1];
-        self.lit(token::Str, Symbol::intern(symbol), None)
-    }
-
-    fn character(&mut self, ch: char) -> Self::Literal {
-        let quoted = format!("{:?}", ch);
-        assert!(quoted.starts_with('\'') && quoted.ends_with('\''));
-        let symbol = &quoted[1..quoted.len() - 1];
-        self.lit(token::Char, Symbol::intern(symbol), None)
-    }
-
-    fn byte_string(&mut self, bytes: &[u8]) -> Self::Literal {
-        let string = bytes
-            .iter()
-            .cloned()
-            .flat_map(ascii::escape_default)
-            .map(Into::<char>::into)
-            .collect::<String>();
-        self.lit(token::ByteStr, Symbol::intern(&string), None)
-    }
-
-    fn span(&mut self, literal: &Self::Literal) -> Self::Span {
-        literal.span
-    }
-
-    fn set_span(&mut self, literal: &mut Self::Literal, span: Self::Span) {
-        literal.span = span;
-    }
-
-    fn subspan(
-        &mut self,
-        literal: &Self::Literal,
-        start: Bound<usize>,
-        end: Bound<usize>,
-    ) -> Option<Self::Span> {
-        let span = literal.span;
-        let length = span.hi().to_usize() - span.lo().to_usize();
-
-        let start = match start {
-            Bound::Included(lo) => lo,
-            Bound::Excluded(lo) => lo.checked_add(1)?,
-            Bound::Unbounded => 0,
-        };
-
-        let end = match end {
-            Bound::Included(hi) => hi.checked_add(1)?,
-            Bound::Excluded(hi) => hi,
-            Bound::Unbounded => length,
-        };
-
-        // Bounds check the values, preventing addition overflow and OOB spans.
-        if start > u32::MAX as usize
-            || end > u32::MAX as usize
-            || (u32::MAX - start as u32) < span.lo().to_u32()
-            || (u32::MAX - end as u32) < span.lo().to_u32()
-            || start >= end
-            || end > length
-        {
-            return None;
-        }
-
-        let new_lo = span.lo() + BytePos::from_usize(start);
-        let new_hi = span.lo() + BytePos::from_usize(end);
-        Some(span.with_lo(new_lo).with_hi(new_hi))
-    }
-}
-
 impl server::SourceFile for Rustc<'_, '_> {
     fn eq(&mut self, file1: &Self::SourceFile, file2: &Self::SourceFile) -> bool {
         Lrc::ptr_eq(file1, file2)
@@ -767,6 +660,42 @@ impl server::Span for Rustc<'_, '_> {
         Some(first.to(second))
     }
 
+    fn subspan(
+        &mut self,
+        span: Self::Span,
+        start: Bound<usize>,
+        end: Bound<usize>,
+    ) -> Option<Self::Span> {
+        let length = span.hi().to_usize() - span.lo().to_usize();
+
+        let start = match start {
+            Bound::Included(lo) => lo,
+            Bound::Excluded(lo) => lo.checked_add(1)?,
+            Bound::Unbounded => 0,
+        };
+
+        let end = match end {
+            Bound::Included(hi) => hi.checked_add(1)?,
+            Bound::Excluded(hi) => hi,
+            Bound::Unbounded => length,
+        };
+
+        // Bounds check the values, preventing addition overflow and OOB spans.
+        if start > u32::MAX as usize
+            || end > u32::MAX as usize
+            || (u32::MAX - start as u32) < span.lo().to_u32()
+            || (u32::MAX - end as u32) < span.lo().to_u32()
+            || start >= end
+            || end > length
+        {
+            return None;
+        }
+
+        let new_lo = span.lo() + BytePos::from_usize(start);
+        let new_hi = span.lo() + BytePos::from_usize(end);
+        Some(span.with_lo(new_lo).with_hi(new_hi))
+    }
+
     fn resolved_at(&mut self, span: Self::Span, at: Self::Span) -> Self::Span {
         span.with_ctxt(at.ctxt())
     }
@@ -812,6 +741,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 +756,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())
+    }
 }