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
17pub(crate) trait DomTreeGraph {
30 fn num_blocks(&self) -> usize;
33
34 fn roots(&self) -> impl Iterator<Item = Block>;
40
41 fn successors(&self, block: Block) -> impl Iterator<Item = Block>;
43
44 fn predecessors(&self, block: Block) -> impl Iterator<Item = Block>;
46}
47
48struct ForwardGraph<'a> {
51 func: &'a Function,
52 cfg: &'a ControlFlowGraph,
53}
54
55impl DomTreeGraph for ForwardGraph<'_> {
56 fn num_blocks(&self) -> usize {
57 self.func.dfg.num_blocks()
58 }
59
60 fn roots(&self) -> impl Iterator<Item = Block> {
61 self.func.layout.entry_block().into_iter()
62 }
63
64 fn successors(&self, block: Block) -> impl Iterator<Item = Block> {
65 self.func.block_successors(block).rev()
73 }
74
75 fn predecessors(&self, block: Block) -> impl Iterator<Item = Block> {
76 self.cfg
77 .pred_iter(block)
78 .map(|pred: BlockPredecessor| pred.block)
79 }
80}
81
82#[derive(Clone, Default)]
84struct SpanningTreeNode {
85 block: PackedOption<Block>,
87 ancestor: u32,
90 label: u32,
94 semi: u32,
96 idom: u32,
99}
100
101const NOT_VISITED: u32 = 0;
103
104#[derive(Clone, Default)]
111struct SpanningTree {
112 nodes: Vec<SpanningTreeNode>,
113}
114
115impl SpanningTree {
116 fn new() -> Self {
117 Self {
119 nodes: vec![Default::default()],
120 }
121 }
122
123 fn with_capacity(capacity: usize) -> Self {
124 let mut nodes = Vec::with_capacity(capacity + 1);
126 nodes.push(Default::default());
127 Self { nodes }
128 }
129
130 fn len(&self) -> usize {
131 self.nodes.len()
132 }
133
134 fn reserve(&mut self, capacity: usize) {
135 self.nodes.reserve(capacity);
137 }
138
139 fn clear(&mut self) {
140 self.nodes.resize(1, Default::default());
141 }
142
143 fn push(&mut self, ancestor: u32, block: Block) -> u32 {
145 debug_assert!(!self.nodes.is_empty());
147
148 let pre_number = self.nodes.len() as u32;
149
150 self.nodes.push(SpanningTreeNode {
151 block: block.into(),
152 ancestor,
153 label: pre_number,
154 semi: pre_number,
155 idom: ancestor,
156 });
157
158 pre_number
159 }
160}
161
162impl core::ops::Index<u32> for SpanningTree {
163 type Output = SpanningTreeNode;
164
165 fn index(&self, idx: u32) -> &Self::Output {
166 &self.nodes[idx as usize]
167 }
168}
169
170impl core::ops::IndexMut<u32> for SpanningTree {
171 fn index_mut(&mut self, idx: u32) -> &mut Self::Output {
172 &mut self.nodes[idx as usize]
173 }
174}
175
176enum TraversalEvent {
180 Enter(u32, Block),
181 Exit(Block),
182}
183
184#[derive(Clone, Default)]
186struct DominatorTreeNode {
187 idom: PackedOption<Block>,
189 pre_number: u32,
191
192 child: PackedOption<Block>,
194
195 sibling: PackedOption<Block>,
198
199 dom_pre_number: u32,
202
203 dom_pre_max: u32,
206}
207
208pub struct DominatorTree {
211 stree: SpanningTree,
213 postorder: Vec<Block>,
215 nodes: SecondaryMap<Block, DominatorTreeNode>,
217
218 dfs_worklist: Vec<TraversalEvent>,
220 eval_worklist: Vec<u32>,
223
224 valid: bool,
225}
226
227impl core::fmt::Debug for DominatorTree {
228 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
229 if !self.is_valid() {
230 return f.write_str("DominatorTree { <invalid> }");
231 }
232
233 let mut s = f.debug_tuple("DominatorTree");
234
235 if let Some((mut root, _)) = self.nodes.iter().find(|n| n.1.pre_number != NOT_VISITED) {
236 loop {
237 if let Some(b) = self.idom(root)
238 && b != root
239 {
240 root = b;
241 } else {
242 break;
243 }
244 }
245
246 fn fmt_block(domtree: &DominatorTree, block: Block) -> impl core::fmt::Debug {
247 core::fmt::from_fn(move |f| {
248 let children = domtree
249 .children(block)
250 .map(|c| fmt_block(domtree, c))
251 .collect::<Vec<_>>();
252 let mut s = f.debug_tuple(&format!("{block}"));
253 if !children.is_empty() {
254 s.field(&children);
255 }
256 s.finish()
257 })
258 }
259
260 s.field(&fmt_block(self, root));
261 }
262
263 s.finish()
264 }
265}
266
267impl DominatorTree {
269 pub fn is_reachable(&self, block: Block) -> bool {
271 self.nodes[block].pre_number != NOT_VISITED
272 }
273
274 pub fn cfg_postorder(&self) -> &[Block] {
279 debug_assert!(self.is_valid());
280 &self.postorder
281 }
282
283 pub fn cfg_rpo(&self) -> impl Iterator<Item = &Block> {
288 debug_assert!(self.is_valid());
289 self.postorder.iter().rev()
290 }
291
292 pub fn idom(&self, block: Block) -> Option<Block> {
303 self.nodes[block].idom.into()
304 }
305
306 pub fn dominates<A, B>(&self, a: A, b: B, layout: &Layout) -> bool
317 where
318 A: Into<ProgramPoint>,
319 B: Into<ProgramPoint>,
320 {
321 let a = a.into();
322 let b = b.into();
323 match a {
324 ProgramPoint::Block(block_a) => match b {
325 ProgramPoint::Block(block_b) => self.block_dominates(block_a, block_b),
326 ProgramPoint::Inst(inst_b) => {
327 let block_b = layout
328 .inst_block(inst_b)
329 .expect("Instruction not in layout.");
330 self.block_dominates(block_a, block_b)
331 }
332 },
333 ProgramPoint::Inst(inst_a) => {
334 let block_a: Block = layout
335 .inst_block(inst_a)
336 .expect("Instruction not in layout.");
337 match b {
338 ProgramPoint::Block(block_b) => {
339 block_a != block_b && self.block_dominates(block_a, block_b)
340 }
341 ProgramPoint::Inst(inst_b) => {
342 let block_b = layout
343 .inst_block(inst_b)
344 .expect("Instruction not in layout.");
345 if block_a == block_b {
346 layout.pp_cmp(a, b) != Ordering::Greater
347 } else {
348 self.block_dominates(block_a, block_b)
349 }
350 }
351 }
352 }
353 }
354 }
355
356 pub fn block_dominates(&self, block_a: Block, block_b: Block) -> bool {
361 let na = &self.nodes[block_a];
362 let nb = &self.nodes[block_b];
363 na.dom_pre_number <= nb.dom_pre_number && na.dom_pre_max >= nb.dom_pre_max
364 }
365
366 pub fn children(&self, block: Block) -> ChildIter<'_> {
371 ChildIter {
372 domtree: self,
373 next: self.nodes[block].child,
374 }
375 }
376}
377
378impl DominatorTree {
379 pub fn new() -> Self {
382 Self {
383 stree: SpanningTree::new(),
384 nodes: SecondaryMap::new(),
385 postorder: Vec::new(),
386 dfs_worklist: Vec::new(),
387 eval_worklist: Vec::new(),
388 valid: false,
389 }
390 }
391
392 pub fn with_function(func: &Function, cfg: &ControlFlowGraph) -> Self {
394 let block_capacity = func.layout.block_capacity();
395 let mut domtree = Self {
396 stree: SpanningTree::with_capacity(block_capacity),
397 nodes: SecondaryMap::with_capacity(block_capacity),
398 postorder: Vec::with_capacity(block_capacity),
399 dfs_worklist: Vec::new(),
400 eval_worklist: Vec::new(),
401 valid: false,
402 };
403 domtree.compute(func, cfg);
404 domtree
405 }
406
407 pub fn compute(&mut self, func: &Function, cfg: &ControlFlowGraph) {
416 let _tt = timing::domtree();
417 debug_assert!(cfg.is_valid());
418 self.compute_from_graph(&ForwardGraph { func, cfg });
419 }
420
421 pub(crate) fn compute_from_graph(&mut self, graph: &impl DomTreeGraph) {
428 self.clear();
429 self.compute_spanning_tree(graph);
430 self.compute_domtree(graph);
431 self.compute_domtree_preorder();
432
433 self.valid = true;
434 }
435
436 pub fn clear(&mut self) {
439 self.stree.clear();
440 self.nodes.clear();
441 self.postorder.clear();
442 self.valid = false;
443 }
444
445 pub fn is_valid(&self) -> bool {
451 self.valid
452 }
453
454 fn compute_spanning_tree(&mut self, graph: &impl DomTreeGraph) {
457 self.nodes.resize(graph.num_blocks());
458 self.stree.reserve(graph.num_blocks());
459
460 for root in graph.roots() {
461 self.dfs_worklist.push(TraversalEvent::Enter(0, root));
462 }
463
464 loop {
465 match self.dfs_worklist.pop() {
466 Some(TraversalEvent::Enter(parent, block)) => {
467 let node = &mut self.nodes[block];
468 if node.pre_number != NOT_VISITED {
469 continue;
470 }
471
472 self.dfs_worklist.push(TraversalEvent::Exit(block));
473
474 let pre_number = self.stree.push(parent, block);
475 node.pre_number = pre_number;
476
477 self.dfs_worklist.extend(
481 graph
482 .successors(block)
483 .filter(|successor| self.nodes[*successor].pre_number == NOT_VISITED)
485 .map(|successor| TraversalEvent::Enter(pre_number, successor)),
486 );
487 }
488 Some(TraversalEvent::Exit(block)) => {
489 self.postorder.push(block);
490 }
491 None => break,
492 }
493 }
494 }
495
496 fn eval(&mut self, v: u32, last_linked: u32) -> u32 {
501 if self.stree[v].ancestor < last_linked {
502 return self.stree[v].label;
503 }
504
505 let mut root = v;
507 loop {
508 self.eval_worklist.push(root);
509 root = self.stree[root].ancestor;
510
511 if self.stree[root].ancestor < last_linked {
512 break;
513 }
514 }
515
516 let mut prev = root;
517 let root = self.stree[prev].ancestor;
518
519 while let Some(curr) = self.eval_worklist.pop() {
522 if self.stree[prev].label < self.stree[curr].label {
523 self.stree[curr].label = self.stree[prev].label;
524 }
525
526 self.stree[curr].ancestor = root;
527 prev = curr;
528 }
529
530 self.stree[v].label
531 }
532
533 fn compute_domtree(&mut self, graph: &impl DomTreeGraph) {
534 for w in (1..self.stree.len() as u32).rev() {
536 let w_node = &mut self.stree[w];
537 let block = w_node.block.expect("Virtual root must have been excluded");
538 let mut semi = w_node.ancestor;
539
540 let last_linked = w + 1;
541
542 for pred in graph.predecessors(block) {
543 if self.nodes[pred].pre_number == NOT_VISITED {
545 continue;
546 }
547
548 let semi_candidate = self.eval(self.nodes[pred].pre_number, last_linked);
549 semi = core::cmp::min(semi, semi_candidate);
550 }
551
552 let w_node = &mut self.stree[w];
553 w_node.label = semi;
554 w_node.semi = semi;
555 }
556
557 for v in 1..self.stree.len() as u32 {
559 let semi = self.stree[v].semi;
560 let block = self.stree[v]
561 .block
562 .expect("Virtual root must have been excluded");
563 let mut idom = self.stree[v].idom;
564
565 while idom > semi {
566 idom = self.stree[idom].idom;
567 }
568
569 self.stree[v].idom = idom;
570
571 self.nodes[block].idom = self.stree[idom].block;
572 }
573 }
574
575 fn compute_domtree_preorder(&mut self) {
579 for &block in &self.postorder {
585 if let Some(idom) = self.idom(block) {
586 let sib = mem::replace(&mut self.nodes[idom].child, block.into());
587 self.nodes[block].sibling = sib;
588 } else {
589 self.dfs_worklist.push(TraversalEvent::Enter(0, block));
593 }
594 }
595
596 let mut n = 0;
604 while let Some(event) = self.dfs_worklist.pop() {
605 if let TraversalEvent::Enter(_, block) = event {
606 n += 1;
607 let node = &mut self.nodes[block];
608 node.dom_pre_number = n;
609 node.dom_pre_max = n;
610 if let Some(sibling) = node.sibling.expand() {
611 self.dfs_worklist.push(TraversalEvent::Enter(0, sibling));
612 }
613 if let Some(child) = node.child.expand() {
614 self.dfs_worklist.push(TraversalEvent::Enter(0, child));
615 }
616 }
617 }
618
619 for &block in &self.postorder {
623 if let Some(idom) = self.idom(block) {
624 let pre_max = cmp::max(self.nodes[block].dom_pre_max, self.nodes[idom].dom_pre_max);
625 self.nodes[idom].dom_pre_max = pre_max;
626 }
627 }
628 }
629}
630
631pub struct ChildIter<'a> {
633 domtree: &'a DominatorTree,
634 next: PackedOption<Block>,
635}
636
637impl<'a> Iterator for ChildIter<'a> {
638 type Item = Block;
639
640 fn next(&mut self) -> Option<Block> {
641 let n = self.next.expand();
642 if let Some(block) = n {
643 self.next = self.domtree.nodes[block].sibling;
644 }
645 n
646 }
647}
648
649#[cfg(test)]
650mod tests {
651 use super::*;
652 use crate::cursor::{Cursor, FuncCursor};
653 use crate::ir::types::*;
654 use crate::ir::{InstBuilder, TrapCode};
655
656 #[test]
657 fn empty() {
658 let func = Function::new();
659 let cfg = ControlFlowGraph::with_function(&func);
660 debug_assert!(cfg.is_valid());
661 let dtree = DominatorTree::with_function(&func, &cfg);
662 assert_eq!(0, dtree.nodes.keys().count());
663 assert_eq!(dtree.cfg_postorder(), &[]);
664 }
665
666 #[test]
667 fn unreachable_node() {
668 let mut func = Function::new();
669 let block0 = func.dfg.make_block();
670 let v0 = func.dfg.append_block_param(block0, I32);
671 let block1 = func.dfg.make_block();
672 let block2 = func.dfg.make_block();
673 let trap_block = func.dfg.make_block();
674
675 let mut cur = FuncCursor::new(&mut func);
676
677 cur.insert_block(block0);
678 cur.ins().brif(v0, block2, &[], trap_block, &[]);
679
680 cur.insert_block(trap_block);
681 cur.ins().trap(TrapCode::unwrap_user(1));
682
683 cur.insert_block(block1);
684 let v1 = cur.ins().iconst(I32, 1);
685 let v2 = cur.ins().iadd(v0, v1);
686 cur.ins().jump(block0, &[v2.into()]);
687
688 cur.insert_block(block2);
689 cur.ins().return_(&[v0]);
690
691 let cfg = ControlFlowGraph::with_function(cur.func);
692 let dt = DominatorTree::with_function(cur.func, &cfg);
693
694 assert_eq!(dt.cfg_postorder(), &[block2, trap_block, block0]);
704
705 let v2_def = cur.func.dfg.value_def(v2).unwrap_inst();
706 assert!(!dt.dominates(v2_def, block0, &cur.func.layout));
707 assert!(!dt.dominates(block0, v2_def, &cur.func.layout));
708
709 assert!(dt.block_dominates(block0, block0));
710 assert!(!dt.block_dominates(block0, block1));
711 assert!(dt.block_dominates(block0, block2));
712 assert!(!dt.block_dominates(block1, block0));
713 assert!(dt.block_dominates(block1, block1));
714 assert!(!dt.block_dominates(block1, block2));
715 assert!(!dt.block_dominates(block2, block0));
716 assert!(!dt.block_dominates(block2, block1));
717 assert!(dt.block_dominates(block2, block2));
718 }
719
720 #[test]
721 fn non_zero_entry_block() {
722 let mut func = Function::new();
723 let block0 = func.dfg.make_block();
724 let block1 = func.dfg.make_block();
725 let block2 = func.dfg.make_block();
726 let block3 = func.dfg.make_block();
727 let cond = func.dfg.append_block_param(block3, I32);
728
729 let mut cur = FuncCursor::new(&mut func);
730
731 cur.insert_block(block3);
732 let jmp_block3_block1 = cur.ins().jump(block1, &[]);
733
734 cur.insert_block(block1);
735 let br_block1_block0_block2 = cur.ins().brif(cond, block0, &[], block2, &[]);
736
737 cur.insert_block(block2);
738 cur.ins().jump(block0, &[]);
739
740 cur.insert_block(block0);
741
742 let cfg = ControlFlowGraph::with_function(cur.func);
743 let dt = DominatorTree::with_function(cur.func, &cfg);
744
745 assert_eq!(dt.cfg_postorder(), &[block0, block2, block1, block3]);
764
765 assert_eq!(cur.func.layout.entry_block().unwrap(), block3);
766 assert_eq!(dt.idom(block3), None);
767 assert_eq!(dt.idom(block1).unwrap(), block3);
768 assert_eq!(dt.idom(block2).unwrap(), block1);
769 assert_eq!(dt.idom(block0).unwrap(), block1);
770
771 assert!(dt.dominates(
772 br_block1_block0_block2,
773 br_block1_block0_block2,
774 &cur.func.layout
775 ));
776 assert!(!dt.dominates(br_block1_block0_block2, jmp_block3_block1, &cur.func.layout));
777 assert!(dt.dominates(jmp_block3_block1, br_block1_block0_block2, &cur.func.layout));
778 }
779
780 #[test]
781 fn backwards_layout() {
782 let mut func = Function::new();
783 let block0 = func.dfg.make_block();
784 let block1 = func.dfg.make_block();
785 let block2 = func.dfg.make_block();
786
787 let mut cur = FuncCursor::new(&mut func);
788
789 cur.insert_block(block0);
790 let jmp02 = cur.ins().jump(block2, &[]);
791
792 cur.insert_block(block1);
793 let trap = cur.ins().trap(TrapCode::unwrap_user(5));
794
795 cur.insert_block(block2);
796 let jmp21 = cur.ins().jump(block1, &[]);
797
798 let cfg = ControlFlowGraph::with_function(cur.func);
799 let dt = DominatorTree::with_function(cur.func, &cfg);
800
801 assert_eq!(cur.func.layout.entry_block(), Some(block0));
802 assert_eq!(dt.idom(block0), None);
803 assert_eq!(dt.idom(block1), Some(block2));
804 assert_eq!(dt.idom(block2), Some(block0));
805
806 assert!(dt.dominates(block0, block0, &cur.func.layout));
807 assert!(dt.dominates(block0, jmp02, &cur.func.layout));
808 assert!(dt.dominates(block0, block1, &cur.func.layout));
809 assert!(dt.dominates(block0, trap, &cur.func.layout));
810 assert!(dt.dominates(block0, block2, &cur.func.layout));
811 assert!(dt.dominates(block0, jmp21, &cur.func.layout));
812
813 assert!(!dt.dominates(jmp02, block0, &cur.func.layout));
814 assert!(dt.dominates(jmp02, jmp02, &cur.func.layout));
815 assert!(dt.dominates(jmp02, block1, &cur.func.layout));
816 assert!(dt.dominates(jmp02, trap, &cur.func.layout));
817 assert!(dt.dominates(jmp02, block2, &cur.func.layout));
818 assert!(dt.dominates(jmp02, jmp21, &cur.func.layout));
819
820 assert!(!dt.dominates(block1, block0, &cur.func.layout));
821 assert!(!dt.dominates(block1, jmp02, &cur.func.layout));
822 assert!(dt.dominates(block1, block1, &cur.func.layout));
823 assert!(dt.dominates(block1, trap, &cur.func.layout));
824 assert!(!dt.dominates(block1, block2, &cur.func.layout));
825 assert!(!dt.dominates(block1, jmp21, &cur.func.layout));
826
827 assert!(!dt.dominates(trap, block0, &cur.func.layout));
828 assert!(!dt.dominates(trap, jmp02, &cur.func.layout));
829 assert!(!dt.dominates(trap, block1, &cur.func.layout));
830 assert!(dt.dominates(trap, trap, &cur.func.layout));
831 assert!(!dt.dominates(trap, block2, &cur.func.layout));
832 assert!(!dt.dominates(trap, jmp21, &cur.func.layout));
833
834 assert!(!dt.dominates(block2, block0, &cur.func.layout));
835 assert!(!dt.dominates(block2, jmp02, &cur.func.layout));
836 assert!(dt.dominates(block2, block1, &cur.func.layout));
837 assert!(dt.dominates(block2, trap, &cur.func.layout));
838 assert!(dt.dominates(block2, block2, &cur.func.layout));
839 assert!(dt.dominates(block2, jmp21, &cur.func.layout));
840
841 assert!(!dt.dominates(jmp21, block0, &cur.func.layout));
842 assert!(!dt.dominates(jmp21, jmp02, &cur.func.layout));
843 assert!(dt.dominates(jmp21, block1, &cur.func.layout));
844 assert!(dt.dominates(jmp21, trap, &cur.func.layout));
845 assert!(!dt.dominates(jmp21, block2, &cur.func.layout));
846 assert!(dt.dominates(jmp21, jmp21, &cur.func.layout));
847 }
848
849 #[test]
850 fn insts_same_block() {
851 let mut func = Function::new();
852 let block0 = func.dfg.make_block();
853
854 let mut cur = FuncCursor::new(&mut func);
855
856 cur.insert_block(block0);
857 let v1 = cur.ins().iconst(I32, 1);
858 let v2 = cur.ins().iadd(v1, v1);
859 let v3 = cur.ins().iadd(v2, v2);
860 cur.ins().return_(&[]);
861
862 let cfg = ControlFlowGraph::with_function(cur.func);
863 let dt = DominatorTree::with_function(cur.func, &cfg);
864
865 let v1_def = cur.func.dfg.value_def(v1).unwrap_inst();
866 let v2_def = cur.func.dfg.value_def(v2).unwrap_inst();
867 let v3_def = cur.func.dfg.value_def(v3).unwrap_inst();
868
869 assert!(dt.dominates(v1_def, v2_def, &cur.func.layout));
870 assert!(dt.dominates(v2_def, v3_def, &cur.func.layout));
871 assert!(dt.dominates(v1_def, v3_def, &cur.func.layout));
872
873 assert!(!dt.dominates(v2_def, v1_def, &cur.func.layout));
874 assert!(!dt.dominates(v3_def, v2_def, &cur.func.layout));
875 assert!(!dt.dominates(v3_def, v1_def, &cur.func.layout));
876
877 assert!(dt.dominates(v2_def, v2_def, &cur.func.layout));
878 assert!(dt.dominates(block0, block0, &cur.func.layout));
879
880 assert!(dt.dominates(block0, v1_def, &cur.func.layout));
881 assert!(dt.dominates(block0, v2_def, &cur.func.layout));
882 assert!(dt.dominates(block0, v3_def, &cur.func.layout));
883
884 assert!(!dt.dominates(v1_def, block0, &cur.func.layout));
885 assert!(!dt.dominates(v2_def, block0, &cur.func.layout));
886 assert!(!dt.dominates(v3_def, block0, &cur.func.layout));
887 }
888}