Rollup merge of rust-lang#118666 - Zalathar:body-closure, r=cjgillot

coverage: Simplify the heuristic for ignoring `async fn` return spans

The code for extracting coverage spans from MIR has a special heuristic for dealing with `async fn`, so that the function's closing brace does not have a confusing double count.

The code implementing that heuristic is currently mixed in with the code for flushing remaining spans after the main refinement loop, making the refinement code harder to understand.

We can solve that by hoisting the heuristic to an earlier stage, after the spans have been extracted and sorted but before they have been processed by the refinement loop.

The coverage tests verify that the heuristic is still effective, so coverage mappings/reports for `async fn` have not changed.

---

This PR also has the side-effect of fixing the `None some_prev` panic that started appearing after rust-lang#118525.

The old code assumed that `prev` would always be present after the refinement loop. That was only true if the list of collected spans was non-empty, but prior to rust-lang#118525 that didn't seem to come up in practice. After that change, the list of collected spans could be empty in some specific circumstances, leading to panics.

The new code uses an `if let` to inspect `prev`, which correctly does nothing if there is no span present.

compiler/rustc_mir_transform/src/coverage/spans.rs

@@ -319,29 +319,16 @@ impl<'a> CoverageSpansGenerator<'a> {
             }
         }
 
-        let prev = self.take_prev();
-        debug!("    AT END, adding last prev={prev:?}");
-
         // Drain any remaining dups into the output.
         for dup in self.pending_dups.drain(..) {
             debug!("    ...adding at least one pending dup={:?}", dup);
             self.refined_spans.push(dup);
         }
 
