coverage: Make unexpansion of closure bodies more precise

This improves the coverage instrumentation of closures declared in macros, as
seen in `closure_macro.rs` and `closure_macro_async.rs`.

7 files changed, 60 insertions, 50 deletions

diff --git a/compiler/rustc_mir_transform/src/coverage/mod.rs b/compiler/rustc_mir_transform/src/coverage/mod.rs
index 5fb72fcf0cf..b8bce7c3fb6 100644
--- a/compiler/rustc_mir_transform/src/coverage/mod.rs
+++ b/compiler/rustc_mir_transform/src/coverage/mod.rs
@@ -408,7 +408,18 @@ fn extract_hir_info<'tcx>(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> ExtractedHir
     let hir_body = tcx.hir().body(fn_body_id);
 
     let is_async_fn = hir_node.fn_sig().is_some_and(|fn_sig| fn_sig.header.is_async());
-    let body_span = get_body_span(tcx, hir_body, def_id);
+
+    let mut body_span = hir_body.value.span;
+
+    use rustc_hir::{Closure, Expr, ExprKind, Node};
+    // Unexpand a closure's body span back to the context of its declaration.
+    // This helps with closure bodies that consist of just a single bang-macro,
+    // and also with closure bodies produced by async desugaring.
+    if let Node::Expr(&Expr { kind: ExprKind::Closure(&Closure { fn_decl_span, .. }), .. }) =
+        hir_node
+    {
+        body_span = body_span.find_ancestor_in_same_ctxt(fn_decl_span).unwrap_or(body_span);
+    }
 
     // The actual signature span is only used if it has the same context and
     // filename as the body, and precedes the body.
@@ -432,23 +443,6 @@ fn extract_hir_info<'tcx>(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> ExtractedHir
     ExtractedHirInfo { function_source_hash, is_async_fn, fn_sig_span, body_span }
 }
 
