Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Sign up
Appearance settings

Commit 88c254f

Browse files
Rollup merge of #145643 - Zalathar:tree, r=SparrowLii
coverage: Build an "expansion tree" and use it to unexpand raw spans Historically and currently, coverage instrumentation assumes that all of a function's spans are in the same file and have the same syntax context. The spans extracted directly from MIR don't satisfy that assumption, so there is an "unexpansion" step that walks up each span's expansion-call-site tree to find a suitable span in the same context as the function's body span. (That unexpansion step is what allows us to have somewhat reasonable coverage instrumentation for macros like `println!`, and for syntax like `for` and `?` that undergo desugaring expansion.) The current unexpansion code mostly works fine in that "flat" single-file single-context world. But it's not suitable for incremental work towards proper expansion-aware coverage instrumentation, which would allow a function's coverage spans to encompass multiple expansion contexts and multiple files. This PR therefore replaces the current unexpansion code with a more sophisticated system that uses the raw MIR spans to reconstruct an "expansion tree", and then uses that tree to help perform most of the unexpansion work. Building the tree is "overkill" for current unexpansion needs (though it does give some minor edge-case improvements), but my hope is that having the explicit tree available will be a big help when taking the next steps towards proper expansion-region support.
2 parents 1d520e2 + c2eb45b commit 88c254f

File tree

12 files changed

+286
-162
lines changed

12 files changed

+286
-162
lines changed
Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
1+
use rustc_data_structures::fx::{FxIndexMap, FxIndexSet, IndexEntry};
2+
use rustc_middle::mir::coverage::BasicCoverageBlock;
3+
use rustc_span::{ExpnId, ExpnKind, Span};
4+
5+
#[derive(Clone, Copy, Debug)]
6+
pub(crate) struct SpanWithBcb {
7+
pub(crate) span: Span,
8+
pub(crate) bcb: BasicCoverageBlock,
9+
}
10+
11+
#[derive(Debug)]
12+
pub(crate) struct ExpnTree {
13+
nodes: FxIndexMap<ExpnId, ExpnNode>,
14+
}
15+
16+
impl ExpnTree {
17+
pub(crate) fn get(&self, expn_id: ExpnId) -> Option<&ExpnNode> {
18+
self.nodes.get(&expn_id)
19+
}
20+
21+
/// Yields the tree node for the given expansion ID (if present), followed
22+
/// by the nodes of all of its descendants in depth-first order.
23+
pub(crate) fn iter_node_and_descendants(
24+
&self,
25+
root_expn_id: ExpnId,
26+
) -> impl Iterator<Item = &ExpnNode> {
27+
gen move {
28+
let Some(root_node) = self.get(root_expn_id) else { return };
29+
yield root_node;
30+
31+
// Stack of child-node-ID iterators that drives the depth-first traversal.
32+
let mut iter_stack = vec![root_node.child_expn_ids.iter()];
33+
34+
while let Some(curr_iter) = iter_stack.last_mut() {
35+
// Pull the next ID from the top of the stack.
36+
let Some(&curr_id) = curr_iter.next() else {
37+
iter_stack.pop();
38+
continue;
39+
};
40+
41+
// Yield this node.
42+
let Some(node) = self.get(curr_id) else { continue };
43+
yield node;
44+
45+
// Push the node's children, to be traversed next.
46+
if !node.child_expn_ids.is_empty() {
47+
iter_stack.push(node.child_expn_ids.iter());
48+
}
49+
}
50+
}
51+
}
52+
}
53+
54+
#[derive(Debug)]
55+
pub(crate) struct ExpnNode {
56+
/// Storing the expansion ID in its own node is not strictly necessary,
57+
/// but is helpful for debugging and might be useful later.
58+
#[expect(dead_code)]
59+
pub(crate) expn_id: ExpnId,
60+
61+
// Useful info extracted from `ExpnData`.
62+
pub(crate) expn_kind: ExpnKind,
63+
/// Non-dummy `ExpnData::call_site` span.
64+
pub(crate) call_site: Option<Span>,
65+
/// Expansion ID of `call_site`, if present.
66+
/// This links an expansion node to its parent in the tree.
67+
pub(crate) call_site_expn_id: Option<ExpnId>,
68+
69+
/// Spans (and their associated BCBs) belonging to this expansion.
70+
pub(crate) spans: Vec<SpanWithBcb>,
71+
/// Expansions whose call-site is in this expansion.
72+
pub(crate) child_expn_ids: FxIndexSet<ExpnId>,
73+
}
74+
75+
impl ExpnNode {
76+
fn new(expn_id: ExpnId) -> Self {
77+
let expn_data = expn_id.expn_data();
78+
79+
let call_site = Some(expn_data.call_site).filter(|sp| !sp.is_dummy());
80+
let call_site_expn_id = try { call_site?.ctxt().outer_expn() };
81+
82+
Self {
83+
expn_id,
84+
85+
expn_kind: expn_data.kind.clone(),
86+
call_site,
87+
call_site_expn_id,
88+
89+
spans: vec![],
90+
child_expn_ids: FxIndexSet::default(),
91+
}
92+
}
93+
}
94+
95+
/// Given a collection of span/BCB pairs from potentially-different syntax contexts,
96+
/// arranges them into an "expansion tree" based on their expansion call-sites.
97+
pub(crate) fn build_expn_tree(spans: impl IntoIterator<Item = SpanWithBcb>) -> ExpnTree {
98+
let mut nodes = FxIndexMap::default();
99+
let new_node = |&expn_id: &ExpnId| ExpnNode::new(expn_id);
100+
101+
for span_with_bcb in spans {
102+
// Create a node for this span's enclosing expansion, and add the span to it.
103+
let expn_id = span_with_bcb.span.ctxt().outer_expn();
104+
let node = nodes.entry(expn_id).or_insert_with_key(new_node);
105+
node.spans.push(span_with_bcb);
106+
107+
// Now walk up the expansion call-site chain, creating nodes and registering children.
108+
let mut prev = expn_id;
109+
let mut curr_expn_id = node.call_site_expn_id;
110+
while let Some(expn_id) = curr_expn_id {
111+
let entry = nodes.entry(expn_id);
112+
let node_existed = matches!(entry, IndexEntry::Occupied(_));
113+
114+
let node = entry.or_insert_with_key(new_node);
115+
node.child_expn_ids.insert(prev);
116+
117+
if node_existed {
118+
break;
119+
}
120+
121+
prev = expn_id;
122+
curr_expn_id = node.call_site_expn_id;
123+
}
124+
}
125+
126+
ExpnTree { nodes }
127+
}

