coverage: Push async special case down into `extract_refined_covspans`

2 files changed, 16 insertions, 14 deletions

diff --git a/compiler/rustc_mir_transform/src/coverage/mappings.rs b/compiler/rustc_mir_transform/src/coverage/mappings.rs
index c79b76d90f2..b0e24cf2bdb 100644
--- a/compiler/rustc_mir_transform/src/coverage/mappings.rs
+++ b/compiler/rustc_mir_transform/src/coverage/mappings.rs
@@ -10,7 +10,7 @@ use rustc_middle::ty::TyCtxt;
 use rustc_span::Span;
 
 use crate::coverage::ExtractedHirInfo;
-use crate::coverage::graph::{BasicCoverageBlock, CoverageGraph, START_BCB};
+use crate::coverage::graph::{BasicCoverageBlock, CoverageGraph};
 use crate::coverage::spans::extract_refined_covspans;
 use crate::coverage::unexpand::unexpand_into_body_span;
 use crate::errors::MCDCExceedsTestVectorLimit;
@@ -82,18 +82,8 @@ pub(super) fn extract_all_mapping_info_from_mir<'tcx>(
     let mut mcdc_degraded_branches = vec![];
     let mut mcdc_mappings = vec![];
 
-    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 {
-            code_mappings.push(CodeMapping { span, bcb: START_BCB });
-        }
-    } else {
-        // Extract coverage spans from MIR statements/terminators as normal.
-        extract_refined_covspans(tcx, mir_body, hir_info, graph, &mut code_mappings);
-    }
+    // Extract ordinary code mappings from MIR statement/terminator spans.
+    extract_refined_covspans(tcx, mir_body, hir_info, graph, &mut code_mappings);
 
     branch_pairs.extend(extract_branch_pairs(mir_body, hir_info, graph));
 
diff --git a/compiler/rustc_mir_transform/src/coverage/spans.rs b/compiler/rustc_mir_transform/src/coverage/spans.rs
index ddeae093df5..0ee42abb195 100644
--- a/compiler/rustc_mir_transform/src/coverage/spans.rs
+++ b/compiler/rustc_mir_transform/src/coverage/spans.rs
@@ -1,5 +1,6 @@
 use rustc_data_structures::fx::FxHashSet;
 use rustc_middle::mir;
+use rustc_middle::mir::coverage::START_BCB;
 use rustc_middle::ty::TyCtxt;
 use rustc_span::source_map::SourceMap;
 use rustc_span::{BytePos, DesugaringKind, ExpnKind, MacroKind, Span};
@@ -16,8 +17,19 @@ pub(super) fn extract_refined_covspans<'tcx>(
     mir_body: &mir::Body<'tcx>,
     hir_info: &ExtractedHirInfo,
     graph: &CoverageGraph,
-    code_mappings: &mut impl Extend<mappings::CodeMapping>,
+    code_mappings: &mut Vec<mappings::CodeMapping>,
 ) {
+    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 {
+            code_mappings.push(mappings::CodeMapping { span, bcb: START_BCB });
+        }
+        return;
+    }
+
     let &ExtractedHirInfo { body_span, .. } = hir_info;
 
     let raw_spans = from_mir::extract_raw_spans_from_mir(mir_body, graph);