macro_rules: Less hacky heuristic for tt macro parameter spans

petrochenkov · petrochenkov · commit a7e862998338 · 2023-12-22T02:17:59.000+03:00
diff --git a/compiler/rustc_ast/src/tokenstream.rs b/compiler/rustc_ast/src/tokenstream.rs
@@ -26,7 +26,7 @@ use rustc_span::{sym, Span, Symbol, DUMMY_SP};
 use smallvec::{smallvec, SmallVec};
 
 use std::borrow::Cow;
-use std::{cmp, fmt, iter, mem};
+use std::{cmp, fmt, iter};
 
 /// When the main Rust parser encounters a syntax-extension invocation, it
 /// parses the arguments to the invocation as a token tree. This is a very
@@ -81,14 +81,6 @@ impl TokenTree {
         }
     }
 
-    /// Modify the `TokenTree`'s span in-place.
-    pub fn set_span(&mut self, span: Span) {
-        match self {
-            TokenTree::Token(token, _) => token.span = span,
-            TokenTree::Delimited(dspan, ..) => *dspan = DelimSpan::from_single(span),
-        }
-    }
-
     /// Create a `TokenTree::Token` with alone spacing.
     pub fn token_alone(kind: TokenKind, span: Span) -> TokenTree {
         TokenTree::Token(Token::new(kind, span), Spacing::Alone)
@@ -461,19 +453,6 @@ impl TokenStream {
         t1.next().is_none() && t2.next().is_none()
     }
 
-    /// Applies the supplied function to each `TokenTree` and its index in `self`, returning a new `TokenStream`
-    ///
-    /// It is equivalent to `TokenStream::new(self.trees().cloned().enumerate().map(|(i, tt)| f(i, tt)).collect())`.
-    pub fn map_enumerated_owned(
-        mut self,
-        mut f: impl FnMut(usize, TokenTree) -> TokenTree,
-    ) -> TokenStream {
-        let owned = Lrc::make_mut(&mut self.0); // clone if necessary
-        // rely on vec's in-place optimizations to avoid another allocation
-        *owned = mem::take(owned).into_iter().enumerate().map(|(i, tree)| f(i, tree)).collect();
-        self
-    }
-
     /// Create a token stream containing a single token with alone spacing. The
     /// spacing used for the final token in a constructed stream doesn't matter
     /// because it's never used. In practice we arbitrarily use
diff --git a/compiler/rustc_expand/src/mbe.rs b/compiler/rustc_expand/src/mbe.rs
@@ -72,7 +72,7 @@ enum TokenTree {
     /// A kleene-style repetition sequence, e.g. `$($e:expr)*` (RHS) or `$($e),*` (LHS).
     Sequence(DelimSpan, SequenceRepetition),
     /// e.g., `$var`.
-    MetaVar(Span, Ident),
+    MetaVar(Span, Ident, bool),
     /// e.g., `$var:expr`. Only appears on the LHS.
     MetaVarDecl(Span, Ident /* name to bind */, Option<NonterminalKind>),
     /// A meta-variable expression inside `${...}`.
@@ -97,7 +97,7 @@ impl TokenTree {
     fn span(&self) -> Span {
         match *self {
             TokenTree::Token(Token { span, .. })
-            | TokenTree::MetaVar(span, _)
+            | TokenTree::MetaVar(span, ..)
             | TokenTree::MetaVarDecl(span, _, _) => span,
             TokenTree::Delimited(span, ..)
             | TokenTree::MetaVarExpr(span, _)
diff --git a/compiler/rustc_expand/src/mbe/macro_check.rs b/compiler/rustc_expand/src/mbe/macro_check.rs
@@ -242,7 +242,7 @@ fn check_binders(
         // the outer macro. See ui/macros/macro-of-higher-order.rs where $y:$fragment in the
         // LHS of the nested macro (and RHS of the outer macro) is parsed as MetaVar(y) Colon
         // MetaVar(fragment) and not as MetaVarDecl(y, fragment).
-        TokenTree::MetaVar(span, name) => {
+        TokenTree::MetaVar(span, name, _) => {
             if macros.is_empty() {
                 sess.dcx.span_bug(span, "unexpected MetaVar in lhs");
             }
@@ -342,7 +342,7 @@ fn check_occurrences(
         TokenTree::MetaVarDecl(span, _name, _kind) => {
             sess.dcx.span_bug(span, "unexpected MetaVarDecl in rhs")
         }
-        TokenTree::MetaVar(span, name) => {
+        TokenTree::MetaVar(span, name, _) => {
             let name = MacroRulesNormalizedIdent::new(name);
             check_ops_is_prefix(sess, node_id, macros, binders, ops, span, name);
         }
diff --git a/compiler/rustc_expand/src/mbe/macro_rules.rs b/compiler/rustc_expand/src/mbe/macro_rules.rs
@@ -10,7 +10,7 @@ use crate::mbe::transcribe::transcribe;
 
 use rustc_ast as ast;
 use rustc_ast::token::{self, Delimiter, NonterminalKind, Token, TokenKind, TokenKind::*};
-use rustc_ast::tokenstream::{DelimSpan, TokenStream, TokenTree};
+use rustc_ast::tokenstream::{DelimSpan, TokenStream};
 use rustc_ast::{NodeId, DUMMY_NODE_ID};
 use rustc_ast_pretty::pprust;
 use rustc_attr::{self as attr, TransparencyError};
@@ -213,45 +213,14 @@ fn expand_macro<'cx>(
             let arm_span = rhses[i].span();
 
             // rhs has holes ( `$id` and `$(...)` that need filled)
-            let mut tts = match transcribe(cx, &named_matches, rhs, rhs_span, transparency) {
+            let tts = match transcribe(cx, &named_matches, rhs, rhs_span, transparency) {
                 Ok(tts) => tts,
                 Err(mut err) => {
                     err.emit();
                     return DummyResult::any(arm_span);
                 }
             };
 
-            // Replace all the tokens for the corresponding positions in the macro, to maintain
-            // proper positions in error reporting, while maintaining the macro_backtrace.
-            if tts.len() == rhs.tts.len() {
-                tts = tts.map_enumerated_owned(|i, mut tt| {
-                    let rhs_tt = &rhs.tts[i];
-                    let ctxt = tt.span().ctxt();
-                    match (&mut tt, rhs_tt) {
-                        // preserve the delim spans if able
-                        (
-                            TokenTree::Delimited(target_sp, ..),
-                            mbe::TokenTree::Delimited(source_sp, ..),
-                        ) => {
-                            target_sp.open = source_sp.open.with_ctxt(ctxt);
-                            target_sp.close = source_sp.close.with_ctxt(ctxt);
-                        }
-                        (
-                            TokenTree::Delimited(target_sp, ..),
-                            mbe::TokenTree::MetaVar(source_sp, ..),
-                        ) => {
-                            target_sp.open = source_sp.with_ctxt(ctxt);
-                            target_sp.close = source_sp.with_ctxt(ctxt).shrink_to_hi();
-                        }
-                        _ => {
-                            let sp = rhs_tt.span().with_ctxt(ctxt);
-                            tt.set_span(sp);
-                        }
-                    }
-                    tt
-                });
-            }
-
             if cx.trace_macros() {
                 let msg = format!("to `{}`", pprust::tts_to_string(&tts));
                 trace_macros_note(&mut cx.expansions, sp, msg);
@@ -1409,7 +1378,7 @@ fn is_in_follow(tok: &mbe::TokenTree, kind: NonterminalKind) -> IsInFollow {
 fn quoted_tt_to_string(tt: &mbe::TokenTree) -> String {
     match tt {
         mbe::TokenTree::Token(token) => pprust::token_to_string(token).into(),
-        mbe::TokenTree::MetaVar(_, name) => format!("${name}"),
+        mbe::TokenTree::MetaVar(_, name, _) => format!("${name}"),
         mbe::TokenTree::MetaVarDecl(_, name, Some(kind)) => format!("${name}:{kind}"),
         mbe::TokenTree::MetaVarDecl(_, name, None) => format!("${name}:"),
         _ => panic!(
diff --git a/compiler/rustc_expand/src/mbe/quoted.rs b/compiler/rustc_expand/src/mbe/quoted.rs
@@ -54,7 +54,7 @@ pub(super) fn parse(
         // parse out the matcher (i.e., in `$id:ident` this would parse the `:` and `ident`).
         let tree = parse_tree(tree, &mut trees, parsing_patterns, sess, node_id, features, edition);
         match tree {
-            TokenTree::MetaVar(start_sp, ident) if parsing_patterns => {
+            TokenTree::MetaVar(start_sp, ident, _) if parsing_patterns => {
                 let span = match trees.next() {
                     Some(&tokenstream::TokenTree::Token(Token { kind: token::Colon, span }, _)) => {
                         match trees.next() {
@@ -202,10 +202,18 @@ fn parse_tree<'a>(
                     // If we didn't find a metavar expression above, then we must have a
                     // repetition sequence in the macro (e.g. `$(pat)*`). Parse the
                     // contents of the sequence itself
-                    let sequence = parse(tts, parsing_patterns, sess, node_id, features, edition);
+                    let mut sequence =
+                        parse(tts, parsing_patterns, sess, node_id, features, edition);
                     // Get the Kleene operator and optional separator
                     let (separator, kleene) =
                         parse_sep_and_kleene_op(&mut trees, delim_span.entire(), sess);
+                    if let [TokenTree::MetaVar(.., undelimited_seq)] = &mut *sequence
+                        && separator.is_none()
+                        && matches!(kleene.op, KleeneOp::ZeroOrMore | KleeneOp::OneOrMore)
+                    {
+                        // See the comment on `fn maybe_use_metavar_location`.
+                        *undelimited_seq = true;
+                    }
                     // Count the number of captured "names" (i.e., named metavars)
                     let num_captures =
                         if parsing_patterns { count_metavar_decls(&sequence) } else { 0 };
@@ -223,7 +231,7 @@ fn parse_tree<'a>(
                     if ident.name == kw::Crate && !is_raw {
                         TokenTree::token(token::Ident(kw::DollarCrate, is_raw), span)
                     } else {
-                        TokenTree::MetaVar(span, ident)
+                        TokenTree::MetaVar(span, ident, false)
                     }
                 }
 
@@ -245,7 +253,7 @@ fn parse_tree<'a>(
                     let msg =
                         format!("expected identifier, found `{}`", pprust::token_to_string(token),);
                     sess.dcx.span_err(token.span, msg);
-                    TokenTree::MetaVar(token.span, Ident::empty())
+                    TokenTree::MetaVar(token.span, Ident::empty(), false)
                 }
 
                 // There are no more tokens. Just return the `$` we already have.
diff --git a/compiler/rustc_expand/src/mbe/transcribe.rs b/compiler/rustc_expand/src/mbe/transcribe.rs
@@ -228,7 +228,7 @@ pub(super) fn transcribe<'a>(
             }
 
             // Replace the meta-var with the matched token tree from the invocation.
-            mbe::TokenTree::MetaVar(mut sp, mut original_ident) => {
+            mbe::TokenTree::MetaVar(mut sp, mut original_ident, undelimited_seq) => {
                 // Find the matched nonterminal from the macro invocation, and use it to replace
                 // the meta-var.
                 let ident = MacroRulesNormalizedIdent::new(original_ident);
@@ -237,7 +237,7 @@ pub(super) fn transcribe<'a>(
                         MatchedTokenTree(tt) => {
                             // `tt`s are emitted into the output stream directly as "raw tokens",
                             // without wrapping them into groups.
-                            result.push(tt.clone());
+                            result.push(maybe_use_metavar_location(cx, tt, sp, *undelimited_seq));
                         }
                         MatchedNonterminal(nt) => {
                             // Other variables are emitted into the output stream as groups with
@@ -302,6 +302,54 @@ pub(super) fn transcribe<'a>(
     }
 }
 
+/// Usually metavariables `$var` produce interpolated tokens, which have an additional place for
+/// keeping both the original span and the metavariable span. For `tt` metavariables that's not the
+/// case however, and there's no place for keeping a second span. So we try to give the single
+/// produced span a location that would be most useful in practice (the hygiene part of the span
+/// must not be changed).
+///
+/// Different locations are useful for different purposes:
+/// - The original location is useful when we need to report a diagnostic for the original token in
+///   isolation, without combining it with any surrounding tokens. This case occurs, but it is not
+///   very common in practice.
+/// - The metavariable location is useful when we need to somehow combine the token span with spans
+///   of its surrounding tokens. This is the most common way to use token spans.
+///
+/// So this function replaces the original location with the metavariable location in all cases
+/// except these two:
+/// - The metavariable is an element of undelimited sequence `$($tt)*`.
+///   These are typically used for passing larger amounts of code, and tokens in that code usually
+///   combine with each other and not with tokens outside of the sequence.
+/// - The metavariable span comes from a different crate, then we prefer the more local span.
+///
+/// FIXME: Find a way to keep both original and metavariable spans for all tokens without
+/// regressing compilation time too much. Several experiments for adding such spans were made in
+/// the past (PR #95580, #118517, #118671) and all showed some regressions.
+fn maybe_use_metavar_location(
+    cx: &ExtCtxt<'_>,
+    orig_tt: &TokenTree,
+    metavar_span: Span,
+    undelimited_seq: bool,
+) -> TokenTree {
+    if undelimited_seq || cx.source_map().is_imported(metavar_span) {
+        return orig_tt.clone();
+    }
+
+    match orig_tt {
+        TokenTree::Token(Token { kind, span }, spacing) => {
+            let span = metavar_span.with_ctxt(span.ctxt());
+            TokenTree::Token(Token { kind: kind.clone(), span }, *spacing)
+        }
+        TokenTree::Delimited(dspan, dspacing, delimiter, tts) => {
+            let dspan = DelimSpan::from_pair(
+                metavar_span.shrink_to_lo().with_ctxt(dspan.open.ctxt()),
+                metavar_span.shrink_to_hi().with_ctxt(dspan.close.ctxt()),
+            );
+            TokenTree::Delimited(dspan, *dspacing, *delimiter, tts.clone())
+        }
+    }
+}
+
 /// Lookup the meta-var named `ident` and return the matched token tree from the invocation using
 /// the set of matches `interpolations`.
 ///
@@ -402,7 +450,7 @@ fn lockstep_iter_size(
                 size.with(lockstep_iter_size(tt, interpolations, repeats))
             })
         }
-        TokenTree::MetaVar(_, name) | TokenTree::MetaVarDecl(_, name, _) => {
+        TokenTree::MetaVar(_, name, _) | TokenTree::MetaVarDecl(_, name, _) => {
             let name = MacroRulesNormalizedIdent::new(*name);
             match lookup_cur_matched(name, interpolations, repeats) {
                 Some(matched) => match matched {
diff --git a/tests/ui/macros/issue-118786.stderr b/tests/ui/macros/issue-118786.stderr
@@ -6,8 +6,8 @@ LL |         macro_rules! $macro_name {
    |
 help: change the delimiters to curly braces
    |
-LL |         macro_rules! {} {
-   |                      ~          +
+LL |         macro_rules! {$macro_name} {
+   |                      +           +
 help: add a semicolon
    |
 LL |         macro_rules! $macro_name; {
diff --git a/tests/ui/parser/issues/issue-68091-unicode-ident-after-if.stderr b/tests/ui/parser/issues/issue-68091-unicode-ident-after-if.stderr
@@ -1,10 +1,10 @@
 error: missing condition for `if` expression
-  --> $DIR/issue-68091-unicode-ident-after-if.rs:3:14
+  --> $DIR/issue-68091-unicode-ident-after-if.rs:3:13
    |
 LL |         $($c)ö* {}
-   |              ^  - if this block is the condition of the `if` expression, then it must be followed by another block
-   |              |
-   |              expected condition here
+   |             ^   - if this block is the condition of the `if` expression, then it must be followed by another block
+   |             |
+   |             expected condition here
 
 error: aborting due to 1 previous error
 
diff --git a/tests/ui/parser/issues/issue-68092-unicode-ident-after-incomplete-expr.stderr b/tests/ui/parser/issues/issue-68092-unicode-ident-after-incomplete-expr.stderr
@@ -1,8 +1,8 @@
 error: macro expansion ends with an incomplete expression: expected expression
-  --> $DIR/issue-68092-unicode-ident-after-incomplete-expr.rs:3:14
+  --> $DIR/issue-68092-unicode-ident-after-incomplete-expr.rs:3:13
    |
 LL |         $($c)ö*
-   |              ^ expected expression
+   |             ^ expected expression
 
 error: aborting due to 1 previous error
 
diff --git a/tests/ui/proc-macro/capture-macro-rules-invoke.stdout b/tests/ui/proc-macro/capture-macro-rules-invoke.stdout
@@ -271,7 +271,7 @@ PRINT-BANG INPUT (DEBUG): TokenStream [
                 span: $DIR/capture-macro-rules-invoke.rs:47:19: 47:20 (#0),
             },
         ],
-        span: $DIR/capture-macro-rules-invoke.rs:47:13: 47:22 (#0),
+        span: $DIR/capture-macro-rules-invoke.rs:15:60: 15:63 (#0),
     },
     Punct {
         ch: ',',
diff --git a/tests/ui/proc-macro/expand-expr.rs b/tests/ui/proc-macro/expand-expr.rs
@@ -37,7 +37,7 @@ expand_expr_is!("hello", stringify!(hello));
 expand_expr_is!("10 + 20", stringify!(10 + 20));
 
 macro_rules! echo_tts {
-    ($($t:tt)*) => { $($t)* };  //~ ERROR: expected expression, found `$`
+    ($($t:tt)*) => { $($t)* };
 }
 
 macro_rules! echo_lit {
@@ -109,7 +109,7 @@ expand_expr_fail!("string"; hello); //~ ERROR: expected one of `.`, `?`, or an o
 
 // Invalid expressions produce errors in addition to returning `Err(())`.
 expand_expr_fail!($); //~ ERROR: expected expression, found `$`
-expand_expr_fail!(echo_tts!($));
+expand_expr_fail!(echo_tts!($)); //~ ERROR: expected expression, found `$`
 expand_expr_fail!(echo_pm!($)); //~ ERROR: expected expression, found `$`
 
 // We get errors reported and recover during macro expansion if the macro
diff --git a/tests/ui/proc-macro/expand-expr.stderr b/tests/ui/proc-macro/expand-expr.stderr
@@ -11,10 +11,10 @@ LL | expand_expr_fail!($);
    |                   ^ expected expression
 
 error: expected expression, found `$`
-  --> $DIR/expand-expr.rs:40:23
+  --> $DIR/expand-expr.rs:112:29
    |
-LL |     ($($t:tt)*) => { $($t)* };
-   |                       ^^^^ expected expression
+LL | expand_expr_fail!(echo_tts!($));
+   |                             ^ expected expression
 
 error: expected expression, found `$`
   --> $DIR/expand-expr.rs:113:28

Original file line number	Diff line number	Diff line change
`@@ -242,7 +242,7 @@ fn check_binders(`
`242`	`242`	`// the outer macro. See ui/macros/macro-of-higher-order.rs where $y:$fragment in the`
`243`	`243`	`// LHS of the nested macro (and RHS of the outer macro) is parsed as MetaVar(y) Colon`
`244`	`244`	`// MetaVar(fragment) and not as MetaVarDecl(y, fragment).`
`245`		`- TokenTree::MetaVar(span, name) => {`
	`245`	`+ TokenTree::MetaVar(span, name, _) => {`
`246`	`246`	`if macros.is_empty() {`
`247`	`247`	`sess.dcx.span_bug(span, "unexpected MetaVar in lhs");`
`248`	`248`	`}`
`@@ -342,7 +342,7 @@ fn check_occurrences(`
`342`	`342`	`TokenTree::MetaVarDecl(span, _name, _kind) => {`
`343`	`343`	`sess.dcx.span_bug(span, "unexpected MetaVarDecl in rhs")`
`344`	`344`	`}`
`345`		`- TokenTree::MetaVar(span, name) => {`
	`345`	`+ TokenTree::MetaVar(span, name, _) => {`
`346`	`346`	`let name = MacroRulesNormalizedIdent::new(name);`
`347`	`347`	`check_ops_is_prefix(sess, node_id, macros, binders, ops, span, name);`
`348`	`348`	`}`
Original file line number	Diff line number	Diff line change
`@@ -6,8 +6,8 @@ LL \| macro_rules! $macro_name {`
`6`	`6`	`\|`
`7`	`7`	`help: change the delimiters to curly braces`
`8`	`8`	`\|`
`9`		`-LL \| macro_rules! {} {`
`10`		`- \| ~ +`
	`9`	`+LL \| macro_rules! {$macro_name} {`
	`10`	`+ \| + +`
`11`	`11`	`help: add a semicolon`
`12`	`12`	`\|`
`13`	`13`	`LL \| macro_rules! $macro_name; {`
Original file line number	Diff line number	Diff line change
`@@ -1,8 +1,8 @@`
`1`	`1`	`error: macro expansion ends with an incomplete expression: expected expression`
`2`		`- --> $DIR/issue-68092-unicode-ident-after-incomplete-expr.rs:3:14`
	`2`	`+ --> $DIR/issue-68092-unicode-ident-after-incomplete-expr.rs:3:13`
`3`	`3`	`\|`
`4`	`4`	`LL \| $($c)ö*`
`5`		`- \| ^ expected expression`
	`5`	`+ \| ^ expected expression`
`6`	`6`
`7`	`7`	`error: aborting due to 1 previous error`
`8`	`8`