rustc_mir_transform/coverage/
expansion.rs

1use rustc_data_structures::fx::{FxIndexMap, FxIndexSet, IndexEntry};
2use rustc_middle::mir::coverage::BasicCoverageBlock;
3use rustc_span::{ExpnId, ExpnKind, Span};
4
5#[derive(Clone, Copy, Debug)]
6pub(crate) struct SpanWithBcb {
7    pub(crate) span: Span,
8    pub(crate) bcb: BasicCoverageBlock,
9}
10
11#[derive(Debug)]
12pub(crate) struct ExpnTree {
13    nodes: FxIndexMap<ExpnId, ExpnNode>,
14}
15
16impl 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)]
55pub(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
75impl 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,
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.
97pub(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}
rustc_mir_transform/coverage/expansion.rs

rustc_mir_transform/coverage/
expansion.rs
