coverage: Improve handling of function/closure spans

compiler/rustc_mir_transform/src/coverage/mod.rs

@@ -394,7 +394,9 @@ fn is_eligible_for_coverage(tcx: TyCtxt<'_>, def_id: LocalDefId) -> bool {
 struct ExtractedHirInfo {
     function_source_hash: u64,
     is_async_fn: bool,
-    fn_sig_span: Span,
+    /// The span of the function's signature, extended to the start of `body_span`.
+    /// Must have the same context and filename as the body span.
+    fn_sig_span_extended: Option<Span>,
     body_span: Span,
 }
 
@@ -407,13 +409,25 @@ fn extract_hir_info<'tcx>(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> ExtractedHir
         hir::map::associated_body(hir_node).expect("HIR node is a function with body");
     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 maybe_fn_sig = hir_node.fn_sig();
+    let is_async_fn = maybe_fn_sig.is_some_and(|fn_sig| fn_sig.header.is_async());
+
+    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.
-    let maybe_fn_sig_span = hir_node.fn_sig().map(|fn_sig| fn_sig.span);
-    let fn_sig_span = maybe_fn_sig_span
+    let fn_sig_span_extended = maybe_fn_sig
+        .map(|fn_sig| fn_sig.span)
         .filter(|&fn_sig_span| {
             let source_map = tcx.sess.source_map();
             let file_idx = |span: Span| source_map.lookup_source_file_idx(span.lo());
@@ -423,30 +437,11 @@ fn extract_hir_info<'tcx>(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> ExtractedHir
                 && file_idx(fn_sig_span) == file_idx(body_span)
         })
         // If so, extend it to the start of the body span.
-        .map(|fn_sig_span| fn_sig_span.with_hi(body_span.lo()))
-        // Otherwise, create a dummy signature span at the start of the body.
-        .unwrap_or_else(|| body_span.shrink_to_lo());
+        .map(|fn_sig_span| fn_sig_span.with_hi(body_span.lo()));
 
     let function_source_hash = hash_mir_source(tcx, hir_body);
 
-    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
+    ExtractedHirInfo { function_source_hash, is_async_fn, fn_sig_span_extended, body_span }
 }
 
 fn hash_mir_source<'tcx>(tcx: TyCtxt<'tcx>, hir_body: &'tcx rustc_hir::Body<'tcx>) -> u64 {

compiler/rustc_mir_transform/src/coverage/spans.rs

@@ -3,7 +3,7 @@ use rustc_index::bit_set::BitSet;
 use rustc_middle::mir;
 use rustc_span::{BytePos, Span, DUMMY_SP};
 
-use super::graph::{BasicCoverageBlock, CoverageGraph};
+use crate::coverage::graph::{BasicCoverageBlock, CoverageGraph, START_BCB};
 use crate::coverage::ExtractedHirInfo;
 
 mod from_mir;
@@ -46,13 +46,26 @@ pub(super) fn generate_coverage_spans(
 ) -> Option<CoverageSpans> {
     let mut mappings = vec![];
 
-    let sorted_spans =
-        from_mir::mir_to_initial_sorted_coverage_spans(mir_body, hir_info, basic_coverage_blocks);
-    let coverage_spans = SpansRefiner::refine_sorted_spans(basic_coverage_blocks, sorted_spans);
-    mappings.extend(coverage_spans.into_iter().map(|CoverageSpan { bcb, span, .. }| {
-        // Each span produced by the generator represents an ordinary code region.
-        BcbMapping { kind: BcbMappingKind::Code(bcb), span }
-    }));
+    if hir_info.is_async_fn {
+        // An async function desugars into a function that returns a future,
+        // with the user code wrapped in a closure. Any spans in the desugared
+        // outer function will be unhelpful, so just keep the signature span
+        // and ignore all of the spans in the MIR body.
+        if let Some(span) = hir_info.fn_sig_span_extended {
+            mappings.push(BcbMapping { kind: BcbMappingKind::Code(START_BCB), span });
+        }
+    } else {
+        let sorted_spans = from_mir::mir_to_initial_sorted_coverage_spans(
+            mir_body,
+            hir_info,
+            basic_coverage_blocks,
+        );
+        let coverage_spans = SpansRefiner::refine_sorted_spans(basic_coverage_blocks, sorted_spans);
+        mappings.extend(coverage_spans.into_iter().map(|CoverageSpan { bcb, span, .. }| {
+            // Each span produced by the generator represents an ordinary code region.
+            BcbMapping { kind: BcbMappingKind::Code(bcb), span }
+        }));
+    }
 
     if mappings.is_empty() {
         return None;

compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs

@@ -23,25 +23,21 @@ pub(super) fn mir_to_initial_sorted_coverage_spans(
     hir_info: &ExtractedHirInfo,
     basic_coverage_blocks: &CoverageGraph,
 ) -> Vec<CoverageSpan> {
-    let &ExtractedHirInfo { is_async_fn, fn_sig_span, body_span, .. } = hir_info;
-
-    let mut initial_spans = vec![SpanFromMir::for_fn_sig(fn_sig_span)];
-
-    if is_async_fn {
-        // An async function desugars into a function that returns a future,
-        // with the user code wrapped in a closure. Any spans in the desugared
-        // outer function will be unhelpful, so just keep the signature span
-        // and ignore all of the spans in the MIR body.
-    } else {
-        for (bcb, bcb_data) in basic_coverage_blocks.iter_enumerated() {
-            initial_spans.extend(bcb_to_initial_coverage_spans(mir_body, body_span, bcb, bcb_data));
-        }
+    let &ExtractedHirInfo { body_span, .. } = hir_info;
 
-        // If no spans were extracted from the body, discard the signature span.
-        // FIXME: This preserves existing behavior; consider getting rid of it.
-        if initial_spans.len() == 1 {
-            initial_spans.clear();
-        }
+    let mut initial_spans = vec![];
+
+    for (bcb, bcb_data) in basic_coverage_blocks.iter_enumerated() {
+        initial_spans.extend(bcb_to_initial_coverage_spans(mir_body, body_span, bcb, bcb_data));
+    }
+
+    // Only add the signature span if we found at least one span in the body.
+    if !initial_spans.is_empty() {
+        // If there is no usable signature span, add a fake one (before refinement)
+        // to avoid an ugly gap between the body start and the first real span.
+        // FIXME: Find a more principled way to solve this problem.
+        let fn_sig_span = hir_info.fn_sig_span_extended.unwrap_or_else(|| body_span.shrink_to_lo());
+        initial_spans.push(SpanFromMir::for_fn_sig(fn_sig_span));
     }
 
     initial_spans.sort_by(|a, b| basic_coverage_blocks.cmp_in_dominator_order(a.bcb, b.bcb));

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))
 

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

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);

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))
 

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

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);

