From ac437169eca16cb4e7595ea4612e0e6f9a1fad2b Mon Sep 17 00:00:00 2001 From: Zalathar Date: Sun, 30 Nov 2025 15:14:47 +1100 Subject: [PATCH] coverage: Store branch spans in the expansion tree --- .../src/coverage/expansion.rs | 17 +++++++++++- .../src/coverage/mappings.rs | 26 +++++++++---------- .../rustc_mir_transform/src/coverage/mod.rs | 1 - .../src/coverage/unexpand.rs | 9 ------- 4 files changed, 29 insertions(+), 24 deletions(-) delete mode 100644 compiler/rustc_mir_transform/src/coverage/unexpand.rs diff --git a/compiler/rustc_mir_transform/src/coverage/expansion.rs b/compiler/rustc_mir_transform/src/coverage/expansion.rs index 328293cfa28b..0288afd95990 100644 --- a/compiler/rustc_mir_transform/src/coverage/expansion.rs +++ b/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, + /// Branch spans (recorded during MIR building) belonging to this expansion. + pub(crate) branch_spans: Vec, + /// Hole spans belonging to this expansion, to be carved out from the /// code spans during span refinement. pub(crate) hole_spans: Vec, @@ -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 } } diff --git a/compiler/rustc_mir_transform/src/coverage/mappings.rs b/compiler/rustc_mir_transform/src/coverage/mappings.rs index 5347e1150e5d..56f2db90ff8c 100644 --- a/compiler/rustc_mir_transform/src/coverage/mappings.rs +++ b/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, ) { 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)?; diff --git a/compiler/rustc_mir_transform/src/coverage/mod.rs b/compiler/rustc_mir_transform/src/coverage/mod.rs index cc9d7c800a96..24a61c9b4a1b 100644 --- a/compiler/rustc_mir_transform/src/coverage/mod.rs +++ b/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 diff --git a/compiler/rustc_mir_transform/src/coverage/unexpand.rs b/compiler/rustc_mir_transform/src/coverage/unexpand.rs deleted file mode 100644 index 922edd3cc4fe..000000000000 --- a/compiler/rustc_mir_transform/src/coverage/unexpand.rs +++ /dev/null @@ -1,9 +0,0 @@ -use rustc_span::Span; - -/// Walks through the expansion ancestors of `original_span` to find a span that -/// is contained in `body_span` and has the same [syntax context] as `body_span`. -pub(crate) fn unexpand_into_body_span(original_span: Span, body_span: Span) -> Option { - // Because we don't need to return any extra ancestor information, - // we can just delegate directly to `find_ancestor_inside_same_ctxt`. - original_span.find_ancestor_inside_same_ctxt(body_span) -}