rustc_mir_transform/coverage/
expansion.rs

1use itertools::Itertools;
2use rustc_data_structures::fx::{FxIndexMap, FxIndexSet, IndexEntry};
3use rustc_middle::mir;
4use rustc_middle::mir::coverage::{BasicCoverageBlock, BranchSpan};
5use rustc_span::{ExpnId, ExpnKind, Span};
6
7use crate::coverage::from_mir;
8use crate::coverage::graph::CoverageGraph;
9use crate::coverage::hir_info::ExtractedHirInfo;
10use crate::coverage::mappings::MappingsError;
11
12#[derive(Clone, Copy, Debug)]
13pub(crate) struct SpanWithBcb {
14    pub(crate) span: Span,
15    pub(crate) bcb: BasicCoverageBlock,
16}
17
18#[derive(Debug)]
19pub(crate) struct ExpnTree {
20    nodes: FxIndexMap<ExpnId, ExpnNode>,
21}
22
23impl ExpnTree {
24    pub(crate) fn get(&self, expn_id: ExpnId) -> Option<&ExpnNode> {
25        self.nodes.get(&expn_id)
26    }
27}
28
29#[derive(Debug)]
30pub(crate) struct ExpnNode {
31    /// Storing the expansion ID in its own node is not strictly necessary,
32    /// but is helpful for debugging and might be useful later.
33    #[expect(dead_code)]
34    pub(crate) expn_id: ExpnId,
35    /// Index of this node in a depth-first traversal from the root.
36    pub(crate) dfs_rank: usize,
37
38    // Useful info extracted from `ExpnData`.
39    pub(crate) expn_kind: ExpnKind,
40    /// Non-dummy `ExpnData::call_site` span.
41    pub(crate) call_site: Option<Span>,
42    /// Expansion ID of `call_site`, if present.
43    /// This links an expansion node to its parent in the tree.
44    pub(crate) call_site_expn_id: Option<ExpnId>,
45
46    /// Holds the function signature span, if it belongs to this expansion.
47    /// Used by special-case code in span refinement.
48    pub(crate) fn_sig_span: Option<Span>,
49    /// Holds the function body span, if it belongs to this expansion.
50    /// Used by special-case code in span refinement.
51    pub(crate) body_span: Option<Span>,
52
53    /// Spans (and their associated BCBs) belonging to this expansion.
54    pub(crate) spans: Vec<SpanWithBcb>,
55    /// Expansions whose call-site is in this expansion.
56    pub(crate) child_expn_ids: FxIndexSet<ExpnId>,
57    /// The "minimum" and "maximum" BCBs (in dominator order) of ordinary spans
58    /// belonging to this tree node and all of its descendants. Used when
59    /// creating a single code mapping representing an entire child expansion.
60    pub(crate) minmax_bcbs: Option<MinMaxBcbs>,
61
62    /// Branch spans (recorded during MIR building) belonging to this expansion.
63    pub(crate) branch_spans: Vec<BranchSpan>,
64
65    /// Hole spans belonging to this expansion, to be carved out from the
66    /// code spans during span refinement.
67    pub(crate) hole_spans: Vec<Span>,
68}
69
70impl ExpnNode {
71    fn new(expn_id: ExpnId) -> Self {
72        let expn_data = expn_id.expn_data();
73
74        let call_site = Some(expn_data.call_site).filter(|sp| !sp.is_dummy());
75        let call_site_expn_id = try { call_site?.ctxt().outer_expn() };
76
77        Self {
78            expn_id,
79            dfs_rank: usize::MAX,
80
81            expn_kind: expn_data.kind,
82            call_site,
83            call_site_expn_id,
84
85            fn_sig_span: None,
86            body_span: None,
87
88            spans: vec![],
89            child_expn_ids: FxIndexSet::default(),
90            minmax_bcbs: None,
91
92            branch_spans: vec![],
93
94            hole_spans: vec![],
95        }
96    }
97}
98
99/// Extracts raw span/BCB pairs from potentially-different syntax contexts, and
100/// arranges them into an "expansion tree" based on their expansion call-sites.
101pub(crate) fn build_expn_tree(
102    mir_body: &mir::Body<'_>,
103    hir_info: &ExtractedHirInfo,
104    graph: &CoverageGraph,
105) -> Result<ExpnTree, MappingsError> {
106    let raw_spans = from_mir::extract_raw_spans_from_mir(mir_body, graph);
107
108    let mut nodes = FxIndexMap::default();
109    let new_node = |&expn_id: &ExpnId| ExpnNode::new(expn_id);
110
111    for from_mir::RawSpanFromMir { raw_span, bcb } in raw_spans {
112        let span_with_bcb = SpanWithBcb { span: raw_span, bcb };
113
114        // Create a node for this span's enclosing expansion, and add the span to it.
115        let expn_id = span_with_bcb.span.ctxt().outer_expn();
116        let node = nodes.entry(expn_id).or_insert_with_key(new_node);
117        node.spans.push(span_with_bcb);
118
119        // Now walk up the expansion call-site chain, creating nodes and registering children.
120        let mut prev = expn_id;
121        let mut curr_expn_id = node.call_site_expn_id;
122        while let Some(expn_id) = curr_expn_id {
123            let entry = nodes.entry(expn_id);
124            let node_existed = matches!(entry, IndexEntry::Occupied(_));
125
126            let node = entry.or_insert_with_key(new_node);
127            node.child_expn_ids.insert(prev);
128
129            if node_existed {
130                break;
131            }
132
133            prev = expn_id;
134            curr_expn_id = node.call_site_expn_id;
135        }
136    }
137
138    // Sort the tree nodes into depth-first order.
139    sort_nodes_depth_first(&mut nodes)?;
140
141    // For each node, determine its "minimum" and "maximum" BCBs, based on its
142    // own spans and its immediate children. This relies on the nodes having
143    // been sorted, so that each node's children are processed before the node
144    // itself.
145    for i in (0..nodes.len()).rev() {
146        // Computing a node's min/max BCBs requires a shared ref to other nodes.
147        let minmax_bcbs = minmax_bcbs_for_expn_tree_node(graph, &nodes, &nodes[i]);
148        // Now we can mutate the current node to set its min/max BCBs.
149        nodes[i].minmax_bcbs = minmax_bcbs;
150    }
151
152    // If we have a span for the function signature, associate it with the
153    // corresponding expansion tree node.
154    if let Some(fn_sig_span) = hir_info.fn_sig_span
155        && let Some(node) = nodes.get_mut(&fn_sig_span.ctxt().outer_expn())
156    {
157        node.fn_sig_span = Some(fn_sig_span);
158    }
159
160    // Also associate the body span with its expansion tree node.
161    let body_span = hir_info.body_span;
162    if let Some(node) = nodes.get_mut(&body_span.ctxt().outer_expn()) {
163        node.body_span = Some(body_span);
164    }
165
166    // Associate each hole span (extracted from HIR) with its corresponding
167    // expansion tree node.
168    for &hole_span in &hir_info.hole_spans {
169        let expn_id = hole_span.ctxt().outer_expn();
170        let Some(node) = nodes.get_mut(&expn_id) else { continue };
171        node.hole_spans.push(hole_span);
172    }
173
174    // Associate each branch span (recorded during MIR building) with its
175    // corresponding expansion tree node.
176    if let Some(coverage_info_hi) = mir_body.coverage_info_hi.as_deref() {
177        for branch_span in &coverage_info_hi.branch_spans {
178            if let Some(node) = nodes.get_mut(&branch_span.span.ctxt().outer_expn()) {
179                node.branch_spans.push(BranchSpan::clone(branch_span));
180            }
181        }
182    }
183
184    Ok(ExpnTree { nodes })
185}
186
187/// Sorts the tree nodes in the map into depth-first order.
188///
189/// This allows subsequent operations to iterate over all nodes, while assuming
190/// that every node occurs before all of its descendants.
191fn sort_nodes_depth_first(nodes: &mut FxIndexMap<ExpnId, ExpnNode>) -> Result<(), MappingsError> {
192    let mut dfs_stack = vec![ExpnId::root()];
193    let mut next_dfs_rank = 0usize;
194    while let Some(expn_id) = dfs_stack.pop() {
195        if let Some(node) = nodes.get_mut(&expn_id) {
196            node.dfs_rank = next_dfs_rank;
197            next_dfs_rank += 1;
198            dfs_stack.extend(node.child_expn_ids.iter().rev().copied());
199        }
200    }
201    nodes.sort_by_key(|_expn_id, node| node.dfs_rank);
202
203    // Verify that the depth-first search visited each node exactly once.
204    for (i, &ExpnNode { dfs_rank, .. }) in nodes.values().enumerate() {
205        if dfs_rank != i {
206            tracing::debug!(dfs_rank, i, "expansion tree node's rank does not match its index");
207            return Err(MappingsError::TreeSortFailure);
208        }
209    }
210
211    Ok(())
212}
213
214#[derive(Clone, Copy, Debug)]
215pub(crate) struct MinMaxBcbs {
216    pub(crate) min: BasicCoverageBlock,
217    pub(crate) max: BasicCoverageBlock,
218}
219
220/// For a single node in the expansion tree, compute its "minimum" and "maximum"
221/// BCBs (in dominator order), from among the BCBs of its immediate spans,
222/// and the min/max of its immediate children.
223fn minmax_bcbs_for_expn_tree_node(
224    graph: &CoverageGraph,
225    nodes: &FxIndexMap<ExpnId, ExpnNode>,
226    node: &ExpnNode,
227) -> Option<MinMaxBcbs> {
228    let immediate_span_bcbs = node.spans.iter().map(|sp: &SpanWithBcb| sp.bcb);
229    let child_minmax_bcbs = node
230        .child_expn_ids
231        .iter()
232        .flat_map(|id| nodes.get(id))
233        .flat_map(|child| child.minmax_bcbs)
234        .flat_map(|MinMaxBcbs { min, max }| [min, max]);
235
236    let (min, max) = Iterator::chain(immediate_span_bcbs, child_minmax_bcbs)
237        .minmax_by(|&a, &b| graph.cmp_in_dominator_order(a, b))
238        .into_option()?;
239    Some(MinMaxBcbs { min, max })
240}
rustc_mir_transform/coverage/expansion.rs

rustc_mir_transform/coverage/
expansion.rs
