coverage: Split out SpanFromMir from CoverageSpan

This draws a clear distinction between the fields/methods that are needed by
initial span extraction and preprocessing, and those that are needed by the
main "refinement" loop.

Author: Zalathar

compiler/rustc_mir_transform/src/coverage/spans.rs

@@ -1,11 +1,9 @@
-use std::cell::OnceCell;
-
 use rustc_data_structures::graph::WithNumNodes;
 use rustc_index::IndexVec;
 use rustc_middle::mir;
-use rustc_span::{BytePos, ExpnKind, MacroKind, Span, Symbol, DUMMY_SP};
+use rustc_span::{BytePos, Span, DUMMY_SP};
 
-use super::graph::{BasicCoverageBlock, CoverageGraph, START_BCB};
+use super::graph::{BasicCoverageBlock, CoverageGraph};
 use crate::coverage::ExtractedHirInfo;
 
 mod from_mir;
@@ -66,8 +64,6 @@ impl CoverageSpans {
 #[derive(Debug, Clone)]
 struct CoverageSpan {
     span: Span,
-    expn_span: Span,
-    current_macro_or_none: OnceCell<Option<Symbol>>,
     bcb: BasicCoverageBlock,
     /// List of all the original spans from MIR that have been merged into this
     /// span. Mainly used to precisely skip over gaps when truncating a span.
@@ -76,24 +72,8 @@ struct CoverageSpan {
 }
 
 impl CoverageSpan {
-    pub fn for_fn_sig(fn_sig_span: Span) -> Self {
-        Self::new(fn_sig_span, fn_sig_span, START_BCB, false)
-    }
-
-    pub(super) fn new(
-        span: Span,
-        expn_span: Span,
-        bcb: BasicCoverageBlock,
-        is_closure: bool,
-    ) -> Self {
-        Self {
-            span,
-            expn_span,
-            current_macro_or_none: Default::default(),
-            bcb,
-            merged_spans: vec![span],
-            is_closure,
-        }
+    fn new(span: Span, bcb: BasicCoverageBlock, is_closure: bool) -> Self {
+        Self { span, bcb, merged_spans: vec![span], is_closure }
     }
 
     pub fn merge_from(&mut self, other: &Self) {
@@ -118,37 +98,6 @@ impl CoverageSpan {
     pub fn is_in_same_bcb(&self, other: &Self) -> bool {
         self.bcb == other.bcb
     }
-
-    /// If the span is part of a macro, returns the macro name symbol.
-    pub fn current_macro(&self) -> Option<Symbol> {
-        self.current_macro_or_none
-            .get_or_init(|| {
-                if let ExpnKind::Macro(MacroKind::Bang, current_macro) =
-                    self.expn_span.ctxt().outer_expn_data().kind
-                {
-                    return Some(current_macro);
-                }
-                None
-            })
-            .map(|symbol| symbol)
-    }
-
-    /// If the span is part of a macro, and the macro is visible (expands directly to the given
-    /// body_span), returns the macro name symbol.
-    pub fn visible_macro(&self, body_span: Span) -> Option<Symbol> {
-        let current_macro = self.current_macro()?;
-        let parent_callsite = self.expn_span.parent_callsite()?;
-
-        // In addition to matching the context of the body span, the parent callsite
-        // must also be the source callsite, i.e. the parent must have no parent.
-        let is_visible_macro =
-            parent_callsite.parent_callsite().is_none() && parent_callsite.eq_ctxt(body_span);
-        is_visible_macro.then_some(current_macro)
-    }
-
-    pub fn is_macro_expansion(&self) -> bool {
-        self.current_macro().is_some()
-    }
 }
 
 /// Converts the initial set of `CoverageSpan`s (one per MIR `Statement` or `Terminator`) into a

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

@@ -4,9 +4,11 @@ use rustc_middle::mir::{
     self, AggregateKind, FakeReadCause, Rvalue, Statement, StatementKind, Terminator,
     TerminatorKind,
 };
-use rustc_span::Span;
+use rustc_span::{ExpnKind, MacroKind, Span, Symbol};
 
-use crate::coverage::graph::{BasicCoverageBlock, BasicCoverageBlockData, CoverageGraph};
+use crate::coverage::graph::{
+    BasicCoverageBlock, BasicCoverageBlockData, CoverageGraph, START_BCB,
+};
 use crate::coverage::spans::CoverageSpan;
 use crate::coverage::ExtractedHirInfo;
 
@@ -17,7 +19,7 @@ pub(super) fn mir_to_initial_sorted_coverage_spans(
 ) -> Vec<CoverageSpan> {
     let &ExtractedHirInfo { is_async_fn, fn_sig_span, body_span, .. } = hir_info;
 
-    let mut initial_spans = vec![CoverageSpan::for_fn_sig(fn_sig_span)];
+    let mut initial_spans = vec![SpanFromMir::for_fn_sig(fn_sig_span)];
 
     if is_async_fn {
         // An async function desugars into a function that returns a future,
@@ -57,7 +59,7 @@ pub(super) fn mir_to_initial_sorted_coverage_spans(
             .then_with(|| Ord::cmp(&a.is_closure, &b.is_closure).reverse())
     });
 
-    initial_spans
+    initial_spans.into_iter().map(SpanFromMir::into_coverage_span).collect::<Vec<_>>()
 }
 
 /// Macros that expand into branches (e.g. `assert!`, `trace!`) tend to generate
@@ -67,7 +69,7 @@ pub(super) fn mir_to_initial_sorted_coverage_spans(
 ///
 /// (The input spans should be sorted in BCB dominator order, so that the
 /// retained "first" span is likely to dominate the others.)
-fn remove_unwanted_macro_spans(initial_spans: &mut Vec<CoverageSpan>) {
+fn remove_unwanted_macro_spans(initial_spans: &mut Vec<SpanFromMir>) {
     let mut seen_spans = FxHashSet::default();
     initial_spans.retain(|covspan| {
         // Ignore (retain) closure spans and non-macro-expansion spans.
@@ -84,7 +86,7 @@ fn remove_unwanted_macro_spans(initial_spans: &mut Vec<CoverageSpan>) {
 /// function body, split it into two parts. The first part covers just the
 /// macro name plus `!`, and the second part covers the rest of the macro
 /// invocation. This seems to give better results for code that uses macros.
-fn split_visible_macro_spans(initial_spans: &mut Vec<CoverageSpan>, hir_info: &ExtractedHirInfo) {
+fn split_visible_macro_spans(initial_spans: &mut Vec<SpanFromMir>, hir_info: &ExtractedHirInfo) {
     let mut extra_spans = vec![];
 
     initial_spans.retain(|covspan| {
@@ -105,8 +107,8 @@ fn split_visible_macro_spans(initial_spans: &mut Vec<CoverageSpan>, hir_info: &E
         }
 
         assert!(!covspan.is_closure);
-        extra_spans.push(CoverageSpan::new(before, covspan.expn_span, covspan.bcb, false));
-        extra_spans.push(CoverageSpan::new(after, covspan.expn_span, covspan.bcb, false));
+        extra_spans.push(SpanFromMir::new(before, covspan.expn_span, covspan.bcb, false));
+        extra_spans.push(SpanFromMir::new(after, covspan.expn_span, covspan.bcb, false));
         false // Discard the original covspan that we just split.
     });
 
@@ -125,22 +127,22 @@ fn bcb_to_initial_coverage_spans<'a, 'tcx>(
     body_span: Span,
     bcb: BasicCoverageBlock,
     bcb_data: &'a BasicCoverageBlockData,
-) -> impl Iterator<Item = CoverageSpan> + Captures<'a> + Captures<'tcx> {
+) -> impl Iterator<Item = SpanFromMir> + 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 = unexpand_into_body_span(expn_span, body_span)?;
 
-            Some(CoverageSpan::new(span, expn_span, bcb, is_closure(statement)))
+            Some(SpanFromMir::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 = unexpand_into_body_span(expn_span, body_span)?;
 
-            Some(CoverageSpan::new(span, expn_span, bcb, false))
+            Some(SpanFromMir::new(span, expn_span, bcb, false))
         });
 
         statement_spans.chain(terminator_span)
@@ -270,3 +272,66 @@ fn unexpand_into_body_span(span: Span, body_span: Span) -> Option<Span> {
     let original_span = original_sp(span, body_span).with_ctxt(body_span.ctxt());
     body_span.contains(original_span).then_some(original_span)
 }
+
+#[derive(Debug)]
+struct SpanFromMir {
+    /// A copy of `expn_span` that has been "un-expanded" back to the current
+    /// function's `body_span`. After various intermediate processing steps,
+    /// this span is emitted as part of the final coverage mappings.
+    ///
+    /// With the exception of `fn_sig_span`, this should always be contained
+    /// within `body_span`.
+    span: Span,
+    /// The actual span that was extracted from MIR, used to look up information
+    /// about macro expansions.
+    expn_span: Span,
+    current_macro_or_none: std::cell::OnceCell<Option<Symbol>>,
+    bcb: BasicCoverageBlock,
+    is_closure: bool,
+}
+
+impl SpanFromMir {
+    fn for_fn_sig(fn_sig_span: Span) -> Self {
+        Self::new(fn_sig_span, fn_sig_span, START_BCB, false)
+    }
+
+    fn new(span: Span, expn_span: Span, bcb: BasicCoverageBlock, is_closure: bool) -> Self {
+        Self { span, expn_span, current_macro_or_none: Default::default(), bcb, is_closure }
+    }
+
+    /// If the span is part of a macro, returns the macro name symbol.
+    fn current_macro(&self) -> Option<Symbol> {
+        self.current_macro_or_none
+            .get_or_init(|| {
+                if let ExpnKind::Macro(MacroKind::Bang, current_macro) =
+                    self.expn_span.ctxt().outer_expn_data().kind
+                {
+                    return Some(current_macro);
+                }
+                None
+            })
+            .map(|symbol| symbol)
+    }
+
+    /// If the span is part of a macro, and the macro is visible (expands directly to the given
+    /// body_span), returns the macro name symbol.
+    fn visible_macro(&self, body_span: Span) -> Option<Symbol> {
+        let current_macro = self.current_macro()?;
+        let parent_callsite = self.expn_span.parent_callsite()?;
+
+        // In addition to matching the context of the body span, the parent callsite
+        // must also be the source callsite, i.e. the parent must have no parent.
+        let is_visible_macro =
+            parent_callsite.parent_callsite().is_none() && parent_callsite.eq_ctxt(body_span);
+        is_visible_macro.then_some(current_macro)
+    }
+
+    fn is_macro_expansion(&self) -> bool {
+        self.current_macro().is_some()
+    }
+
+    fn into_coverage_span(self) -> CoverageSpan {
+        let Self { span, expn_span: _, current_macro_or_none: _, bcb, is_closure } = self;
+        CoverageSpan::new(span, bcb, is_closure)
+    }
+}