tests/coverage/coverage_attr_closure.cov-map

@@ -0,0 +1,34 @@
+Function name: coverage_attr_closure::GLOBAL_CLOSURE_ON::{closure#0}
+Raw bytes (9): 0x[01, 01, 00, 01, 01, 06, 0f, 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 + 6, 15) to (start + 2, 2)
+
+Function name: coverage_attr_closure::contains_closures_off::{closure#0} (unused)
+Raw bytes (9): 0x[01, 01, 00, 01, 00, 1d, 13, 02, 06]
+Number of files: 1
+- file 0 => global file 1
+Number of expressions: 0
+Number of file 0 mappings: 1
+- Code(Zero) at (prev + 29, 19) to (start + 2, 6)
+
+Function name: coverage_attr_closure::contains_closures_on
+Raw bytes (19): 0x[01, 01, 00, 03, 01, 0f, 01, 02, 05, 01, 04, 06, 02, 05, 01, 04, 06, 01, 02]
+Number of files: 1
+- file 0 => global file 1
+Number of expressions: 0
+Number of file 0 mappings: 3
+- Code(Counter(0)) at (prev + 15, 1) to (start + 2, 5)
+- Code(Counter(0)) at (prev + 4, 6) to (start + 2, 5)
+- Code(Counter(0)) at (prev + 4, 6) to (start + 1, 2)
+
+Function name: coverage_attr_closure::contains_closures_on::{closure#0} (unused)
+Raw bytes (9): 0x[01, 01, 00, 01, 00, 11, 13, 02, 06]
+Number of files: 1
+- file 0 => global file 1
+Number of expressions: 0
+Number of file 0 mappings: 1
+- Code(Zero) at (prev + 17, 19) to (start + 2, 6)
+

tests/coverage/coverage_attr_closure.coverage

@@ -0,0 +1,43 @@
+   LL|       |#![feature(coverage_attribute, stmt_expr_attributes)]
+   LL|       |#![allow(dead_code)]
+   LL|       |// edition: 2021
+   LL|       |
+   LL|       |static GLOBAL_CLOSURE_ON: fn(&str) = #[coverage(on)]
+   LL|      0||input: &str| {
+   LL|      0|    println!("{input}");
+   LL|      0|};
+   LL|       |static GLOBAL_CLOSURE_OFF: fn(&str) = #[coverage(off)]
+   LL|       ||input: &str| {
+   LL|       |    println!("{input}");
+   LL|       |};
+   LL|       |
+   LL|       |#[coverage(on)]
+   LL|      1|fn contains_closures_on() {
+   LL|      1|    let _local_closure_on = #[coverage(on)]
+   LL|      1|    |input: &str| {
+   LL|      0|        println!("{input}");
+   LL|      1|    };
+   LL|      1|    let _local_closure_off = #[coverage(off)]
+   LL|      1|    |input: &str| {
+   LL|       |        println!("{input}");
+   LL|      1|    };
+   LL|      1|}
+   LL|       |
+   LL|       |#[coverage(off)]
+   LL|       |fn contains_closures_off() {
+   LL|       |    let _local_closure_on = #[coverage(on)]
+   LL|      0|    |input: &str| {
+   LL|      0|        println!("{input}");
+   LL|      0|    };
+   LL|       |    let _local_closure_off = #[coverage(off)]
+   LL|       |    |input: &str| {
+   LL|       |        println!("{input}");
+   LL|       |    };
+   LL|       |}
+   LL|       |
+   LL|       |#[coverage(off)]
+   LL|       |fn main() {
+   LL|       |    contains_closures_on();
+   LL|       |    contains_closures_off();
+   LL|       |}
+