sonatina_codegen/
domtree.rs

1//! This module contains dominantor tree related structs.
2//!
3//! The algorithm is based on Keith D. Cooper., Timothy J. Harvey., and Ken Kennedy.: A Simple, Fast Dominance Algorithm:  
4//! <https://www.cs.rice.edu/~keith/EMBED/dom.pdf>
5
6use std::collections::BTreeSet;
7
8use cranelift_entity::{packed_option::PackedOption, SecondaryMap};
9
10use sonatina_ir::Block;
11
12use super::cfg::ControlFlowGraph;
13
14#[derive(Default, Debug)]
15pub struct DomTree {
16    doms: SecondaryMap<Block, PackedOption<Block>>,
17    rpo: Vec<Block>,
18}
19
20impl DomTree {
21    pub fn new() -> Self {
22        Self::default()
23    }
24
25    pub fn clear(&mut self) {
26        self.doms.clear();
27        self.rpo.clear();
28    }
29
30    /// Returns the immediate dominator of the `block`.
31    /// Returns None if the `block` is unreachable from the entry block, or the `block` is the entry block itself.
32    pub fn idom_of(&self, block: Block) -> Option<Block> {
33        if self.rpo[0] == block {
34            return None;
35        }
36        self.doms[block].expand()
37    }
38
39    /// Returns `true` if block1 strictly dominates block2.
40    pub fn strictly_dominates(&self, block1: Block, block2: Block) -> bool {
41        let mut current_block = block2;
42        while let Some(block) = self.idom_of(current_block) {
43            if block == block1 {
44                return true;
45            }
46            current_block = block;
47        }
48
49        false
50    }
51
52    /// Returns `true` if block1 dominates block2.
53    pub fn dominates(&self, block1: Block, block2: Block) -> bool {
54        if block1 == block2 {
55            return true;
56        }
57
58        self.strictly_dominates(block1, block2)
59    }
60
61    pub fn compute(&mut self, cfg: &ControlFlowGraph) {
62        self.clear();
63
64        self.rpo = cfg.post_order().collect();
65        self.rpo.reverse();
66
67        let block_num = self.rpo.len();
68
69        if self.doms.capacity() < block_num {
70            self.doms = SecondaryMap::with_capacity(block_num);
71        } else {
72            self.doms.clear();
73        }
74
75        let mut rpo_nums = SecondaryMap::with_capacity(block_num);
76        for (i, &block) in self.rpo.iter().enumerate() {
77            rpo_nums[block] = (block_num - i) as u32;
78        }
79
80        match self.rpo.first() {
81            Some(&entry) => self.doms[entry] = entry.into(),
82            None => return,
83        }
84
85        let mut changed = true;
86        while changed {
87            changed = false;
88            for &block in self.rpo.iter().skip(1) {
89                let processed_pred =
90                    match cfg.preds_of(block).find(|&&pred| self.doms[pred].is_some()) {
91                        Some(pred) => *pred,
92                        _ => continue,
93                    };
94                let mut new_dom = processed_pred;
95
96                for &pred in cfg.preds_of(block) {
97                    if pred != processed_pred && self.doms[pred].is_some() {
98                        new_dom = self.intersect(new_dom, pred, &rpo_nums);
99                    }
100                }
101                if Some(new_dom) != self.doms[block].expand() {
102                    changed = true;
103                    self.doms[block] = new_dom.into();
104                }
105            }
106        }
107    }
108
109    /// Compute dominance frontiers of each blocks.
110    pub fn compute_df(&self, cfg: &ControlFlowGraph) -> DFSet {
111        let mut df = DFSet::default();
112
113        for &block in &self.rpo {
114            if cfg.pred_num_of(block) < 2 {
115                continue;
116            }
117            for pred in cfg.preds_of(block) {
118                let mut runner = *pred;
119                while PackedOption::from(runner) != self.doms[block] && self.is_reachable(runner) {
120                    df.0[runner].insert(block);
121                    runner = self.doms[runner].unwrap();
122                }
123            }
124        }
125
126        df
127    }
128
129    /// Returns `true` if block is reachable from the entry block.
130    pub fn is_reachable(&self, block: Block) -> bool {
131        self.idom_of(block).is_some()
132    }
133
134    /// Returns blocks in RPO.
135    pub fn rpo(&self) -> &[Block] {
136        &self.rpo
137    }
138
139    fn intersect(
140        &self,
141        mut b1: Block,
142        mut b2: Block,
143        rpo_nums: &SecondaryMap<Block, u32>,
144    ) -> Block {
145        while b1 != b2 {
146            while rpo_nums[b1] < rpo_nums[b2] {
147                b1 = self.doms[b1].unwrap();
148            }
149            while rpo_nums[b2] < rpo_nums[b1] {
150                b2 = self.doms[b2].unwrap();
151            }
152        }
153
154        b1
155    }
156}
157
158/// Dominance frontiers of each blocks.
159#[derive(Default, Debug)]
160pub struct DFSet(SecondaryMap<Block, BTreeSet<Block>>);
161
162impl DFSet {
163    pub fn frontiers(&self, block: Block) -> impl Iterator<Item = &Block> {
164        self.0[block].iter()
165    }
166
167    pub fn in_frontier_of(&self, block: Block, of: Block) -> bool {
168        self.0[of].contains(&block)
169    }
170
171    pub fn frontier_num_of(&self, of: Block) -> usize {
172        self.0[of].len()
173    }
174
175    pub fn clear(&mut self) {
176        self.0.clear()
177    }
178}
179
180#[derive(Default)]
181pub struct DominatorTreeTraversable {
182    children: SecondaryMap<Block, Vec<Block>>,
183}
184
185impl DominatorTreeTraversable {
186    pub fn compute(&mut self, domtree: &DomTree) {
187        for &block in domtree.rpo() {
188            if let Some(idom) = domtree.idom_of(block) {
189                self.children[idom].push(block)
190            }
191        }
192    }
193
194    pub fn children_of(&self, block: Block) -> &[Block] {
195        &self.children[block]
196    }
197
198    pub fn clear(&mut self) {
199        self.children.clear();
200    }
201}
202
203#[cfg(test)]
204mod tests {
205    #![allow(clippy::many_single_char_names)]
206
207    use super::*;
208
209    use sonatina_ir::{builder::test_util::*, Function, Type};
210
211    fn calc_dom(func: &Function) -> (DomTree, DFSet) {
212        let mut cfg = ControlFlowGraph::default();
213        cfg.compute(func);
214        let mut dom_tree = DomTree::default();
215        dom_tree.compute(&cfg);
216        let df = dom_tree.compute_df(&cfg);
217        (dom_tree, df)
218    }
219
220    fn test_df(df: &DFSet, of: Block, frontiers: &[Block]) -> bool {
221        if df.frontier_num_of(of) != frontiers.len() {
222            return false;
223        }
224
225        for &block in frontiers {
226            if !df.in_frontier_of(block, of) {
227                return false;
228            }
229        }
230        true
231    }
232
233    #[test]
234    fn dom_tree_if_else() {
235        let mut test_module_builder = TestModuleBuilder::new();
236        let mut builder = test_module_builder.func_builder(&[], Type::Void);
237
238        let entry_block = builder.append_block();
239        let then_block = builder.append_block();
240        let else_block = builder.append_block();
241        let merge_block = builder.append_block();
242
243        builder.switch_to_block(entry_block);
244        let v0 = builder.make_imm_value(true);
245        builder.br(v0, else_block, then_block);
246
247        builder.switch_to_block(then_block);
248        builder.jump(merge_block);
249
250        builder.switch_to_block(else_block);
251        builder.jump(merge_block);
252
253        builder.switch_to_block(merge_block);
254        builder.ret(None);
255
256        builder.seal_all();
257        let func_ref = builder.finish();
258
259        let module = test_module_builder.build();
260        let func = &module.funcs[func_ref];
261
262        let (dom_tree, df) = calc_dom(func);
263        assert_eq!(dom_tree.idom_of(entry_block), None);
264        assert_eq!(dom_tree.idom_of(then_block), Some(entry_block));
265        assert_eq!(dom_tree.idom_of(else_block), Some(entry_block));
266        assert_eq!(dom_tree.idom_of(merge_block), Some(entry_block));
267
268        assert!(test_df(&df, entry_block, &[]));
269        assert!(test_df(&df, then_block, &[merge_block]));
270        assert!(test_df(&df, else_block, &[merge_block]));
271        assert!(test_df(&df, merge_block, &[]));
272    }
273
274    #[test]
275    fn unreachable_edge() {
276        let mut test_module_builder = TestModuleBuilder::new();
277        let mut builder = test_module_builder.func_builder(&[], Type::Void);
278
279        let a = builder.append_block();
280        let b = builder.append_block();
281        let c = builder.append_block();
282        let d = builder.append_block();
283        let e = builder.append_block();
284
285        builder.switch_to_block(a);
286        let v0 = builder.make_imm_value(true);
287        builder.br(v0, b, c);
288
289        builder.switch_to_block(b);
290        builder.jump(e);
291
292        builder.switch_to_block(c);
293        builder.jump(e);
294
295        builder.switch_to_block(d);
296        builder.jump(e);
297
298        builder.switch_to_block(e);
299        builder.ret(None);
300
301        builder.seal_all();
302        let func_ref = builder.finish();
303
304        let module = test_module_builder.build();
305        let func = &module.funcs[func_ref];
306
307        let (dom_tree, df) = calc_dom(func);
308        assert_eq!(dom_tree.idom_of(a), None);
309        assert_eq!(dom_tree.idom_of(b), Some(a));
310        assert_eq!(dom_tree.idom_of(c), Some(a));
311        assert_eq!(dom_tree.idom_of(d), None);
312        assert!(!dom_tree.is_reachable(d));
313        assert_eq!(dom_tree.idom_of(e), Some(a));
314
315        assert!(test_df(&df, a, &[]));
316        assert!(test_df(&df, b, &[e]));
317        assert!(test_df(&df, c, &[e]));
318        assert!(test_df(&df, d, &[]));
319        assert!(test_df(&df, e, &[]));
320    }
321
322    #[test]
323    fn dom_tree_complex() {
324        let mut test_module_builder = TestModuleBuilder::new();
325        let mut builder = test_module_builder.func_builder(&[], Type::Void);
326
327        let a = builder.append_block();
328        let b = builder.append_block();
329        let c = builder.append_block();
330        let d = builder.append_block();
331        let e = builder.append_block();
332        let f = builder.append_block();
333        let g = builder.append_block();
334        let h = builder.append_block();
335        let i = builder.append_block();
336        let j = builder.append_block();
337        let k = builder.append_block();
338        let l = builder.append_block();
339        let m = builder.append_block();
340
341        builder.switch_to_block(a);
342        let v0 = builder.make_imm_value(true);
343        builder.br(v0, c, b);
344
345        builder.switch_to_block(b);
346        builder.br(v0, g, d);
347
348        builder.switch_to_block(c);
349        builder.br(v0, h, e);
350
351        builder.switch_to_block(d);
352        builder.br(v0, g, f);
353
354        builder.switch_to_block(e);
355        builder.br(v0, h, c);
356
357        builder.switch_to_block(f);
358        builder.br(v0, k, i);
359
360        builder.switch_to_block(g);
361        builder.jump(j);
362
363        builder.switch_to_block(h);
364        builder.jump(m);
365
366        builder.switch_to_block(i);
367        builder.jump(l);
368
369        builder.switch_to_block(j);
370        builder.jump(i);
371
372        builder.switch_to_block(k);
373        builder.jump(l);
374
375        builder.switch_to_block(l);
376        builder.br(v0, m, b);
377
378        builder.switch_to_block(m);
379        builder.ret(None);
380
381        builder.seal_all();
382        let func_ref = builder.finish();
383
384        let module = test_module_builder.build();
385        let func = &module.funcs[func_ref];
386
387        let (dom_tree, df) = calc_dom(func);
388        assert_eq!(dom_tree.idom_of(a), None);
389        assert_eq!(dom_tree.idom_of(b), Some(a));
390        assert_eq!(dom_tree.idom_of(c), Some(a));
391        assert_eq!(dom_tree.idom_of(d), Some(b));
392        assert_eq!(dom_tree.idom_of(e), Some(c));
393        assert_eq!(dom_tree.idom_of(f), Some(d));
394        assert_eq!(dom_tree.idom_of(g), Some(b));
395        assert_eq!(dom_tree.idom_of(h), Some(c));
396        assert_eq!(dom_tree.idom_of(i), Some(b));
397        assert_eq!(dom_tree.idom_of(j), Some(g));
398        assert_eq!(dom_tree.idom_of(k), Some(f));
399
400        assert!(test_df(&df, a, &[]));
401        assert!(test_df(&df, b, &[b, m]));
402        assert!(test_df(&df, c, &[c, m]));
403        assert!(test_df(&df, d, &[g, i, l]));
404        assert!(test_df(&df, e, &[c, h]));
405        assert!(test_df(&df, f, &[i, l]));
406        assert!(test_df(&df, g, &[i]));
407        assert!(test_df(&df, h, &[m]));
408        assert!(test_df(&df, i, &[l]));
409        assert!(test_df(&df, j, &[i]));
410        assert!(test_df(&df, k, &[l]));
411        assert!(test_df(&df, l, &[b, m]));
412        assert!(test_df(&df, m, &[]));
413    }
414
415    #[test]
416    fn dom_tree_br_table() {
417        let mut test_module_builder = TestModuleBuilder::new();
418        let mut builder = test_module_builder.func_builder(&[], Type::Void);
419
420        let a = builder.append_block();
421        let b = builder.append_block();
422        let c = builder.append_block();
423        let d = builder.append_block();
424        let e = builder.append_block();
425        let f = builder.append_block();
426
427        builder.switch_to_block(a);
428        let v0 = builder.make_imm_value(0i32);
429        let v1 = builder.make_imm_value(1i32);
430        let v2 = builder.make_imm_value(2i32);
431        builder.br_table(v0, Some(b), &[(v1, c), (v2, d)]);
432
433        builder.switch_to_block(b);
434        let v3 = builder.make_imm_value(true);
435        builder.br(v3, a, e);
436
437        builder.switch_to_block(c);
438        builder.jump(f);
439
440        builder.switch_to_block(d);
441        builder.jump(f);
442
443        builder.switch_to_block(e);
444        builder.ret(None);
445
446        builder.switch_to_block(f);
447        builder.ret(None);
448
449        builder.seal_all();
450        let func_ref = builder.finish();
451
452        let module = test_module_builder.build();
453        let func = &module.funcs[func_ref];
454
455        let (dom_tree, df) = calc_dom(func);
456        assert_eq!(dom_tree.idom_of(a), None);
457        assert_eq!(dom_tree.idom_of(b), Some(a));
458        assert_eq!(dom_tree.idom_of(c), Some(a));
459        assert_eq!(dom_tree.idom_of(d), Some(a));
460        assert_eq!(dom_tree.idom_of(e), Some(b));
461        assert_eq!(dom_tree.idom_of(f), Some(a));
462
463        assert!(test_df(&df, a, &[]));
464        assert!(test_df(&df, b, &[]));
465        assert!(test_df(&df, c, &[f]));
466        assert!(test_df(&df, d, &[f]));
467        assert!(test_df(&df, e, &[]));
468        assert!(test_df(&df, f, &[]));
469    }
470}
sonatina_codegen/domtree.rs

sonatina_codegen/
domtree.rs
