Rollup merge of rust-lang#119155 - Zalathar:async-fn, r=compiler-errors

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

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 look back at the original definition and check whether the current function is actually an `async fn`.

In addition to being generally nicer, this also gets rid of the one piece of code that specifically cares about `CoverageSpan::is_closure` representing an actual closure. All remaining code that inspects that field just uses it as an indication that the span is a hole that should be carved out of other spans, and then discarded.

That opens up the possibility of introducing other kinds of “hole” spans, e.g. for nested functions/types/macros, and having them all behave uniformly.

---

`@rustbot` label +A-code-coverage

Author: GuillaumeGomez

compiler/rustc_mir_transform/src/coverage/mod.rs

@@ -68,41 +68,29 @@ impl<'tcx> MirPass<'tcx> for InstrumentCoverage {
 struct Instrumentor<'a, 'tcx> {
     tcx: TyCtxt<'tcx>,
     mir_body: &'a mut mir::Body<'tcx>,
-    fn_sig_span: Span,
-    body_span: Span,
-    function_source_hash: u64,
+    hir_info: ExtractedHirInfo,
     basic_coverage_blocks: CoverageGraph,
     coverage_counters: CoverageCounters,
 }
 
 impl<'a, 'tcx> Instrumentor<'a, 'tcx> {
     fn new(tcx: TyCtxt<'tcx>, mir_body: &'a mut mir::Body<'tcx>) -> Self {
-        let hir_info @ ExtractedHirInfo { function_source_hash, fn_sig_span, body_span } =
-            extract_hir_info(tcx, mir_body.source.def_id().expect_local());
+        let hir_info = extract_hir_info(tcx, mir_body.source.def_id().expect_local());
 
         debug!(?hir_info, "instrumenting {:?}", mir_body.source.def_id());
 
         let basic_coverage_blocks = CoverageGraph::from_mir(mir_body);
         let coverage_counters = CoverageCounters::new(&basic_coverage_blocks);
 
-        Self {
-            tcx,
-            mir_body,
-            fn_sig_span,
-            body_span,
-            function_source_hash,
-            basic_coverage_blocks,
-            coverage_counters,
-        }
+        Self { tcx, mir_body, hir_info, basic_coverage_blocks, coverage_counters }
     }
 
     fn inject_counters(&'a mut self) {
         ////////////////////////////////////////////////////
         // Compute coverage spans from the `CoverageGraph`.
         let Some(coverage_spans) = CoverageSpans::generate_coverage_spans(
             self.mir_body,
-            self.fn_sig_span,
-            self.body_span,
+            &self.hir_info,
             &self.basic_coverage_blocks,
         ) else {
             // No relevant spans were found in MIR, so skip instrumenting this function.
@@ -121,7 +109,7 @@ impl<'a, 'tcx> Instrumentor<'a, 'tcx> {
         let mappings = self.create_mappings_and_inject_coverage_statements(&coverage_spans);
 
         self.mir_body.function_coverage_info = Some(Box::new(FunctionCoverageInfo {
-            function_source_hash: self.function_source_hash,
+            function_source_hash: self.hir_info.function_source_hash,
             num_counters: self.coverage_counters.num_counters(),
             expressions: self.coverage_counters.take_expressions(),
             mappings,
@@ -136,7 +124,7 @@ impl<'a, 'tcx> Instrumentor<'a, 'tcx> {
         coverage_spans: &CoverageSpans,
     ) -> Vec<Mapping> {
         let source_map = self.tcx.sess.source_map();
-        let body_span = self.body_span;
+        let body_span = self.hir_info.body_span;
 
         let source_file = source_map.lookup_source_file(body_span.lo());
         use rustc_session::RemapFileNameExt;
@@ -311,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,
 }
@@ -324,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
@@ -345,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>(

compiler/rustc_mir_transform/src/coverage/spans.rs

@@ -6,6 +6,7 @@ use rustc_middle::mir;
 use rustc_span::{BytePos, ExpnKind, MacroKind, Span, Symbol, DUMMY_SP};
 
 use super::graph::{BasicCoverageBlock, CoverageGraph, START_BCB};
+use crate::coverage::ExtractedHirInfo;
 
 mod from_mir;
 
@@ -21,14 +22,12 @@ impl CoverageSpans {
     /// Returns `None` if no coverage-relevant spans could be extracted.
     pub(super) fn generate_coverage_spans(
         mir_body: &mir::Body<'_>,
-        fn_sig_span: Span,
-        body_span: Span,
+        hir_info: &ExtractedHirInfo,
         basic_coverage_blocks: &CoverageGraph,
     ) -> Option<Self> {
         let coverage_spans = CoverageSpansGenerator::generate_coverage_spans(
             mir_body,
-            fn_sig_span,
-            body_span,
+            hir_info,
             basic_coverage_blocks,
         );
 
@@ -230,19 +229,17 @@ impl<'a> CoverageSpansGenerator<'a> {
     /// to be).
     pub(super) fn generate_coverage_spans(
         mir_body: &mir::Body<'_>,
-        fn_sig_span: Span, // Ensured to be same SourceFile and SyntaxContext as `body_span`
-        body_span: Span,
+        hir_info: &ExtractedHirInfo,
         basic_coverage_blocks: &'a CoverageGraph,
     ) -> Vec<CoverageSpan> {
         let sorted_spans = from_mir::mir_to_initial_sorted_coverage_spans(
             mir_body,
-            fn_sig_span,
-            body_span,
+            hir_info,
             basic_coverage_blocks,
         );
 
         let coverage_spans = Self {
-            body_span,
+            body_span: hir_info.body_span,
             basic_coverage_blocks,
             sorted_spans_iter: sorted_spans.into_iter(),
             some_curr: None,

compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs

@@ -7,13 +7,22 @@ use rustc_span::Span;
 
 use crate::coverage::graph::{BasicCoverageBlock, BasicCoverageBlockData, CoverageGraph};
 use crate::coverage::spans::CoverageSpan;
+use crate::coverage::ExtractedHirInfo;
 
 pub(super) fn mir_to_initial_sorted_coverage_spans(
     mir_body: &mir::Body<'_>,
-    fn_sig_span: Span,
-    body_span: Span,
+    hir_info: &ExtractedHirInfo,
     basic_coverage_blocks: &CoverageGraph,
 ) -> Vec<CoverageSpan> {
+    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() {
         initial_spans.extend(bcb_to_initial_coverage_spans(mir_body, body_span, bcb, bcb_data));
@@ -44,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
 }
 

tests/coverage/async_block.cov-map

@@ -0,0 +1,32 @@
+Function name: async_block::main
+Raw bytes (38): 0x[01, 01, 02, 01, 05, 03, 05, 06, 01, 05, 01, 00, 0b, 05, 01, 09, 00, 0a, 03, 00, 0e, 00, 13, 05, 00, 14, 01, 16, 05, 07, 0a, 02, 06, 06, 03, 01, 00, 02]
+Number of files: 1
+- file 0 => global file 1
+Number of expressions: 2
+- expression 0 operands: lhs = Counter(0), rhs = Counter(1)
+- expression 1 operands: lhs = Expression(0, Add), rhs = Counter(1)
+Number of file 0 mappings: 6
+- Code(Counter(0)) at (prev + 5, 1) to (start + 0, 11)
+- Code(Counter(1)) at (prev + 1, 9) to (start + 0, 10)
+- Code(Expression(0, Add)) at (prev + 0, 14) to (start + 0, 19)
+    = (c0 + c1)
+- Code(Counter(1)) at (prev + 0, 20) to (start + 1, 22)
+- Code(Counter(1)) at (prev + 7, 10) to (start + 2, 6)
+- Code(Expression(1, Sub)) at (prev + 3, 1) to (start + 0, 2)
+    = ((c0 + c1) - c1)
+
+Function name: async_block::main::{closure#0}
+Raw bytes (28): 0x[01, 01, 02, 01, 05, 05, 02, 04, 01, 07, 1c, 01, 17, 05, 01, 18, 02, 0e, 02, 02, 14, 02, 0e, 07, 03, 09, 00, 0a]
+Number of files: 1
+- file 0 => global file 1
+Number of expressions: 2
+- expression 0 operands: lhs = Counter(0), rhs = Counter(1)
+- expression 1 operands: lhs = Counter(1), rhs = Expression(0, Sub)
+Number of file 0 mappings: 4
+- Code(Counter(0)) at (prev + 7, 28) to (start + 1, 23)
+- Code(Counter(1)) at (prev + 1, 24) to (start + 2, 14)
+- Code(Expression(0, Sub)) at (prev + 2, 20) to (start + 2, 14)
+    = (c0 - c1)
+- Code(Expression(1, Add)) at (prev + 3, 9) to (start + 0, 10)
+    = (c1 + (c0 - c1))
+

tests/coverage/async_block.coverage

@@ -0,0 +1,37 @@
+   LL|       |#![feature(coverage_attribute)]
+   LL|       |#![feature(noop_waker)]
+   LL|       |// edition: 2021
+   LL|       |
+   LL|      1|fn main() {
+   LL|     17|    for i in 0..16 {
+                      ^16
+   LL|     16|        let future = async {
+   LL|     16|            if i >= 12 {
+   LL|      4|                println!("big");
+   LL|     12|            } else {
+   LL|     12|                println!("small");
+   LL|     12|            }
+   LL|     16|        };
+   LL|     16|        executor::block_on(future);
+   LL|     16|    }
+   LL|      1|}
+   LL|       |
+   LL|       |mod executor {
+   LL|       |    use core::future::Future;
+   LL|       |    use core::pin::pin;
+   LL|       |    use core::task::{Context, Poll, Waker};
+   LL|       |
+   LL|       |    #[coverage(off)]
+   LL|       |    pub fn block_on<F: Future>(mut future: F) -> F::Output {
+   LL|       |        let mut future = pin!(future);
+   LL|       |        let waker = Waker::noop();
+   LL|       |        let mut context = Context::from_waker(&waker);
+   LL|       |
+   LL|       |        loop {
+   LL|       |            if let Poll::Ready(val) = future.as_mut().poll(&mut context) {
+   LL|       |                break val;
+   LL|       |            }
+   LL|       |        }
+   LL|       |    }
+   LL|       |}
+

tests/coverage/async_block.rs

@@ -0,0 +1,35 @@
+#![feature(coverage_attribute)]
+#![feature(noop_waker)]
+// edition: 2021
+
+fn main() {
+    for i in 0..16 {
+        let future = async {
+            if i >= 12 {
+                println!("big");
+            } else {
+                println!("small");
+            }
+        };
+        executor::block_on(future);
+    }
+}
+
+mod executor {
+    use core::future::Future;
+    use core::pin::pin;
+    use core::task::{Context, Poll, Waker};
+
+    #[coverage(off)]
+    pub fn block_on<F: Future>(mut future: F) -> F::Output {
+        let mut future = pin!(future);
+        let waker = Waker::noop();
+        let mut context = Context::from_waker(&waker);
+
+        loop {
+            if let Poll::Ready(val) = future.as_mut().poll(&mut context) {
+                break val;
+            }
+        }
+    }
+}