Rollup merge of #148812 - Zalathar:expansions, r=JonathanBrouwer

coverage: Associate hole spans with expansion tree nodes

This PR is another incremental step towards expansion region support in coverage instrumentation.

When preparing coverage mappings for a function, we extract “raw” spans from MIR, and then use “hole” spans extracted from HIR to avoid overlap with nested items and closures. That hole-carving process was historically built around the assumption that there would be one set of spans and holes per function, but expansion region support will need to invalidate that assumption.

Therefore, to be more friendly to future work on expansion regions, this PR associates each hole span with its corresponding node in the expansion tree, and makes the span refinement step obtain holes from the current node.

There should be no change to compiler output.

compiler/rustc_mir_transform/src/coverage/expansion.rs

@@ -1,7 +1,12 @@
 use rustc_data_structures::fx::{FxIndexMap, FxIndexSet, IndexEntry};
+use rustc_middle::mir;
 use rustc_middle::mir::coverage::BasicCoverageBlock;
 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 {
     pub(crate) span: Span,
@@ -70,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 {
@@ -88,17 +97,27 @@ impl ExpnNode {
 
             spans: vec![],
             child_expn_ids: FxIndexSet::default(),
+
+            hole_spans: vec![],
         }
     }
 }
 
-/// Given a collection of span/BCB pairs from potentially-different syntax contexts,
+/// 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(spans: impl IntoIterator<Item = SpanWithBcb>) -> 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();
     let new_node = |&expn_id: &ExpnId| ExpnNode::new(expn_id);
 
-    for span_with_bcb in spans {
+    for from_mir::RawSpanFromMir { raw_span, bcb } in raw_spans {
+        let span_with_bcb = SpanWithBcb { span: raw_span, bcb };
+
         // Create a node for this span's enclosing expansion, and add the span to it.
         let expn_id = span_with_bcb.span.ctxt().outer_expn();
         let node = nodes.entry(expn_id).or_insert_with_key(new_node);
@@ -123,5 +142,13 @@ pub(crate) fn build_expn_tree(spans: impl IntoIterator<Item = SpanWithBcb>) -> E
         }
     }
 
+    // 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

@@ -5,6 +5,7 @@ use rustc_middle::mir::coverage::{
 use rustc_middle::mir::{self, BasicBlock, StatementKind};
 use rustc_middle::ty::TyCtxt;
 
+use crate::coverage::expansion;
 use crate::coverage::graph::CoverageGraph;
 use crate::coverage::hir_info::ExtractedHirInfo;
 use crate::coverage::spans::extract_refined_covspans;
@@ -23,10 +24,12 @@ pub(crate) fn extract_mappings_from_mir<'tcx>(
     hir_info: &ExtractedHirInfo,
     graph: &CoverageGraph,
 ) -> ExtractedMappings {
+    let expn_tree = expansion::build_expn_tree(mir_body, hir_info, graph);
+
     let mut mappings = vec![];
 
     // Extract ordinary code mappings from MIR statement/terminator spans.
-    extract_refined_covspans(tcx, mir_body, hir_info, graph, &mut mappings);
+    extract_refined_covspans(tcx, hir_info, graph, &expn_tree, &mut mappings);
 
     extract_branch_mappings(mir_body, hir_info, graph, &mut mappings);
 

compiler/rustc_mir_transform/src/coverage/mod.rs

@@ -9,6 +9,7 @@ use crate::coverage::mappings::ExtractedMappings;
 
 mod counters;
 mod expansion;
+mod from_mir;
 mod graph;
 mod hir_info;
 mod mappings;

compiler/rustc_mir_transform/src/coverage/spans.rs

@@ -1,22 +1,18 @@
-use rustc_middle::mir;
 use rustc_middle::mir::coverage::{Mapping, MappingKind, START_BCB};
 use rustc_middle::ty::TyCtxt;
 use rustc_span::source_map::SourceMap;
 use rustc_span::{BytePos, DesugaringKind, ExpnId, ExpnKind, MacroKind, Span};
 use tracing::instrument;
 
-use crate::coverage::expansion::{self, ExpnTree, SpanWithBcb};
+use crate::coverage::expansion::{ExpnTree, SpanWithBcb};
 use crate::coverage::graph::{BasicCoverageBlock, CoverageGraph};
 use crate::coverage::hir_info::ExtractedHirInfo;
-use crate::coverage::spans::from_mir::{Hole, RawSpanFromMir};
-
-mod from_mir;
 
 pub(super) fn extract_refined_covspans<'tcx>(
     tcx: TyCtxt<'tcx>,
-    mir_body: &mir::Body<'tcx>,
     hir_info: &ExtractedHirInfo,
     graph: &CoverageGraph,
+    expn_tree: &ExpnTree,
     mappings: &mut Vec<Mapping>,
 ) {
     if hir_info.is_async_fn {
@@ -32,22 +28,32 @@ pub(super) fn extract_refined_covspans<'tcx>(
 
     let &ExtractedHirInfo { body_span, .. } = hir_info;
 
-    let raw_spans = from_mir::extract_raw_spans_from_mir(mir_body, graph);
-    // Use the raw spans to build a tree of expansions for this function.
-    let expn_tree = expansion::build_expn_tree(
-        raw_spans
-            .into_iter()
-            .map(|RawSpanFromMir { raw_span, bcb }| SpanWithBcb { span: raw_span, bcb }),
-    );
+    // If there somehow isn't an expansion tree node corresponding to the
+    // body span, return now and don't create any mappings.
+    let Some(node) = expn_tree.get(body_span.ctxt().outer_expn()) else { return };
 
     let mut covspans = vec![];
-    let mut push_covspan = |covspan: Covspan| {
+
+    for &SpanWithBcb { span, bcb } in &node.spans {
+        covspans.push(Covspan { span, bcb });
+    }
+
+    // For each expansion with its call-site in the body span, try to
+    // distill a corresponding covspan.
+    for &child_expn_id in &node.child_expn_ids {
+        if let Some(covspan) = single_covspan_for_child_expn(tcx, graph, &expn_tree, child_expn_id)
+        {
+            covspans.push(covspan);
+        }
+    }
+
+    covspans.retain(|covspan: &Covspan| {
         let covspan_span = covspan.span;
         // Discard any spans not contained within the function body span.
         // Also discard any spans that fill the entire body, because they tend
         // to represent compiler-inserted code, e.g. implicitly returning `()`.
         if !body_span.contains(covspan_span) || body_span.source_equal(covspan_span) {
-            return;
+            return false;
         }
 
         // Each pushed covspan should have the same context as the body span.
@@ -57,27 +63,11 @@ pub(super) fn extract_refined_covspans<'tcx>(
                 false,
                 "span context mismatch: body_span={body_span:?}, covspan.span={covspan_span:?}"
             );
-            return;
+            return false;
         }
 
-        covspans.push(covspan);
-    };
-
-    if let Some(node) = expn_tree.get(body_span.ctxt().outer_expn()) {
-        for &SpanWithBcb { span, bcb } in &node.spans {
-            push_covspan(Covspan { span, bcb });
-        }
-
-        // For each expansion with its call-site in the body span, try to
-        // distill a corresponding covspan.
-        for &child_expn_id in &node.child_expn_ids {
-            if let Some(covspan) =
-                single_covspan_for_child_expn(tcx, graph, &expn_tree, child_expn_id)
-            {
-                push_covspan(covspan);
-            }
-        }
-    }
+        true
+    });
 
     // Only proceed if we found at least one usable span.
     if covspans.is_empty() {
@@ -107,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));
 
@@ -295,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
+    }
+}