macro_rules: Preserve all metavariable spans in a global side table

3 files changed, 105 insertions, 24 deletions

diff --git a/compiler/rustc_expand/src/lib.rs b/compiler/rustc_expand/src/lib.rs
index 4da86d77dc8..0e73abc9ed8 100644
--- a/compiler/rustc_expand/src/lib.rs
+++ b/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)]
diff --git a/compiler/rustc_expand/src/mbe/transcribe.rs b/compiler/rustc_expand/src/mbe/transcribe.rs
index 434891ebc76..519e4a634d8 100644
--- a/compiler/rustc_expand/src/mbe/transcribe.rs
+++ b/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())
         }
diff --git a/compiler/rustc_span/src/lib.rs b/compiler/rustc_span/src/lib.rs
index 49ae9f16c8e..616a7ccc7c6 100644
--- a/compiler/rustc_span/src/lib.rs
+++ b/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.