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::{ExpnKind, Span, SyntaxContext};
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<SyntaxContext, ExpnNode>,
21}
22
23impl ExpnTree {
24    pub(crate) fn get(&self, context: SyntaxContext) -> Option<&ExpnNode> {
25        self.nodes.get(&context)
26    }
27}
28
29#[derive(Debug)]
30pub(crate) struct ExpnNode {
31    /// Storing the syntax context 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) context: SyntaxContext,
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    /// Syntax context of `call_site`, if present.
43    /// This links an expansion node to its parent in the tree.
44    pub(crate) call_site_context: Option<SyntaxContext>,
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_contexts: FxIndexSet<SyntaxContext>,
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 for_context(context: SyntaxContext) -> Self {
72        let expn_data = context.outer_expn_data();
73
74        let call_site = Some(expn_data.call_site).filter(|sp| !sp.is_dummy());
75        let call_site_context = try { call_site?.ctxt() };
76
77        Self {
78            context,
79            dfs_rank: usize::MAX,
80
81            expn_kind: expn_data.kind,
82            call_site,
83            call_site_context,
84
85            fn_sig_span: None,
86            body_span: None,
87
88            spans: vec![],
89            child_contexts: 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 = |&context: &SyntaxContext| ExpnNode::for_context(context);
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 context = span_with_bcb.span.ctxt();
116        let node = nodes.entry(context).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 = context;
121        let mut curr_context = node.call_site_context;
122        while let Some(context) = curr_context {
123            let entry = nodes.entry(context);
124            let node_existed = matches!(entry, IndexEntry::Occupied(_));
125
126            let node = entry.or_insert_with_key(new_node);
127            node.child_contexts.insert(prev);
128
129            if node_existed {
130                break;
131            }
132
133            prev = context;
134            curr_context = node.call_site_context;
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())
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()) {
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 context = hole_span.ctxt();
170        let Some(node) = nodes.get_mut(&context) 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()) {
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(
192    nodes: &mut FxIndexMap<SyntaxContext, ExpnNode>,
193) -> Result<(), MappingsError> {
194    let mut dfs_stack = vec![SyntaxContext::root()];
195    let mut next_dfs_rank = 0usize;
196    while let Some(context) = dfs_stack.pop() {
197        if let Some(node) = nodes.get_mut(&context) {
198            node.dfs_rank = next_dfs_rank;
199            next_dfs_rank += 1;
200            dfs_stack.extend(node.child_contexts.iter().rev().copied());
201        }
202    }
203    nodes.sort_by_key(|_context, node| node.dfs_rank);
204
205    // Verify that the depth-first search visited each node exactly once.
206    for (i, &ExpnNode { dfs_rank, .. }) in nodes.values().enumerate() {
207        if dfs_rank != i {
208            tracing::debug!(dfs_rank, i, "expansion tree node's rank does not match its index");
209            return Err(MappingsError::TreeSortFailure);
210        }
211    }
212
213    Ok(())
214}
215
216#[derive(Clone, Copy, Debug)]
217pub(crate) struct MinMaxBcbs {
218    pub(crate) min: BasicCoverageBlock,
219    pub(crate) max: BasicCoverageBlock,
220}
221
222/// For a single node in the expansion tree, compute its "minimum" and "maximum"
223/// BCBs (in dominator order), from among the BCBs of its immediate spans,
224/// and the min/max of its immediate children.
225fn minmax_bcbs_for_expn_tree_node(
226    graph: &CoverageGraph,
227    nodes: &FxIndexMap<SyntaxContext, ExpnNode>,
228    node: &ExpnNode,
229) -> Option<MinMaxBcbs> {
230    let immediate_span_bcbs = node.spans.iter().map(|sp: &SpanWithBcb| sp.bcb);
231    let child_minmax_bcbs = node
232        .child_contexts
233        .iter()
234        .flat_map(|id| nodes.get(id))
235        .flat_map(|child| child.minmax_bcbs)
236        .flat_map(|MinMaxBcbs { min, max }| [min, max]);
237
238    let (min, max) = Iterator::chain(immediate_span_bcbs, child_minmax_bcbs)
239        .minmax_by(|&a, &b| graph.cmp_in_dominator_order(a, b))
240        .into_option()?;
241    Some(MinMaxBcbs { min, max })
242}
rustc_mir_transform/coverage/expansion.rs

rustc_mir_transform/coverage/
expansion.rs
