1use fxhash::FxHashMap;
2
3use crate::ir::{BasicBlockId, FunctionBody, InstId};
4
5#[derive(Debug, Clone, PartialEq, Eq)]
6pub struct ControlFlowGraph {
7 entry: BasicBlockId,
8 blocks: FxHashMap<BasicBlockId, BlockNode>,
9 pub(super) exits: Vec<BasicBlockId>,
10}
11
12impl ControlFlowGraph {
13 pub fn compute(func: &FunctionBody) -> Self {
14 let entry = func.order.entry();
15 let mut cfg = Self {
16 entry,
17 blocks: FxHashMap::default(),
18 exits: vec![],
19 };
20
21 for block in func.order.iter_block() {
22 let terminator = func
23 .order
24 .terminator(&func.store, block)
25 .expect("a block must have terminator");
26 cfg.analyze_terminator(func, terminator);
27 }
28
29 cfg
30 }
31
32 pub fn entry(&self) -> BasicBlockId {
33 self.entry
34 }
35
36 pub fn preds(&self, block: BasicBlockId) -> &[BasicBlockId] {
37 self.blocks[&block].preds()
38 }
39
40 pub fn succs(&self, block: BasicBlockId) -> &[BasicBlockId] {
41 self.blocks[&block].succs()
42 }
43
44 pub fn post_order(&self) -> CfgPostOrder {
45 CfgPostOrder::new(self)
46 }
47
48 pub(super) fn add_edge(&mut self, from: BasicBlockId, to: BasicBlockId) {
49 self.node_mut(to).push_pred(from);
50 self.node_mut(from).push_succ(to);
51 }
52
53 pub(super) fn reverse_edge(&mut self, new_entry: BasicBlockId, new_exits: Vec<BasicBlockId>) {
54 for (_, block) in self.blocks.iter_mut() {
55 block.reverse_edge()
56 }
57
58 self.entry = new_entry;
59 self.exits = new_exits;
60 }
61
62 fn analyze_terminator(&mut self, func: &FunctionBody, terminator: InstId) {
63 let block = func.order.inst_block(terminator);
64 let branch_info = func.store.branch_info(terminator);
65 if branch_info.is_not_a_branch() {
66 self.node_mut(block);
67 self.exits.push(block)
68 } else {
69 for dest in branch_info.block_iter() {
70 self.add_edge(block, dest)
71 }
72 }
73 }
74
75 fn node_mut(&mut self, block: BasicBlockId) -> &mut BlockNode {
76 self.blocks.entry(block).or_default()
77 }
78}
79
80#[derive(Default, Clone, Debug, PartialEq, Eq)]
81struct BlockNode {
82 preds: Vec<BasicBlockId>,
83 succs: Vec<BasicBlockId>,
84}
85
86impl BlockNode {
87 fn push_pred(&mut self, pred: BasicBlockId) {
88 self.preds.push(pred);
89 }
90
91 fn push_succ(&mut self, succ: BasicBlockId) {
92 self.succs.push(succ);
93 }
94
95 fn preds(&self) -> &[BasicBlockId] {
96 &self.preds
97 }
98
99 fn succs(&self) -> &[BasicBlockId] {
100 &self.succs
101 }
102
103 fn reverse_edge(&mut self) {
104 std::mem::swap(&mut self.preds, &mut self.succs)
105 }
106}
107
108pub struct CfgPostOrder<'a> {
109 cfg: &'a ControlFlowGraph,
110 node_state: FxHashMap<BasicBlockId, NodeState>,
111 stack: Vec<BasicBlockId>,
112}
113
114impl<'a> CfgPostOrder<'a> {
115 fn new(cfg: &'a ControlFlowGraph) -> Self {
116 let stack = vec![cfg.entry()];
117
118 Self {
119 cfg,
120 node_state: FxHashMap::default(),
121 stack,
122 }
123 }
124}
125
126impl<'a> Iterator for CfgPostOrder<'a> {
127 type Item = BasicBlockId;
128
129 fn next(&mut self) -> Option<BasicBlockId> {
130 while let Some(&block) = self.stack.last() {
131 let node_state = self.node_state.entry(block).or_default();
132 if *node_state == NodeState::Unvisited {
133 *node_state = NodeState::Visited;
134 for &succ in self.cfg.succs(block) {
135 let pred_state = self.node_state.entry(succ).or_default();
136 if *pred_state == NodeState::Unvisited {
137 self.stack.push(succ);
138 }
139 }
140 } else {
141 self.stack.pop().unwrap();
142 if *node_state != NodeState::Finished {
143 *node_state = NodeState::Finished;
144 return Some(block);
145 }
146 }
147 }
148
149 None
150 }
151}
152
153#[derive(Debug, Clone, Copy, PartialEq, Eq)]
154enum NodeState {
155 Unvisited,
156 Visited,
157 Finished,
158}
159
160impl Default for NodeState {
161 fn default() -> Self {
162 Self::Unvisited
163 }
164}