Auto merge of #77271 - petrochenkov:notokenexp, r=Aaron1011

Expand `NtExpr` tokens only in key-value attributes

Implement the experiment described in #55414 (comment)

This PR also removes some customization points and token visiting functionality from AST visitors.
Read-only visitor no longer visits tokens, mutable visitor visits tokens only when specifically enabled, mutable token visiting is restricted to its single intended use case.

I haven't changed the representation of `MacArgs::Eq` yet, but it potentially can use a `TokenTree` or a `Token` instead of `TokenStream`.
It's hard to get rid of `Nonterminal::NtExpr` there (and e.g. replace it with `ast::Expr`) due to the dual nature of key-value attributes (the value is both an expression and a token stream, depending on context), and `Nonterminal` has all the machinery for maintaining both representations in sync.

Fixes #55414
Fixes #43860

Author: bors

compiler/rustc_ast/src/mut_visit.rs

@@ -34,6 +34,13 @@ impl<A: Array> ExpectOne<A> for SmallVec<A> {
 }
 
 pub trait MutVisitor: Sized {
+    /// Mutable token visiting only exists for the `macro_rules` token marker and should not be
+    /// used otherwise. Token visitor would be entirely separate from the regular visitor if
+    /// the marker didn't have to visit AST fragments in nonterminal tokens.
+    fn token_visiting_enabled(&self) -> bool {
+        false
+    }
+
     // Methods in this trait have one of three forms:
     //
     //   fn visit_t(&mut self, t: &mut T);                      // common
@@ -246,22 +253,6 @@ pub trait MutVisitor: Sized {
         noop_flat_map_generic_param(param, self)
     }
 
-    fn visit_tt(&mut self, tt: &mut TokenTree) {
-        noop_visit_tt(tt, self);
-    }
-
-    fn visit_tts(&mut self, tts: &mut TokenStream) {
-        noop_visit_tts(tts, self);
-    }
-
-    fn visit_token(&mut self, t: &mut Token) {
-        noop_visit_token(t, self);
-    }
-
-    fn visit_interpolated(&mut self, nt: &mut token::Nonterminal) {
-        noop_visit_interpolated(nt, self);
-    }
-
     fn visit_param_bound(&mut self, tpb: &mut GenericBound) {
         noop_visit_param_bound(tpb, self);
     }
@@ -375,11 +366,30 @@ pub fn visit_mac_args<T: MutVisitor>(args: &mut MacArgs, vis: &mut T) {
         MacArgs::Empty => {}
         MacArgs::Delimited(dspan, _delim, tokens) => {
             visit_delim_span(dspan, vis);
-            vis.visit_tts(tokens);
+            visit_tts(tokens, vis);
         }
         MacArgs::Eq(eq_span, tokens) => {
             vis.visit_span(eq_span);
-            vis.visit_tts(tokens);
+            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() {
+                if let Some(TokenTree::Token(token)) = tokens.trees_ref().next() {
+                    if let token::Interpolated(..) = token.kind {
+                        // ^^ Do not `make_mut` unless we have to.
+                        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),
+                                },
+                                t => panic!("unexpected token in key-value attribute: {:?}", t),
+                            },
+                            t => panic!("unexpected token in key-value attribute: {:?}", t),
+                        }
+                    }
+                }
+            }
         }
     }
 }
@@ -626,28 +636,33 @@ pub fn noop_flat_map_param<T: MutVisitor>(mut param: Param, vis: &mut T) -> Smal
     smallvec![param]
 }
 
