coverage: Hoist special handling of async function spans

This sidesteps the normal span refinement code in cases where we know that we
are only dealing with the special signature span that represents having called
an async function.

Author: Zalathar

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,25 @@ 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.
+        let span = hir_info.fn_sig_span;
+        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,17 @@ 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 { fn_sig_span, 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() {
+        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));