Auto merge of rust-lang#80441 - petrochenkov:kwtok, r=Aaron1011

ast: Remove some indirection layers from values in key-value attributes

Trying to address some perf regressions from rust-lang#78837 (comment).

compiler/rustc_ast/src/ast.rs

@@ -23,8 +23,8 @@ pub use GenericArgs::*;
 pub use UnsafeSource::*;
 
 use crate::ptr::P;
-use crate::token::{self, CommentKind, DelimToken};
-use crate::tokenstream::{DelimSpan, LazyTokenStream, TokenStream};
+use crate::token::{self, CommentKind, DelimToken, Token};
+use crate::tokenstream::{DelimSpan, LazyTokenStream, TokenStream, TokenTree};
 
 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
 use rustc_data_structures::stack::ensure_sufficient_stack;
@@ -1464,8 +1464,8 @@ pub enum MacArgs {
     Eq(
         /// Span of the `=` token.
         Span,
-        /// Token stream of the "value".
-        TokenStream,
+        /// "value" as a nonterminal token.
+        Token,
     ),
 }
 
@@ -1478,10 +1478,10 @@ impl MacArgs {
     }
 
     pub fn span(&self) -> Option<Span> {
-        match *self {
+        match self {
             MacArgs::Empty => None,
             MacArgs::Delimited(dspan, ..) => Some(dspan.entire()),
-            MacArgs::Eq(eq_span, ref tokens) => Some(eq_span.to(tokens.span().unwrap_or(eq_span))),
+            MacArgs::Eq(eq_span, token) => Some(eq_span.to(token.span)),
         }
     }
 
@@ -1490,7 +1490,8 @@ impl MacArgs {
     pub fn inner_tokens(&self) -> TokenStream {
         match self {
             MacArgs::Empty => TokenStream::default(),
-            MacArgs::Delimited(.., tokens) | MacArgs::Eq(.., tokens) => tokens.clone(),
+            MacArgs::Delimited(.., tokens) => tokens.clone(),
+            MacArgs::Eq(.., token) => TokenTree::Token(token.clone()).into(),
         }
     }
 

compiler/rustc_ast/src/attr/mod.rs

@@ -476,7 +476,7 @@ impl MetaItemKind {
     pub fn mac_args(&self, span: Span) -> MacArgs {
         match self {
             MetaItemKind::Word => MacArgs::Empty,
-            MetaItemKind::NameValue(lit) => MacArgs::Eq(span, lit.token_tree().into()),
+            MetaItemKind::NameValue(lit) => MacArgs::Eq(span, lit.to_token()),
             MetaItemKind::List(list) => {
                 let mut tts = Vec::new();
                 for (i, item) in list.iter().enumerate() {
@@ -498,7 +498,10 @@ impl MetaItemKind {
         match *self {
             MetaItemKind::Word => vec![],
             MetaItemKind::NameValue(ref lit) => {
-                vec![TokenTree::token(token::Eq, span).into(), lit.token_tree().into()]
+                vec![
+                    TokenTree::token(token::Eq, span).into(),
+                    TokenTree::Token(lit.to_token()).into(),
+                ]
             }
             MetaItemKind::List(ref list) => {
                 let mut tokens = Vec::new();
@@ -554,10 +557,7 @@ impl MetaItemKind {
                 MetaItemKind::list_from_tokens(tokens.clone())
             }
             MacArgs::Delimited(..) => None,
-            MacArgs::Eq(_, tokens) => {
-                assert!(tokens.len() == 1);
-                MetaItemKind::name_value_from_tokens(&mut tokens.trees())
-            }
+            MacArgs::Eq(_, token) => Lit::from_token(token).ok().map(MetaItemKind::NameValue),
             MacArgs::Empty => Some(MetaItemKind::Word),
         }
     }
@@ -592,7 +592,7 @@ impl NestedMetaItem {
     fn token_trees_and_spacings(&self) -> Vec<TreeAndSpacing> {
         match *self {
             NestedMetaItem::MetaItem(ref item) => item.token_trees_and_spacings(),
-            NestedMetaItem::Literal(ref lit) => vec![lit.token_tree().into()],
+            NestedMetaItem::Literal(ref lit) => vec![TokenTree::Token(lit.to_token()).into()],
         }
     }
 

compiler/rustc_ast/src/mut_visit.rs

@@ -365,18 +365,16 @@ pub fn visit_mac_args<T: MutVisitor>(args: &mut MacArgs, vis: &mut T) {
             visit_delim_span(dspan, vis);
             visit_tts(tokens, vis);
         }
-        MacArgs::Eq(eq_span, tokens) => {
+        MacArgs::Eq(eq_span, token) => {
             vis.visit_span(eq_span);
-            visit_tts(tokens, vis);
-            // The value in `#[key = VALUE]` must be visited as an expression for backward
-            // compatibility, so that macros can be expanded in that position.
-            if !vis.token_visiting_enabled() {
-                match Lrc::make_mut(&mut tokens.0).get_mut(0) {
-                    Some((TokenTree::Token(token), _spacing)) => match &mut token.kind {
-                        token::Interpolated(nt) => match Lrc::make_mut(nt) {
-                            token::NtExpr(expr) => vis.visit_expr(expr),
-                            t => panic!("unexpected token in key-value attribute: {:?}", t),
-                        },
+            if vis.token_visiting_enabled() {
+                visit_token(token, vis);
+            } else {
+                // The value in `#[key = VALUE]` must be visited as an expression for backward
+                // compatibility, so that macros can be expanded in that position.
+                match &mut token.kind {
+                    token::Interpolated(nt) => match Lrc::make_mut(nt) {
+                        token::NtExpr(expr) => vis.visit_expr(expr),
                         t => panic!("unexpected token in key-value attribute: {:?}", t),
                     },
                     t => panic!("unexpected token in key-value attribute: {:?}", t),

compiler/rustc_ast/src/util/literal.rs

@@ -2,7 +2,6 @@
 
 use crate::ast::{self, Lit, LitKind};
 use crate::token::{self, Token};
-use crate::tokenstream::TokenTree;
 
 use rustc_lexer::unescape::{unescape_byte, unescape_char};
 use rustc_lexer::unescape::{unescape_byte_literal, unescape_literal, Mode};
@@ -225,13 +224,13 @@ impl Lit {
         Lit { token: kind.to_lit_token(), kind, span }
     }
 
-    /// Losslessly convert an AST literal into a token stream.
-    pub fn token_tree(&self) -> TokenTree {
-        let token = match self.token.kind {
+    /// Losslessly convert an AST literal into a token.
+    pub fn to_token(&self) -> Token {
+        let kind = match self.token.kind {
             token::Bool => token::Ident(self.token.symbol, false),
             _ => token::Literal(self.token),
         };
-        TokenTree::token(token, self.span)
+        Token::new(kind, self.span)
     }
 }
 

compiler/rustc_ast/src/visit.rs

@@ -15,7 +15,6 @@
 
 use crate::ast::*;
 use crate::token;
-use crate::tokenstream::TokenTree;
 
 use rustc_span::symbol::{Ident, Symbol};
 use rustc_span::Span;
@@ -905,12 +904,9 @@ pub fn walk_mac_args<'a, V: Visitor<'a>>(visitor: &mut V, args: &'a MacArgs) {
         MacArgs::Delimited(_dspan, _delim, _tokens) => {}
         // The value in `#[key = VALUE]` must be visited as an expression for backward
         // compatibility, so that macros can be expanded in that position.
-        MacArgs::Eq(_eq_span, tokens) => match tokens.trees_ref().next() {
-            Some(TokenTree::Token(token)) => match &token.kind {
-                token::Interpolated(nt) => match &**nt {
-                    token::NtExpr(expr) => visitor.visit_expr(expr),
-                    t => panic!("unexpected token in key-value attribute: {:?}", t),
-                },
+        MacArgs::Eq(_eq_span, token) => match &token.kind {
+            token::Interpolated(nt) => match &**nt {
+                token::NtExpr(expr) => visitor.visit_expr(expr),
                 t => panic!("unexpected token in key-value attribute: {:?}", t),
             },
             t => panic!("unexpected token in key-value attribute: {:?}", t),

compiler/rustc_ast_lowering/src/lib.rs

@@ -1013,10 +1013,36 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
             // This is an inert key-value attribute - it will never be visible to macros
             // after it gets lowered to HIR. Therefore, we can synthesize tokens with fake
             // spans to handle nonterminals in `#[doc]` (e.g. `#[doc = $e]`).
-            MacArgs::Eq(eq_span, ref tokens) => MacArgs::Eq(
-                eq_span,
-                self.lower_token_stream(tokens.clone(), CanSynthesizeMissingTokens::Yes),
-            ),
+            MacArgs::Eq(eq_span, ref token) => {
+                // In valid code the value is always representable as a single literal token.
+                fn unwrap_single_token(sess: &Session, tokens: TokenStream, span: Span) -> Token {
+                    if tokens.len() != 1 {
+                        sess.diagnostic()
+                            .delay_span_bug(span, "multiple tokens in key-value attribute's value");
+                    }
+                    match tokens.into_trees().next() {
+                        Some(TokenTree::Token(token)) => token,
+                        Some(TokenTree::Delimited(_, delim, tokens)) => {
+                            if delim != token::NoDelim {
+                                sess.diagnostic().delay_span_bug(
+                                    span,
+                                    "unexpected delimiter in key-value attribute's value",
+                                )
+                            }
+                            unwrap_single_token(sess, tokens, span)
+                        }
+                        None => Token::dummy(),
+                    }
+                }
+
+                let tokens = TokenStreamLowering {
+                    parse_sess: &self.sess.parse_sess,
+                    synthesize_tokens: CanSynthesizeMissingTokens::Yes,
+                    nt_to_tokenstream: self.nt_to_tokenstream,
+                }
+                .lower_token(token.clone());
+                MacArgs::Eq(eq_span, unwrap_single_token(self.sess, tokens, token.span))
+            }
         }
     }
 

compiler/rustc_ast_pretty/src/pprust/state.rs

@@ -445,10 +445,11 @@ pub trait PrintState<'a>: std::ops::Deref<Target = pp::Printer> + std::ops::Dere
             ),
             MacArgs::Empty | MacArgs::Eq(..) => {
                 self.print_path(&item.path, false, 0);
-                if let MacArgs::Eq(_, tokens) = &item.args {
+                if let MacArgs::Eq(_, token) = &item.args {
                     self.space();
                     self.word_space("=");
-                    self.print_tts(tokens, true);
+                    let token_str = self.token_to_string_ext(token, true);
+                    self.word(token_str);
                 }
             }
         }

compiler/rustc_parse/src/parser/mod.rs

@@ -983,8 +983,8 @@ impl<'a> Parser<'a> {
                         _ => self.sess.gated_spans.gate(sym::extended_key_value_attributes, span),
                     }
 
-                    let token = token::Interpolated(Lrc::new(token::NtExpr(expr)));
-                    MacArgs::Eq(eq_span, TokenTree::token(token, span).into())
+                    let token_kind = token::Interpolated(Lrc::new(token::NtExpr(expr)));
+                    MacArgs::Eq(eq_span, Token::new(token_kind, span))
                 } else {
                     MacArgs::Empty
                 }

compiler/rustc_parse/src/validate_attr.rs

@@ -2,7 +2,7 @@
 
 use crate::parse_in;
 
-use rustc_ast::tokenstream::DelimSpan;
+use rustc_ast::tokenstream::{DelimSpan, TokenTree};
 use rustc_ast::{self as ast, Attribute, MacArgs, MacDelimiter, MetaItem, MetaItemKind};
 use rustc_errors::{Applicability, PResult};
 use rustc_feature::{AttributeTemplate, BUILTIN_ATTRIBUTE_MAP};
@@ -45,7 +45,8 @@ pub fn parse_meta<'a>(sess: &'a ParseSess, attr: &Attribute) -> PResult<'a, Meta
         kind: match &item.args {
             MacArgs::Empty => MetaItemKind::Word,
             MacArgs::Eq(_, t) => {
-                let v = parse_in(sess, t.clone(), "name value", |p| p.parse_unsuffixed_lit())?;
+                let t = TokenTree::Token(t.clone()).into();
+                let v = parse_in(sess, t, "name value", |p| p.parse_unsuffixed_lit())?;
                 MetaItemKind::NameValue(v)
             }
             MacArgs::Delimited(dspan, delim, t) => {

src/test/ui/ast-json/ast-json-noexpand-output.stdout

@@ -1 +1 @@
-{"module":{"inner":{"lo":0,"hi":0},"unsafety":"No","items":[{"attrs":[],"id":0,"span":{"lo":0,"hi":0},"vis":{"kind":"Inherited","span":{"lo":0,"hi":0},"tokens":null},"ident":{"name":"core","span":{"lo":0,"hi":0}},"kind":{"variant":"ExternCrate","fields":[null]},"tokens":null}],"inline":true},"attrs":[{"kind":{"variant":"Normal","fields":[{"path":{"span":{"lo":0,"hi":0},"segments":[{"ident":{"name":"crate_type","span":{"lo":0,"hi":0}},"id":0,"args":null}],"tokens":null},"args":{"variant":"Eq","fields":[{"lo":0,"hi":0},{"0":[[{"variant":"Token","fields":[{"kind":{"variant":"Interpolated","fields":[{"variant":"NtExpr","fields":[{"id":0,"kind":{"variant":"Lit","fields":[{"token":{"kind":"Str","symbol":"lib","suffix":null},"kind":{"variant":"Str","fields":["lib","Cooked"]},"span":{"lo":0,"hi":0}}]},"span":{"lo":0,"hi":0},"attrs":{"0":null},"tokens":null}]}]},"span":{"lo":0,"hi":0}}]},"Alone"]]}]},"tokens":null},{"0":[[{"variant":"Token","fields":[{"kind":"Pound","span":{"lo":0,"hi":0}}]},"Joint"],[{"variant":"Token","fields":[{"kind":"Not","span":{"lo":0,"hi":0}}]},"Alone"],[{"variant":"Delimited","fields":[{"open":{"lo":0,"hi":0},"close":{"lo":0,"hi":0}},"Bracket",{"0":[[{"variant":"Token","fields":[{"kind":{"variant":"Ident","fields":["crate_type",false]},"span":{"lo":0,"hi":0}}]},"Alone"],[{"variant":"Token","fields":[{"kind":"Eq","span":{"lo":0,"hi":0}}]},"Alone"],[{"variant":"Token","fields":[{"kind":{"variant":"Literal","fields":[{"kind":"Str","symbol":"lib","suffix":null}]},"span":{"lo":0,"hi":0}}]},"Alone"]]}]},"Alone"]]}]},"id":null,"style":"Inner","span":{"lo":0,"hi":0}}],"span":{"lo":0,"hi":0},"proc_macros":[]}
+{"module":{"inner":{"lo":0,"hi":0},"unsafety":"No","items":[{"attrs":[],"id":0,"span":{"lo":0,"hi":0},"vis":{"kind":"Inherited","span":{"lo":0,"hi":0},"tokens":null},"ident":{"name":"core","span":{"lo":0,"hi":0}},"kind":{"variant":"ExternCrate","fields":[null]},"tokens":null}],"inline":true},"attrs":[{"kind":{"variant":"Normal","fields":[{"path":{"span":{"lo":0,"hi":0},"segments":[{"ident":{"name":"crate_type","span":{"lo":0,"hi":0}},"id":0,"args":null}],"tokens":null},"args":{"variant":"Eq","fields":[{"lo":0,"hi":0},{"kind":{"variant":"Interpolated","fields":[{"variant":"NtExpr","fields":[{"id":0,"kind":{"variant":"Lit","fields":[{"token":{"kind":"Str","symbol":"lib","suffix":null},"kind":{"variant":"Str","fields":["lib","Cooked"]},"span":{"lo":0,"hi":0}}]},"span":{"lo":0,"hi":0},"attrs":{"0":null},"tokens":null}]}]},"span":{"lo":0,"hi":0}}]},"tokens":null},{"0":[[{"variant":"Token","fields":[{"kind":"Pound","span":{"lo":0,"hi":0}}]},"Joint"],[{"variant":"Token","fields":[{"kind":"Not","span":{"lo":0,"hi":0}}]},"Alone"],[{"variant":"Delimited","fields":[{"open":{"lo":0,"hi":0},"close":{"lo":0,"hi":0}},"Bracket",{"0":[[{"variant":"Token","fields":[{"kind":{"variant":"Ident","fields":["crate_type",false]},"span":{"lo":0,"hi":0}}]},"Alone"],[{"variant":"Token","fields":[{"kind":"Eq","span":{"lo":0,"hi":0}}]},"Alone"],[{"variant":"Token","fields":[{"kind":{"variant":"Literal","fields":[{"kind":"Str","symbol":"lib","suffix":null}]},"span":{"lo":0,"hi":0}}]},"Alone"]]}]},"Alone"]]}]},"id":null,"style":"Inner","span":{"lo":0,"hi":0}}],"span":{"lo":0,"hi":0},"proc_macros":[]}