-pub fn noop_visit_tt<T: MutVisitor>(tt: &mut TokenTree, vis: &mut T) {
+// No `noop_` prefix because there isn't a corresponding method in `MutVisitor`.
+pub fn visit_tt<T: MutVisitor>(tt: &mut TokenTree, vis: &mut T) {
     match tt {
         TokenTree::Token(token) => {
-            vis.visit_token(token);
+            visit_token(token, vis);
         }
         TokenTree::Delimited(DelimSpan { open, close }, _delim, tts) => {
             vis.visit_span(open);
             vis.visit_span(close);
-            vis.visit_tts(tts);
+            visit_tts(tts, vis);
         }
     }
 }
 
-pub fn noop_visit_tts<T: MutVisitor>(TokenStream(tts): &mut TokenStream, vis: &mut T) {
-    let tts = Lrc::make_mut(tts);
-    visit_vec(tts, |(tree, _is_joint)| vis.visit_tt(tree));
+// No `noop_` prefix because there isn't a corresponding method in `MutVisitor`.
+pub fn visit_tts<T: MutVisitor>(TokenStream(tts): &mut TokenStream, vis: &mut T) {
+    if vis.token_visiting_enabled() {
+        let tts = Lrc::make_mut(tts);
+        visit_vec(tts, |(tree, _is_joint)| visit_tt(tree, vis));
+    }
 }
 
+// No `noop_` prefix because there isn't a corresponding method in `MutVisitor`.
 // Applies ident visitor if it's an ident; applies other visits to interpolated nodes.
 // In practice the ident part is not actually used by specific visitors right now,
 // but there's a test below checking that it works.
-pub fn noop_visit_token<T: MutVisitor>(t: &mut Token, vis: &mut T) {
+pub fn visit_token<T: MutVisitor>(t: &mut Token, vis: &mut T) {
     let Token { kind, span } = t;
     match kind {
         token::Ident(name, _) | token::Lifetime(name) => {
@@ -659,13 +674,14 @@ pub fn noop_visit_token<T: MutVisitor>(t: &mut Token, vis: &mut T) {
         }
         token::Interpolated(nt) => {
             let mut nt = Lrc::make_mut(nt);
-            vis.visit_interpolated(&mut nt);
+            visit_interpolated(&mut nt, vis);
         }
         _ => {}
     }
     vis.visit_span(span);
 }
 
+// No `noop_` prefix because there isn't a corresponding method in `MutVisitor`.
 /// Applies the visitor to elements of interpolated nodes.
 //
 // N.B., this can occur only when applying a visitor to partially expanded
@@ -689,7 +705,7 @@ pub fn noop_visit_token<T: MutVisitor>(t: &mut Token, vis: &mut T) {
 // contain multiple items, but decided against it when I looked at
 // `parse_item_or_view_item` and tried to figure out what I would do with
 // multiple items there....
-pub fn noop_visit_interpolated<T: MutVisitor>(nt: &mut token::Nonterminal, vis: &mut T) {
+pub fn visit_interpolated<T: MutVisitor>(nt: &mut token::Nonterminal, vis: &mut T) {
     match nt {
         token::NtItem(item) => visit_clobber(item, |item| {
             // This is probably okay, because the only visitors likely to
@@ -714,7 +730,7 @@ pub fn noop_visit_interpolated<T: MutVisitor>(nt: &mut token::Nonterminal, vis:
             visit_mac_args(args, vis);
         }
         token::NtPath(path) => vis.visit_path(path),
-        token::NtTT(tt) => vis.visit_tt(tt),
+        token::NtTT(tt) => visit_tt(tt, vis),
         token::NtVis(visib) => vis.visit_vis(visib),
     }
 }

compiler/rustc_ast/src/tokenstream.rs

@@ -318,6 +318,10 @@ impl TokenStream {
         }
     }
 
+    pub fn trees_ref(&self) -> CursorRef<'_> {
+        CursorRef::new(self)
+    }
+
     pub fn trees(&self) -> Cursor {
         self.clone().into_trees()
     }
@@ -408,6 +412,36 @@ impl TokenStreamBuilder {
     }
 }
 
+/// By-reference iterator over a `TokenStream`.
+#[derive(Clone)]
+pub struct CursorRef<'t> {
+    stream: &'t TokenStream,
+    index: usize,
+}
+
+impl<'t> CursorRef<'t> {
+    fn new(stream: &TokenStream) -> CursorRef<'_> {
+        CursorRef { stream, index: 0 }
+    }
+
+    fn next_with_spacing(&mut self) -> Option<&'t TreeAndSpacing> {
+        self.stream.0.get(self.index).map(|tree| {
+            self.index += 1;
+            tree
+        })
+    }
+}
+
+impl<'t> Iterator for CursorRef<'t> {
+    type Item = &'t TokenTree;
+
+    fn next(&mut self) -> Option<&'t TokenTree> {
+        self.next_with_spacing().map(|(tree, _)| tree)
+    }
+}
+
+/// Owning by-value iterator over a `TokenStream`.
+/// FIXME: Many uses of this can be replaced with by-reference iterator to avoid clones.
 #[derive(Clone)]
 pub struct Cursor {
     pub stream: TokenStream,

compiler/rustc_ast/src/visit.rs

@@ -14,8 +14,8 @@
 //! those that are created by the expansion of a macro.
 
 use crate::ast::*;
-use crate::token::Token;
-use crate::tokenstream::{TokenStream, TokenTree};
+use crate::token;
+use crate::tokenstream::TokenTree;
 
 use rustc_span::symbol::{Ident, Symbol};
 use rustc_span::Span;
@@ -208,14 +208,6 @@ pub trait Visitor<'ast>: Sized {
     fn visit_attribute(&mut self, attr: &'ast Attribute) {
         walk_attribute(self, attr)
     }
-    fn visit_tt(&mut self, tt: TokenTree) {
-        walk_tt(self, tt)
-    }
-    fn visit_tts(&mut self, tts: TokenStream) {
-        walk_tts(self, tts)
-    }
-    fn visit_token(&mut self, _t: Token) {}
-    // FIXME: add `visit_interpolated` and `walk_interpolated`
     fn visit_vis(&mut self, vis: &'ast Visibility) {
         walk_vis(self, vis)
     }
@@ -902,20 +894,19 @@ pub fn walk_attribute<'a, V: Visitor<'a>>(visitor: &mut V, attr: &'a Attribute)
 pub fn walk_mac_args<'a, V: Visitor<'a>>(visitor: &mut V, args: &'a MacArgs) {
     match args {
         MacArgs::Empty => {}
-        MacArgs::Delimited(_dspan, _delim, tokens) => visitor.visit_tts(tokens.clone()),
-        MacArgs::Eq(_eq_span, tokens) => visitor.visit_tts(tokens.clone()),
-    }
-}
-
-pub fn walk_tt<'a, V: Visitor<'a>>(visitor: &mut V, tt: TokenTree) {
-    match tt {
-        TokenTree::Token(token) => visitor.visit_token(token),
-        TokenTree::Delimited(_, _, tts) => visitor.visit_tts(tts),
-    }
-}
-
-pub fn walk_tts<'a, V: Visitor<'a>>(visitor: &mut V, tts: TokenStream) {
-    for tt in tts.trees() {
-        visitor.visit_tt(tt);
+        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),
+                },
+                token::Literal(..) | token::Ident(..) => {}
+                t => panic!("unexpected token in key-value attribute: {:?}", t),
+            },
+            t => panic!("unexpected token in key-value attribute: {:?}", t),
+        },
     }
 }

