Rollup merge of #116974 - Zalathar:signature-spans, r=oli-obk,cjgillot

coverage: Fix inconsistent handling of function signature spans

While doing some more cleanup of `spans`, I noticed a strange inconsistency in how function signatures are handled. Normally the function signature span is treated as though it were executable as part of the start of the function, but in some cases the signature span disappears entirely from coverage, for no obvious reason.

This is caused by the fact that spans created by `CoverageSpan::for_fn_sig` don't add the span to their `merged_spans` field (unlike normal statement/terminator spans). In cases where the span-processing code looks at those merged spans, it thinks the signature span is no longer visible and deletes it.

Adding the signature span to `merged_spans` resolves the inconsistency.

(Prior to #116409 this wouldn't have been possible, because there was no case in the old `CoverageStatement` enum representing a signature. Now that `merged_spans` is just a list of spans, that's no longer an obstacle.)

9 files changed, 233 insertions, 67 deletions

diff --git a/compiler/rustc_mir_transform/src/coverage/spans.rs b/compiler/rustc_mir_transform/src/coverage/spans.rs
index dde16cbd782..704eea413e1 100644
--- a/compiler/rustc_mir_transform/src/coverage/spans.rs
+++ b/compiler/rustc_mir_transform/src/coverage/spans.rs
@@ -2,7 +2,7 @@ use std::cell::OnceCell;
 
 use rustc_data_structures::graph::WithNumNodes;
 use rustc_index::IndexVec;
-use rustc_middle::mir::{self, AggregateKind, Rvalue, Statement, StatementKind};
+use rustc_middle::mir;
 use rustc_span::{BytePos, ExpnKind, MacroKind, Span, Symbol, DUMMY_SP};
 
 use super::graph::{BasicCoverageBlock, CoverageGraph, START_BCB};
@@ -70,29 +70,15 @@ struct CoverageSpan {
 
 impl CoverageSpan {
     pub fn for_fn_sig(fn_sig_span: Span) -> Self {
-        Self {
-            span: fn_sig_span,
-            expn_span: fn_sig_span,
-            current_macro_or_none: Default::default(),
-            bcb: START_BCB,
-            merged_spans: vec![],
-            is_closure: false,
-        }
+        Self::new(fn_sig_span, fn_sig_span, START_BCB, false)
     }
 
-    pub fn for_statement(
-        statement: &Statement<'_>,
+    pub(super) fn new(
         span: Span,
         expn_span: Span,
         bcb: BasicCoverageBlock,
+        is_closure: bool,
     ) -> Self {
-        let is_closure = match statement.kind {
-            StatementKind::Assign(box (_, Rvalue::Aggregate(box ref kind, _))) => {
-                matches!(kind, AggregateKind::Closure(_, _) | AggregateKind::Coroutine(_, _, _))
-            }
-            _ => false,
-        };
-
         Self {
             span,
             expn_span,
@@ -103,17 +89,6 @@ impl CoverageSpan {
         }
     }
 
-    pub fn for_terminator(span: Span, expn_span: Span, bcb: BasicCoverageBlock) -> Self {
-        Self {
-            span,
-            expn_span,
-            current_macro_or_none: Default::default(),
-            bcb,
-            merged_spans: vec![span],
-            is_closure: false,
-        }
-    }
-
     pub fn merge_from(&mut self, mut other: CoverageSpan) {
         debug_assert!(self.is_mergeable(&other));
         self.span = self.span.to(other.span);
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 02e2cf6b05e..6189e5379ea 100644
--- a/compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs
+++ b/compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs
@@ -1,5 +1,7 @@
+use rustc_data_structures::captures::Captures;
 use rustc_middle::mir::{
-    self, FakeReadCause, Statement, StatementKind, Terminator, TerminatorKind,
+    self, AggregateKind, FakeReadCause, Rvalue, Statement, StatementKind, Terminator,
+    TerminatorKind,
 };
 use rustc_span::Span;
 
@@ -12,7 +14,7 @@ pub(super) fn mir_to_initial_sorted_coverage_spans(
     body_span: Span,
     basic_coverage_blocks: &CoverageGraph,
 ) -> Vec<CoverageSpan> {
-    let mut initial_spans = Vec::<CoverageSpan>::with_capacity(mir_body.basic_blocks.len() * 2);
+    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));
     }
@@ -50,34 +52,41 @@ pub(super) fn mir_to_initial_sorted_coverage_spans(
 // for each `Statement` and `Terminator`. (Note that subsequent stages of coverage analysis will
 // merge some `CoverageSpan`s, at which point a `CoverageSpan` may represent multiple
 // `Statement`s and/or `Terminator`s.)
-fn bcb_to_initial_coverage_spans(
-    mir_body: &mir::Body<'_>,
+fn bcb_to_initial_coverage_spans<'a, 'tcx>(
+    mir_body: &'a mir::Body<'tcx>,
     body_span: Span,
     bcb: BasicCoverageBlock,
-    bcb_data: &BasicCoverageBlockData,
-) -> Vec<CoverageSpan> {
-    bcb_data
-        .basic_blocks
-        .iter()
-        .flat_map(|&bb| {
-            let data = &mir_body[bb];
-            data.statements
-                .iter()
-                .filter_map(move |statement| {
-                    filtered_statement_span(statement).map(|span| {
-                        CoverageSpan::for_statement(
-                            statement,
-                            function_source_span(span, body_span),
-                            span,
-                            bcb,
-                        )
-                    })
-                })
-                .chain(filtered_terminator_span(data.terminator()).map(|span| {
-                    CoverageSpan::for_terminator(function_source_span(span, body_span), span, bcb)
-                }))
-        })
-        .collect()
+    bcb_data: &'a BasicCoverageBlockData,
+) -> impl Iterator<Item = CoverageSpan> + Captures<'a> + Captures<'tcx> {
+    bcb_data.basic_blocks.iter().flat_map(move |&bb| {
+        let data = &mir_body[bb];
+
+        let statement_spans = data.statements.iter().filter_map(move |statement| {
+            let expn_span = filtered_statement_span(statement)?;
+            let span = function_source_span(expn_span, body_span);
+
+            Some(CoverageSpan::new(span, expn_span, bcb, is_closure(statement)))
+        });
+
+        let terminator_span = Some(data.terminator()).into_iter().filter_map(move |terminator| {
+            let expn_span = filtered_terminator_span(terminator)?;
+            let span = function_source_span(expn_span, body_span);
+
+            Some(CoverageSpan::new(span, expn_span, bcb, false))
+        });
+
+        statement_spans.chain(terminator_span)
+    })
+}
+
+fn is_closure(statement: &Statement<'_>) -> bool {
+    match statement.kind {
+        StatementKind::Assign(box (_, Rvalue::Aggregate(box ref agg_kind, _))) => match agg_kind {
+            AggregateKind::Closure(_, _) | AggregateKind::Coroutine(_, _, _) => true,
+            _ => false,
+        },
+        _ => false,
+    }
 }
 
 /// If the MIR `Statement` has a span contributive to computing coverage spans,
diff --git a/tests/coverage-map/fn_sig_into_try.cov-map b/tests/coverage-map/fn_sig_into_try.cov-map
new file mode 100644
index 00000000000..4672e7c1ce9
--- /dev/null
+++ b/tests/coverage-map/fn_sig_into_try.cov-map
@@ -0,0 +1,53 @@
+Function name: fn_sig_into_try::a
+Raw bytes (9): 0x[01, 01, 00, 01, 01, 0a, 01, 04, 02]
+Number of files: 1
+- file 0 => global file 1
+Number of expressions: 0
+Number of file 0 mappings: 1
+- Code(Counter(0)) at (prev + 10, 1) to (start + 4, 2)
+
+Function name: fn_sig_into_try::b
+Raw bytes (28): 0x[01, 01, 02, 01, 05, 05, 02, 04, 01, 10, 01, 02, 0f, 00, 02, 0f, 00, 10, 02, 01, 05, 00, 0c, 07, 01, 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 = Counter(1), rhs = Expression(0, Sub)
+Number of file 0 mappings: 4
+- Code(Counter(0)) at (prev + 16, 1) to (start + 2, 15)
+- Code(Zero) at (prev + 2, 15) to (start + 0, 16)
+- Code(Expression(0, Sub)) at (prev + 1, 5) to (start + 0, 12)
+    = (c0 - c1)
+- Code(Expression(1, Add)) at (prev + 1, 1) to (start + 0, 2)
+    = (c1 + (c0 - c1))
+
+Function name: fn_sig_into_try::c
+Raw bytes (28): 0x[01, 01, 02, 01, 05, 05, 02, 04, 01, 16, 01, 02, 17, 00, 02, 17, 00, 18, 02, 01, 05, 00, 0c, 07, 01, 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 = Counter(1), rhs = Expression(0, Sub)
+Number of file 0 mappings: 4
+- Code(Counter(0)) at (prev + 22, 1) to (start + 2, 23)
+- Code(Zero) at (prev + 2, 23) to (start + 0, 24)
+- Code(Expression(0, Sub)) at (prev + 1, 5) to (start + 0, 12)
+    = (c0 - c1)
+- Code(Expression(1, Add)) at (prev + 1, 1) to (start + 0, 2)
+    = (c1 + (c0 - c1))
+
+Function name: fn_sig_into_try::d
+Raw bytes (28): 0x[01, 01, 02, 01, 05, 05, 02, 04, 01, 1c, 01, 03, 0f, 00, 03, 0f, 00, 10, 02, 01, 05, 00, 0c, 07, 01, 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 = Counter(1), rhs = Expression(0, Sub)
+Number of file 0 mappings: 4
+- Code(Counter(0)) at (prev + 28, 1) to (start + 3, 15)
+- Code(Zero) at (prev + 3, 15) to (start + 0, 16)
+- Code(Expression(0, Sub)) at (prev + 1, 5) to (start + 0, 12)
+    = (c0 - c1)
+- Code(Expression(1, Add)) at (prev + 1, 1) to (start + 0, 2)
+    = (c1 + (c0 - c1))
+
diff --git a/tests/coverage-map/fn_sig_into_try.rs b/tests/coverage-map/fn_sig_into_try.rs
new file mode 100644
index 00000000000..92850c8a188
--- /dev/null
+++ b/tests/coverage-map/fn_sig_into_try.rs
@@ -0,0 +1,41 @@
+#![feature(coverage_attribute)]
+// compile-flags: --edition=2021
+
+// Regression test for inconsistent handling of function signature spans that
+// are followed by code using the `?` operator.
+//
+// For each of these similar functions, the line containing the function
+// signature should be handled in the same way.
+
+fn a() -> Option<i32>
+{
+    Some(7i32);
+    Some(0)
+}
+
+fn b() -> Option<i32>
+{
+    Some(7i32)?;
+    Some(0)
+}
+
+fn c() -> Option<i32>
+{
+    let _ = Some(7i32)?;
+    Some(0)
+}
+
+fn d() -> Option<i32>
+{
+    let _: () = ();
+    Some(7i32)?;
+    Some(0)
+}
+
+#[coverage(off)]
+fn main() {
+    a();
+    b();
+    c();
+    d();
+}
diff --git a/tests/coverage-map/status-quo/inline-dead.cov-map b/tests/coverage-map/status-quo/inline-dead.cov-map
index 483f7ef79c6..06b64da5723 100644
--- a/tests/coverage-map/status-quo/inline-dead.cov-map
+++ b/tests/coverage-map/status-quo/inline-dead.cov-map
@@ -31,13 +31,15 @@ Number of file 0 mappings: 2
 - Code(Counter(0)) at (prev + 7, 6) to (start + 2, 2)
 
 Function name: inline_dead::main::{closure#0}
-Raw bytes (16): 0x[01, 01, 01, 01, 05, 02, 00, 09, 0d, 00, 0e, 03, 02, 05, 00, 06]
+Raw bytes (23): 0x[01, 01, 02, 09, 06, 01, 05, 03, 01, 07, 17, 00, 18, 00, 02, 0d, 00, 0e, 03, 02, 05, 00, 06]
 Number of files: 1
 - file 0 => global file 1
-Number of expressions: 1
-- expression 0 operands: lhs = Counter(0), rhs = Counter(1)
-Number of file 0 mappings: 2
-- Code(Zero) at (prev + 9, 13) to (start + 0, 14)
+Number of expressions: 2
+- expression 0 operands: lhs = Counter(2), rhs = Expression(1, Sub)
+- expression 1 operands: lhs = Counter(0), rhs = Counter(1)
+Number of file 0 mappings: 3
+- Code(Counter(0)) at (prev + 7, 23) to (start + 0, 24)
+- Code(Zero) at (prev + 2, 13) to (start + 0, 14)
 - Code(Expression(0, Add)) at (prev + 2, 5) to (start + 0, 6)
-    = (c0 + c1)
+    = (c2 + (c0 - c1))
 
diff --git a/tests/coverage-map/status-quo/issue-84561.cov-map b/tests/coverage-map/status-quo/issue-84561.cov-map
index 01fa7ec573c..76340b1a78c 100644
--- a/tests/coverage-map/status-quo/issue-84561.cov-map
+++ b/tests/coverage-map/status-quo/issue-84561.cov-map
@@ -7,15 +7,15 @@ Number of file 0 mappings: 1
 - Code(Counter(0)) at (prev + 4, 10) to (start + 0, 19)
 
 Function name: <issue_84561::Foo as core::fmt::Debug>::fmt
-Raw bytes (29): 0x[01, 01, 02, 01, 05, 05, 02, 04, 01, 89, 01, 09, 00, 25, 05, 00, 25, 00, 26, 02, 01, 09, 00, 0f, 07, 01, 05, 00, 06]
+Raw bytes (29): 0x[01, 01, 02, 01, 05, 05, 02, 04, 01, 88, 01, 05, 01, 25, 05, 01, 25, 00, 26, 02, 01, 09, 00, 0f, 07, 01, 05, 00, 06]
 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 + 137, 9) to (start + 0, 37)
-- Code(Counter(1)) at (prev + 0, 37) to (start + 0, 38)
+- Code(Counter(0)) at (prev + 136, 5) to (start + 1, 37)
+- Code(Counter(1)) at (prev + 1, 37) to (start + 0, 38)
 - Code(Expression(0, Sub)) at (prev + 1, 9) to (start + 0, 15)
     = (c0 - c1)
 - Code(Expression(1, Add)) at (prev + 1, 5) to (start + 0, 6)
diff --git a/tests/run-coverage/fn_sig_into_try.coverage b/tests/run-coverage/fn_sig_into_try.coverage
new file mode 100644
index 00000000000..f1ddb1da780
--- /dev/null
+++ b/tests/run-coverage/fn_sig_into_try.coverage
@@ -0,0 +1,45 @@
+   LL|       |#![feature(coverage_attribute)]
+   LL|       |// compile-flags: --edition=2021
+   LL|       |
+   LL|       |// Regression test for inconsistent handling of function signature spans that
+   LL|       |// are followed by code using the `?` operator.
+   LL|       |//
+   LL|       |// For each of these similar functions, the line containing the function
+   LL|       |// signature should be handled in the same way.
+   LL|       |
+   LL|      1|fn a() -> Option<i32>
+   LL|      1|{
+   LL|      1|    Some(7i32);
+   LL|      1|    Some(0)
+   LL|      1|}
+   LL|       |
+   LL|      1|fn b() -> Option<i32>
+   LL|      1|{
+   LL|      1|    Some(7i32)?;
+                            ^0
+   LL|      1|    Some(0)
+   LL|      1|}
+   LL|       |
+   LL|      1|fn c() -> Option<i32>
+   LL|      1|{
+   LL|      1|    let _ = Some(7i32)?;
+                                    ^0
+   LL|      1|    Some(0)
+   LL|      1|}
+   LL|       |
+   LL|      1|fn d() -> Option<i32>
+   LL|      1|{
+   LL|      1|    let _: () = ();
+   LL|      1|    Some(7i32)?;
+                            ^0
+   LL|      1|    Some(0)
+   LL|      1|}
+   LL|       |
+   LL|       |#[coverage(off)]
+   LL|       |fn main() {
+   LL|       |    a();
+   LL|       |    b();
+   LL|       |    c();
+   LL|       |    d();
+   LL|       |}
+
diff --git a/tests/run-coverage/fn_sig_into_try.rs b/tests/run-coverage/fn_sig_into_try.rs
new file mode 100644
index 00000000000..92850c8a188
--- /dev/null
+++ b/tests/run-coverage/fn_sig_into_try.rs
@@ -0,0 +1,41 @@
+#![feature(coverage_attribute)]
+// compile-flags: --edition=2021
+
+// Regression test for inconsistent handling of function signature spans that
+// are followed by code using the `?` operator.
+//
+// For each of these similar functions, the line containing the function
+// signature should be handled in the same way.
+
+fn a() -> Option<i32>
+{
+    Some(7i32);
+    Some(0)
+}
+
+fn b() -> Option<i32>
+{
+    Some(7i32)?;
+    Some(0)
+}
+
+fn c() -> Option<i32>
+{
+    let _ = Some(7i32)?;
+    Some(0)
+}
+
+fn d() -> Option<i32>
+{
+    let _: () = ();
+    Some(7i32)?;
+    Some(0)
+}
+
+#[coverage(off)]
+fn main() {
+    a();
+    b();
+    c();
+    d();
+}
diff --git a/tests/run-coverage/issue-84561.coverage b/tests/run-coverage/issue-84561.coverage
index 222f877d36a..e693866e277 100644
--- a/tests/run-coverage/issue-84561.coverage
+++ b/tests/run-coverage/issue-84561.coverage
@@ -135,7 +135,7 @@
    LL|      0|}
    LL|       |
    LL|       |impl std::fmt::Debug for Foo {
-   LL|       |    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+   LL|      7|    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
    LL|      7|        write!(f, "try and succeed")?;
                                                   ^0
    LL|      7|        Ok(())