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`.

1 files changed, 12 insertions, 18 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;