coverage: Overhaul how "visible macros" are determined

4 files changed, 61 insertions, 69 deletions

diff --git a/compiler/rustc_mir_transform/src/coverage/spans.rs b/compiler/rustc_mir_transform/src/coverage/spans.rs
index ed091752187..6c0c62a1b76 100644
--- a/compiler/rustc_mir_transform/src/coverage/spans.rs
+++ b/compiler/rustc_mir_transform/src/coverage/spans.rs
@@ -1,9 +1,7 @@
-use std::cell::OnceCell;
-
 use rustc_data_structures::graph::WithNumNodes;
 use rustc_index::IndexVec;
 use rustc_middle::mir;
-use rustc_span::{BytePos, ExpnKind, MacroKind, Span, Symbol, DUMMY_SP};
+use rustc_span::{BytePos, Span, Symbol, DUMMY_SP};
 
 use super::graph::{BasicCoverageBlock, CoverageGraph, START_BCB};
 use crate::coverage::ExtractedHirInfo;
@@ -71,8 +69,7 @@ impl CoverageSpans {
 #[derive(Debug, Clone)]
 struct CoverageSpan {
     pub span: Span,
-    pub expn_span: Span,
-    pub current_macro_or_none: OnceCell<Option<Symbol>>,
+    pub visible_macro: Option<Symbol>,
     pub bcb: BasicCoverageBlock,
     /// List of all the original spans from MIR that have been merged into this
     /// span. Mainly used to precisely skip over gaps when truncating a span.
@@ -82,23 +79,16 @@ struct CoverageSpan {
 
 impl CoverageSpan {
     pub fn for_fn_sig(fn_sig_span: Span) -> Self {
-        Self::new(fn_sig_span, fn_sig_span, START_BCB, false)
+        Self::new(fn_sig_span, None, START_BCB, false)
     }
 
     pub(super) fn new(
         span: Span,
-        expn_span: Span,
+        visible_macro: Option<Symbol>,
         bcb: BasicCoverageBlock,
         is_closure: bool,
     ) -> Self {
-        Self {
-            span,
-            expn_span,
-            current_macro_or_none: Default::default(),
-            bcb,
-            merged_spans: vec![span],
-            is_closure,
-        }
+        Self { span, visible_macro, bcb, merged_spans: vec![span], is_closure }
     }
 
     pub fn merge_from(&mut self, other: &Self) {
@@ -123,37 +113,6 @@ impl CoverageSpan {
     pub fn is_in_same_bcb(&self, other: &Self) -> bool {
         self.bcb == other.bcb
     }
-
-    /// If the span is part of a macro, returns the macro name symbol.
-    pub fn current_macro(&self) -> Option<Symbol> {
-        self.current_macro_or_none
-            .get_or_init(|| {
-                if let ExpnKind::Macro(MacroKind::Bang, current_macro) =
-                    self.expn_span.ctxt().outer_expn_data().kind
-                {
-                    return Some(current_macro);
-                }
-                None
-            })
-            .map(|symbol| symbol)
-    }
-
-    /// If the span is part of a macro, and the macro is visible (expands directly to the given
-    /// body_span), returns the macro name symbol.
-    pub fn visible_macro(&self, body_span: Span) -> Option<Symbol> {
-        let current_macro = self.current_macro()?;
-        let parent_callsite = self.expn_span.parent_callsite()?;
-
-        // In addition to matching the context of the body span, the parent callsite
-        // must also be the source callsite, i.e. the parent must have no parent.
-        let is_visible_macro =
-            parent_callsite.parent_callsite().is_none() && parent_callsite.eq_ctxt(body_span);
-        is_visible_macro.then_some(current_macro)
-    }
-
-    pub fn is_macro_expansion(&self) -> bool {
-        self.current_macro().is_some()
-    }
 }
 
 /// Converts the initial set of `CoverageSpan`s (one per MIR `Statement` or `Terminator`) into a
@@ -164,10 +123,6 @@ impl CoverageSpan {
 ///    execution
 ///  * Carve out (leave uncovered) any span that will be counted by another MIR (notably, closures)
 struct CoverageSpansGenerator<'a> {
-    /// A `Span` covering the function body of the MIR (typically from left curly brace to right
-    /// curly brace).
-    body_span: Span,
-
     /// The BasicCoverageBlock Control Flow Graph (BCB CFG).
     basic_coverage_blocks: &'a CoverageGraph,
 
@@ -244,7 +199,6 @@ impl<'a> CoverageSpansGenerator<'a> {
         );
 
         let coverage_spans = Self {
-            body_span: hir_info.body_span,
             basic_coverage_blocks,
             sorted_spans_iter: sorted_spans.into_iter(),
             some_curr: None,
@@ -303,7 +257,7 @@ impl<'a> CoverageSpansGenerator<'a> {
                 // **originally** the same as the original span of `prev()`. The original spans
                 // reflect their original sort order, and for equal spans, conveys a partial
                 // ordering based on CFG dominator priority.
-                if prev.is_macro_expansion() && curr.is_macro_expansion() {
+                if prev.visible_macro.is_some() && curr.visible_macro.is_some() {
                     // Macros that expand to include branching (such as
                     // `assert_eq!()`, `assert_ne!()`, `info!()`, `debug!()`, or
                     // `trace!()`) typically generate callee spans with identical
@@ -365,12 +319,7 @@ impl<'a> CoverageSpansGenerator<'a> {
     fn maybe_push_macro_name_span(&mut self) {
         let curr = self.curr();
 
-        let Some(visible_macro) = curr.visible_macro(self.body_span) else { return };
-        if let Some(prev) = &self.some_prev
-            && prev.expn_span.eq_ctxt(curr.expn_span)
-        {
-            return;
-        }
+        let Some(visible_macro) = curr.visible_macro else { return };
 
         // The split point is relative to `curr_original_span`,
         // because `curr.span` may have been merged with preceding spans.
diff --git a/compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs b/compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs
index 913819ef33d..fde25ad994f 100644
--- a/compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs
+++ b/compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs
@@ -3,7 +3,7 @@ use rustc_middle::mir::{
     self, AggregateKind, FakeReadCause, Rvalue, Statement, StatementKind, Terminator,
     TerminatorKind,
 };
-use rustc_span::Span;
+use rustc_span::{ExpnKind, MacroKind, Span, Symbol};
 
 use crate::coverage::graph::{BasicCoverageBlock, BasicCoverageBlockData, CoverageGraph};
 use crate::coverage::spans::CoverageSpan;
@@ -71,16 +71,18 @@ fn bcb_to_initial_coverage_spans<'a, 'tcx>(
 
         let statement_spans = data.statements.iter().filter_map(move |statement| {
             let expn_span = filtered_statement_span(statement)?;
-            let span = unexpand_into_body_span(expn_span, body_span)?;
+            let (span, visible_macro) =
+                unexpand_into_body_span_with_visible_macro(expn_span, body_span)?;
 
-            Some(CoverageSpan::new(span, expn_span, bcb, is_closure_or_coroutine(statement)))
+            Some(CoverageSpan::new(span, visible_macro, bcb, is_closure_or_coroutine(statement)))
         });
 
         let terminator_span = Some(data.terminator()).into_iter().filter_map(move |terminator| {
             let expn_span = filtered_terminator_span(terminator)?;
-            let span = unexpand_into_body_span(expn_span, body_span)?;
+            let (span, visible_macro) =
+                unexpand_into_body_span_with_visible_macro(expn_span, body_span)?;
 
-            Some(CoverageSpan::new(span, expn_span, bcb, false))
+            Some(CoverageSpan::new(span, visible_macro, bcb, false))
         });
 
         statement_spans.chain(terminator_span)
@@ -201,7 +203,48 @@ fn filtered_terminator_span(terminator: &Terminator<'_>) -> Option<Span> {
 ///
 /// [^1]Expansions result from Rust syntax including macros, syntactic sugar,
 /// etc.).
-#[inline]
-fn unexpand_into_body_span(span: Span, body_span: Span) -> Option<Span> {
-    span.find_ancestor_inside_same_ctxt(body_span)
+fn unexpand_into_body_span_with_visible_macro(
+    original_span: Span,
+    body_span: Span,
+) -> Option<(Span, Option<Symbol>)> {
+    let (span, prev) = unexpand_into_body_span_with_prev(original_span, body_span)?;
+
+    let visible_macro = prev
+        .map(|prev| match prev.ctxt().outer_expn_data().kind {
+            ExpnKind::Macro(MacroKind::Bang, name) => Some(name),
+            _ => None,
+        })
+        .flatten();
+
+    Some((span, visible_macro))
+}
+
+/// Walks through the expansion ancestors of `original_span` to find a span that
+/// is contained in `body_span` and has the same [`SyntaxContext`] as `body_span`.
+/// The ancestor that was traversed just before the matching span (if any) is
+/// also returned.
+///
+/// For example, a return value of `Some((ancestor, Some(prev))` means that:
+/// - `ancestor == original_span.find_ancestor_inside_same_ctxt(body_span)`
+/// - `ancestor == prev.parent_callsite()`
+///
+/// [`SyntaxContext`]: rustc_span::SyntaxContext
+fn unexpand_into_body_span_with_prev(
+    original_span: Span,
+    body_span: Span,
+) -> Option<(Span, Option<Span>)> {
+    let mut prev = None;
+    let mut curr = original_span;
+
+    while !body_span.contains(curr) || !curr.eq_ctxt(body_span) {
+        prev = Some(curr);
+        curr = curr.parent_callsite()?;
+    }
+
+    debug_assert_eq!(Some(curr), original_span.find_ancestor_in_same_ctxt(body_span));
+    if let Some(prev) = prev {
+        debug_assert_eq!(Some(curr), prev.parent_callsite());
+    }
+
+    Some((curr, prev))
 }
diff --git a/tests/coverage/inline-dead.cov-map b/tests/coverage/inline-dead.cov-map
index 958b423f24c..ab04e746b91 100644
--- a/tests/coverage/inline-dead.cov-map
+++ b/tests/coverage/inline-dead.cov-map
@@ -31,14 +31,14 @@ Number of file 0 mappings: 2
 - Code(Counter(0)) at (prev + 7, 6) to (start + 2, 2)
 
 Function name: inline_dead::main::{closure#0}
-Raw bytes (23): 0x[01, 01, 02, 00, 06, 01, 00, 03, 01, 07, 17, 00, 18, 00, 02, 0d, 00, 0e, 03, 02, 05, 00, 06]
+Raw bytes (23): 0x[01, 01, 02, 00, 06, 01, 00, 03, 01, 07, 17, 01, 16, 00, 02, 0d, 00, 0e, 03, 02, 05, 00, 06]
 Number of files: 1
 - file 0 => global file 1
 Number of expressions: 2
 - expression 0 operands: lhs = Zero, rhs = Expression(1, Sub)
 - expression 1 operands: lhs = Counter(0), rhs = Zero
 Number of file 0 mappings: 3
-- Code(Counter(0)) at (prev + 7, 23) to (start + 0, 24)
+- Code(Counter(0)) at (prev + 7, 23) to (start + 1, 22)
 - Code(Zero) at (prev + 2, 13) to (start + 0, 14)
 - Code(Expression(0, Add)) at (prev + 2, 5) to (start + 0, 6)
     = (Zero + (c0 - Zero))
diff --git a/tests/coverage/inline-dead.coverage b/tests/coverage/inline-dead.coverage
index de96aa17acd..7c201f482db 100644
--- a/tests/coverage/inline-dead.coverage
+++ b/tests/coverage/inline-dead.coverage
@@ -5,7 +5,7 @@
    LL|      1|    println!("{}", live::<false>());
    LL|      1|
    LL|      1|    let f = |x: bool| {
-   LL|       |        debug_assert!(
+   LL|      1|        debug_assert!(
    LL|      0|            x
    LL|       |        );
    LL|      1|    };