Rollup merge of #138776 - Zalathar:unexpand, r=oli-obk

coverage: Separate span-extraction from unexpansion

Historically, coverage instrumentation has relied on eagerly “unexpanding” MIR spans back to ancestor spans that have the same context as the function body, and lie within that body. Doing so makes several subsequent operations more straightforward.

In order to support expansion regions, we need to stop doing that, and handle layers of macro-expansion more explicitly. This PR takes a step in that direction, by deferring some of the unexpansion steps, and concentrating them in one place (`spans::extract_refined_covspans`).

Unexpansion still takes place as before, but these changes will make it easier to experiment with expansion-aware coverage instrumentation.

3 files changed, 81 insertions, 90 deletions

diff --git a/compiler/rustc_mir_transform/src/coverage/mod.rs b/compiler/rustc_mir_transform/src/coverage/mod.rs
index 572dbae8fd2..aa4c0ef1e1f 100644
--- a/compiler/rustc_mir_transform/src/coverage/mod.rs
+++ b/compiler/rustc_mir_transform/src/coverage/mod.rs
@@ -273,8 +273,9 @@ struct ExtractedHirInfo {
     /// Must have the same context and filename as the body span.
     fn_sig_span_extended: Option<Span>,
     body_span: Span,
-    /// "Holes" are regions within the body span that should not be included in
-    /// coverage spans for this function (e.g. closures and nested items).
+    /// "Holes" are regions within the function body (or its expansions) that
+    /// should not be included in coverage spans for this function
+    /// (e.g. closures and nested items).
     hole_spans: Vec<Span>,
 }
 
@@ -323,7 +324,7 @@ fn extract_hir_info<'tcx>(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> ExtractedHir
 
     let function_source_hash = hash_mir_source(tcx, hir_body);
 
-    let hole_spans = extract_hole_spans_from_hir(tcx, body_span, hir_body);
+    let hole_spans = extract_hole_spans_from_hir(tcx, hir_body);
 
     ExtractedHirInfo {
         function_source_hash,
@@ -340,14 +341,9 @@ fn hash_mir_source<'tcx>(tcx: TyCtxt<'tcx>, hir_body: &'tcx hir::Body<'tcx>) ->
     tcx.hir_owner_nodes(owner).opt_hash_including_bodies.unwrap().to_smaller_hash().as_u64()
 }
 
-fn extract_hole_spans_from_hir<'tcx>(
-    tcx: TyCtxt<'tcx>,
-    body_span: Span, // Usually `hir_body.value.span`, but not always
-    hir_body: &hir::Body<'tcx>,
-) -> Vec<Span> {
+fn extract_hole_spans_from_hir<'tcx>(tcx: TyCtxt<'tcx>, hir_body: &hir::Body<'tcx>) -> Vec<Span> {
     struct HolesVisitor<'tcx> {
         tcx: TyCtxt<'tcx>,
-        body_span: Span,
         hole_spans: Vec<Span>,
     }
 
@@ -387,14 +383,11 @@ fn extract_hole_spans_from_hir<'tcx>(
     }
     impl HolesVisitor<'_> {
         fn visit_hole_span(&mut self, hole_span: Span) {
-            // Discard any holes that aren't directly visible within the body span.
-            if self.body_span.contains(hole_span) && self.body_span.eq_ctxt(hole_span) {
-                self.hole_spans.push(hole_span);
-            }
+            self.hole_spans.push(hole_span);
         }
     }
 
-    let mut visitor = HolesVisitor { tcx, body_span, hole_spans: vec![] };
+    let mut visitor = HolesVisitor { tcx, hole_spans: vec![] };
 
     visitor.visit_body(hir_body);
     visitor.hole_spans
diff --git a/compiler/rustc_mir_transform/src/coverage/spans.rs b/compiler/rustc_mir_transform/src/coverage/spans.rs
index b9ed6984ddb..8befe9c5d8d 100644
--- a/compiler/rustc_mir_transform/src/coverage/spans.rs
+++ b/compiler/rustc_mir_transform/src/coverage/spans.rs
@@ -6,10 +6,8 @@ use rustc_span::{DesugaringKind, ExpnKind, MacroKind, Span};
 use tracing::{debug, debug_span, instrument};
 
 use crate::coverage::graph::{BasicCoverageBlock, CoverageGraph};
-use crate::coverage::spans::from_mir::{
-    ExtractedCovspans, Hole, SpanFromMir, extract_covspans_from_mir,
-};
-use crate::coverage::{ExtractedHirInfo, mappings};
+use crate::coverage::spans::from_mir::{Hole, RawSpanFromMir, SpanFromMir};
+use crate::coverage::{ExtractedHirInfo, mappings, unexpand};
 
 mod from_mir;
 
