rapx/graphs/
cfg.rs

1use crate::compat::FxHashSet;
2use crate::graphs::scc::{Scc, SccExit, SccInfo};
3use rustc_middle::{
4    mir::{BasicBlock, Terminator},
5    ty::TyCtxt,
6};
7use rustc_span::def_id::DefId;
8
9/// Reusable CFG block structure shared by analyses built over MIR.
10///
11/// Each `CfgBlock` corresponds to a MIR basic block and stores:
12/// - its block index,
13/// - whether it is a cleanup block,
14/// - its outgoing CFG edges,
15/// - and SCC metadata for loop/cycle-aware traversal.
16///
17/// Terminator data is intentionally not cached here; use
18/// [`ControlFlowGraph::terminator`] to retrieve it on demand from MIR.
19#[derive(Debug, Clone)]
20pub struct CfgBlock {
21    /// Index of this block in the CFG block list (global, unique across an
22    /// inlined multi-function CFG).
23    pub index: usize,
24    /// The function whose MIR contains this block.
25    pub def_id: DefId,
26    /// This block's index within `def_id`'s MIR basic blocks.
27    pub local_index: usize,
28    /// Whether this block belongs to MIR cleanup/unwind control flow.
29    pub is_cleanup: bool,
30    /// Outgoing successor block indices.
31    pub next: FxHashSet<usize>,
32    /// SCC information for this block.
33    ///
34    /// For non-root blocks inside an SCC, `enter` points to the SCC root.
35    /// For SCC roots, this field also stores member nodes, exits, and back edges.
36    pub scc: SccInfo,
37}
38
39impl CfgBlock {
40    /// Create a new CFG block for `def_id` at MIR basic block `index`.
41    pub fn new(def_id: DefId, index: usize, is_cleanup: bool) -> Self {
42        Self {
43            index,
44            def_id,
45            local_index: index,
46            is_cleanup,
47            next: FxHashSet::default(),
48            scc: SccInfo::new(index),
49        }
50    }
51
52    /// Add a successor edge from this block to `index`.
53    pub fn add_next(&mut self, index: usize) {
54        self.next.insert(index);
55    }
56}
57
58/// Generic MIR control-flow graph container.
59///
60/// This structure intentionally keeps only generic CFG shape and SCC metadata.
61#[derive(Clone)]
62pub struct ControlFlowGraph<'tcx> {
63    /// Definition being analyzed.
64    pub def_id: DefId,
65    /// Type context from the Rust compiler.
66    pub tcx: TyCtxt<'tcx>,
67    /// All CFG blocks for the current body.
68    pub blocks: Vec<CfgBlock>,
69}
70
71impl<'tcx> ControlFlowGraph<'tcx> {
72    /// Construct a control-flow graph wrapper from prebuilt blocks.
73    pub fn new(def_id: DefId, tcx: TyCtxt<'tcx>, blocks: Vec<CfgBlock>) -> Self {
74        Self {
75            def_id,
76            tcx,
77            blocks,
78        }
79    }
80
81    /// Get an immutable reference to a block by index.
82    pub fn block(&self, index: usize) -> &CfgBlock {
83        &self.blocks[index]
84    }
85
86    /// Get a mutable reference to a block by index.
87    pub fn block_mut(&mut self, index: usize) -> &mut CfgBlock {
88        &mut self.blocks[index]
89    }
90
91    /// Retrieve the MIR terminator for the block at `index` on demand.
92    ///
93    /// Returns `None` only for blocks whose terminator has not yet been
94    /// elaborated (which is unusual for optimized MIR).
95    pub fn terminator(&self, index: usize) -> Option<&Terminator<'tcx>> {
96        let block = self.blocks.get(index)?;
97        let body = self.tcx.optimized_mir(block.def_id);
98        body.basic_blocks
99            .get(BasicBlock::from(block.local_index))
100            .and_then(|bb| bb.terminator.as_ref())
101    }
102}
103
104/// Record exits from the SCC root to blocks outside the SCC.
105fn record_root_exits<'tcx>(
106    graph: &mut ControlFlowGraph<'tcx>,
107    root: usize,
108    scc_components: &[usize],
109) {
110    let nexts = graph.block(root).next.clone();
111    for next in nexts {
112        if !scc_components.contains(&next) {
113            graph
114                .block_mut(root)
115                .scc
116                .exits
117                .insert(SccExit::new(root, next));
118        }
119    }
120}
121
122/// Record membership, exit edges, and back edges for all non-root SCC members.
123fn record_member_nodes<'tcx>(
124    graph: &mut ControlFlowGraph<'tcx>,
125    root: usize,
126    scc_components: &[usize],
127) {
128    for &node in &scc_components[1..] {
129        // Record membership under the root SCC.
130        graph.block_mut(root).scc.nodes.insert(node);
131        // Make each member point to the SCC root.
132        graph.block_mut(node).scc.enter = root;
133
134        let nexts = graph.block(node).next.clone();
135        for next in nexts {
136            // Any edge leaving the SCC is an SCC exit.
137            if !scc_components.contains(&next) {
138                graph
139                    .block_mut(root)
140                    .scc
141                    .exits
142                    .insert(SccExit::new(node, next));
143            }
144            // Any edge back to the root is tracked as a back edge.
145            if next == root && !graph.block(root).scc.backedges.contains(&(node, root)) {
146                graph.block_mut(root).scc.backedges.push((node, root));
147            }
148        }
149    }
150}
151
152/// Re-run SCC discovery with a temporarily reduced graph to discover nested SCCs.
153///
154/// Isolates the SCC rooted at `root` by:
155/// 1. Redirecting block 0 to point only to `root`.
156/// 2. Removing back edges to `root`.
157/// 3. Cutting all outgoing edges from SCC exit targets.
158///
159/// After re-running SCC discovery, all edges are restored.
160fn rerun_scc_in_isolation<'tcx>(
161    graph: &mut ControlFlowGraph<'tcx>,
162    root: usize,
163) {
164    let scc_exits = graph.block(root).scc.exits.clone();
165    let backedges = graph.block(root).scc.backedges.clone();
166    let mut backups: Vec<(usize, FxHashSet<usize>)> = Vec::new();
167
168    // Temporarily redirect entry block 0 to this SCC root only.
169    // This helps isolate SCC structure for the recursive `find_scc()` call.
170    let block0 = graph.block_mut(0);
171    backups.push((0, block0.next.clone()));
172    block0.next.clear();
173    block0.next.insert(root);
174
175    // Temporarily remove back edges to the root.
176    for &(node, target) in &backedges {
177        if target != root {
178            continue;
179        }
180        let block = graph.block_mut(node);
181        backups.push((node, block.next.clone()));
182        block.next.remove(&root);
183    }
184
185    // Temporarily cut all outgoing edges from SCC exit targets.
186    for exit in &scc_exits {
187        let block_to = graph.block_mut(exit.to);
188        backups.push((exit.to, block_to.next.clone()));
189        block_to.next.clear();
190    }
191
192    // Re-run SCC discovery on the transformed graph.
193    graph.find_scc();
194
195    // Restore all modified edges.
196    for (idx, saved_next) in backups {
197        graph.block_mut(idx).next = saved_next;
198    }
199}
200
201/// Handle a newly discovered SCC: mark the root, collect membership and edge metadata,
202/// then re-run SCC discovery on an isolated subgraph to populate nested SCC structure.
203fn scc_handler<'tcx>(graph: &mut ControlFlowGraph<'tcx>, root: usize, scc_components: &[usize]) {
204    rap_debug!(
205        "Scc found: root = {}, components = {:?}",
206        root,
207        scc_components
208    );
209
210    // The SCC root always points to itself.
211    graph.block_mut(root).scc.enter = root;
212
213    // A single-node SCC is trivial; nothing else needs to be recorded.
214    if scc_components.len() <= 1 {
215        return;
216    }
217
218    record_root_exits(graph, root, scc_components);
219    record_member_nodes(graph, root, scc_components);
220
221    rap_debug!("Scc Info: {:?}", graph.block(root).scc);
222    rerun_scc_in_isolation(graph, root);
223}
224
225impl<'tcx> Scc for ControlFlowGraph<'tcx> {
226    /// Callback invoked when an SCC is discovered.
227    fn on_scc_found(&mut self, root: usize, scc_components: &[usize]) {
228        scc_handler(self, root, scc_components);
229    }
230
231    /// Return the outgoing successors of a node.
232    fn get_next(&mut self, root: usize) -> FxHashSet<usize> {
233        self.block(root).next.clone()
234    }
235
236    /// Return the total number of CFG blocks.
237    fn get_size(&mut self) -> usize {
238        self.blocks.len()
239    }
240}
rapx/graphs/cfg.rs

rapx/graphs/
cfg.rs