compiler/rustc_expand/src/mbe/transcribe.rs

@@ -20,6 +20,10 @@ use std::mem;
 struct Marker(ExpnId, Transparency);
 
 impl MutVisitor for Marker {
+    fn token_visiting_enabled(&self) -> bool {
+        true
+    }
+
     fn visit_span(&mut self, span: &mut Span) {
         *span = span.apply_mark(self.0, self.1)
     }
@@ -277,7 +281,7 @@ pub(super) fn transcribe<'a>(
             // preserve syntax context.
             mbe::TokenTree::Token(token) => {
                 let mut tt = TokenTree::Token(token);
-                marker.visit_tt(&mut tt);
+                mut_visit::visit_tt(&mut tt, &mut marker);
                 result.push(tt.into());
             }
 

compiler/rustc_expand/src/mut_visit/tests.rs

@@ -15,6 +15,9 @@ fn fake_print_crate(s: &mut pprust::State<'_>, krate: &ast::Crate) {
 struct ToZzIdentMutVisitor;
 
 impl MutVisitor for ToZzIdentMutVisitor {
+    fn token_visiting_enabled(&self) -> bool {
+        true
+    }
     fn visit_ident(&mut self, ident: &mut Ident) {
         *ident = Ident::from_str("zz");
     }

compiler/rustc_resolve/src/build_reduced_graph.rs

@@ -15,7 +15,6 @@ use crate::{
 };
 use crate::{Module, ModuleData, ModuleKind, NameBinding, NameBindingKind, Segment, ToNameBinding};
 
-use rustc_ast::token::{self, Token};
 use rustc_ast::visit::{self, AssocCtxt, Visitor};
 use rustc_ast::{self as ast, Block, ForeignItem, ForeignItemKind, Item, ItemKind, NodeId};
 use rustc_ast::{AssocItem, AssocItemKind, MetaItemKind, StmtKind};
@@ -1395,16 +1394,6 @@ impl<'a, 'b> Visitor<'b> for BuildReducedGraphVisitor<'a, 'b> {
         visit::walk_assoc_item(self, item, ctxt);
     }
 
-    fn visit_token(&mut self, t: Token) {
-        if let token::Interpolated(nt) = t.kind {
-            if let token::NtExpr(ref expr) = *nt {
-                if let ast::ExprKind::MacCall(..) = expr.kind {
-                    self.visit_invoc(expr.id);
-                }
-            }
-        }
-    }
-
     fn visit_attribute(&mut self, attr: &'b ast::Attribute) {
         if !attr.is_doc_comment() && attr::is_builtin_attr(attr) {
             self.r

compiler/rustc_resolve/src/def_collector.rs

@@ -1,5 +1,4 @@
 use crate::Resolver;
-use rustc_ast::token::{self, Token};
 use rustc_ast::visit::{self, FnKind};
 use rustc_ast::walk_list;
 use rustc_ast::*;
@@ -256,16 +255,6 @@ impl<'a, 'b> visit::Visitor<'a> for DefCollector<'a, 'b> {
         }
     }
 
-    fn visit_token(&mut self, t: Token) {
-        if let token::Interpolated(nt) = t.kind {
-            if let token::NtExpr(ref expr) = *nt {
-                if let ExprKind::MacCall(..) = expr.kind {
-                    self.visit_macro_invoc(expr.id);
-                }
-            }
-        }
-    }
-
     fn visit_arm(&mut self, arm: &'a Arm) {
         if arm.is_placeholder { self.visit_macro_invoc(arm.id) } else { visit::walk_arm(self, arm) }
     }