macro_rules: Preserve all metavariable spans in a global side table

compiler/rustc_expand/src/lib.rs

@@ -6,6 +6,7 @@
 #![feature(if_let_guard)]
 #![feature(let_chains)]
 #![feature(macro_metavar_expr)]
+#![feature(map_try_insert)]
 #![feature(proc_macro_diagnostic)]
 #![feature(proc_macro_internals)]
 #![feature(proc_macro_span)]

compiler/rustc_expand/src/mbe/transcribe.rs

@@ -13,7 +13,7 @@ use rustc_errors::DiagnosticBuilder;
 use rustc_errors::{pluralize, PResult};
 use rustc_span::hygiene::{LocalExpnId, Transparency};
 use rustc_span::symbol::{sym, Ident, MacroRulesNormalizedIdent};
-use rustc_span::{Span, SyntaxContext};
+use rustc_span::{with_metavar_spans, Span, SyntaxContext};
 
 use smallvec::{smallvec, SmallVec};
 use std::mem;
@@ -254,7 +254,8 @@ 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(maybe_use_metavar_location(cx, &stack, sp, tt));
+                            let tt = maybe_use_metavar_location(cx, &stack, sp, tt, &mut marker);
+                            result.push(tt);
                         }
                         MatchedNonterminal(nt) => {
                             // Other variables are emitted into the output stream as groups with
@@ -319,6 +320,17 @@ pub(super) fn transcribe<'a>(
     }
 }
 
+/// Store the metavariable span for this original span into a side table.
+/// FIXME: Try to put the metavariable span into `SpanData` instead of a side table (#118517).
+/// An optimal encoding for inlined spans will need to be selected to minimize regressions.
+/// The side table approach is relatively good, but not perfect due to collisions.
+/// In particular, collisions happen when token is passed as an argument through several macro
+/// calls, like in recursive macros.
+/// The old heuristic below is used to improve spans in case of collisions, but diagnostics are
+/// still degraded sometimes in those cases.
+///
+/// The old heuristic:
+///
 /// 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
@@ -338,15 +350,12 @@ pub(super) fn transcribe<'a>(
 ///   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<'_>,
     stack: &[Frame<'_>],
-    metavar_span: Span,
+    mut metavar_span: Span,
     orig_tt: &TokenTree,
+    marker: &mut Marker,
 ) -> TokenTree {
     let undelimited_seq = matches!(
         stack.last(),
@@ -357,18 +366,44 @@ fn maybe_use_metavar_location(
             ..
         })
     );
