1use crate::types::BlockId;
2use std::collections::{HashMap, HashSet, VecDeque};
3
4use super::dominators::DominatorTree;
5use super::paths::Path;
6use super::types::Loop;
7
8#[derive(Clone, Debug)]
9pub struct TestCfg {
10 pub entry: BlockId,
11 pub exits: HashSet<BlockId>,
12 pub error_blocks: HashSet<BlockId>,
13 pub successors: HashMap<BlockId, Vec<BlockId>>,
14 pub predecessors: HashMap<BlockId, Vec<BlockId>>,
15}
16
17impl TestCfg {
18 pub fn new(entry: BlockId) -> Self {
19 Self {
20 entry,
21 exits: HashSet::new(),
22 error_blocks: HashSet::new(),
23 successors: HashMap::new(),
24 predecessors: HashMap::new(),
25 }
26 }
27
28 pub fn add_edge(&mut self, from: BlockId, to: BlockId) -> &mut Self {
29 self.successors.entry(from).or_default().push(to);
30 self.predecessors.entry(to).or_default().push(from);
31 self
32 }
33
34 pub fn add_exit(&mut self, block: BlockId) -> &mut Self {
35 self.exits.insert(block);
36 self
37 }
38
39 pub fn add_error(&mut self, block: BlockId) -> &mut Self {
40 self.error_blocks.insert(block);
41 self
42 }
43
44 pub fn chain(start: i64, count: usize) -> Self {
45 let mut cfg = Self::new(BlockId(start));
46 for i in start..(start + count as i64 - 1) {
47 cfg.add_edge(BlockId(i), BlockId(i + 1));
48 }
49 cfg.add_exit(BlockId(start + count as i64 - 1));
50 cfg
51 }
52
53 pub fn if_else() -> Self {
54 let mut cfg = Self::new(BlockId(0));
55 cfg.add_edge(BlockId(0), BlockId(1))
56 .add_edge(BlockId(0), BlockId(2))
57 .add_edge(BlockId(1), BlockId(3))
58 .add_edge(BlockId(2), BlockId(3))
59 .add_exit(BlockId(3));
60 cfg
61 }
62
63 pub fn simple_loop() -> Self {
64 let mut cfg = Self::new(BlockId(0));
65 cfg.add_edge(BlockId(0), BlockId(1))
66 .add_edge(BlockId(1), BlockId(2))
67 .add_edge(BlockId(2), BlockId(1))
68 .add_edge(BlockId(1), BlockId(3))
69 .add_exit(BlockId(3));
70 cfg
71 }
72
73 pub fn enumerate_paths(&self) -> Vec<Path> {
74 let mut paths = Vec::new();
75 let mut current = vec![self.entry];
76 let mut visited = HashSet::new();
77 self.dfs(&mut paths, &mut current, &mut visited, self.entry);
78 paths
79 }
80
81 fn dfs(
82 &self,
83 paths: &mut Vec<Path>,
84 current: &mut Vec<BlockId>,
85 visited: &mut HashSet<BlockId>,
86 block: BlockId,
87 ) {
88 if self.exits.contains(&block) {
89 paths.push(Path::new(current.clone()));
90 return;
91 }
92 if visited.contains(&block) {
93 return;
94 }
95 visited.insert(block);
96 if let Some(successors) = self.successors.get(&block) {
97 for &succ in successors {
98 current.push(succ);
99 self.dfs(paths, current, visited, succ);
100 current.pop();
101 }
102 }
103 visited.remove(&block);
104 }
105
106 pub fn compute_dominators(&self) -> DominatorTree {
107 let mut blocks: HashSet<BlockId> = HashSet::new();
108 blocks.insert(self.entry);
109 for (from, tos) in &self.successors {
110 blocks.insert(*from);
111 for to in tos {
112 blocks.insert(*to);
113 }
114 }
115
116 if blocks.is_empty() {
117 return DominatorTree::new(self.entry);
118 }
119
120 let block_list: Vec<BlockId> = blocks.iter().copied().collect();
121 let mut dom: HashMap<BlockId, HashSet<BlockId>> = HashMap::new();
122
123 for &block in &block_list {
124 if block == self.entry {
125 dom.insert(block, HashSet::from([self.entry]));
126 } else {
127 dom.insert(block, blocks.clone());
128 }
129 }
130
131 let mut changed = true;
132 while changed {
133 changed = false;
134 for &block in &block_list {
135 if block == self.entry {
136 continue;
137 }
138 let preds = self.predecessors.get(&block);
139 if preds.is_none() || preds.unwrap().is_empty() {
140 continue;
141 }
142 let mut new_dom: HashSet<BlockId> =
143 dom.get(&preds.unwrap()[0]).cloned().unwrap_or_default();
144 for pred in &preds.unwrap()[1..] {
145 if let Some(pred_dom) = dom.get(pred) {
146 new_dom = new_dom.intersection(pred_dom).copied().collect();
147 }
148 }
149 new_dom.insert(block);
150 if dom.get(&block) != Some(&new_dom) {
151 dom.insert(block, new_dom);
152 changed = true;
153 }
154 }
155 }
156
157 let mut idom: HashMap<BlockId, BlockId> = HashMap::new();
158 for &block in &block_list {
159 if block == self.entry {
160 continue;
161 }
162 if let Some(doms) = dom.get(&block) {
163 let mut best_candidate: Option<BlockId> = None;
164 let mut best_size = 0;
165
166 for &candidate in doms {
167 if candidate == block {
168 continue;
169 }
170 if let Some(candidate_doms) = dom.get(&candidate) {
171 if candidate_doms.len() > best_size {
172 best_size = candidate_doms.len();
173 best_candidate = Some(candidate);
174 }
175 }
176 }
177
178 if let Some(candidate) = best_candidate {
179 idom.insert(block, candidate);
180 }
181 }
182 }
183
184 DominatorTree {
185 root: self.entry,
186 dominators: idom,
187 }
188 }
189
190 pub fn detect_loops(&self) -> Vec<Loop> {
191 let dom = self.compute_dominators();
192 let mut loops = Vec::new();
193
194 for (from, tos) in &self.successors {
195 for to in tos {
196 if dom.dominates(*to, *from) {
197 let header = *to;
198 let mut loop_blocks = HashSet::new();
199 loop_blocks.insert(header);
200 let mut worklist = VecDeque::new();
201 worklist.push_back(*from);
202
203 while let Some(block) = worklist.pop_front() {
204 if loop_blocks.contains(&block) {
205 continue;
206 }
207 if dom.dominates(header, block) || block == header {
208 loop_blocks.insert(block);
209 if let Some(preds) = self.predecessors.get(&block) {
210 for &pred in preds {
211 if !loop_blocks.contains(&pred) {
212 worklist.push_back(pred);
213 }
214 }
215 }
216 }
217 }
218
219 let mut blocks: Vec<BlockId> =
220 loop_blocks.into_iter().filter(|&b| b != header).collect();
221 blocks.sort();
222 loops.push(Loop::with_depth(header, blocks, 0));
223 }
224 }
225 }
226
227 loops
228 }
229}
230
231#[cfg(test)]
232mod tests {
233 use super::*;
234
235 #[test]
236 fn test_test_cfg_chain() {
237 let cfg = TestCfg::chain(0, 5);
238
239 assert_eq!(cfg.entry, BlockId(0));
240 assert!(cfg.exits.contains(&BlockId(4)));
241
242 assert_eq!(cfg.successors.get(&BlockId(0)), Some(&vec![BlockId(1)]));
243 assert_eq!(cfg.successors.get(&BlockId(1)), Some(&vec![BlockId(2)]));
244 assert_eq!(cfg.successors.get(&BlockId(2)), Some(&vec![BlockId(3)]));
245 assert_eq!(cfg.successors.get(&BlockId(3)), Some(&vec![BlockId(4)]));
246 }
247
248 #[test]
249 fn test_test_cfg_if_else() {
250 let cfg = TestCfg::if_else();
251
252 assert_eq!(cfg.entry, BlockId(0));
253 assert!(cfg.exits.contains(&BlockId(3)));
254
255 let succ0 = cfg.successors.get(&BlockId(0)).unwrap();
256 assert!(succ0.contains(&BlockId(1)));
257 assert!(succ0.contains(&BlockId(2)));
258 assert_eq!(cfg.successors.get(&BlockId(1)), Some(&vec![BlockId(3)]));
259 assert_eq!(cfg.successors.get(&BlockId(2)), Some(&vec![BlockId(3)]));
260 }
261
262 #[test]
263 fn test_paths_simple_chain() {
264 let cfg = TestCfg::chain(0, 4);
265 let paths = cfg.enumerate_paths();
266
267 assert_eq!(paths.len(), 1);
268 assert_eq!(
269 paths[0].blocks,
270 vec![BlockId(0), BlockId(1), BlockId(2), BlockId(3)]
271 );
272 assert!(paths[0].is_normal());
273 }
274
275 #[test]
276 fn test_paths_if_else() {
277 let cfg = TestCfg::if_else();
278 let paths = cfg.enumerate_paths();
279
280 assert_eq!(paths.len(), 2);
281
282 assert_eq!(paths[0].entry(), Some(BlockId(0)));
283 assert_eq!(paths[0].exit(), Some(BlockId(3)));
284 assert_eq!(paths[1].entry(), Some(BlockId(0)));
285 assert_eq!(paths[1].exit(), Some(BlockId(3)));
286
287 let paths_set: HashSet<_> = paths.iter().map(|p| p.blocks.clone()).collect();
288
289 assert!(paths_set.contains(&vec![BlockId(0), BlockId(1), BlockId(3)]));
290 assert!(paths_set.contains(&vec![BlockId(0), BlockId(2), BlockId(3)]));
291 }
292
293 #[test]
294 fn test_dominators_chain() {
295 let cfg = TestCfg::chain(0, 5);
296 let dom = cfg.compute_dominators();
297
298 assert_eq!(dom.root, BlockId(0));
299 assert_eq!(dom.immediate_dominator(BlockId(1)), Some(BlockId(0)));
300 assert_eq!(dom.immediate_dominator(BlockId(2)), Some(BlockId(1)));
301 assert_eq!(dom.immediate_dominator(BlockId(3)), Some(BlockId(2)));
302 assert_eq!(dom.immediate_dominator(BlockId(4)), Some(BlockId(3)));
303 }
304
305 #[test]
306 fn test_dominators_if_else() {
307 let cfg = TestCfg::if_else();
308 let dom = cfg.compute_dominators();
309
310 assert!(dom.dominates(BlockId(0), BlockId(0)));
311 assert!(dom.dominates(BlockId(0), BlockId(1)));
312 assert!(dom.dominates(BlockId(0), BlockId(2)));
313 assert!(dom.dominates(BlockId(0), BlockId(3)));
314 assert_eq!(dom.immediate_dominator(BlockId(3)), Some(BlockId(0)));
315 }
316
317 #[test]
318 fn test_loops_simple_loop() {
319 let cfg = TestCfg::simple_loop();
320 let loops = cfg.detect_loops();
321
322 assert_eq!(loops.len(), 1);
323 assert_eq!(loops[0].header, BlockId(1));
324 assert!(!loops[0].blocks.is_empty());
325 assert!(loops[0].blocks.contains(&BlockId(2)));
326 }
327
328 #[test]
329 fn test_loops_none_in_chain() {
330 let cfg = TestCfg::chain(0, 5);
331 let loops = cfg.detect_loops();
332
333 assert_eq!(loops.len(), 0);
334 }
335
336 #[test]
337 fn test_loops_none_in_if_else() {
338 let cfg = TestCfg::if_else();
339 let loops = cfg.detect_loops();
340
341 assert_eq!(loops.len(), 0);
342 }
343}