‎compiler/rustc_mir_transform/src/coverage/mod.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ use crate::coverage::graph::CoverageGraph;
88
use crate::coverage::mappings::ExtractedMappings;
99

1010
mod counters;
11+
mod expansion;
1112
mod graph;
1213
mod hir_info;
1314
mod mappings;

‎compiler/rustc_mir_transform/src/coverage/spans.rs

Lines changed: 80 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,14 @@
1-
use rustc_data_structures::fx::FxHashSet;
21
use rustc_middle::mir;
32
use rustc_middle::mir::coverage::{Mapping, MappingKind, START_BCB};
43
use rustc_middle::ty::TyCtxt;
54
use rustc_span::source_map::SourceMap;
6-
use rustc_span::{BytePos, DesugaringKind, ExpnKind, MacroKind, Span};
5+
use rustc_span::{BytePos, DesugaringKind, ExpnId,ExpnKind, MacroKind, Span};
76
use tracing::instrument;
87

8+
use crate::coverage::expansion::{self, ExpnTree, SpanWithBcb};
99
use crate::coverage::graph::{BasicCoverageBlock, CoverageGraph};
1010
use crate::coverage::hir_info::ExtractedHirInfo;
11-
use crate::coverage::spans::from_mir::{Hole, RawSpanFromMir, SpanFromMir};
12-
use crate::coverage::unexpand;
11+
use crate::coverage::spans::from_mir::{Hole, RawSpanFromMir};
1312

1413
mod from_mir;
1514

@@ -34,19 +33,51 @@ pub(super) fn extract_refined_covspans<'tcx>(
3433
let &ExtractedHirInfo { body_span, .. } = hir_info;
3534