@@ -19,7 +17,35 @@ pub(super) fn extract_refined_covspans(
     graph: &CoverageGraph,
     code_mappings: &mut impl Extend<mappings::CodeMapping>,
 ) {
-    let ExtractedCovspans { mut covspans } = extract_covspans_from_mir(mir_body, hir_info, graph);
+    let &ExtractedHirInfo { body_span, .. } = hir_info;
+
+    let raw_spans = from_mir::extract_raw_spans_from_mir(mir_body, graph);
+    let mut covspans = raw_spans
+        .into_iter()
+        .filter_map(|RawSpanFromMir { raw_span, bcb }| try {
+            let (span, expn_kind) =
+                unexpand::unexpand_into_body_span_with_expn_kind(raw_span, body_span)?;
+            // Discard any spans that fill the entire body, because they tend
+            // to represent compiler-inserted code, e.g. implicitly returning `()`.
+            if span.source_equal(body_span) {
+                return None;
+            };
+            SpanFromMir { span, expn_kind, bcb }
+        })
+        .collect::<Vec<_>>();
+
+    // Only proceed if we found at least one usable span.
+    if covspans.is_empty() {
+        return;
+    }
+
+    // Also add the adjusted function signature span, if available.
+    // Otherwise, add a fake span at the start of the body, to avoid an ugly
+    // gap between the start of the body and the first real span.
+    // FIXME: Find a more principled way to solve this problem.
+    covspans.push(SpanFromMir::for_fn_sig(
+        hir_info.fn_sig_span_extended.unwrap_or_else(|| body_span.shrink_to_lo()),
+    ));
 
     // First, perform the passes that need macro information.
     covspans.sort_by(|a, b| graph.cmp_in_dominator_order(a.bcb, b.bcb));
@@ -43,7 +69,14 @@ pub(super) fn extract_refined_covspans(
     covspans.dedup_by(|b, a| a.span.source_equal(b.span));
 
     // Sort the holes, and merge overlapping/adjacent holes.
-    let mut holes = hir_info.hole_spans.iter().map(|&span| Hole { span }).collect::<Vec<_>>();
+    let mut holes = hir_info
+        .hole_spans
+        .iter()
+        .copied()
+        // Discard any holes that aren't directly visible within the body span.
+        .filter(|&hole_span| body_span.contains(hole_span) && body_span.eq_ctxt(hole_span))
+        .map(|span| Hole { span })
+        .collect::<Vec<_>>();
     holes.sort_by(|a, b| compare_spans(a.span, b.span));
     holes.dedup_by(|b, a| a.merge_if_overlapping_or_adjacent(b));
 
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 73b68d7155c..1faa2171c0b 100644
--- a/compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs
+++ b/compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs
@@ -1,3 +1,5 @@
+use std::iter;
+
 use rustc_middle::bug;
 use rustc_middle::mir::coverage::CoverageKind;
 use rustc_middle::mir::{
@@ -5,87 +7,50 @@ use rustc_middle::mir::{
 };
 use rustc_span::{ExpnKind, Span};
 
-use crate::coverage::ExtractedHirInfo;
-use crate::coverage::graph::{
-    BasicCoverageBlock, BasicCoverageBlockData, CoverageGraph, START_BCB,
-};
+use crate::coverage::graph::{BasicCoverageBlock, CoverageGraph, START_BCB};
 use crate::coverage::spans::Covspan;
-use crate::coverage::unexpand::unexpand_into_body_span_with_expn_kind;
 
-pub(crate) struct ExtractedCovspans {
-    pub(crate) covspans: Vec<SpanFromMir>,
+#[derive(Debug)]
+pub(crate) struct RawSpanFromMir {
+    /// A span that has been extracted from a MIR statement/terminator, but
+    /// hasn't been "unexpanded", so it might not lie within the function body
+    /// span and might be part of an expansion with a different context.
+    pub(crate) raw_span: Span,
+    pub(crate) bcb: BasicCoverageBlock,
 }
 
-/// Traverses the MIR body to produce an initial collection of coverage-relevant
-/// spans, each associated with a node in the coverage graph (BCB) and possibly
-/// other metadata.
-pub(crate) fn extract_covspans_from_mir(
-    mir_body: &mir::Body<'_>,
-    hir_info: &ExtractedHirInfo,
+/// Generates an initial set of coverage spans from the statements and
+/// terminators in the function's MIR body, each associated with its
+/// corresponding node in the coverage graph.
+///
+/// This is necessarily an inexact process, because MIR isn't designed to
+/// capture source spans at the level of detail we would want for coverage,
+/// but it's good enough to be better than nothing.
+pub(crate) fn extract_raw_spans_from_mir<'tcx>(
+    mir_body: &mir::Body<'tcx>,
     graph: &CoverageGraph,
-) -> ExtractedCovspans {
-    let &ExtractedHirInfo { body_span, .. } = hir_info;
-
-    let mut covspans = vec![];
+) -> Vec<RawSpanFromMir> {
+    let mut raw_spans = vec![];
 
+    // We only care about blocks that are part of the coverage graph.
     for (bcb, bcb_data) in graph.iter_enumerated() {
-        bcb_to_initial_coverage_spans(mir_body, body_span, bcb, bcb_data, &mut covspans);
-    }
+        let make_raw_span = |raw_span: Span| RawSpanFromMir { raw_span, bcb };
 
-    // Only add the signature span if we found at least one span in the body.
-    if !covspans.is_empty() {
-        // If there is no usable signature span, add a fake one (before refinement)
-        // to avoid an ugly gap between the body start and the first real span.
-        // FIXME: Find a more principled way to solve this problem.
-        let fn_sig_span = hir_info.fn_sig_span_extended.unwrap_or_else(|| body_span.shrink_to_lo());
-        covspans.push(SpanFromMir::for_fn_sig(fn_sig_span));
-    }
+        // A coverage graph node can consist of multiple basic blocks.
+        for &bb in &bcb_data.basic_blocks {
+            let bb_data = &mir_body[bb];
 
-    ExtractedCovspans { covspans }
-}
+            let statements = bb_data.statements.iter();
+            raw_spans.extend(statements.filter_map(filtered_statement_span).map(make_raw_span));
 
-// Generate a set of coverage spans from the filtered set of `Statement`s and `Terminator`s of
-// the `BasicBlock`(s) in the given `BasicCoverageBlockData`. One coverage span is generated
-// for each `Statement` and `Terminator`. (Note that subsequent stages of coverage analysis will
-// merge some coverage spans, at which point a coverage span may represent multiple
-// `Statement`s and/or `Terminator`s.)
-fn bcb_to_initial_coverage_spans<'a, 'tcx>(
-    mir_body: &'a mir::Body<'tcx>,
-    body_span: Span,
-    bcb: BasicCoverageBlock,
-    bcb_data: &'a BasicCoverageBlockData,
-    initial_covspans: &mut Vec<SpanFromMir>,
-) {
-    for &bb in &bcb_data.basic_blocks {
-        let data = &mir_body[bb];
-
-        let unexpand = move |expn_span| {
-            unexpand_into_body_span_with_expn_kind(expn_span, body_span)
-                // Discard any spans that fill the entire body, because they tend
-                // to represent compiler-inserted code, e.g. implicitly returning `()`.
-                .filter(|(span, _)| !span.source_equal(body_span))
-        };
-
-        let mut extract_statement_span = |statement| {
-            let expn_span = filtered_statement_span(statement)?;
-            let (span, expn_kind) = unexpand(expn_span)?;
-
-            initial_covspans.push(SpanFromMir::new(span, expn_kind, bcb));
-            Some(())
-        };
-        for statement in data.statements.iter() {
-            extract_statement_span(statement);
+            // There's only one terminator, but wrap it in an iterator to
+            // mirror the handling of statements.
+            let terminator = iter::once(bb_data.terminator());
+            raw_spans.extend(terminator.filter_map(filtered_terminator_span).map(make_raw_span));
         }
-
-        let mut extract_terminator_span = |terminator| {
-            let expn_span = filtered_terminator_span(terminator)?;
-            let (span, expn_kind) = unexpand(expn_span)?;
-
-            initial_covspans.push(SpanFromMir::new(span, expn_kind, bcb));
-            Some(())
-        };
-        extract_terminator_span(data.terminator());
     }
+
+    raw_spans
 }
 
 /// If the MIR `Statement` has a span contributive to computing coverage spans,
@@ -219,7 +184,7 @@ pub(crate) struct SpanFromMir {
 }
 
 impl SpanFromMir {
-    fn for_fn_sig(fn_sig_span: Span) -> Self {
+    pub(crate) fn for_fn_sig(fn_sig_span: Span) -> Self {
         Self::new(fn_sig_span, None, START_BCB)
     }