From fde1702db86a8de4849caca12787d74525d06344 Mon Sep 17 00:00:00 2001
From: Zalathar <Zalathar@users.noreply.github.com>
Date: Tue, 16 Jan 2024 22:07:21 +1100
Subject: [PATCH] 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.
---
 .../rustc_mir_transform/src/coverage/spans.rs | 28 +++++++++++++------
 .../src/coverage/spans/from_mir.rs            | 24 ++++++----------
 2 files changed, 28 insertions(+), 24 deletions(-)

diff --git a/compiler/rustc_mir_transform/src/coverage/spans.rs b/compiler/rustc_mir_transform/src/coverage/spans.rs
index dee6a3b7143..6db16262a70 100644
--- a/compiler/rustc_mir_transform/src/coverage/spans.rs
+++ b/compiler/rustc_mir_transform/src/coverage/spans.rs
@@ -3,7 +3,7 @@
 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;
diff --git a/compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs b/compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs
index 8d8e8e61327..63371a39f5a 100644
--- a/compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs
+++ b/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 &ExtractedHirInfo { fn_sig_span, body_span, .. } = hir_info;
 
-    let mut initial_spans = vec![SpanFromMir::for_fn_sig(fn_sig_span)];
+    let mut initial_spans = vec![];
 
-    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));
-        }
+    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));
+    }
 
-        // 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();
-        }
+    // 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));