coverage: Treat await similar to a macro

Currently `await` is only counted towards coverage if the containing
function is suspended and resumed at least once. This is because it
expands to code which contains a branch on the discriminant of `Poll`.

By treating it like a branching macro (e.g. `assert!`), these
implementation details will be hidden from the coverage results.

8 files changed, 67 insertions, 71 deletions

diff --git a/compiler/rustc_mir_transform/src/coverage/spans.rs b/compiler/rustc_mir_transform/src/coverage/spans.rs
index fcc774503f4..b904b0d2599 100644
--- a/compiler/rustc_mir_transform/src/coverage/spans.rs
+++ b/compiler/rustc_mir_transform/src/coverage/spans.rs
@@ -3,7 +3,7 @@ use std::collections::VecDeque;
 use rustc_data_structures::captures::Captures;
 use rustc_data_structures::fx::FxHashSet;
 use rustc_middle::mir;
-use rustc_span::Span;
+use rustc_span::{DesugaringKind, ExpnKind, MacroKind, Span};
 use tracing::{debug, debug_span, instrument};
 
 use crate::coverage::graph::{BasicCoverageBlock, CoverageGraph};
@@ -25,7 +25,7 @@ pub(super) fn extract_refined_covspans(
 
     // First, perform the passes that need macro information.
     covspans.sort_by(|a, b| basic_coverage_blocks.cmp_in_dominator_order(a.bcb, b.bcb));
-    remove_unwanted_macro_spans(&mut covspans);
+    remove_unwanted_expansion_spans(&mut covspans);
     split_visible_macro_spans(&mut covspans);
 
     // We no longer need the extra information in `SpanFromMir`, so convert to `Covspan`.
@@ -76,18 +76,24 @@ pub(super) fn extract_refined_covspans(
 /// invocation, which is unhelpful. Keeping only the first such span seems to
 /// give better mappings, so remove the others.
 ///
+/// Similarly, `await` expands to a branch on the discriminant of `Poll`, which
+/// leads to incorrect coverage if the `Future` is immediately ready (#98712).
+///
 /// (The input spans should be sorted in BCB dominator order, so that the
 /// retained "first" span is likely to dominate the others.)
-fn remove_unwanted_macro_spans(covspans: &mut Vec<SpanFromMir>) {
-    let mut seen_macro_spans = FxHashSet::default();
+fn remove_unwanted_expansion_spans(covspans: &mut Vec<SpanFromMir>) {
+    let mut deduplicated_spans = FxHashSet::default();
+
     covspans.retain(|covspan| {
-        // Ignore (retain) non-macro-expansion spans.
-        if covspan.visible_macro.is_none() {
-            return true;
+        match covspan.expn_kind {
+            // Retain only the first await-related or macro-expanded covspan with this span.
+            Some(ExpnKind::Desugaring(kind)) if kind == DesugaringKind::Await => {
+                deduplicated_spans.insert(covspan.span)
+            }
+            Some(ExpnKind::Macro(MacroKind::Bang, _)) => deduplicated_spans.insert(covspan.span),
+            // Ignore (retain) other spans.
+            _ => true,
         }
-
-        // Retain only the first macro-expanded covspan with this span.
-        seen_macro_spans.insert(covspan.span)
     });
 }
 
@@ -99,7 +105,9 @@ fn split_visible_macro_spans(covspans: &mut Vec<SpanFromMir>) {
     let mut extra_spans = vec![];
 
     covspans.retain(|covspan| {
-        let Some(visible_macro) = covspan.visible_macro else { return true };
+        let Some(ExpnKind::Macro(MacroKind::Bang, visible_macro)) = covspan.expn_kind else {
+            return true;
+        };
 
         let split_len = visible_macro.as_str().len() as u32 + 1;
         let (before, after) = covspan.span.split_at(split_len);
@@ -111,8 +119,8 @@ fn split_visible_macro_spans(covspans: &mut Vec<SpanFromMir>) {
             return true;
         }
 
-        extra_spans.push(SpanFromMir::new(before, covspan.visible_macro, covspan.bcb));
-        extra_spans.push(SpanFromMir::new(after, covspan.visible_macro, covspan.bcb));
+        extra_spans.push(SpanFromMir::new(before, covspan.expn_kind.clone(), covspan.bcb));
+        extra_spans.push(SpanFromMir::new(after, covspan.expn_kind.clone(), covspan.bcb));
         false // Discard the original covspan that we just split.
     });
 
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 32bd25bf4b9..7f5765c9462 100644
--- a/compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs
+++ b/compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs
@@ -3,13 +3,13 @@ use rustc_middle::mir::coverage::CoverageKind;
 use rustc_middle::mir::{
     self, FakeReadCause, Statement, StatementKind, Terminator, TerminatorKind,
 };
-use rustc_span::{Span, Symbol};
+use rustc_span::{ExpnKind, Span};
 
 use crate::coverage::graph::{
     BasicCoverageBlock, BasicCoverageBlockData, CoverageGraph, START_BCB,
 };
 use crate::coverage::spans::Covspan;
-use crate::coverage::unexpand::unexpand_into_body_span_with_visible_macro;
+use crate::coverage::unexpand::unexpand_into_body_span_with_expn_kind;
 use crate::coverage::ExtractedHirInfo;
 
 pub(crate) struct ExtractedCovspans {
@@ -60,7 +60,7 @@ fn bcb_to_initial_coverage_spans<'a, 'tcx>(
         let data = &mir_body[bb];
 
         let unexpand = move |expn_span| {
-            unexpand_into_body_span_with_visible_macro(expn_span, body_span)
+            unexpand_into_body_span_with_expn_kind(expn_span, body_span)
                 // Discard any spans that fill the entire body, because they tend
                 // to represent compiler-inserted code, e.g. implicitly returning `()`.
                 .filter(|(span, _)| !span.source_equal(body_span))
@@ -68,9 +68,9 @@ fn bcb_to_initial_coverage_spans<'a, 'tcx>(
 
         let mut extract_statement_span = |statement| {
             let expn_span = filtered_statement_span(statement)?;
-            let (span, visible_macro) = unexpand(expn_span)?;
+            let (span, expn_kind) = unexpand(expn_span)?;
 
-            initial_covspans.push(SpanFromMir::new(span, visible_macro, bcb));
+            initial_covspans.push(SpanFromMir::new(span, expn_kind, bcb));
             Some(())
         };
         for statement in data.statements.iter() {
@@ -79,9 +79,9 @@ fn bcb_to_initial_coverage_spans<'a, 'tcx>(
 
         let mut extract_terminator_span = |terminator| {
             let expn_span = filtered_terminator_span(terminator)?;
-            let (span, visible_macro) = unexpand(expn_span)?;
+            let (span, expn_kind) = unexpand(expn_span)?;
 
-            initial_covspans.push(SpanFromMir::new(span, visible_macro, bcb));
+            initial_covspans.push(SpanFromMir::new(span, expn_kind, bcb));
             Some(())
         };
         extract_terminator_span(data.terminator());
@@ -214,7 +214,7 @@ pub(crate) struct SpanFromMir {
     /// With the exception of `fn_sig_span`, this should always be contained
     /// within `body_span`.
     pub(crate) span: Span,
-    pub(crate) visible_macro: Option<Symbol>,
+    pub(crate) expn_kind: Option<ExpnKind>,
     pub(crate) bcb: BasicCoverageBlock,
 }
 
@@ -223,12 +223,12 @@ impl SpanFromMir {
         Self::new(fn_sig_span, None, START_BCB)
     }
 
-    pub(crate) fn new(span: Span, visible_macro: Option<Symbol>, bcb: BasicCoverageBlock) -> Self {
-        Self { span, visible_macro, bcb }
+    pub(crate) fn new(span: Span, expn_kind: Option<ExpnKind>, bcb: BasicCoverageBlock) -> Self {
+        Self { span, expn_kind, bcb }
     }
 
     pub(crate) fn into_covspan(self) -> Covspan {
-        let Self { span, visible_macro: _, bcb } = self;
+        let Self { span, expn_kind: _, bcb } = self;
         Covspan { span, bcb }
     }
 }
diff --git a/compiler/rustc_mir_transform/src/coverage/unexpand.rs b/compiler/rustc_mir_transform/src/coverage/unexpand.rs
index 8cde291b907..cb861544736 100644
--- a/compiler/rustc_mir_transform/src/coverage/unexpand.rs
+++ b/compiler/rustc_mir_transform/src/coverage/unexpand.rs
@@ -1,4 +1,4 @@
-use rustc_span::{ExpnKind, MacroKind, Span, Symbol};
+use rustc_span::{ExpnKind, Span};
 
 /// Walks through the expansion ancestors of `original_span` to find a span that
 /// is contained in `body_span` and has the same [syntax context] as `body_span`.
@@ -13,20 +13,15 @@ pub(crate) fn unexpand_into_body_span(original_span: Span, body_span: Span) -> O
 ///
 /// If the returned span represents a bang-macro invocation (e.g. `foo!(..)`),
 /// the returned symbol will be the name of that macro (e.g. `foo`).
-pub(crate) fn unexpand_into_body_span_with_visible_macro(
+pub(crate) fn unexpand_into_body_span_with_expn_kind(
     original_span: Span,
     body_span: Span,
-) -> Option<(Span, Option<Symbol>)> {
+) -> Option<(Span, Option<ExpnKind>)> {
     let (span, prev) = unexpand_into_body_span_with_prev(original_span, body_span)?;
 
-    let visible_macro = prev
-        .map(|prev| match prev.ctxt().outer_expn_data().kind {
-            ExpnKind::Macro(MacroKind::Bang, name) => Some(name),
-            _ => None,
-        })
-        .flatten();
+    let expn_kind = prev.map(|prev| prev.ctxt().outer_expn_data().kind);
 
-    Some((span, visible_macro))
+    Some((span, expn_kind))
 }
 
 /// Walks through the expansion ancestors of `original_span` to find a span that
diff --git a/tests/coverage/async.cov-map b/tests/coverage/async.cov-map
index 9e5a4bdc60f..1ba165f1e49 100644
--- a/tests/coverage/async.cov-map
+++ b/tests/coverage/async.cov-map
@@ -92,20 +92,18 @@ Number of file 0 mappings: 1
 - Code(Counter(0)) at (prev + 25, 1) to (start + 0, 23)
 
 Function name: async::g::{closure#0} (unused)
-Raw bytes (69): 0x[01, 01, 00, 0d, 00, 19, 17, 01, 0c, 00, 02, 09, 00, 0a, 00, 00, 0e, 00, 11, 00, 00, 12, 00, 17, 00, 00, 1b, 00, 1c, 00, 00, 20, 00, 22, 00, 01, 09, 00, 0a, 00, 00, 0e, 00, 11, 00, 00, 12, 00, 17, 00, 00, 1b, 00, 1c, 00, 00, 20, 00, 22, 00, 01, 0e, 00, 10, 00, 02, 01, 00, 02]
+Raw bytes (59): 0x[01, 01, 00, 0b, 00, 19, 17, 01, 0c, 00, 02, 09, 00, 0a, 00, 00, 0e, 00, 17, 00, 00, 1b, 00, 1c, 00, 00, 20, 00, 22, 00, 01, 09, 00, 0a, 00, 00, 0e, 00, 17, 00, 00, 1b, 00, 1c, 00, 00, 20, 00, 22, 00, 01, 0e, 00, 10, 00, 02, 01, 00, 02]
 Number of files: 1
 - file 0 => global file 1
 Number of expressions: 0
-Number of file 0 mappings: 13
+Number of file 0 mappings: 11
 - Code(Zero) at (prev + 25, 23) to (start + 1, 12)
 - Code(Zero) at (prev + 2, 9) to (start + 0, 10)
-- Code(Zero) at (prev + 0, 14) to (start + 0, 17)
-- Code(Zero) at (prev + 0, 18) to (start + 0, 23)
+- Code(Zero) at (prev + 0, 14) to (start + 0, 23)
 - Code(Zero) at (prev + 0, 27) to (start + 0, 28)
 - Code(Zero) at (prev + 0, 32) to (start + 0, 34)
 - Code(Zero) at (prev + 1, 9) to (start + 0, 10)
-- Code(Zero) at (prev + 0, 14) to (start + 0, 17)
-- Code(Zero) at (prev + 0, 18) to (start + 0, 23)
+- Code(Zero) at (prev + 0, 14) to (start + 0, 23)
 - Code(Zero) at (prev + 0, 27) to (start + 0, 28)
 - Code(Zero) at (prev + 0, 32) to (start + 0, 34)
 - Code(Zero) at (prev + 1, 14) to (start + 0, 16)
@@ -120,15 +118,14 @@ Number of file 0 mappings: 1
 - Code(Counter(0)) at (prev + 33, 1) to (start + 0, 22)
 
 Function name: async::h::{closure#0} (unused)
-Raw bytes (44): 0x[01, 01, 00, 08, 00, 21, 16, 03, 0c, 00, 04, 09, 00, 0a, 00, 00, 0e, 00, 13, 00, 00, 14, 00, 19, 00, 00, 1a, 00, 1b, 00, 00, 20, 00, 22, 00, 01, 0e, 00, 10, 00, 02, 01, 00, 02]
+Raw bytes (39): 0x[01, 01, 00, 07, 00, 21, 16, 03, 0c, 00, 04, 09, 00, 0a, 00, 00, 0e, 00, 19, 00, 00, 1a, 00, 1b, 00, 00, 20, 00, 22, 00, 01, 0e, 00, 10, 00, 02, 01, 00, 02]
 Number of files: 1
 - file 0 => global file 1
 Number of expressions: 0
-Number of file 0 mappings: 8
+Number of file 0 mappings: 7
 - Code(Zero) at (prev + 33, 22) to (start + 3, 12)
 - Code(Zero) at (prev + 4, 9) to (start + 0, 10)
-- Code(Zero) at (prev + 0, 14) to (start + 0, 19)
-- Code(Zero) at (prev + 0, 20) to (start + 0, 25)
+- Code(Zero) at (prev + 0, 14) to (start + 0, 25)
 - Code(Zero) at (prev + 0, 26) to (start + 0, 27)
 - Code(Zero) at (prev + 0, 32) to (start + 0, 34)
 - Code(Zero) at (prev + 1, 14) to (start + 0, 16)
@@ -143,28 +140,25 @@ Number of file 0 mappings: 1
 - Code(Counter(0)) at (prev + 42, 1) to (start + 0, 19)
 
 Function name: async::i::{closure#0}
-Raw bytes (78): 0x[01, 01, 02, 07, 21, 19, 1d, 0e, 01, 2a, 13, 04, 0c, 0d, 05, 09, 00, 0a, 01, 00, 0e, 00, 12, 05, 00, 13, 00, 18, 09, 00, 1c, 00, 21, 0d, 00, 27, 00, 2a, 15, 00, 2b, 00, 30, 1d, 01, 09, 00, 0a, 11, 00, 0e, 00, 11, 25, 00, 12, 00, 17, 29, 00, 1b, 00, 20, 1d, 00, 24, 00, 26, 21, 01, 0e, 00, 10, 03, 02, 01, 00, 02]
+Raw bytes (63): 0x[01, 01, 02, 07, 19, 11, 15, 0b, 01, 2a, 13, 04, 0c, 09, 05, 09, 00, 0a, 01, 00, 0e, 00, 18, 05, 00, 1c, 00, 21, 09, 00, 27, 00, 30, 15, 01, 09, 00, 0a, 0d, 00, 0e, 00, 17, 1d, 00, 1b, 00, 20, 15, 00, 24, 00, 26, 19, 01, 0e, 00, 10, 03, 02, 01, 00, 02]
 Number of files: 1
 - file 0 => global file 1
 Number of expressions: 2
-- expression 0 operands: lhs = Expression(1, Add), rhs = Counter(8)
-- expression 1 operands: lhs = Counter(6), rhs = Counter(7)
-Number of file 0 mappings: 14
+- expression 0 operands: lhs = Expression(1, Add), rhs = Counter(6)
+- expression 1 operands: lhs = Counter(4), rhs = Counter(5)
+Number of file 0 mappings: 11
 - Code(Counter(0)) at (prev + 42, 19) to (start + 4, 12)
-- Code(Counter(3)) at (prev + 5, 9) to (start + 0, 10)
-- Code(Counter(0)) at (prev + 0, 14) to (start + 0, 18)
-- Code(Counter(1)) at (prev + 0, 19) to (start + 0, 24)
-- Code(Counter(2)) at (prev + 0, 28) to (start + 0, 33)
-- Code(Counter(3)) at (prev + 0, 39) to (start + 0, 42)
-- Code(Counter(5)) at (prev + 0, 43) to (start + 0, 48)
-- Code(Counter(7)) at (prev + 1, 9) to (start + 0, 10)
-- Code(Counter(4)) at (prev + 0, 14) to (start + 0, 17)
-- Code(Counter(9)) at (prev + 0, 18) to (start + 0, 23)
-- Code(Counter(10)) at (prev + 0, 27) to (start + 0, 32)
-- Code(Counter(7)) at (prev + 0, 36) to (start + 0, 38)
-- Code(Counter(8)) at (prev + 1, 14) to (start + 0, 16)
+- Code(Counter(2)) at (prev + 5, 9) to (start + 0, 10)
+- Code(Counter(0)) at (prev + 0, 14) to (start + 0, 24)
+- Code(Counter(1)) at (prev + 0, 28) to (start + 0, 33)
+- Code(Counter(2)) at (prev + 0, 39) to (start + 0, 48)
+- Code(Counter(5)) at (prev + 1, 9) to (start + 0, 10)
+- Code(Counter(3)) at (prev + 0, 14) to (start + 0, 23)
+- Code(Counter(7)) at (prev + 0, 27) to (start + 0, 32)
+- Code(Counter(5)) at (prev + 0, 36) to (start + 0, 38)
+- Code(Counter(6)) at (prev + 1, 14) to (start + 0, 16)
 - Code(Expression(0, Add)) at (prev + 2, 1) to (start + 0, 2)
-    = ((c6 + c7) + c8)
+    = ((c4 + c5) + c6)
 
 Function name: async::j
 Raw bytes (58): 0x[01, 01, 02, 07, 0d, 05, 09, 0a, 01, 35, 01, 00, 0d, 01, 0b, 0b, 00, 0c, 05, 01, 09, 00, 0a, 01, 00, 0e, 00, 1b, 05, 00, 1f, 00, 27, 09, 01, 09, 00, 0a, 11, 00, 0e, 00, 1a, 09, 00, 1e, 00, 20, 0d, 01, 0e, 00, 10, 03, 02, 01, 00, 02]
diff --git a/tests/coverage/async.coverage b/tests/coverage/async.coverage
index f5473829b02..995674257c4 100644
--- a/tests/coverage/async.coverage
+++ b/tests/coverage/async.coverage
@@ -45,9 +45,9 @@
    LL|      1|                    // executed asynchronously.
    LL|      1|    match x {
    LL|      1|        y if c(x).await == y + 1 => { d().await; }
-                      ^0        ^0                  ^0  ^0
+                      ^0                            ^0
    LL|      1|        y if f().await == y + 1 => (),
-                      ^0       ^0                ^0
+                      ^0                         ^0
    LL|      1|        _ => (),
    LL|       |    }
    LL|      1|}
diff --git a/tests/coverage/await_ready.cov-map b/tests/coverage/await_ready.cov-map
index c6dfc01e861..0c9f2ae29a8 100644
--- a/tests/coverage/await_ready.cov-map
+++ b/tests/coverage/await_ready.cov-map
@@ -7,14 +7,13 @@ Number of file 0 mappings: 1
 - Code(Counter(0)) at (prev + 10, 1) to (start + 0, 30)
 
 Function name: await_ready::await_ready::{closure#0}
-Raw bytes (19): 0x[01, 01, 00, 03, 01, 0a, 1e, 02, 0c, 05, 03, 0a, 00, 0f, 09, 01, 01, 00, 02]
+Raw bytes (14): 0x[01, 01, 00, 02, 01, 0a, 1e, 03, 0f, 05, 04, 01, 00, 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 + 10, 30) to (start + 2, 12)
-- Code(Counter(1)) at (prev + 3, 10) to (start + 0, 15)
-- Code(Counter(2)) at (prev + 1, 1) to (start + 0, 2)
+Number of file 0 mappings: 2
+- Code(Counter(0)) at (prev + 10, 30) to (start + 3, 15)
+- Code(Counter(1)) at (prev + 4, 1) to (start + 0, 2)
 
 Function name: await_ready::main
 Raw bytes (9): 0x[01, 01, 00, 01, 01, 10, 01, 03, 02]
diff --git a/tests/coverage/await_ready.coverage b/tests/coverage/await_ready.coverage
index bddaf15cf80..0075f09426e 100644
--- a/tests/coverage/await_ready.coverage
+++ b/tests/coverage/await_ready.coverage
@@ -8,9 +8,9 @@
    LL|       |async fn ready() -> u8 { 1 }
    LL|       |
    LL|      1|async fn await_ready() -> u8 {
-   LL|      1|    // FIXME(#98712): await is only covered if the function yields
+   LL|      1|    // await should be covered even if the function never yields
    LL|      1|    ready()
-   LL|      0|        .await
+   LL|      1|        .await
    LL|      1|}
    LL|       |
    LL|      1|fn main() {
diff --git a/tests/coverage/await_ready.rs b/tests/coverage/await_ready.rs
index 884ff16fff6..9212a4ba705 100644
--- a/tests/coverage/await_ready.rs
+++ b/tests/coverage/await_ready.rs
@@ -8,7 +8,7 @@
 async fn ready() -> u8 { 1 }
 
 async fn await_ready() -> u8 {
-    // FIXME(#98712): await is only covered if the function yields
+    // await should be covered even if the function never yields
     ready()
         .await
 }