coverage: Fix inconsistent handling of function signature spans

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};
@@ -75,29 +75,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,
@@ -108,17 +94,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);

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,

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))
+

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();
+}

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))
 

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)

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|       |}
+

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();
+}

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(())