coverage: Associate hole spans with expansion tree nodes

This will make it easier to perform span refinement for child expansions.

Author: Zalathar

compiler/rustc_mir_transform/src/coverage/expansion.rs

@@ -5,6 +5,7 @@ use rustc_span::{ExpnId, ExpnKind, Span};
 
 use crate::coverage::from_mir;
 use crate::coverage::graph::CoverageGraph;
+use crate::coverage::hir_info::ExtractedHirInfo;
 
 #[derive(Clone, Copy, Debug)]
 pub(crate) struct SpanWithBcb {
@@ -74,6 +75,10 @@ pub(crate) struct ExpnNode {
     pub(crate) spans: Vec<SpanWithBcb>,
     /// Expansions whose call-site is in this expansion.
     pub(crate) child_expn_ids: FxIndexSet<ExpnId>,
+
+    /// Hole spans belonging to this expansion, to be carved out from the
+    /// code spans during span refinement.
+    pub(crate) hole_spans: Vec<Span>,
 }
 
 impl ExpnNode {
@@ -92,13 +97,19 @@ impl ExpnNode {
 
             spans: vec![],
             child_expn_ids: FxIndexSet::default(),
+
+            hole_spans: vec![],
         }
     }
 }
 
 /// Extracts raw span/BCB pairs from potentially-different syntax contexts, and
 /// arranges them into an "expansion tree" based on their expansion call-sites.
-pub(crate) fn build_expn_tree(mir_body: &mir::Body<'_>, graph: &CoverageGraph) -> ExpnTree {
+pub(crate) fn build_expn_tree(
+    mir_body: &mir::Body<'_>,
+    hir_info: &ExtractedHirInfo,
+    graph: &CoverageGraph,
+) -> ExpnTree {
     let raw_spans = from_mir::extract_raw_spans_from_mir(mir_body, graph);
 
     let mut nodes = FxIndexMap::default();
@@ -131,5 +142,13 @@ pub(crate) fn build_expn_tree(mir_body: &mir::Body<'_>, graph: &CoverageGraph) -
         }
     }
 
+    // Associate each hole span (extracted from HIR) with its corresponding
+    // expansion tree node.
+    for &hole_span in &hir_info.hole_spans {
+        let expn_id = hole_span.ctxt().outer_expn();
+        let Some(node) = nodes.get_mut(&expn_id) else { continue };
+        node.hole_spans.push(hole_span);
+    }
+
     ExpnTree { nodes }
 }

compiler/rustc_mir_transform/src/coverage/from_mir.rs

@@ -142,19 +142,3 @@ fn filtered_terminator_span(terminator: &Terminator<'_>) -> Option<Span> {
         | TerminatorKind::InlineAsm { .. } => Some(terminator.source_info.span),
     }
 }
-
-#[derive(Debug)]
-pub(crate) struct Hole {
-    pub(crate) span: Span,
-}
-
-impl Hole {
-    pub(crate) fn merge_if_overlapping_or_adjacent(&mut self, other: &mut Self) -> bool {
-        if !self.span.overlaps_or_adjacent(other.span) {
-            return false;
-        }
-
-        self.span = self.span.to(other.span);
-        true
-    }
-}

compiler/rustc_mir_transform/src/coverage/mappings.rs

@@ -24,7 +24,7 @@ pub(crate) fn extract_mappings_from_mir<'tcx>(
     hir_info: &ExtractedHirInfo,
     graph: &CoverageGraph,
 ) -> ExtractedMappings {
-    let expn_tree = expansion::build_expn_tree(mir_body, graph);
+    let expn_tree = expansion::build_expn_tree(mir_body, hir_info, graph);
 
     let mut mappings = vec![];
 

compiler/rustc_mir_transform/src/coverage/spans.rs

@@ -5,7 +5,6 @@ use rustc_span::{BytePos, DesugaringKind, ExpnId, ExpnKind, MacroKind, Span};
 use tracing::instrument;
 
 use crate::coverage::expansion::{ExpnTree, SpanWithBcb};
-use crate::coverage::from_mir::Hole;
 use crate::coverage::graph::{BasicCoverageBlock, CoverageGraph};
 use crate::coverage::hir_info::ExtractedHirInfo;
 
@@ -98,14 +97,8 @@ pub(super) fn extract_refined_covspans<'tcx>(
     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()
-        .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<_>>();
+    let mut holes = node.hole_spans.iter().copied().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));
 
@@ -286,3 +279,19 @@ fn ensure_non_empty_span(source_map: &SourceMap, span: Span) -> Option<Span> {
         })
         .ok()?
 }
+
+#[derive(Debug)]
+struct Hole {
+    span: Span,
+}
+
+impl Hole {
+    fn merge_if_overlapping_or_adjacent(&mut self, other: &mut Self) -> bool {
+        if !self.span.overlaps_or_adjacent(other.span) {
+            return false;
+        }
+
+        self.span = self.span.to(other.span);
+        true
+    }
+}