fe_mir/analysis/
cfg.rs

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}