coverage: Store branch spans in the expansion tree

Author: Zalathar

compiler/rustc_mir_transform/src/coverage/expansion.rs

@@ -1,6 +1,6 @@
 use rustc_data_structures::fx::{FxIndexMap, FxIndexSet, IndexEntry};
 use rustc_middle::mir;
-use rustc_middle::mir::coverage::BasicCoverageBlock;
+use rustc_middle::mir::coverage::{BasicCoverageBlock, BranchSpan};
 use rustc_span::{ExpnId, ExpnKind, Span};
 
 use crate::coverage::from_mir;
@@ -83,6 +83,9 @@ pub(crate) struct ExpnNode {
     /// Expansions whose call-site is in this expansion.
     pub(crate) child_expn_ids: FxIndexSet<ExpnId>,
 
+    /// Branch spans (recorded during MIR building) belonging to this expansion.
+    pub(crate) branch_spans: Vec<BranchSpan>,
+
     /// Hole spans belonging to this expansion, to be carved out from the
     /// code spans during span refinement.
     pub(crate) hole_spans: Vec<Span>,
@@ -108,6 +111,8 @@ impl ExpnNode {
             spans: vec![],
             child_expn_ids: FxIndexSet::default(),
 
+            branch_spans: vec![],
+
             hole_spans: vec![],
         }
     }
@@ -174,5 +179,15 @@ pub(crate) fn build_expn_tree(
         node.hole_spans.push(hole_span);
     }
 
+    // Associate each branch span (recorded during MIR building) with its
+    // corresponding expansion tree node.
+    if let Some(coverage_info_hi) = mir_body.coverage_info_hi.as_deref() {
+        for branch_span in &coverage_info_hi.branch_spans {
+            if let Some(node) = nodes.get_mut(&branch_span.span.ctxt().outer_expn()) {
+                node.branch_spans.push(BranchSpan::clone(branch_span));
+            }
+        }
+    }
+
     ExpnTree { nodes }
 }

compiler/rustc_mir_transform/src/coverage/mappings.rs

@@ -4,12 +4,12 @@ use rustc_middle::mir::coverage::{
 };
 use rustc_middle::mir::{self, BasicBlock, StatementKind};
 use rustc_middle::ty::TyCtxt;
+use rustc_span::ExpnKind;
 
-use crate::coverage::expansion;
+use crate::coverage::expansion::{self, ExpnTree};
 use crate::coverage::graph::CoverageGraph;
 use crate::coverage::hir_info::ExtractedHirInfo;
 use crate::coverage::spans::extract_refined_covspans;
-use crate::coverage::unexpand::unexpand_into_body_span;
 
 #[derive(Default)]
 pub(crate) struct ExtractedMappings {
@@ -31,7 +31,7 @@ pub(crate) fn extract_mappings_from_mir<'tcx>(
     // Extract ordinary code mappings from MIR statement/terminator spans.
     extract_refined_covspans(tcx, hir_info, graph, &expn_tree, &mut mappings);
 
-    extract_branch_mappings(mir_body, hir_info, graph, &mut mappings);
+    extract_branch_mappings(mir_body, hir_info, graph, &expn_tree, &mut mappings);
 
     ExtractedMappings { mappings }
 }
@@ -57,25 +57,25 @@ fn resolve_block_markers(
     block_markers
 }
 
-pub(super) fn extract_branch_mappings(
+fn extract_branch_mappings(
     mir_body: &mir::Body<'_>,
     hir_info: &ExtractedHirInfo,
     graph: &CoverageGraph,
+    expn_tree: &ExpnTree,
     mappings: &mut Vec<Mapping>,
 ) {
     let Some(coverage_info_hi) = mir_body.coverage_info_hi.as_deref() else { return };
-
     let block_markers = resolve_block_markers(coverage_info_hi, mir_body);
 
-    mappings.extend(coverage_info_hi.branch_spans.iter().filter_map(
-        |&BranchSpan { span: raw_span, true_marker, false_marker }| try {
-            // For now, ignore any branch span that was introduced by
-            // expansion. This makes things like assert macros less noisy.
-            if !raw_span.ctxt().outer_expn_data().is_root() {
-                return None;
-            }
-            let span = unexpand_into_body_span(raw_span, hir_info.body_span)?;
+    // For now, ignore any branch span that was introduced by
+    // expansion. This makes things like assert macros less noisy.
+    let Some(node) = expn_tree.get(hir_info.body_span.ctxt().outer_expn()) else { return };
+    if node.expn_kind != ExpnKind::Root {
+        return;
+    }
 
+    mappings.extend(node.branch_spans.iter().filter_map(
+        |&BranchSpan { span, true_marker, false_marker }| try {
             let bcb_from_marker = |marker: BlockMarkerId| graph.bcb_from_bb(block_markers[marker]?);
 
             let true_bcb = bcb_from_marker(true_marker)?;

compiler/rustc_mir_transform/src/coverage/mod.rs

@@ -17,7 +17,6 @@ pub(super) mod query;
 mod spans;
 #[cfg(test)]
 mod tests;
-mod unexpand;
 
 /// Inserts `StatementKind::Coverage` statements that either instrument the binary with injected
 /// counters, via intrinsic `llvm.instrprof.increment`, and/or inject metadata used during codegen