1use crate::entity::SecondaryMap;
4use crate::flowgraph::{BlockPredecessor, ControlFlowGraph};
5use crate::ir::{Block, Function, Inst, Layout, ProgramPoint};
6use crate::packed_option::PackedOption;
7use crate::timing;
8use alloc::vec::Vec;
9use core::cmp;
10use core::cmp::Ordering;
11use core::mem;
12
13mod simple;
14
15pub use simple::SimpleDominatorTree;
16
17#[derive(Clone, Default)]
19struct SpanningTreeNode {
20 block: PackedOption<Block>,
22 ancestor: u32,
25 label: u32,
29 semi: u32,
31 idom: u32,
34}
35
36const NOT_VISITED: u32 = 0;
38
39#[derive(Clone, Default)]
46struct SpanningTree {
47 nodes: Vec<SpanningTreeNode>,
48}
49
50impl SpanningTree {
51 fn new() -> Self {
52 Self {
54 nodes: vec![Default::default()],
55 }
56 }
57
58 fn with_capacity(capacity: usize) -> Self {
59 let mut nodes = Vec::with_capacity(capacity + 1);
61 nodes.push(Default::default());
62 Self { nodes }
63 }
64
65 fn len(&self) -> usize {
66 self.nodes.len()
67 }
68
69 fn reserve(&mut self, capacity: usize) {
70 self.nodes.reserve(capacity);
72 }
73
74 fn clear(&mut self) {
75 self.nodes.resize(1, Default::default());
76 }
77
78 fn push(&mut self, ancestor: u32, block: Block) -> u32 {
80 debug_assert!(!self.nodes.is_empty());
82
83 let pre_number = self.nodes.len() as u32;
84
85 self.nodes.push(SpanningTreeNode {
86 block: block.into(),
87 ancestor,
88 label: pre_number,
89 semi: pre_number,
90 idom: ancestor,
91 });
92
93 pre_number
94 }
95}
96
97impl std::ops::Index<u32> for SpanningTree {
98 type Output = SpanningTreeNode;
99
100 fn index(&self, idx: u32) -> &Self::Output {
101 &self.nodes[idx as usize]
102 }
103}
104
105impl std::ops::IndexMut<u32> for SpanningTree {
106 fn index_mut(&mut self, idx: u32) -> &mut Self::Output {
107 &mut self.nodes[idx as usize]
108 }
109}
110
111enum TraversalEvent {
115 Enter(u32, Block),
116 Exit(Block),
117}
118
119#[derive(Clone, Default)]
121struct DominatorTreeNode {
122 idom: PackedOption<Block>,
124 pre_number: u32,
126}
127
128pub struct DominatorTree {
131 stree: SpanningTree,
133 postorder: Vec<Block>,
135 nodes: SecondaryMap<Block, DominatorTreeNode>,
137
138 dfs_worklist: Vec<TraversalEvent>,
140 eval_worklist: Vec<u32>,
143
144 valid: bool,
145}
146
147impl DominatorTree {
149 pub fn is_reachable(&self, block: Block) -> bool {
151 self.nodes[block].pre_number != NOT_VISITED
152 }
153
154 pub fn cfg_postorder(&self) -> &[Block] {
159 debug_assert!(self.is_valid());
160 &self.postorder
161 }
162
163 pub fn cfg_rpo(&self) -> impl Iterator<Item = &Block> {
168 debug_assert!(self.is_valid());
169 self.postorder.iter().rev()
170 }
171
172 pub fn idom(&self, block: Block) -> Option<Block> {
183 self.nodes[block].idom.into()
184 }
185
186 pub fn dominates<A, B>(&self, a: A, b: B, layout: &Layout) -> bool
197 where
198 A: Into<ProgramPoint>,
199 B: Into<ProgramPoint>,
200 {
201 let a = a.into();
202 let b = b.into();
203 match a {
204 ProgramPoint::Block(block_a) => match b {
205 ProgramPoint::Block(block_b) => self.block_dominates(block_a, block_b),
206 ProgramPoint::Inst(inst_b) => {
207 let block_b = layout
208 .inst_block(inst_b)
209 .expect("Instruction not in layout.");
210 self.block_dominates(block_a, block_b)
211 }
212 },
213 ProgramPoint::Inst(inst_a) => {
214 let block_a: Block = layout
215 .inst_block(inst_a)
216 .expect("Instruction not in layout.");
217 match b {
218 ProgramPoint::Block(block_b) => {
219 block_a != block_b && self.block_dominates(block_a, block_b)
220 }
221 ProgramPoint::Inst(inst_b) => {
222 let block_b = layout
223 .inst_block(inst_b)
224 .expect("Instruction not in layout.");
225 if block_a == block_b {
226 layout.pp_cmp(a, b) != Ordering::Greater
227 } else {
228 self.block_dominates(block_a, block_b)
229 }
230 }
231 }
232 }
233 }
234 }
235
236 fn block_dominates(&self, block_a: Block, mut block_b: Block) -> bool {
240 let pre_a = self.nodes[block_a].pre_number;
241
242 while pre_a < self.nodes[block_b].pre_number {
245 let idom = match self.idom(block_b) {
246 Some(idom) => idom,
247 None => return false, };
249 block_b = idom;
250 }
251
252 block_a == block_b
253 }
254}
255
256impl DominatorTree {
257 pub fn new() -> Self {
260 Self {
261 stree: SpanningTree::new(),
262 nodes: SecondaryMap::new(),
263 postorder: Vec::new(),
264 dfs_worklist: Vec::new(),
265 eval_worklist: Vec::new(),
266 valid: false,
267 }
268 }
269
270 pub fn with_function(func: &Function, cfg: &ControlFlowGraph) -> Self {
272 let block_capacity = func.layout.block_capacity();
273 let mut domtree = Self {
274 stree: SpanningTree::with_capacity(block_capacity),
275 nodes: SecondaryMap::with_capacity(block_capacity),
276 postorder: Vec::with_capacity(block_capacity),
277 dfs_worklist: Vec::new(),
278 eval_worklist: Vec::new(),
279 valid: false,
280 };
281 domtree.compute(func, cfg);
282 domtree
283 }
284
285 pub fn compute(&mut self, func: &Function, cfg: &ControlFlowGraph) {
294 let _tt = timing::domtree();
295 debug_assert!(cfg.is_valid());
296
297 self.clear();
298 self.compute_spanning_tree(func);
299 self.compute_domtree(cfg);
300
301 self.valid = true;
302 }
303
304 pub fn clear(&mut self) {
307 self.stree.clear();
308 self.nodes.clear();
309 self.postorder.clear();
310 self.valid = false;
311 }
312
313 pub fn is_valid(&self) -> bool {
319 self.valid
320 }
321
322 fn compute_spanning_tree(&mut self, func: &Function) {
325 self.nodes.resize(func.dfg.num_blocks());
326 self.stree.reserve(func.dfg.num_blocks());
327
328 if let Some(block) = func.layout.entry_block() {
329 self.dfs_worklist.push(TraversalEvent::Enter(0, block));
330 }
331
332 loop {
333 match self.dfs_worklist.pop() {
334 Some(TraversalEvent::Enter(parent, block)) => {
335 let node = &mut self.nodes[block];
336 if node.pre_number != NOT_VISITED {
337 continue;
338 }
339
340 self.dfs_worklist.push(TraversalEvent::Exit(block));
341
342 let pre_number = self.stree.push(parent, block);
343 node.pre_number = pre_number;
344
345 self.dfs_worklist.extend(
347 func.block_successors(block)
348 .rev()
358 .filter(|successor| self.nodes[*successor].pre_number == NOT_VISITED)
360 .map(|successor| TraversalEvent::Enter(pre_number, successor)),
361 );
362 }
363 Some(TraversalEvent::Exit(block)) => self.postorder.push(block),
364 None => break,
365 }
366 }
367 }
368
369 fn eval(&mut self, v: u32, last_linked: u32) -> u32 {
374 if self.stree[v].ancestor < last_linked {
375 return self.stree[v].label;
376 }
377
378 let mut root = v;
380 loop {
381 self.eval_worklist.push(root);
382 root = self.stree[root].ancestor;
383
384 if self.stree[root].ancestor < last_linked {
385 break;
386 }
387 }
388
389 let mut prev = root;
390 let root = self.stree[prev].ancestor;
391
392 while let Some(curr) = self.eval_worklist.pop() {
395 if self.stree[prev].label < self.stree[curr].label {
396 self.stree[curr].label = self.stree[prev].label;
397 }
398
399 self.stree[curr].ancestor = root;
400 prev = curr;
401 }
402
403 self.stree[v].label
404 }
405
406 fn compute_domtree(&mut self, cfg: &ControlFlowGraph) {
407 for w in (1..self.stree.len() as u32).rev() {
409 let w_node = &mut self.stree[w];
410 let block = w_node.block.expect("Virtual root must have been excluded");
411 let mut semi = w_node.ancestor;
412
413 let last_linked = w + 1;
414
415 for pred in cfg
416 .pred_iter(block)
417 .map(|pred: BlockPredecessor| pred.block)
418 {
419 if self.nodes[pred].pre_number == NOT_VISITED {
421 continue;
422 }
423
424 let semi_candidate = self.eval(self.nodes[pred].pre_number, last_linked);
425 semi = std::cmp::min(semi, semi_candidate);
426 }
427
428 let w_node = &mut self.stree[w];
429 w_node.label = semi;
430 w_node.semi = semi;
431 }
432
433 for v in 1..self.stree.len() as u32 {
435 let semi = self.stree[v].semi;
436 let block = self.stree[v]
437 .block
438 .expect("Virtual root must have been excluded");
439 let mut idom = self.stree[v].idom;
440
441 while idom > semi {
442 idom = self.stree[idom].idom;
443 }
444
445 self.stree[v].idom = idom;
446
447 self.nodes[block].idom = self.stree[idom].block;
448 }
449 }
450}
451
452pub struct DominatorTreePreorder {
463 nodes: SecondaryMap<Block, ExtraNode>,
464
465 stack: Vec<Block>,
467}
468
469#[derive(Default, Clone)]
470struct ExtraNode {
471 child: PackedOption<Block>,
473
474 sibling: PackedOption<Block>,
476
477 pre_number: u32,
480
481 pre_max: u32,
484}
485
486impl DominatorTreePreorder {
488 pub fn new() -> Self {
490 Self {
491 nodes: SecondaryMap::new(),
492 stack: Vec::new(),
493 }
494 }
495
496 pub fn compute(&mut self, domtree: &DominatorTree) {
498 self.nodes.clear();
499
500 for &block in domtree.cfg_postorder() {
505 if let Some(idom) = domtree.idom(block) {
506 let sib = mem::replace(&mut self.nodes[idom].child, block.into());
507 self.nodes[block].sibling = sib;
508 } else {
509 self.stack.push(block);
511 }
512 }
513
514 debug_assert!(self.stack.len() <= 1);
516 let mut n = 0;
517 while let Some(block) = self.stack.pop() {
518 n += 1;
519 let node = &mut self.nodes[block];
520 node.pre_number = n;
521 node.pre_max = n;
522 if let Some(n) = node.sibling.expand() {
523 self.stack.push(n);
524 }
525 if let Some(n) = node.child.expand() {
526 self.stack.push(n);
527 }
528 }
529
530 for &block in domtree.cfg_postorder() {
534 if let Some(idom) = domtree.idom(block) {
535 let pre_max = cmp::max(self.nodes[block].pre_max, self.nodes[idom].pre_max);
536 self.nodes[idom].pre_max = pre_max;
537 }
538 }
539 }
540}
541
542pub struct ChildIter<'a> {
544 dtpo: &'a DominatorTreePreorder,
545 next: PackedOption<Block>,
546}
547
548impl<'a> Iterator for ChildIter<'a> {
549 type Item = Block;
550
551 fn next(&mut self) -> Option<Block> {
552 let n = self.next.expand();
553 if let Some(block) = n {
554 self.next = self.dtpo.nodes[block].sibling;
555 }
556 n
557 }
558}
559
560impl DominatorTreePreorder {
562 pub fn children(&self, block: Block) -> ChildIter<'_> {
567 ChildIter {
568 dtpo: self,
569 next: self.nodes[block].child,
570 }
571 }
572
573 pub fn dominates(&self, a: Block, b: Block) -> bool {
581 let na = &self.nodes[a];
582 let nb = &self.nodes[b];
583 na.pre_number <= nb.pre_number && na.pre_max >= nb.pre_max
584 }
585
586 pub fn dominates_inst(&self, a: Inst, b: Inst, layout: &Layout) -> bool {
588 match (layout.inst_block(a), layout.inst_block(b)) {
589 (Some(block_a), Some(block_b)) => {
590 if block_a == block_b {
591 layout.pp_cmp(a, b) != Ordering::Greater
592 } else {
593 self.dominates(block_a, block_b)
594 }
595 }
596 _ => false,
597 }
598 }
599
600 pub fn pre_cmp_block(&self, a: Block, b: Block) -> Ordering {
602 self.nodes[a].pre_number.cmp(&self.nodes[b].pre_number)
603 }
604
605 pub fn pre_cmp<A, B>(&self, a: A, b: B, layout: &Layout) -> Ordering
610 where
611 A: Into<ProgramPoint>,
612 B: Into<ProgramPoint>,
613 {
614 let a = a.into();
615 let b = b.into();
616 self.pre_cmp_block(layout.pp_block(a), layout.pp_block(b))
617 .then_with(|| layout.pp_cmp(a, b))
618 }
619}
620
621#[cfg(test)]
622mod tests {
623 use super::*;
624 use crate::cursor::{Cursor, FuncCursor};
625 use crate::ir::types::*;
626 use crate::ir::{InstBuilder, TrapCode};
627
628 #[test]
629 fn empty() {
630 let func = Function::new();
631 let cfg = ControlFlowGraph::with_function(&func);
632 debug_assert!(cfg.is_valid());
633 let dtree = DominatorTree::with_function(&func, &cfg);
634 assert_eq!(0, dtree.nodes.keys().count());
635 assert_eq!(dtree.cfg_postorder(), &[]);
636
637 let mut dtpo = DominatorTreePreorder::new();
638 dtpo.compute(&dtree);
639 }
640
641 #[test]
642 fn unreachable_node() {
643 let mut func = Function::new();
644 let block0 = func.dfg.make_block();
645 let v0 = func.dfg.append_block_param(block0, I32);
646 let block1 = func.dfg.make_block();
647 let block2 = func.dfg.make_block();
648 let trap_block = func.dfg.make_block();
649
650 let mut cur = FuncCursor::new(&mut func);
651
652 cur.insert_block(block0);
653 cur.ins().brif(v0, block2, &[], trap_block, &[]);
654
655 cur.insert_block(trap_block);
656 cur.ins().trap(TrapCode::unwrap_user(1));
657
658 cur.insert_block(block1);
659 let v1 = cur.ins().iconst(I32, 1);
660 let v2 = cur.ins().iadd(v0, v1);
661 cur.ins().jump(block0, &[v2.into()]);
662
663 cur.insert_block(block2);
664 cur.ins().return_(&[v0]);
665
666 let cfg = ControlFlowGraph::with_function(cur.func);
667 let dt = DominatorTree::with_function(cur.func, &cfg);
668
669 assert_eq!(dt.cfg_postorder(), &[block2, trap_block, block0]);
679
680 let v2_def = cur.func.dfg.value_def(v2).unwrap_inst();
681 assert!(!dt.dominates(v2_def, block0, &cur.func.layout));
682 assert!(!dt.dominates(block0, v2_def, &cur.func.layout));
683
684 let mut dtpo = DominatorTreePreorder::new();
685 dtpo.compute(&dt);
686 assert!(dtpo.dominates(block0, block0));
687 assert!(!dtpo.dominates(block0, block1));
688 assert!(dtpo.dominates(block0, block2));
689 assert!(!dtpo.dominates(block1, block0));
690 assert!(dtpo.dominates(block1, block1));
691 assert!(!dtpo.dominates(block1, block2));
692 assert!(!dtpo.dominates(block2, block0));
693 assert!(!dtpo.dominates(block2, block1));
694 assert!(dtpo.dominates(block2, block2));
695 }
696
697 #[test]
698 fn non_zero_entry_block() {
699 let mut func = Function::new();
700 let block0 = func.dfg.make_block();
701 let block1 = func.dfg.make_block();
702 let block2 = func.dfg.make_block();
703 let block3 = func.dfg.make_block();
704 let cond = func.dfg.append_block_param(block3, I32);
705
706 let mut cur = FuncCursor::new(&mut func);
707
708 cur.insert_block(block3);
709 let jmp_block3_block1 = cur.ins().jump(block1, &[]);
710
711 cur.insert_block(block1);
712 let br_block1_block0_block2 = cur.ins().brif(cond, block0, &[], block2, &[]);
713
714 cur.insert_block(block2);
715 cur.ins().jump(block0, &[]);
716
717 cur.insert_block(block0);
718
719 let cfg = ControlFlowGraph::with_function(cur.func);
720 let dt = DominatorTree::with_function(cur.func, &cfg);
721
722 assert_eq!(dt.cfg_postorder(), &[block0, block2, block1, block3]);
741
742 assert_eq!(cur.func.layout.entry_block().unwrap(), block3);
743 assert_eq!(dt.idom(block3), None);
744 assert_eq!(dt.idom(block1).unwrap(), block3);
745 assert_eq!(dt.idom(block2).unwrap(), block1);
746 assert_eq!(dt.idom(block0).unwrap(), block1);
747
748 assert!(dt.dominates(
749 br_block1_block0_block2,
750 br_block1_block0_block2,
751 &cur.func.layout
752 ));
753 assert!(!dt.dominates(br_block1_block0_block2, jmp_block3_block1, &cur.func.layout));
754 assert!(dt.dominates(jmp_block3_block1, br_block1_block0_block2, &cur.func.layout));
755 }
756
757 #[test]
758 fn backwards_layout() {
759 let mut func = Function::new();
760 let block0 = func.dfg.make_block();
761 let block1 = func.dfg.make_block();
762 let block2 = func.dfg.make_block();
763
764 let mut cur = FuncCursor::new(&mut func);
765
766 cur.insert_block(block0);
767 let jmp02 = cur.ins().jump(block2, &[]);
768
769 cur.insert_block(block1);
770 let trap = cur.ins().trap(TrapCode::unwrap_user(5));
771
772 cur.insert_block(block2);
773 let jmp21 = cur.ins().jump(block1, &[]);
774
775 let cfg = ControlFlowGraph::with_function(cur.func);
776 let dt = DominatorTree::with_function(cur.func, &cfg);
777
778 assert_eq!(cur.func.layout.entry_block(), Some(block0));
779 assert_eq!(dt.idom(block0), None);
780 assert_eq!(dt.idom(block1), Some(block2));
781 assert_eq!(dt.idom(block2), Some(block0));
782
783 assert!(dt.dominates(block0, block0, &cur.func.layout));
784 assert!(dt.dominates(block0, jmp02, &cur.func.layout));
785 assert!(dt.dominates(block0, block1, &cur.func.layout));
786 assert!(dt.dominates(block0, trap, &cur.func.layout));
787 assert!(dt.dominates(block0, block2, &cur.func.layout));
788 assert!(dt.dominates(block0, jmp21, &cur.func.layout));
789
790 assert!(!dt.dominates(jmp02, block0, &cur.func.layout));
791 assert!(dt.dominates(jmp02, jmp02, &cur.func.layout));
792 assert!(dt.dominates(jmp02, block1, &cur.func.layout));
793 assert!(dt.dominates(jmp02, trap, &cur.func.layout));
794 assert!(dt.dominates(jmp02, block2, &cur.func.layout));
795 assert!(dt.dominates(jmp02, jmp21, &cur.func.layout));
796
797 assert!(!dt.dominates(block1, block0, &cur.func.layout));
798 assert!(!dt.dominates(block1, jmp02, &cur.func.layout));
799 assert!(dt.dominates(block1, block1, &cur.func.layout));
800 assert!(dt.dominates(block1, trap, &cur.func.layout));
801 assert!(!dt.dominates(block1, block2, &cur.func.layout));
802 assert!(!dt.dominates(block1, jmp21, &cur.func.layout));
803
804 assert!(!dt.dominates(trap, block0, &cur.func.layout));
805 assert!(!dt.dominates(trap, jmp02, &cur.func.layout));
806 assert!(!dt.dominates(trap, block1, &cur.func.layout));
807 assert!(dt.dominates(trap, trap, &cur.func.layout));
808 assert!(!dt.dominates(trap, block2, &cur.func.layout));
809 assert!(!dt.dominates(trap, jmp21, &cur.func.layout));
810
811 assert!(!dt.dominates(block2, block0, &cur.func.layout));
812 assert!(!dt.dominates(block2, jmp02, &cur.func.layout));
813 assert!(dt.dominates(block2, block1, &cur.func.layout));
814 assert!(dt.dominates(block2, trap, &cur.func.layout));
815 assert!(dt.dominates(block2, block2, &cur.func.layout));
816 assert!(dt.dominates(block2, jmp21, &cur.func.layout));
817
818 assert!(!dt.dominates(jmp21, block0, &cur.func.layout));
819 assert!(!dt.dominates(jmp21, jmp02, &cur.func.layout));
820 assert!(dt.dominates(jmp21, block1, &cur.func.layout));
821 assert!(dt.dominates(jmp21, trap, &cur.func.layout));
822 assert!(!dt.dominates(jmp21, block2, &cur.func.layout));
823 assert!(dt.dominates(jmp21, jmp21, &cur.func.layout));
824 }
825
826 #[test]
827 fn insts_same_block() {
828 let mut func = Function::new();
829 let block0 = func.dfg.make_block();
830
831 let mut cur = FuncCursor::new(&mut func);
832
833 cur.insert_block(block0);
834 let v1 = cur.ins().iconst(I32, 1);
835 let v2 = cur.ins().iadd(v1, v1);
836 let v3 = cur.ins().iadd(v2, v2);
837 cur.ins().return_(&[]);
838
839 let cfg = ControlFlowGraph::with_function(cur.func);
840 let dt = DominatorTree::with_function(cur.func, &cfg);
841
842 let v1_def = cur.func.dfg.value_def(v1).unwrap_inst();
843 let v2_def = cur.func.dfg.value_def(v2).unwrap_inst();
844 let v3_def = cur.func.dfg.value_def(v3).unwrap_inst();
845
846 assert!(dt.dominates(v1_def, v2_def, &cur.func.layout));
847 assert!(dt.dominates(v2_def, v3_def, &cur.func.layout));
848 assert!(dt.dominates(v1_def, v3_def, &cur.func.layout));
849
850 assert!(!dt.dominates(v2_def, v1_def, &cur.func.layout));
851 assert!(!dt.dominates(v3_def, v2_def, &cur.func.layout));
852 assert!(!dt.dominates(v3_def, v1_def, &cur.func.layout));
853
854 assert!(dt.dominates(v2_def, v2_def, &cur.func.layout));
855 assert!(dt.dominates(block0, block0, &cur.func.layout));
856
857 assert!(dt.dominates(block0, v1_def, &cur.func.layout));
858 assert!(dt.dominates(block0, v2_def, &cur.func.layout));
859 assert!(dt.dominates(block0, v3_def, &cur.func.layout));
860
861 assert!(!dt.dominates(v1_def, block0, &cur.func.layout));
862 assert!(!dt.dominates(v2_def, block0, &cur.func.layout));
863 assert!(!dt.dominates(v3_def, block0, &cur.func.layout));
864 }
865}