-        // Async functions wrap a closure that implements the body to be executed. The enclosing
-        // function is called and returns an `impl Future` without initially executing any of the
-        // body. To avoid showing the return from the enclosing function as a "covered" return from
-        // the closure, the enclosing function's `TerminatorKind::Return`s `CoverageSpan` is
-        // excluded. The closure's `Return` is the only one that will be counted. This provides
-        // adequate coverage, and more intuitive counts. (Avoids double-counting the closing brace
-        // of the function body.)
-        let body_ends_with_closure = if let Some(last_covspan) = self.refined_spans.last() {
-            last_covspan.is_closure && last_covspan.span.hi() == self.body_span.hi()
-        } else {
-            false
-        };
-
-        if !body_ends_with_closure {
+        // There is usually a final span remaining in `prev` after the loop ends,
+        // so add it to the output as well.
+        if let Some(prev) = self.some_prev.take() {
+            debug!("    AT END, adding last prev={prev:?}");
             self.refined_spans.push(prev);
         }
 
@@ -398,38 +385,36 @@ impl<'a> CoverageSpansGenerator<'a> {
         self.refined_spans.push(macro_name_cov);
     }
 
+    #[track_caller]
     fn curr(&self) -> &CoverageSpan {
-        self.some_curr
-            .as_ref()
-            .unwrap_or_else(|| bug!("invalid attempt to unwrap a None some_curr"))
+        self.some_curr.as_ref().unwrap_or_else(|| bug!("some_curr is None (curr)"))
     }
 
+    #[track_caller]
     fn curr_mut(&mut self) -> &mut CoverageSpan {
-        self.some_curr
-            .as_mut()
-            .unwrap_or_else(|| bug!("invalid attempt to unwrap a None some_curr"))
+        self.some_curr.as_mut().unwrap_or_else(|| bug!("some_curr is None (curr_mut)"))
     }
 
     /// If called, then the next call to `next_coverage_span()` will *not* update `prev` with the
     /// `curr` coverage span.
+    #[track_caller]
     fn take_curr(&mut self) -> CoverageSpan {
-        self.some_curr.take().unwrap_or_else(|| bug!("invalid attempt to unwrap a None some_curr"))
+        self.some_curr.take().unwrap_or_else(|| bug!("some_curr is None (take_curr)"))
     }
 
+    #[track_caller]
     fn prev(&self) -> &CoverageSpan {
-        self.some_prev
-            .as_ref()
-            .unwrap_or_else(|| bug!("invalid attempt to unwrap a None some_prev"))
+        self.some_prev.as_ref().unwrap_or_else(|| bug!("some_prev is None (prev)"))
     }
 
+    #[track_caller]
     fn prev_mut(&mut self) -> &mut CoverageSpan {
-        self.some_prev
-            .as_mut()
-            .unwrap_or_else(|| bug!("invalid attempt to unwrap a None some_prev"))
+        self.some_prev.as_mut().unwrap_or_else(|| bug!("some_prev is None (prev_mut)"))
     }
 
+    #[track_caller]
     fn take_prev(&mut self) -> CoverageSpan {
-        self.some_prev.take().unwrap_or_else(|| bug!("invalid attempt to unwrap a None some_prev"))
+        self.some_prev.take().unwrap_or_else(|| bug!("some_prev is None (take_prev)"))
     }
 
     /// If there are `pending_dups` but `prev` is not a matching dup (`prev.span` doesn't match the

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

@@ -44,6 +44,16 @@ 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/no_spans.cov-map

@@ -0,0 +1,8 @@
+Function name: no_spans::affected_function::{closure#0}
+Raw bytes (9): 0x[01, 01, 00, 01, 01, 1b, 0c, 00, 0e]
+Number of files: 1
+- file 0 => global file 1
+Number of expressions: 0
+Number of file 0 mappings: 1
+- Code(Counter(0)) at (prev + 27, 12) to (start + 0, 14)
+

tests/coverage/no_spans.coverage

@@ -0,0 +1,30 @@
+   LL|       |#![feature(coverage_attribute)]
+   LL|       |// edition: 2021
+   LL|       |
+   LL|       |// If the span extractor can't find any relevant spans for a function, the
+   LL|       |// refinement loop will terminate with nothing in its `prev` slot. If the
+   LL|       |// subsequent code tries to unwrap `prev`, it will panic.
+   LL|       |//
+   LL|       |// This scenario became more likely after #118525 started discarding spans that
+   LL|       |// can't be un-expanded back to within the function body.
+   LL|       |//
+   LL|       |// Regression test for "invalid attempt to unwrap a None some_prev", as seen
+   LL|       |// in issues such as #118643 and #118662.
+   LL|       |
+   LL|       |#[coverage(off)]
+   LL|       |fn main() {
+   LL|       |    affected_function()();
+   LL|       |}
+   LL|       |
+   LL|       |macro_rules! macro_that_defines_a_function {
+   LL|       |    (fn $name:ident () $body:tt) => {
+   LL|       |        fn $name () -> impl Fn() $body
+   LL|       |    }
+   LL|       |}
+   LL|       |
+   LL|       |macro_that_defines_a_function! {
+   LL|       |    fn affected_function() {
+   LL|      1|        || ()
+   LL|       |    }
+   LL|       |}
+

tests/coverage/no_spans.rs

@@ -0,0 +1,29 @@
+#![feature(coverage_attribute)]
+// edition: 2021
+
+// If the span extractor can't find any relevant spans for a function, the
+// refinement loop will terminate with nothing in its `prev` slot. If the
+// subsequent code tries to unwrap `prev`, it will panic.
+//
+// This scenario became more likely after #118525 started discarding spans that
+// can't be un-expanded back to within the function body.
+//
+// Regression test for "invalid attempt to unwrap a None some_prev", as seen
+// in issues such as #118643 and #118662.
+
+#[coverage(off)]
+fn main() {
+    affected_function()();
+}
+
+macro_rules! macro_that_defines_a_function {
+    (fn $name:ident () $body:tt) => {
+        fn $name () -> impl Fn() $body
+    }
+}
+
+macro_that_defines_a_function! {
+    fn affected_function() {
+        || ()
+    }
+}