-    if undelimited_seq || cx.source_map().is_imported(metavar_span) {
+    if undelimited_seq {
+        // Do not record metavar spans for tokens from undelimited sequences, for perf reasons.
+        return orig_tt.clone();
+    }
+
+    let insert = |mspans: &mut FxHashMap<_, _>, s, ms| match mspans.try_insert(s, ms) {
+        Ok(_) => true,
+        Err(err) => *err.entry.get() == ms, // Tried to insert the same span, still success
+    };
+    marker.visit_span(&mut metavar_span);
+    let no_collision = match orig_tt {
+        TokenTree::Token(token, ..) => {
+            with_metavar_spans(|mspans| insert(mspans, token.span, metavar_span))
+        }
+        TokenTree::Delimited(dspan, ..) => with_metavar_spans(|mspans| {
+            insert(mspans, dspan.open, metavar_span)
+                && insert(mspans, dspan.close, metavar_span)
+                && insert(mspans, dspan.entire(), metavar_span)
+        }),
+    };
+    if no_collision || cx.source_map().is_imported(metavar_span) {
         return orig_tt.clone();
     }
 
+    // Setting metavar spans for the heuristic spans gives better opportunities for combining them
+    // with neighboring spans even despite their different syntactic contexts.
     match orig_tt {
         TokenTree::Token(Token { kind, span }, spacing) => {
             let span = metavar_span.with_ctxt(span.ctxt());
+            with_metavar_spans(|mspans| insert(mspans, span, metavar_span));
             TokenTree::Token(Token { kind: kind.clone(), span }, *spacing)
         }
         TokenTree::Delimited(dspan, dspacing, delimiter, tts) => {
-            let open = metavar_span.shrink_to_lo().with_ctxt(dspan.open.ctxt());
-            let close = metavar_span.shrink_to_hi().with_ctxt(dspan.close.ctxt());
+            let open = metavar_span.with_ctxt(dspan.open.ctxt());
+            let close = metavar_span.with_ctxt(dspan.close.ctxt());
+            with_metavar_spans(|mspans| {
+                insert(mspans, open, metavar_span) && insert(mspans, close, metavar_span)
+            });
             let dspan = DelimSpan::from_pair(open, close);
             TokenTree::Delimited(dspan, *dspacing, *delimiter, tts.clone())
         }

compiler/rustc_span/src/lib.rs

@@ -72,6 +72,7 @@ pub mod fatal_error;
 
 pub mod profiling;
 
+use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::stable_hasher::{Hash128, Hash64, HashStable, StableHasher};
 use rustc_data_structures::sync::{FreezeLock, FreezeWriteGuard, Lock, Lrc};
 
@@ -98,6 +99,9 @@ mod tests;
 pub struct SessionGlobals {
     symbol_interner: symbol::Interner,
     span_interner: Lock<span_encoding::SpanInterner>,
+    /// Maps a macro argument token into use of the corresponding metavariable in the macro body.
+    /// Collisions are possible and processed in `maybe_use_metavar_location` on best effort basis.
+    metavar_spans: Lock<FxHashMap<Span, Span>>,
     hygiene_data: Lock<hygiene::HygieneData>,
 
     /// A reference to the source map in the `Session`. It's an `Option`
@@ -115,6 +119,7 @@ impl SessionGlobals {
         SessionGlobals {
             symbol_interner: symbol::Interner::fresh(),
             span_interner: Lock::new(span_encoding::SpanInterner::default()),
+            metavar_spans: Default::default(),
             hygiene_data: Lock::new(hygiene::HygieneData::new(edition)),
             source_map: Lock::new(None),
         }
@@ -168,6 +173,11 @@ pub fn create_default_session_globals_then<R>(f: impl FnOnce() -> R) -> R {
 // deserialization.
 scoped_tls::scoped_thread_local!(static SESSION_GLOBALS: SessionGlobals);
 
+#[inline]
+pub fn with_metavar_spans<R>(f: impl FnOnce(&mut FxHashMap<Span, Span>) -> R) -> R {
+    with_session_globals(|session_globals| f(&mut session_globals.metavar_spans.lock()))
+}
+
 // FIXME: We should use this enum or something like it to get rid of the
 // use of magic `/rust/1.x/...` paths across the board.
 #[derive(Debug, Eq, PartialEq, Clone, Ord, PartialOrd, Decodable)]
@@ -824,29 +834,64 @@ impl Span {
         )
     }
 
+    /// Check if you can select metavar spans for the given spans to get matching contexts.
+    fn try_metavars(a: SpanData, b: SpanData, a_orig: Span, b_orig: Span) -> (SpanData, SpanData) {
+        let get = |mspans: &FxHashMap<_, _>, s| mspans.get(&s).copied();
+        match with_metavar_spans(|mspans| (get(mspans, a_orig), get(mspans, b_orig))) {
+            (None, None) => {}
+            (Some(meta_a), None) => {
+                let meta_a = meta_a.data();
+                if meta_a.ctxt == b.ctxt {
+                    return (meta_a, b);
+                }
+            }
+            (None, Some(meta_b)) => {
+                let meta_b = meta_b.data();
+                if a.ctxt == meta_b.ctxt {
+                    return (a, meta_b);
+                }
+            }
+            (Some(meta_a), Some(meta_b)) => {
+                let meta_b = meta_b.data();
+                if a.ctxt == meta_b.ctxt {
+                    return (a, meta_b);
+                }
+                let meta_a = meta_a.data();
+                if meta_a.ctxt == b.ctxt {
+                    return (meta_a, b);
+                } else if meta_a.ctxt == meta_b.ctxt {
+                    return (meta_a, meta_b);
+                }
+            }
+        }
+
+        (a, b)
+    }
+
     /// Prepare two spans to a combine operation like `to` or `between`.
-    /// FIXME: consider using declarative macro metavariable spans for the given spans if they are
-    /// better suitable for combining (#119412).
     fn prepare_to_combine(
         a_orig: Span,
         b_orig: Span,
     ) -> Result<(SpanData, SpanData, Option<LocalDefId>), Span> {
         let (a, b) = (a_orig.data(), b_orig.data());
+        if a.ctxt == b.ctxt {
+            return Ok((a, b, if a.parent == b.parent { a.parent } else { None }));
+        }
 
-        if a.ctxt != b.ctxt {
-            // Context mismatches usually happen when procedural macros combine spans copied from
-            // the macro input with spans produced by the macro (`Span::*_site`).
-            // In that case we consider the combined span to be produced by the macro and return
-            // the original macro-produced span as the result.
-            // Otherwise we just fall back to returning the first span.
-            // Combining locations typically doesn't make sense in case of context mismatches.
-            // `is_root` here is a fast path optimization.
-            let a_is_callsite = a.ctxt.is_root() || a.ctxt == b.span().source_callsite().ctxt();
-            return Err(if a_is_callsite { b_orig } else { a_orig });
+        let (a, b) = Span::try_metavars(a, b, a_orig, b_orig);
+        if a.ctxt == b.ctxt {
+            return Ok((a, b, if a.parent == b.parent { a.parent } else { None }));
         }
 
-        let parent = if a.parent == b.parent { a.parent } else { None };
-        Ok((a, b, parent))
+        // Context mismatches usually happen when procedural macros combine spans copied from
+        // the macro input with spans produced by the macro (`Span::*_site`).
+        // In that case we consider the combined span to be produced by the macro and return
+        // the original macro-produced span as the result.
+        // Otherwise we just fall back to returning the first span.
+        // Combining locations typically doesn't make sense in case of context mismatches.
+        // `is_root` here is a fast path optimization.
+        let a_is_callsite = a.ctxt.is_root() || a.ctxt == b.span().source_callsite().ctxt();
+        Err(if a_is_callsite { b_orig } else { a_orig })
     }
 
     /// This span, but in a larger context, may switch to the metavariable span if suitable.

tests/coverage/no_spans.cov-map

@@ -1,3 +1,11 @@
+Function name: no_spans::affected_function
+Raw bytes (9): 0x[01, 01, 00, 01, 01, 1a, 1c, 00, 1d]
+Number of files: 1
+- file 0 => global file 1
+Number of expressions: 0
+Number of file 0 mappings: 1
+- Code(Counter(0)) at (prev + 26, 28) to (start + 0, 29)
+
 Function name: no_spans::affected_function::{closure#0}
 Raw bytes (9): 0x[01, 01, 00, 01, 01, 1b, 0c, 00, 0e]
 Number of files: 1

tests/coverage/no_spans.coverage

@@ -23,7 +23,7 @@
    LL|       |}
    LL|       |
    LL|       |macro_that_defines_a_function! {
-   LL|       |    fn affected_function() {
+   LL|      1|    fn affected_function() {
    LL|      1|        || ()
    LL|       |    }
    LL|       |}

tests/coverage/no_spans_if_not.cov-map

@@ -1,3 +1,15 @@
+Function name: no_spans_if_not::affected_function
+Raw bytes (21): 0x[01, 01, 01, 01, 00, 03, 01, 16, 1c, 01, 12, 02, 02, 0d, 00, 0f, 00, 02, 0d, 00, 0f]
+Number of files: 1
+- file 0 => global file 1
+Number of expressions: 1
+- expression 0 operands: lhs = Counter(0), rhs = Zero
+Number of file 0 mappings: 3
+- Code(Counter(0)) at (prev + 22, 28) to (start + 1, 18)
+- Code(Expression(0, Sub)) at (prev + 2, 13) to (start + 0, 15)
+    = (c0 - Zero)
+- Code(Zero) at (prev + 2, 13) to (start + 0, 15)
+
 Function name: no_spans_if_not::main
 Raw bytes (9): 0x[01, 01, 00, 01, 01, 0b, 01, 02, 02]
 Number of files: 1

tests/coverage/no_spans_if_not.coverage

@@ -19,11 +19,11 @@
    LL|       |}
    LL|       |
    LL|       |macro_that_defines_a_function! {
-   LL|       |    fn affected_function() {
-   LL|       |        if !false {
-   LL|       |            ()
+   LL|      1|    fn affected_function() {
+   LL|      1|        if !false {
+   LL|      1|            ()
    LL|       |        } else {
-   LL|       |            ()
+   LL|      0|            ()
    LL|       |        }
    LL|       |    }
    LL|       |}

tests/ui/anon-params/anon-params-edition-hygiene.rs

@@ -1,3 +1,4 @@
+//@ check-pass
 //@ edition:2018
 //@ aux-build:anon-params-edition-hygiene.rs
 
@@ -8,7 +9,6 @@
 extern crate anon_params_edition_hygiene;
 
 generate_trait_2015_ident!(u8);
-// FIXME: Edition hygiene doesn't work correctly with `tt`s in this case.
-generate_trait_2015_tt!(u8); //~ ERROR expected one of `:`, `@`, or `|`, found `)`
+generate_trait_2015_tt!(u8);
 
 fn main() {}

tests/ui/editions/edition-keywords-2018-2018-parsing.rs

@@ -16,7 +16,7 @@ macro_rules! local_passes_ident {
     ($i: ident) => ($i) //~ ERROR macro expansion ends with an incomplete expression
 }
 macro_rules! local_passes_tt {
-    ($i: tt) => ($i) //~ ERROR macro expansion ends with an incomplete expression
+    ($i: tt) => ($i)
 }
 
 pub fn check_async() {
@@ -34,7 +34,7 @@ pub fn check_async() {
     if passes_tt!(r#async) == 1 {} // OK
     if local_passes_ident!(async) == 1 {} // Error reported above in the macro
     if local_passes_ident!(r#async) == 1 {} // OK
-    if local_passes_tt!(async) == 1 {} // Error reported above in the macro
+    if local_passes_tt!(async) == 1 {} //~ ERROR macro expansion ends with an incomplete expression
     if local_passes_tt!(r#async) == 1 {} // OK
     module::async(); //~ ERROR expected identifier, found keyword `async`
     module::r#async(); // OK

tests/ui/editions/edition-keywords-2018-2018-parsing.stderr

@@ -68,10 +68,10 @@ LL |     ($i: ident) => ($i)
    |                       ^ expected one of `move`, `|`, or `||`
 
 error: macro expansion ends with an incomplete expression: expected one of `move`, `|`, or `||`
-  --> $DIR/edition-keywords-2018-2018-parsing.rs:19:20
+  --> $DIR/edition-keywords-2018-2018-parsing.rs:37:30
    |
-LL |     ($i: tt) => ($i)
-   |                    ^ expected one of `move`, `|`, or `||`
+LL |     if local_passes_tt!(async) == 1 {}
+   |                              ^ expected one of `move`, `|`, or `||`
 
 error[E0308]: mismatched types
   --> $DIR/edition-keywords-2018-2018-parsing.rs:42:33

tests/ui/fmt/format-args-capture-first-literal-is-macro.rs

@@ -6,7 +6,6 @@ extern crate format_string_proc_macro;
 macro_rules! identity_mbe {
     ($tt:tt) => {
         $tt
-        //~^ ERROR there is no argument named `a`
     };
 }
 
@@ -16,6 +15,7 @@ fn main() {
     format!(identity_pm!("{a}"));
     //~^ ERROR there is no argument named `a`
     format!(identity_mbe!("{a}"));
+    //~^ ERROR there is no argument named `a`
     format!(concat!("{a}"));
     //~^ ERROR there is no argument named `a`
 }

tests/ui/fmt/format-args-capture-first-literal-is-macro.stderr

@@ -1,5 +1,5 @@
 error: there is no argument named `a`
-  --> $DIR/format-args-capture-first-literal-is-macro.rs:16:26
+  --> $DIR/format-args-capture-first-literal-is-macro.rs:15:26
    |
 LL |     format!(identity_pm!("{a}"));
    |                          ^^^^^
@@ -8,10 +8,10 @@ LL |     format!(identity_pm!("{a}"));
    = note: to avoid ambiguity, `format_args!` cannot capture variables when the format string is expanded from a macro
 
 error: there is no argument named `a`
-  --> $DIR/format-args-capture-first-literal-is-macro.rs:8:9
+  --> $DIR/format-args-capture-first-literal-is-macro.rs:17:27
    |
-LL |         $tt
-   |         ^^^
+LL |     format!(identity_mbe!("{a}"));
+   |                           ^^^^^
    |
    = note: did you intend to capture a variable `a` from the surrounding scope?
    = note: to avoid ambiguity, `format_args!` cannot capture variables when the format string is expanded from a macro