rustc_mir_transform/coverage/
expansion.rs

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

rustc_mir_transform/coverage/
expansion.rs