-fn get_body_span<'tcx>(
-    tcx: TyCtxt<'tcx>,
-    hir_body: &rustc_hir::Body<'tcx>,
-    def_id: LocalDefId,
-) -> Span {
-    let mut body_span = hir_body.value.span;
-
-    if tcx.is_closure_or_coroutine(def_id.to_def_id()) {
-        // If the current function is a closure, and its "body" span was created
-        // by macro expansion or compiler desugaring, try to walk backwards to
-        // the pre-expansion call site or body.
-        body_span = body_span.source_callsite();
-    }
-
-    body_span
-}
-
 fn hash_mir_source<'tcx>(tcx: TyCtxt<'tcx>, hir_body: &'tcx rustc_hir::Body<'tcx>) -> u64 {
     // FIXME(cjgillot) Stop hashing HIR manually here.
     let owner = hir_body.id().hir_id.owner;
diff --git a/tests/coverage/closure_macro.cov-map b/tests/coverage/closure_macro.cov-map
index 323d6e3c8aa..571e5564b65 100644
--- a/tests/coverage/closure_macro.cov-map
+++ b/tests/coverage/closure_macro.cov-map
@@ -1,20 +1,20 @@
 Function name: closure_macro::load_configuration_files
-Raw bytes (9): 0x[01, 01, 00, 01, 01, 1e, 01, 02, 02]
+Raw bytes (9): 0x[01, 01, 00, 01, 01, 1d, 01, 02, 02]
 Number of files: 1
 - file 0 => global file 1
 Number of expressions: 0
 Number of file 0 mappings: 1
-- Code(Counter(0)) at (prev + 30, 1) to (start + 2, 2)
+- Code(Counter(0)) at (prev + 29, 1) to (start + 2, 2)
 
 Function name: closure_macro::main
-Raw bytes (43): 0x[01, 01, 02, 01, 05, 05, 02, 07, 01, 22, 01, 01, 21, 02, 02, 09, 00, 0f, 05, 00, 12, 00, 13, 02, 00, 12, 00, 13, 05, 00, 54, 00, 55, 02, 02, 09, 02, 0b, 07, 03, 01, 00, 02]
+Raw bytes (43): 0x[01, 01, 02, 01, 05, 05, 02, 07, 01, 21, 01, 01, 21, 02, 02, 09, 00, 0f, 05, 00, 12, 00, 13, 02, 00, 12, 00, 13, 05, 00, 54, 00, 55, 02, 02, 09, 02, 0b, 07, 03, 01, 00, 02]
 Number of files: 1
 - file 0 => global file 1
 Number of expressions: 2
 - expression 0 operands: lhs = Counter(0), rhs = Counter(1)
 - expression 1 operands: lhs = Counter(1), rhs = Expression(0, Sub)
 Number of file 0 mappings: 7
-- Code(Counter(0)) at (prev + 34, 1) to (start + 1, 33)
+- Code(Counter(0)) at (prev + 33, 1) to (start + 1, 33)
 - Code(Expression(0, Sub)) at (prev + 2, 9) to (start + 0, 15)
     = (c0 - c1)
 - Code(Counter(1)) at (prev + 0, 18) to (start + 0, 19)
@@ -27,10 +27,19 @@ Number of file 0 mappings: 7
     = (c1 + (c0 - c1))
 
 Function name: closure_macro::main::{closure#0}
-Raw bytes (9): 0x[01, 01, 00, 01, 01, 24, 12, 00, 54]
+Raw bytes (35): 0x[01, 01, 03, 01, 05, 05, 0b, 09, 00, 05, 01, 10, 1c, 03, 21, 05, 04, 11, 01, 27, 02, 03, 11, 00, 16, 00, 00, 17, 00, 1e, 07, 02, 09, 00, 0a]
 Number of files: 1
 - file 0 => global file 1
-Number of expressions: 0
-Number of file 0 mappings: 1
-- Code(Counter(0)) at (prev + 36, 18) to (start + 0, 84)
+Number of expressions: 3
+- expression 0 operands: lhs = Counter(0), rhs = Counter(1)
+- expression 1 operands: lhs = Counter(1), rhs = Expression(2, Add)
+- expression 2 operands: lhs = Counter(2), rhs = Zero
+Number of file 0 mappings: 5
+- Code(Counter(0)) at (prev + 16, 28) to (start + 3, 33)
+- Code(Counter(1)) at (prev + 4, 17) to (start + 1, 39)
+- Code(Expression(0, Sub)) at (prev + 3, 17) to (start + 0, 22)
+    = (c0 - c1)
+- Code(Zero) at (prev + 0, 23) to (start + 0, 30)
+- Code(Expression(1, Add)) at (prev + 2, 9) to (start + 0, 10)
+    = (c1 + (c2 + Zero))
 
diff --git a/tests/coverage/closure_macro.coverage b/tests/coverage/closure_macro.coverage
index ab152a91356..716d75cb8d7 100644
--- a/tests/coverage/closure_macro.coverage
+++ b/tests/coverage/closure_macro.coverage
@@ -1,4 +1,3 @@
-   LL|       |#![feature(coverage_attribute)]
    LL|       |// edition: 2018
    LL|       |
    LL|       |macro_rules! bail {
@@ -14,16 +13,16 @@
    LL|       |
    LL|       |macro_rules! on_error {
    LL|       |    ($value:expr, $error_message:expr) => {
-   LL|       |        $value.or_else(|e| {
-   LL|       |            // FIXME(85000): no coverage in closure macros
-   LL|       |            let message = format!($error_message, e);
-   LL|       |            if message.len() > 0 {
-   LL|       |                println!("{}", message);
-   LL|       |                Ok(String::from("ok"))
+   LL|      0|        $value.or_else(|e| {
+   LL|      0|            // This closure, which is declared in a macro, should be instrumented.
+   LL|      0|            let message = format!($error_message, e);
+   LL|      0|            if message.len() > 0 {
+   LL|      0|                println!("{}", message);
+   LL|      0|                Ok(String::from("ok"))
    LL|       |            } else {
-   LL|       |                bail!("error");
+   LL|      0|                bail!("error");
    LL|       |            }
-   LL|       |        })
+   LL|      0|        })
    LL|       |    };
    LL|       |}
    LL|       |
diff --git a/tests/coverage/closure_macro.rs b/tests/coverage/closure_macro.rs
index 38a81ea434b..6fe1212de8d 100644
--- a/tests/coverage/closure_macro.rs
+++ b/tests/coverage/closure_macro.rs
@@ -1,4 +1,3 @@
-#![feature(coverage_attribute)]
 // edition: 2018
 
 macro_rules! bail {
@@ -15,7 +14,7 @@ macro_rules! bail {
 macro_rules! on_error {
     ($value:expr, $error_message:expr) => {
         $value.or_else(|e| {
-            // FIXME(85000): no coverage in closure macros
+            // This closure, which is declared in a macro, should be instrumented.
             let message = format!($error_message, e);
             if message.len() > 0 {
                 println!("{}", message);
diff --git a/tests/coverage/closure_macro_async.cov-map b/tests/coverage/closure_macro_async.cov-map
index 5f5e0644278..49ec767eab3 100644
--- a/tests/coverage/closure_macro_async.cov-map
+++ b/tests/coverage/closure_macro_async.cov-map
@@ -35,10 +35,19 @@ Number of file 0 mappings: 7
     = (c1 + (c0 - c1))
 
 Function name: closure_macro_async::test::{closure#0}::{closure#0}
-Raw bytes (9): 0x[01, 01, 00, 01, 01, 25, 12, 00, 54]
+Raw bytes (35): 0x[01, 01, 03, 01, 05, 05, 0b, 09, 00, 05, 01, 12, 1c, 03, 21, 05, 04, 11, 01, 27, 02, 03, 11, 00, 16, 00, 00, 17, 00, 1e, 07, 02, 09, 00, 0a]
 Number of files: 1
 - file 0 => global file 1
-Number of expressions: 0
-Number of file 0 mappings: 1
-- Code(Counter(0)) at (prev + 37, 18) to (start + 0, 84)
+Number of expressions: 3
+- expression 0 operands: lhs = Counter(0), rhs = Counter(1)
+- expression 1 operands: lhs = Counter(1), rhs = Expression(2, Add)
+- expression 2 operands: lhs = Counter(2), rhs = Zero
+Number of file 0 mappings: 5
+- Code(Counter(0)) at (prev + 18, 28) to (start + 3, 33)
+- Code(Counter(1)) at (prev + 4, 17) to (start + 1, 39)
+- Code(Expression(0, Sub)) at (prev + 3, 17) to (start + 0, 22)
+    = (c0 - c1)
+- Code(Zero) at (prev + 0, 23) to (start + 0, 30)
+- Code(Expression(1, Add)) at (prev + 2, 9) to (start + 0, 10)
+    = (c1 + (c2 + Zero))
 
diff --git a/tests/coverage/closure_macro_async.coverage b/tests/coverage/closure_macro_async.coverage
index d67c2ed524b..1032e027cd9 100644
--- a/tests/coverage/closure_macro_async.coverage
+++ b/tests/coverage/closure_macro_async.coverage
@@ -15,16 +15,16 @@
    LL|       |
    LL|       |macro_rules! on_error {
    LL|       |    ($value:expr, $error_message:expr) => {
-   LL|       |        $value.or_else(|e| {
-   LL|       |            // FIXME(85000): no coverage in closure macros
-   LL|       |            let message = format!($error_message, e);
-   LL|       |            if message.len() > 0 {
-   LL|       |                println!("{}", message);
-   LL|       |                Ok(String::from("ok"))
+   LL|      0|        $value.or_else(|e| {
+   LL|      0|            // This closure, which is declared in a macro, should be instrumented.
+   LL|      0|            let message = format!($error_message, e);
+   LL|      0|            if message.len() > 0 {
+   LL|      0|                println!("{}", message);
+   LL|      0|                Ok(String::from("ok"))
    LL|       |            } else {
-   LL|       |                bail!("error");
+   LL|      0|                bail!("error");
    LL|       |            }
-   LL|       |        })
+   LL|      0|        })
    LL|       |    };
    LL|       |}
    LL|       |
diff --git a/tests/coverage/closure_macro_async.rs b/tests/coverage/closure_macro_async.rs
index 116cb72b0f3..db656fca198 100644
--- a/tests/coverage/closure_macro_async.rs
+++ b/tests/coverage/closure_macro_async.rs
@@ -16,7 +16,7 @@ macro_rules! bail {
 macro_rules! on_error {
     ($value:expr, $error_message:expr) => {
         $value.or_else(|e| {
-            // FIXME(85000): no coverage in closure macros
+            // This closure, which is declared in a macro, should be instrumented.
             let message = format!($error_message, e);
             if message.len() > 0 {
                 println!("{}", message);