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

rapx/graphs/
cfg.rs
