coverage: Check for async fn explicitly, without needing a heuristic

Zalathar · Zalathar · commit cf6dc7adb3b9 · 2023-12-20T18:48:04.000+11:00
The old code used a heuristic to detect async functions and adjust their
coverage spans to produce better output. But there's no need to resort to a
heuristic when we can just check whether the current function is actually an
`async fn`.
diff --git a/compiler/rustc_mir_transform/src/coverage/mod.rs b/compiler/rustc_mir_transform/src/coverage/mod.rs
@@ -299,6 +299,7 @@ fn is_eligible_for_coverage(tcx: TyCtxt<'_>, def_id: LocalDefId) -> bool {
 #[derive(Debug)]
 struct ExtractedHirInfo {
     function_source_hash: u64,
+    is_async_fn: bool,
     fn_sig_span: Span,
     body_span: Span,
 }
@@ -312,6 +313,7 @@ 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);
 
     // The actual signature span is only used if it has the same context and
@@ -333,7 +335,7 @@ fn extract_hir_info<'tcx>(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> ExtractedHir
 
     let function_source_hash = hash_mir_source(tcx, hir_body);
 
-    ExtractedHirInfo { function_source_hash, fn_sig_span, body_span }
+    ExtractedHirInfo { function_source_hash, is_async_fn, fn_sig_span, body_span }
 }
 
 fn get_body_span<'tcx>(
diff --git a/compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs b/compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs
@@ -14,7 +14,14 @@ pub(super) fn mir_to_initial_sorted_coverage_spans(
     hir_info: &ExtractedHirInfo,
     basic_coverage_blocks: &CoverageGraph,
 ) -> Vec<CoverageSpan> {
-    let &ExtractedHirInfo { fn_sig_span, body_span, .. } = hir_info;
+    let &ExtractedHirInfo { is_async_fn, fn_sig_span, body_span, .. } = hir_info;
+    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 produce a single span
+        // associating the function signature with its entry BCB.
+        return vec![CoverageSpan::for_fn_sig(fn_sig_span)];
+    }
 
     let mut initial_spans = Vec::with_capacity(mir_body.basic_blocks.len() * 2);
     for (bcb, bcb_data) in basic_coverage_blocks.iter_enumerated() {
@@ -46,16 +53,6 @@ pub(super) fn mir_to_initial_sorted_coverage_spans(
             .then_with(|| Ord::cmp(&a.is_closure, &b.is_closure).reverse())
     });
 
-    // The desugaring of an async function includes a closure containing the
-    // original function body, and a terminator that returns the `impl Future`.
-    // That terminator will cause a confusing coverage count for the function's
-    // closing brace, so discard everything after the body closure span.
-    if let Some(body_closure_index) =
-        initial_spans.iter().rposition(|covspan| covspan.is_closure && covspan.span == body_span)
-    {
-        initial_spans.truncate(body_closure_index + 1);
-    }
-
     initial_spans
 }
 
