addressed review feedback

3 files changed, 78 insertions, 18 deletions

diff --git a/compiler/rustc_mir/src/transform/coverage/spans.rs b/compiler/rustc_mir/src/transform/coverage/spans.rs
index e1db5191c89..e0c09d9ed0d 100644
--- a/compiler/rustc_mir/src/transform/coverage/spans.rs
+++ b/compiler/rustc_mir/src/transform/coverage/spans.rs
@@ -13,6 +13,7 @@ use rustc_middle::ty::TyCtxt;
 use rustc_span::source_map::original_sp;
 use rustc_span::{BytePos, ExpnKind, MacroKind, Span, Symbol};
 
+use std::cell::RefCell;
 use std::cmp::Ordering;
 
 #[derive(Debug, Copy, Clone)]
@@ -68,6 +69,7 @@ impl CoverageStatement {
 pub(super) struct CoverageSpan {
     pub span: Span,
     pub expn_span: Span,
+    pub current_macro_or_none: RefCell<Option<Option<Symbol>>>,
     pub bcb: BasicCoverageBlock,
     pub coverage_statements: Vec<CoverageStatement>,
     pub is_closure: bool,
@@ -78,6 +80,7 @@ impl CoverageSpan {
         Self {
             span: fn_sig_span,
             expn_span: fn_sig_span,
+            current_macro_or_none: Default::default(),
             bcb: START_BCB,
             coverage_statements: vec![],
             is_closure: false,
@@ -103,6 +106,7 @@ impl CoverageSpan {
         Self {
             span,
             expn_span,
+            current_macro_or_none: Default::default(),
             bcb,
             coverage_statements: vec![CoverageStatement::Statement(bb, span, stmt_index)],
             is_closure,
@@ -118,6 +122,7 @@ impl CoverageSpan {
         Self {
             span,
             expn_span,
+            current_macro_or_none: Default::default(),
             bcb,
             coverage_statements: vec![CoverageStatement::Terminator(bb, span)],
             is_closure: false,
@@ -174,15 +179,19 @@ impl CoverageSpan {
             .join("\n")
     }
 
-    /// 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.
+    /// If the span is part of a macro, returns the macro name symbol.
     pub fn current_macro(&self) -> Option<Symbol> {
-        if let ExpnKind::Macro(MacroKind::Bang, current_macro) =
-            self.expn_span.ctxt().outer_expn_data().kind
-        {
-            return Some(current_macro);
-        }
-        None
+        self.current_macro_or_none
+            .borrow_mut()
+            .get_or_insert_with(|| {
+                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
@@ -474,8 +483,8 @@ impl<'a, 'tcx> CoverageSpans<'a, 'tcx> {
                 self.curr().expn_span.ctxt() != prev_expn_span.ctxt()
             }) {
                 let merged_prefix_len = self.curr_original_span.lo() - self.curr().span.lo();
-                let after_macro_bang = merged_prefix_len
-                    + BytePos(visible_macro.to_string().bytes().count() as u32 + 1);
+                let after_macro_bang =
+                    merged_prefix_len + BytePos(visible_macro.as_str().bytes().count() as u32 + 1);
                 let mut macro_name_cov = self.curr().clone();
                 self.curr_mut().span =
                     self.curr().span.with_lo(self.curr().span.lo() + after_macro_bang);
@@ -766,6 +775,9 @@ impl<'a, 'tcx> CoverageSpans<'a, 'tcx> {
     }
 }
 
+/// See `function_source_span()` for a description of the two returned spans.
+/// If the MIR `Statement` is not contributive to computing coverage spans,
+/// returns `None`.
 pub(super) fn filtered_statement_span(
     statement: &'a Statement<'tcx>,
     body_span: Span,
@@ -811,6 +823,9 @@ pub(super) fn filtered_statement_span(
     }
 }
 
+/// See `function_source_span()` for a description of the two returned spans.
+/// If the MIR `Terminator` is not contributive to computing coverage spans,
+/// returns `None`.
 pub(super) fn filtered_terminator_span(
     terminator: &'a Terminator<'tcx>,
     body_span: Span,
@@ -854,8 +869,21 @@ pub(super) fn filtered_terminator_span(
     }
 }
 
-/// Returns the span within the function source body, and the given span, which will be different
-/// if the given span is an expansion (macro, syntactic sugar, etc.).
+/// Returns two spans from the given span (the span associated with a
+/// `Statement` or `Terminator`):
+///
+///   1. An extrapolated span (pre-expansion[^1]) corresponding to a range within
+///      the function's body source. This span is guaranteed to be contained
+///      within, or equal to, the `body_span`. If the extrapolated span is not
+///      contained within the `body_span`, the `body_span` is returned.
+///   2. The actual `span` value from the `Statement`, before expansion.
+///
+/// Only the first span is used when computing coverage code regions. The second
+/// span is useful if additional expansion data is needed (such as to look up
+/// the macro name for a composed span within that macro).)
+///
+/// [^1]Expansions result from Rust syntax including macros, syntactic
+/// sugar, etc.).
 #[inline]
 fn function_source_span(span: Span, body_span: Span) -> (Span, Span) {
     let original_span = original_sp(span, body_span).with_ctxt(body_span.ctxt());
diff --git a/src/test/run-make-fulldeps/coverage-reports/expected_show_coverage.issue-84561.txt b/src/test/run-make-fulldeps/coverage-reports/expected_show_coverage.issue-84561.txt
index 5d266b5db15..8256daa1419 100644
--- a/src/test/run-make-fulldeps/coverage-reports/expected_show_coverage.issue-84561.txt
+++ b/src/test/run-make-fulldeps/coverage-reports/expected_show_coverage.issue-84561.txt
@@ -10,7 +10,7 @@
   | Unexecuted instantiation: <issue_84561::Foo as core::cmp::PartialEq>::ne
   ------------------
     5|       |struct Foo(u32);
-    6|      1|fn test2() {
+    6|      1|fn test3() {
     7|      1|    let is_true = std::env::args().len() == 1;
     8|      1|    let bar = Foo(1);
     9|      1|    assert_eq!(bar, Foo(1));
@@ -173,8 +173,24 @@
   160|      1|    debug!("debug is enabled");
   161|      1|}
   162|       |
-  163|      1|fn main() {
-  164|      1|    test1();
-  165|      1|    test2();
-  166|      1|}
+  163|       |macro_rules! call_debug {
+  164|       |    ($($arg:tt)+) => (
+  165|      1|        fn call_print(s: &str) {
+  166|      1|            print!("{}", s);
+  167|      1|        }
+  168|       |
+  169|       |        call_print("called from call_debug: ");
+  170|       |        debug!($($arg)+);
+  171|       |    );
+  172|       |}
+  173|       |
+  174|      1|fn test2() {
+  175|      1|    call_debug!("debug is enabled");
+  176|      1|}
+  177|       |
+  178|      1|fn main() {
+  179|      1|    test1();
+  180|      1|    test2();
+  181|      1|    test3();
+  182|      1|}
 
diff --git a/src/test/run-make-fulldeps/coverage/issue-84561.rs b/src/test/run-make-fulldeps/coverage/issue-84561.rs
index 5c8fd0b7cae..b39a289c45e 100644
--- a/src/test/run-make-fulldeps/coverage/issue-84561.rs
+++ b/src/test/run-make-fulldeps/coverage/issue-84561.rs
@@ -3,7 +3,7 @@
 // expect-exit-status-101
 #[derive(PartialEq, Eq)]
 struct Foo(u32);
-fn test2() {
+fn test3() {
     let is_true = std::env::args().len() == 1;
     let bar = Foo(1);
     assert_eq!(bar, Foo(1));
@@ -160,7 +160,23 @@ fn test1() {
     debug!("debug is enabled");
 }
 
+macro_rules! call_debug {
+    ($($arg:tt)+) => (
+        fn call_print(s: &str) {
+            print!("{}", s);
+        }
+
+        call_print("called from call_debug: ");
+        debug!($($arg)+);
+    );
+}
+
+fn test2() {
+    call_debug!("debug is enabled");
+}
+
 fn main() {
     test1();
     test2();
+    test3();
 }