sonatina_codegen/
cfg.rs

1use std::collections::BTreeSet;
2
3use cranelift_entity::{packed_option::PackedOption, SecondaryMap};
4
5use sonatina_ir::{Block, Function, Insn};
6
7#[derive(Default, Debug, Clone, PartialEq, Eq)]
8pub struct ControlFlowGraph {
9    entry: PackedOption<Block>,
10    blocks: SecondaryMap<Block, BlockNode>,
11    pub(super) exits: smallvec::SmallVec<[Block; 8]>,
12}
13
14impl ControlFlowGraph {
15    pub fn new() -> Self {
16        Self::default()
17    }
18
19    pub fn compute(&mut self, func: &Function) {
20        self.clear();
21
22        self.entry = func.layout.entry_block().into();
23
24        for block in func.layout.iter_block() {
25            if let Some(last_insn) = func.layout.last_insn_of(block) {
26                self.analyze_insn(func, last_insn);
27            }
28        }
29    }
30
31    pub fn preds_of(&self, block: Block) -> impl Iterator<Item = &Block> {
32        self.blocks[block].preds()
33    }
34
35    pub fn succs_of(&self, block: Block) -> impl Iterator<Item = &Block> {
36        self.blocks[block].succs()
37    }
38
39    pub fn pred_num_of(&self, block: Block) -> usize {
40        self.blocks[block].pred_num()
41    }
42
43    pub fn succ_num_of(&self, block: Block) -> usize {
44        self.blocks[block].succ_num()
45    }
46
47    pub fn entry(&self) -> Option<Block> {
48        self.entry.expand()
49    }
50
51    pub fn post_order(&self) -> CfgPostOrder {
52        CfgPostOrder::new(self)
53    }
54
55    pub fn add_edge(&mut self, from: Block, to: Block) {
56        self.blocks[to].push_pred(from);
57        self.blocks[from].push_succ(to);
58    }
59
60    pub fn remove_edge(&mut self, from: Block, to: Block) {
61        self.blocks[to].remove_pred(from);
62        self.blocks[from].remove_succ(to);
63    }
64
65    pub(super) fn reverse_edges(&mut self, new_entry: Block, new_exits: &[Block]) {
66        for node in self.blocks.values_mut() {
67            node.reverse_edge();
68        }
69        self.entry = new_entry.into();
70        self.exits = new_exits.into();
71    }
72
73    pub fn clear(&mut self) {
74        self.entry = None.into();
75        self.blocks.clear();
76        self.exits.clear();
77    }
78
79    fn analyze_insn(&mut self, func: &Function, insn: Insn) {
80        if func.dfg.is_return(insn) {
81            let exit = func.layout.insn_block(insn);
82            self.exits.push(exit);
83        }
84
85        for dest in func.dfg.analyze_branch(insn).iter_dests() {
86            let block = func.layout.insn_block(insn);
87            self.add_edge(block, dest);
88        }
89    }
90}
91
92#[derive(Default, Clone, Debug, PartialEq, Eq)]
93struct BlockNode {
94    preds: BTreeSet<Block>,
95    succs: BTreeSet<Block>,
96}
97
98impl BlockNode {
99    fn push_pred(&mut self, pred: Block) {
100        self.preds.insert(pred);
101    }
102
103    fn push_succ(&mut self, succ: Block) {
104        self.succs.insert(succ);
105    }
106
107    fn remove_pred(&mut self, pred: Block) {
108        self.preds.remove(&pred);
109    }
110
111    fn remove_succ(&mut self, succ: Block) {
112        self.succs.remove(&succ);
113    }
114
115    fn preds(&self) -> impl Iterator<Item = &Block> {
116        self.preds.iter()
117    }
118
119    fn succs(&self) -> impl Iterator<Item = &Block> {
120        self.succs.iter()
121    }
122
123    fn pred_num(&self) -> usize {
124        self.preds.len()
125    }
126
127    fn succ_num(&self) -> usize {
128        self.succs.len()
129    }
130
131    fn reverse_edge(&mut self) {
132        std::mem::swap(&mut self.preds, &mut self.succs);
133    }
134}
135
136pub struct CfgPostOrder<'a> {
137    cfg: &'a ControlFlowGraph,
138    node_state: SecondaryMap<Block, NodeState>,
139    stack: Vec<Block>,
140}
141
142impl<'a> CfgPostOrder<'a> {
143    fn new(cfg: &'a ControlFlowGraph) -> Self {
144        let mut stack = Vec::new();
145
146        if let Some(entry) = cfg.entry() {
147            stack.push(entry);
148        }
149
150        Self {
151            cfg,
152            node_state: SecondaryMap::default(),
153            stack,
154        }
155    }
156}
157
158impl<'a> Iterator for CfgPostOrder<'a> {
159    type Item = Block;
160
161    fn next(&mut self) -> Option<Block> {
162        while let Some(&block) = self.stack.last() {
163            if self.node_state[block].is_unvisited() {
164                self.node_state[block].set_visited();
165                for &pred in self.cfg.succs_of(block) {
166                    if self.node_state[pred].is_unvisited() {
167                        self.stack.push(pred);
168                    }
169                }
170            } else {
171                self.stack.pop().unwrap();
172                if !self.node_state[block].has_finished() {
173                    self.node_state[block].set_finished();
174                    return Some(block);
175                }
176            }
177        }
178
179        None
180    }
181}
182
183#[derive(Default, Debug, Clone, Copy)]
184struct NodeState(u8);
185
186impl NodeState {
187    fn is_unvisited(self) -> bool {
188        self.0 == 0
189    }
190
191    fn has_finished(self) -> bool {
192        self.0 == 2
193    }
194
195    fn set_visited(&mut self) {
196        self.0 = 1;
197    }
198
199    fn set_finished(&mut self) {
200        self.0 = 2;
201    }
202}
sonatina_codegen/cfg.rs

sonatina_codegen/
cfg.rs
