coverage: Simplify initial creation of coverage spans

Zalathar · Zalathar · commit 319693a927f3 · 2023-10-21T11:56:15.000+11:00
diff --git 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};
@@ -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![fn_sig_span],
-            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);
diff --git a/compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs b/compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs
@@ -1,6 +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;
 
@@ -59,24 +60,35 @@ fn bcb_to_initial_coverage_spans<'a, 'tcx>(
 ) -> impl Iterator<Item = CoverageSpan> + Captures<'a> + Captures<'tcx> {
     bcb_data.basic_blocks.iter().flat_map(move |&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)
-            }))
+
+        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,
 /// return it; otherwise return `None`.
 fn filtered_statement_span(statement: &Statement<'_>) -> Option<Span> {
