1use crate::entity::SecondaryMap;
4use crate::flowgraph::{BlockPredecessor, ControlFlowGraph};
5use crate::ir::{Block, Function, 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: 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 pre_cmp_block(&self, a: Block, b: Block) -> Ordering {
588 self.nodes[a].pre_number.cmp(&self.nodes[b].pre_number)
589 }
590
591 pub fn pre_cmp<A, B>(&self, a: A, b: B, layout: &Layout) -> Ordering
596 where
597 A: Into<ProgramPoint>,
598 B: Into<ProgramPoint>,
599 {
600 let a = a.into();
601 let b = b.into();
602 self.pre_cmp_block(layout.pp_block(a), layout.pp_block(b))
603 .then_with(|| layout.pp_cmp(a, b))
604 }
605}
606
607#[cfg(test)]
608mod tests {
609 use super::*;
610 use crate::cursor::{Cursor, FuncCursor};
611 use crate::ir::types::*;
612 use crate::ir::{InstBuilder, TrapCode};
613
614 #[test]
615 fn empty() {
616 let func = Function::new();
617 let cfg = ControlFlowGraph::with_function(&func);
618 debug_assert!(cfg.is_valid());
619 let dtree = DominatorTree::with_function(&func, &cfg);
620 assert_eq!(0, dtree.nodes.keys().count());
621 assert_eq!(dtree.cfg_postorder(), &[]);
622
623 let mut dtpo = DominatorTreePreorder::new();
624 dtpo.compute(&dtree);
625 }
626
627 #[test]
628 fn unreachable_node() {
629 let mut func = Function::new();
630 let block0 = func.dfg.make_block();
631 let v0 = func.dfg.append_block_param(block0, I32);
632 let block1 = func.dfg.make_block();
633 let block2 = func.dfg.make_block();
634 let trap_block = func.dfg.make_block();
635
636 let mut cur = FuncCursor::new(&mut func);
637
638 cur.insert_block(block0);
639 cur.ins().brif(v0, block2, &[], trap_block, &[]);
640
641 cur.insert_block(trap_block);
642 cur.ins().trap(TrapCode::unwrap_user(1));
643
644 cur.insert_block(block1);
645 let v1 = cur.ins().iconst(I32, 1);
646 let v2 = cur.ins().iadd(v0, v1);
647 cur.ins().jump(block0, &[v2]);
648
649 cur.insert_block(block2);
650 cur.ins().return_(&[v0]);
651
652 let cfg = ControlFlowGraph::with_function(cur.func);
653 let dt = DominatorTree::with_function(cur.func, &cfg);
654
655 assert_eq!(dt.cfg_postorder(), &[block2, trap_block, block0]);
665
666 let v2_def = cur.func.dfg.value_def(v2).unwrap_inst();
667 assert!(!dt.dominates(v2_def, block0, &cur.func.layout));
668 assert!(!dt.dominates(block0, v2_def, &cur.func.layout));
669
670 let mut dtpo = DominatorTreePreorder::new();
671 dtpo.compute(&dt);
672 assert!(dtpo.dominates(block0, block0));
673 assert!(!dtpo.dominates(block0, block1));
674 assert!(dtpo.dominates(block0, block2));
675 assert!(!dtpo.dominates(block1, block0));
676 assert!(dtpo.dominates(block1, block1));
677 assert!(!dtpo.dominates(block1, block2));
678 assert!(!dtpo.dominates(block2, block0));
679 assert!(!dtpo.dominates(block2, block1));
680 assert!(dtpo.dominates(block2, block2));
681 }
682
683 #[test]
684 fn non_zero_entry_block() {
685 let mut func = Function::new();
686 let block0 = func.dfg.make_block();
687 let block1 = func.dfg.make_block();
688 let block2 = func.dfg.make_block();
689 let block3 = func.dfg.make_block();
690 let cond = func.dfg.append_block_param(block3, I32);
691
692 let mut cur = FuncCursor::new(&mut func);
693
694 cur.insert_block(block3);
695 let jmp_block3_block1 = cur.ins().jump(block1, &[]);
696
697 cur.insert_block(block1);
698 let br_block1_block0_block2 = cur.ins().brif(cond, block0, &[], block2, &[]);
699
700 cur.insert_block(block2);
701 cur.ins().jump(block0, &[]);
702
703 cur.insert_block(block0);
704
705 let cfg = ControlFlowGraph::with_function(cur.func);
706 let dt = DominatorTree::with_function(cur.func, &cfg);
707
708 assert_eq!(dt.cfg_postorder(), &[block0, block2, block1, block3]);
727
728 assert_eq!(cur.func.layout.entry_block().unwrap(), block3);
729 assert_eq!(dt.idom(block3), None);
730 assert_eq!(dt.idom(block1).unwrap(), block3);
731 assert_eq!(dt.idom(block2).unwrap(), block1);
732 assert_eq!(dt.idom(block0).unwrap(), block1);
733
734 assert!(dt.dominates(
735 br_block1_block0_block2,
736 br_block1_block0_block2,
737 &cur.func.layout
738 ));
739 assert!(!dt.dominates(br_block1_block0_block2, jmp_block3_block1, &cur.func.layout));
740 assert!(dt.dominates(jmp_block3_block1, br_block1_block0_block2, &cur.func.layout));
741 }
742
743 #[test]
744 fn backwards_layout() {
745 let mut func = Function::new();
746 let block0 = func.dfg.make_block();
747 let block1 = func.dfg.make_block();
748 let block2 = func.dfg.make_block();
749
750 let mut cur = FuncCursor::new(&mut func);
751
752 cur.insert_block(block0);
753 let jmp02 = cur.ins().jump(block2, &[]);
754
755 cur.insert_block(block1);
756 let trap = cur.ins().trap(TrapCode::unwrap_user(5));
757
758 cur.insert_block(block2);
759 let jmp21 = cur.ins().jump(block1, &[]);
760
761 let cfg = ControlFlowGraph::with_function(cur.func);
762 let dt = DominatorTree::with_function(cur.func, &cfg);
763
764 assert_eq!(cur.func.layout.entry_block(), Some(block0));
765 assert_eq!(dt.idom(block0), None);
766 assert_eq!(dt.idom(block1), Some(block2));
767 assert_eq!(dt.idom(block2), Some(block0));
768
769 assert!(dt.dominates(block0, block0, &cur.func.layout));
770 assert!(dt.dominates(block0, jmp02, &cur.func.layout));
771 assert!(dt.dominates(block0, block1, &cur.func.layout));
772 assert!(dt.dominates(block0, trap, &cur.func.layout));
773 assert!(dt.dominates(block0, block2, &cur.func.layout));
774 assert!(dt.dominates(block0, jmp21, &cur.func.layout));
775
776 assert!(!dt.dominates(jmp02, block0, &cur.func.layout));
777 assert!(dt.dominates(jmp02, jmp02, &cur.func.layout));
778 assert!(dt.dominates(jmp02, block1, &cur.func.layout));
779 assert!(dt.dominates(jmp02, trap, &cur.func.layout));
780 assert!(dt.dominates(jmp02, block2, &cur.func.layout));
781 assert!(dt.dominates(jmp02, jmp21, &cur.func.layout));
782
783 assert!(!dt.dominates(block1, block0, &cur.func.layout));
784 assert!(!dt.dominates(block1, jmp02, &cur.func.layout));
785 assert!(dt.dominates(block1, block1, &cur.func.layout));
786 assert!(dt.dominates(block1, trap, &cur.func.layout));
787 assert!(!dt.dominates(block1, block2, &cur.func.layout));
788 assert!(!dt.dominates(block1, jmp21, &cur.func.layout));
789
790 assert!(!dt.dominates(trap, block0, &cur.func.layout));
791 assert!(!dt.dominates(trap, jmp02, &cur.func.layout));
792 assert!(!dt.dominates(trap, block1, &cur.func.layout));
793 assert!(dt.dominates(trap, trap, &cur.func.layout));
794 assert!(!dt.dominates(trap, block2, &cur.func.layout));
795 assert!(!dt.dominates(trap, jmp21, &cur.func.layout));
796
797 assert!(!dt.dominates(block2, block0, &cur.func.layout));
798 assert!(!dt.dominates(block2, jmp02, &cur.func.layout));
799 assert!(dt.dominates(block2, block1, &cur.func.layout));
800 assert!(dt.dominates(block2, trap, &cur.func.layout));
801 assert!(dt.dominates(block2, block2, &cur.func.layout));
802 assert!(dt.dominates(block2, jmp21, &cur.func.layout));
803
804 assert!(!dt.dominates(jmp21, block0, &cur.func.layout));
805 assert!(!dt.dominates(jmp21, jmp02, &cur.func.layout));
806 assert!(dt.dominates(jmp21, block1, &cur.func.layout));
807 assert!(dt.dominates(jmp21, trap, &cur.func.layout));
808 assert!(!dt.dominates(jmp21, block2, &cur.func.layout));
809 assert!(dt.dominates(jmp21, jmp21, &cur.func.layout));
810 }
811
812 #[test]
813 fn insts_same_block() {
814 let mut func = Function::new();
815 let block0 = func.dfg.make_block();
816
817 let mut cur = FuncCursor::new(&mut func);
818
819 cur.insert_block(block0);
820 let v1 = cur.ins().iconst(I32, 1);
821 let v2 = cur.ins().iadd(v1, v1);
822 let v3 = cur.ins().iadd(v2, v2);
823 cur.ins().return_(&[]);
824
825 let cfg = ControlFlowGraph::with_function(cur.func);
826 let dt = DominatorTree::with_function(cur.func, &cfg);
827
828 let v1_def = cur.func.dfg.value_def(v1).unwrap_inst();
829 let v2_def = cur.func.dfg.value_def(v2).unwrap_inst();
830 let v3_def = cur.func.dfg.value_def(v3).unwrap_inst();
831
832 assert!(dt.dominates(v1_def, v2_def, &cur.func.layout));
833 assert!(dt.dominates(v2_def, v3_def, &cur.func.layout));
834 assert!(dt.dominates(v1_def, v3_def, &cur.func.layout));
835
836 assert!(!dt.dominates(v2_def, v1_def, &cur.func.layout));
837 assert!(!dt.dominates(v3_def, v2_def, &cur.func.layout));
838 assert!(!dt.dominates(v3_def, v1_def, &cur.func.layout));
839
840 assert!(dt.dominates(v2_def, v2_def, &cur.func.layout));
841 assert!(dt.dominates(block0, block0, &cur.func.layout));
842
843 assert!(dt.dominates(block0, v1_def, &cur.func.layout));
844 assert!(dt.dominates(block0, v2_def, &cur.func.layout));
845 assert!(dt.dominates(block0, v3_def, &cur.func.layout));
846
847 assert!(!dt.dominates(v1_def, block0, &cur.func.layout));
848 assert!(!dt.dominates(v2_def, block0, &cur.func.layout));
849 assert!(!dt.dominates(v3_def, block0, &cur.func.layout));
850 }
851}