3635
let raw_spans = from_mir::extract_raw_spans_from_mir(mir_body, graph);
37-
let mut covspans = raw_spans
38-
.into_iter()
39-
.filter_map(|RawSpanFromMir { raw_span, bcb }| try {
40-
let (span, expn_kind) =
41-
unexpand::unexpand_into_body_span_with_expn_kind(raw_span, body_span)?;
42-
// Discard any spans that fill the entire body, because they tend
43-
// to represent compiler-inserted code, e.g. implicitly returning `()`.
44-
if span.source_equal(body_span) {
45-
return None;
46-
};
47-
SpanFromMir { span, expn_kind, bcb }
48-
})
49-
.collect::<Vec<_>>();
36+
// Use the raw spans to build a tree of expansions for this function.
37+
let expn_tree = expansion::build_expn_tree(
38+
raw_spans
39+
.into_iter()
40+
.map(|RawSpanFromMir { raw_span, bcb }| SpanWithBcb { span: raw_span, bcb }),
41+
);
42+
43+
let mut covspans = vec![];
44+
let mut push_covspan = |covspan: Covspan| {
45+
let covspan_span = covspan.span;
46+
// Discard any spans not contained within the function body span.
47+
// Also discard any spans that fill the entire body, because they tend
48+
// to represent compiler-inserted code, e.g. implicitly returning `()`.
49+
if !body_span.contains(covspan_span) || body_span.source_equal(covspan_span) {
50+
return;
51+
}
52+
53+
// Each pushed covspan should have the same context as the body span.
54+
// If it somehow doesn't, discard the covspan, or panic in debug builds.
55+
if !body_span.eq_ctxt(covspan_span) {
56+
debug_assert!(
57+
false,
58+
"span context mismatch: body_span={body_span:?}, covspan.span={covspan_span:?}"
59+
);
60+
return;
61+
}
62+
63+
covspans.push(covspan);
64+
};
65+
66+
if let Some(node) = expn_tree.get(body_span.ctxt().outer_expn()) {
67+
for &SpanWithBcb { span, bcb } in &node.spans {
68+
push_covspan(Covspan { span, bcb });
69+
}
70+
71+
// For each expansion with its call-site in the body span, try to
72+
// distill a corresponding covspan.
73+
for &child_expn_id in &node.child_expn_ids {
74+
if let Some(covspan) =
75+
single_covspan_for_child_expn(tcx, graph, &expn_tree, child_expn_id)
76+
{
77+
push_covspan(covspan);
78+
}
79+
}
80+
}
5081

5182
// Only proceed if we found at least one usable span.
5283
if covspans.is_empty() {
@@ -57,17 +88,10 @@ pub(super) fn extract_refined_covspans<'tcx>(
5788
// Otherwise, add a fake span at the start of the body, to avoid an ugly
5889
// gap between the start of the body and the first real span.
5990
// FIXME: Find a more principled way to solve this problem.
60-
covspans.push(SpanFromMir::for_fn_sig(
61-
hir_info.fn_sig_span.unwrap_or_else(|| body_span.shrink_to_lo()),
62-
));
63-
64-
// First, perform the passes that need macro information.
65-
covspans.sort_by(|a, b| graph.cmp_in_dominator_order(a.bcb, b.bcb));
66-
remove_unwanted_expansion_spans(&mut covspans);
67-
shrink_visible_macro_spans(tcx, &mut covspans);
68-
69-
// We no longer need the extra information in `SpanFromMir`, so convert to `Covspan`.
70-
let mut covspans = covspans.into_iter().map(SpanFromMir::into_covspan).collect::<Vec<_>>();
91+
covspans.push(Covspan {
92+
span: hir_info.fn_sig_span.unwrap_or_else(|| body_span.shrink_to_lo()),
93+
bcb: START_BCB,
94+
});
7195

7296
let compare_covspans = |a: &Covspan, b: &Covspan| {
7397
compare_spans(a.span, b.span)
@@ -117,43 +141,37 @@ pub(super) fn extract_refined_covspans<'tcx>(
117141
}));
118142
}
119143

120-
/// Macros that expand into branches (e.g. `assert!`, `trace!`) tend to generate
121-
/// multiple condition/consequent blocks that have the span of the whole macro
122-
/// invocation, which is unhelpful. Keeping only the first such span seems to
123-
/// give better mappings, so remove the others.
124-
///
125-
/// Similarly, `await` expands to a branch on the discriminant of `Poll`, which
126-
/// leads to incorrect coverage if the `Future` is immediately ready (#98712).
127-
///
128-
/// (The input spans should be sorted in BCB dominator order, so that the
129-
/// retained "first" span is likely to dominate the others.)
130-
fn remove_unwanted_expansion_spans(covspans: &mut Vec<SpanFromMir>) {
131-
let mut deduplicated_spans = FxHashSet::default();
132-
133-
covspans.retain(|covspan| {
134-
match covspan.expn_kind {
135-
// Retain only the first await-related or macro-expanded covspan with this span.
136-
Some(ExpnKind::Desugaring(DesugaringKind::Await)) => {
137-
deduplicated_spans.insert(covspan.span)
138-
}
139-
Some(ExpnKind::Macro(MacroKind::Bang, _)) => deduplicated_spans.insert(covspan.span),
140-
// Ignore (retain) other spans.
141-
_ => true,
144+
/// For a single child expansion, try to distill it into a single span+BCB mapping.
145+
fn single_covspan_for_child_expn(
146+
tcx: TyCtxt<'_>,
147+
graph: &CoverageGraph,
148+
expn_tree: &ExpnTree,
149+
expn_id: ExpnId,
150+
) -> Option<Covspan> {
151+
let node = expn_tree.get(expn_id)?;
152+
153+
let bcbs =
154+
expn_tree.iter_node_and_descendants(expn_id).flat_map(|n| n.spans.iter().map(|s| s.bcb));
155+
156+
let bcb = match node.expn_kind {
157+
// For bang-macros (e.g. `assert!`, `trace!`) and for `await`, taking
158+
// the "first" BCB in dominator order seems to give good results.
159+
ExpnKind::Macro(MacroKind::Bang, _) | ExpnKind::Desugaring(DesugaringKind::Await) => {
160+
bcbs.min_by(|&a, &b| graph.cmp_in_dominator_order(a, b))?
142161
}
143-
});
144-
}
145-
146-
/// When a span corresponds to a macro invocation that is visible from the
147-
/// function body, truncate it to just the macro name plus `!`.
148-
/// This seems to give better results for code that uses macros.
149-
fn shrink_visible_macro_spans(tcx: TyCtxt<'_>, covspans: &mut Vec<SpanFromMir>) {
150-
let source_map = tcx.sess.source_map();
162+
// For other kinds of expansion, taking the "last" (most-dominated) BCB
163+
// seems to give good results.
164+
_ => bcbs.max_by(|&a, &b| graph.cmp_in_dominator_order(a, b))?,
165+
};
151166

152-
for covspan in covspans {
153-
if matches!(covspan.expn_kind, Some(ExpnKind::Macro(MacroKind::Bang, _))) {
154-
covspan.span = source_map.span_through_char(covspan.span, '!');
155-
}
167+
// For bang-macro expansions, limit the call-site span to just the macro
168+
// name plus `!`, excluding the macro arguments.
169+
let mut span = node.call_site?;
170+
if matches!(node.expn_kind, ExpnKind::Macro(MacroKind::Bang, _)) {
171+
span = tcx.sess.source_map().span_through_char(span, '!');
156172
}
173+
174+
Some(Covspan { span, bcb })
157175
}
158176

159177
/// Discard all covspans that overlap a hole.

‎compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs

Lines changed: 2 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,9 @@ use rustc_middle::mir::coverage::CoverageKind;
55
use rustc_middle::mir::{
66
self, FakeReadCause, Statement, StatementKind, Terminator, TerminatorKind,
77
};
8-
use rustc_span::{ExpnKind,Span};
8+
use rustc_span::Span;
99

10-
use crate::coverage::graph::{BasicCoverageBlock, CoverageGraph, START_BCB};
11-
use crate::coverage::spans::Covspan;
10+
use crate::coverage::graph::{BasicCoverageBlock, CoverageGraph};
1211

1312
#[derive(Debug)]
1413
pub(crate) struct RawSpanFromMir {
@@ -160,32 +159,3 @@ impl Hole {
160159
true
161160
}
162161
}
163-
164-
#[derive(Debug)]
165-
pub(crate) struct SpanFromMir {
166-
/// A span that has been extracted from MIR and then "un-expanded" back to
167-
/// within the current function's `body_span`. After various intermediate
168-
/// processing steps, this span is emitted as part of the final coverage
169-
/// mappings.
170-
///
171-
/// With the exception of `fn_sig_span`, this should always be contained
172-
/// within `body_span`.
173-
pub(crate) span: Span,
174-
pub(crate) expn_kind: Option<ExpnKind>,
175-
pub(crate) bcb: BasicCoverageBlock,
176-
}
177-
178-
impl SpanFromMir {
179-
pub(crate) fn for_fn_sig(fn_sig_span: Span) -> Self {
180-
Self::new(fn_sig_span, None, START_BCB)
181-
}
182-
183-
pub(crate) fn new(span: Span, expn_kind: Option<ExpnKind>, bcb: BasicCoverageBlock) -> Self {
184-
Self { span, expn_kind, bcb }
185-
}
186-
187-
pub(crate) fn into_covspan(self) -> Covspan {
188-
let Self { span, expn_kind: _, bcb } = self;
189-
Covspan { span, bcb }
190-
}
191-
}

0 commit comments

Comments
(0)

AltStyle によって変換されたページ (->オリジナル) /