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,
88 label: pre_number,
89 semi: pre_number,
90 idom: ancestor,
91 });
92
93 pre_number
94 }
95}
96
97impl core::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 core::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 child: PackedOption<Block>,
129
130 sibling: PackedOption<Block>,
133
134 dom_pre_number: u32,
137
138 dom_pre_max: u32,
141}
142
143pub struct DominatorTree {
146 stree: SpanningTree,
148 postorder: Vec<Block>,
150 nodes: SecondaryMap<Block, DominatorTreeNode>,
152
153 dfs_worklist: Vec<TraversalEvent>,
155 eval_worklist: Vec<u32>,
158
159 valid: bool,
160}
161
162impl core::fmt::Debug for DominatorTree {
163 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
164 if !self.is_valid() {
165 return f.write_str("DominatorTree { <invalid> }");
166 }
167
168 let mut s = f.debug_tuple("DominatorTree");
169
170 if let Some((mut root, _)) = self.nodes.iter().find(|n| n.1.pre_number != NOT_VISITED) {
171 loop {
172 if let Some(b) = self.idom(root)
173 && b != root
174 {
175 root = b;
176 } else {
177 break;
178 }
179 }
180
181 fn fmt_block(domtree: &DominatorTree, block: Block) -> impl core::fmt::Debug {
182 core::fmt::from_fn(move |f| {
183 let children = domtree
184 .children(block)
185 .map(|c| fmt_block(domtree, c))
186 .collect::<Vec<_>>();
187 let mut s = f.debug_tuple(&format!("{block}"));
188 if !children.is_empty() {
189 s.field(&children);
190 }
191 s.finish()
192 })
193 }
194
195 s.field(&fmt_block(self, root));
196 }
197
198 s.finish()
199 }
200}
201
202impl DominatorTree {
204 pub fn is_reachable(&self, block: Block) -> bool {
206 self.nodes[block].pre_number != NOT_VISITED
207 }
208
209 pub fn cfg_postorder(&self) -> &[Block] {
214 debug_assert!(self.is_valid());
215 &self.postorder
216 }
217
218 pub fn cfg_rpo(&self) -> impl Iterator<Item = &Block> {
223 debug_assert!(self.is_valid());
224 self.postorder.iter().rev()
225 }
226
227 pub fn idom(&self, block: Block) -> Option<Block> {
238 self.nodes[block].idom.into()
239 }
240
241 pub fn dominates<A, B>(&self, a: A, b: B, layout: &Layout) -> bool
252 where
253 A: Into<ProgramPoint>,
254 B: Into<ProgramPoint>,
255 {
256 let a = a.into();
257 let b = b.into();
258 match a {
259 ProgramPoint::Block(block_a) => match b {
260 ProgramPoint::Block(block_b) => self.block_dominates(block_a, block_b),
261 ProgramPoint::Inst(inst_b) => {
262 let block_b = layout
263 .inst_block(inst_b)
264 .expect("Instruction not in layout.");
265 self.block_dominates(block_a, block_b)
266 }
267 },
268 ProgramPoint::Inst(inst_a) => {
269 let block_a: Block = layout
270 .inst_block(inst_a)
271 .expect("Instruction not in layout.");
272 match b {
273 ProgramPoint::Block(block_b) => {
274 block_a != block_b && self.block_dominates(block_a, block_b)
275 }
276 ProgramPoint::Inst(inst_b) => {
277 let block_b = layout
278 .inst_block(inst_b)
279 .expect("Instruction not in layout.");
280 if block_a == block_b {
281 layout.pp_cmp(a, b) != Ordering::Greater
282 } else {
283 self.block_dominates(block_a, block_b)
284 }
285 }
286 }
287 }
288 }
289 }
290
291 pub fn block_dominates(&self, block_a: Block, block_b: Block) -> bool {
296 let na = &self.nodes[block_a];
297 let nb = &self.nodes[block_b];
298 na.dom_pre_number <= nb.dom_pre_number && na.dom_pre_max >= nb.dom_pre_max
299 }
300
301 pub fn children(&self, block: Block) -> ChildIter<'_> {
306 ChildIter {
307 domtree: self,
308 next: self.nodes[block].child,
309 }
310 }
311}
312
313impl DominatorTree {
314 pub fn new() -> Self {
317 Self {
318 stree: SpanningTree::new(),
319 nodes: SecondaryMap::new(),
320 postorder: Vec::new(),
321 dfs_worklist: Vec::new(),
322 eval_worklist: Vec::new(),
323 valid: false,
324 }
325 }
326
327 pub fn with_function(func: &Function, cfg: &ControlFlowGraph) -> Self {
329 let block_capacity = func.layout.block_capacity();
330 let mut domtree = Self {
331 stree: SpanningTree::with_capacity(block_capacity),
332 nodes: SecondaryMap::with_capacity(block_capacity),
333 postorder: Vec::with_capacity(block_capacity),
334 dfs_worklist: Vec::new(),
335 eval_worklist: Vec::new(),
336 valid: false,
337 };
338 domtree.compute(func, cfg);
339 domtree
340 }
341
342 pub fn compute(&mut self, func: &Function, cfg: &ControlFlowGraph) {
351 let _tt = timing::domtree();
352 debug_assert!(cfg.is_valid());
353
354 self.clear();
355 self.compute_spanning_tree(func);
356 self.compute_domtree(cfg);
357 self.compute_domtree_preorder();
358
359 self.valid = true;
360 }
361
362 pub fn clear(&mut self) {
365 self.stree.clear();
366 self.nodes.clear();
367 self.postorder.clear();
368 self.valid = false;
369 }
370
371 pub fn is_valid(&self) -> bool {
377 self.valid
378 }
379
380 fn compute_spanning_tree(&mut self, func: &Function) {
383 self.nodes.resize(func.dfg.num_blocks());
384 self.stree.reserve(func.dfg.num_blocks());
385
386 if let Some(block) = func.layout.entry_block() {
387 self.dfs_worklist.push(TraversalEvent::Enter(0, block));
388 }
389
390 loop {
391 match self.dfs_worklist.pop() {
392 Some(TraversalEvent::Enter(parent, block)) => {
393 let node = &mut self.nodes[block];
394 if node.pre_number != NOT_VISITED {
395 continue;
396 }
397
398 self.dfs_worklist.push(TraversalEvent::Exit(block));
399
400 let pre_number = self.stree.push(parent, block);
401 node.pre_number = pre_number;
402
403 self.dfs_worklist.extend(
405 func.block_successors(block)
406 .rev()
416 .filter(|successor| self.nodes[*successor].pre_number == NOT_VISITED)
418 .map(|successor| TraversalEvent::Enter(pre_number, successor)),
419 );
420 }
421 Some(TraversalEvent::Exit(block)) => {
422 self.postorder.push(block);
423 }
424 None => break,
425 }
426 }
427 }
428
429 fn eval(&mut self, v: u32, last_linked: u32) -> u32 {
434 if self.stree[v].ancestor < last_linked {
435 return self.stree[v].label;
436 }
437
438 let mut root = v;
440 loop {
441 self.eval_worklist.push(root);
442 root = self.stree[root].ancestor;
443
444 if self.stree[root].ancestor < last_linked {
445 break;
446 }
447 }
448
449 let mut prev = root;
450 let root = self.stree[prev].ancestor;
451
452 while let Some(curr) = self.eval_worklist.pop() {
455 if self.stree[prev].label < self.stree[curr].label {
456 self.stree[curr].label = self.stree[prev].label;
457 }
458
459 self.stree[curr].ancestor = root;
460 prev = curr;
461 }
462
463 self.stree[v].label
464 }
465
466 fn compute_domtree(&mut self, cfg: &ControlFlowGraph) {
467 for w in (1..self.stree.len() as u32).rev() {
469 let w_node = &mut self.stree[w];
470 let block = w_node.block.expect("Virtual root must have been excluded");
471 let mut semi = w_node.ancestor;
472
473 let last_linked = w + 1;
474
475 for pred in cfg
476 .pred_iter(block)
477 .map(|pred: BlockPredecessor| pred.block)
478 {
479 if self.nodes[pred].pre_number == NOT_VISITED {
481 continue;
482 }
483
484 let semi_candidate = self.eval(self.nodes[pred].pre_number, last_linked);
485 semi = core::cmp::min(semi, semi_candidate);
486 }
487
488 let w_node = &mut self.stree[w];
489 w_node.label = semi;
490 w_node.semi = semi;
491 }
492
493 for v in 1..self.stree.len() as u32 {
495 let semi = self.stree[v].semi;
496 let block = self.stree[v]
497 .block
498 .expect("Virtual root must have been excluded");
499 let mut idom = self.stree[v].idom;
500
501 while idom > semi {
502 idom = self.stree[idom].idom;
503 }
504
505 self.stree[v].idom = idom;
506
507 self.nodes[block].idom = self.stree[idom].block;
508 }
509 }
510
511 fn compute_domtree_preorder(&mut self) {
515 for &block in &self.postorder {
521 if let Some(idom) = self.idom(block) {
522 let sib = mem::replace(&mut self.nodes[idom].child, block.into());
523 self.nodes[block].sibling = sib;
524 } else {
525 self.dfs_worklist.push(TraversalEvent::Enter(0, block));
527 }
528 }
529
530 debug_assert!(self.dfs_worklist.len() <= 1);
532 let mut n = 0;
533 while let Some(event) = self.dfs_worklist.pop() {
534 if let TraversalEvent::Enter(_, block) = event {
535 n += 1;
536 let node = &mut self.nodes[block];
537 node.dom_pre_number = n;
538 node.dom_pre_max = n;
539 if let Some(sibling) = node.sibling.expand() {
540 self.dfs_worklist.push(TraversalEvent::Enter(0, sibling));
541 }
542 if let Some(child) = node.child.expand() {
543 self.dfs_worklist.push(TraversalEvent::Enter(0, child));
544 }
545 }
546 }
547
548 for &block in &self.postorder {
552 if let Some(idom) = self.idom(block) {
553 let pre_max = cmp::max(self.nodes[block].dom_pre_max, self.nodes[idom].dom_pre_max);
554 self.nodes[idom].dom_pre_max = pre_max;
555 }
556 }
557 }
558}
559
560pub struct ChildIter<'a> {
562 domtree: &'a DominatorTree,
563 next: PackedOption<Block>,
564}
565
566impl<'a> Iterator for ChildIter<'a> {
567 type Item = Block;
568
569 fn next(&mut self) -> Option<Block> {
570 let n = self.next.expand();
571 if let Some(block) = n {
572 self.next = self.domtree.nodes[block].sibling;
573 }
574 n
575 }
576}
577
578#[cfg(test)]
579mod tests {
580 use super::*;
581 use crate::cursor::{Cursor, FuncCursor};
582 use crate::ir::types::*;
583 use crate::ir::{InstBuilder, TrapCode};
584
585 #[test]
586 fn empty() {
587 let func = Function::new();
588 let cfg = ControlFlowGraph::with_function(&func);
589 debug_assert!(cfg.is_valid());
590 let dtree = DominatorTree::with_function(&func, &cfg);
591 assert_eq!(0, dtree.nodes.keys().count());
592 assert_eq!(dtree.cfg_postorder(), &[]);
593 }
594
595 #[test]
596 fn unreachable_node() {
597 let mut func = Function::new();
598 let block0 = func.dfg.make_block();
599 let v0 = func.dfg.append_block_param(block0, I32);
600 let block1 = func.dfg.make_block();
601 let block2 = func.dfg.make_block();
602 let trap_block = func.dfg.make_block();
603
604 let mut cur = FuncCursor::new(&mut func);
605
606 cur.insert_block(block0);
607 cur.ins().brif(v0, block2, &[], trap_block, &[]);
608
609 cur.insert_block(trap_block);
610 cur.ins().trap(TrapCode::unwrap_user(1));
611
612 cur.insert_block(block1);
613 let v1 = cur.ins().iconst(I32, 1);
614 let v2 = cur.ins().iadd(v0, v1);
615 cur.ins().jump(block0, &[v2.into()]);
616
617 cur.insert_block(block2);
618 cur.ins().return_(&[v0]);
619
620 let cfg = ControlFlowGraph::with_function(cur.func);
621 let dt = DominatorTree::with_function(cur.func, &cfg);
622
623 assert_eq!(dt.cfg_postorder(), &[block2, trap_block, block0]);
633
634 let v2_def = cur.func.dfg.value_def(v2).unwrap_inst();
635 assert!(!dt.dominates(v2_def, block0, &cur.func.layout));
636 assert!(!dt.dominates(block0, v2_def, &cur.func.layout));
637
638 assert!(dt.block_dominates(block0, block0));
639 assert!(!dt.block_dominates(block0, block1));
640 assert!(dt.block_dominates(block0, block2));
641 assert!(!dt.block_dominates(block1, block0));
642 assert!(dt.block_dominates(block1, block1));
643 assert!(!dt.block_dominates(block1, block2));
644 assert!(!dt.block_dominates(block2, block0));
645 assert!(!dt.block_dominates(block2, block1));
646 assert!(dt.block_dominates(block2, block2));
647 }
648
649 #[test]
650 fn non_zero_entry_block() {
651 let mut func = Function::new();
652 let block0 = func.dfg.make_block();
653 let block1 = func.dfg.make_block();
654 let block2 = func.dfg.make_block();
655 let block3 = func.dfg.make_block();
656 let cond = func.dfg.append_block_param(block3, I32);
657
658 let mut cur = FuncCursor::new(&mut func);
659
660 cur.insert_block(block3);
661 let jmp_block3_block1 = cur.ins().jump(block1, &[]);
662
663 cur.insert_block(block1);
664 let br_block1_block0_block2 = cur.ins().brif(cond, block0, &[], block2, &[]);
665
666 cur.insert_block(block2);
667 cur.ins().jump(block0, &[]);
668
669 cur.insert_block(block0);
670
671 let cfg = ControlFlowGraph::with_function(cur.func);
672 let dt = DominatorTree::with_function(cur.func, &cfg);
673
674 assert_eq!(dt.cfg_postorder(), &[block0, block2, block1, block3]);
693
694 assert_eq!(cur.func.layout.entry_block().unwrap(), block3);
695 assert_eq!(dt.idom(block3), None);
696 assert_eq!(dt.idom(block1).unwrap(), block3);
697 assert_eq!(dt.idom(block2).unwrap(), block1);
698 assert_eq!(dt.idom(block0).unwrap(), block1);
699
700 assert!(dt.dominates(
701 br_block1_block0_block2,
702 br_block1_block0_block2,
703 &cur.func.layout
704 ));
705 assert!(!dt.dominates(br_block1_block0_block2, jmp_block3_block1, &cur.func.layout));
706 assert!(dt.dominates(jmp_block3_block1, br_block1_block0_block2, &cur.func.layout));
707 }
708
709 #[test]
710 fn backwards_layout() {
711 let mut func = Function::new();
712 let block0 = func.dfg.make_block();
713 let block1 = func.dfg.make_block();
714 let block2 = func.dfg.make_block();
715
716 let mut cur = FuncCursor::new(&mut func);
717
718 cur.insert_block(block0);
719 let jmp02 = cur.ins().jump(block2, &[]);
720
721 cur.insert_block(block1);
722 let trap = cur.ins().trap(TrapCode::unwrap_user(5));
723
724 cur.insert_block(block2);
725 let jmp21 = cur.ins().jump(block1, &[]);
726
727 let cfg = ControlFlowGraph::with_function(cur.func);
728 let dt = DominatorTree::with_function(cur.func, &cfg);
729
730 assert_eq!(cur.func.layout.entry_block(), Some(block0));
731 assert_eq!(dt.idom(block0), None);
732 assert_eq!(dt.idom(block1), Some(block2));
733 assert_eq!(dt.idom(block2), Some(block0));
734
735 assert!(dt.dominates(block0, block0, &cur.func.layout));
736 assert!(dt.dominates(block0, jmp02, &cur.func.layout));
737 assert!(dt.dominates(block0, block1, &cur.func.layout));
738 assert!(dt.dominates(block0, trap, &cur.func.layout));
739 assert!(dt.dominates(block0, block2, &cur.func.layout));
740 assert!(dt.dominates(block0, jmp21, &cur.func.layout));
741
742 assert!(!dt.dominates(jmp02, block0, &cur.func.layout));
743 assert!(dt.dominates(jmp02, jmp02, &cur.func.layout));
744 assert!(dt.dominates(jmp02, block1, &cur.func.layout));
745 assert!(dt.dominates(jmp02, trap, &cur.func.layout));
746 assert!(dt.dominates(jmp02, block2, &cur.func.layout));
747 assert!(dt.dominates(jmp02, jmp21, &cur.func.layout));
748
749 assert!(!dt.dominates(block1, block0, &cur.func.layout));
750 assert!(!dt.dominates(block1, jmp02, &cur.func.layout));
751 assert!(dt.dominates(block1, block1, &cur.func.layout));
752 assert!(dt.dominates(block1, trap, &cur.func.layout));
753 assert!(!dt.dominates(block1, block2, &cur.func.layout));
754 assert!(!dt.dominates(block1, jmp21, &cur.func.layout));
755
756 assert!(!dt.dominates(trap, block0, &cur.func.layout));
757 assert!(!dt.dominates(trap, jmp02, &cur.func.layout));
758 assert!(!dt.dominates(trap, block1, &cur.func.layout));
759 assert!(dt.dominates(trap, trap, &cur.func.layout));
760 assert!(!dt.dominates(trap, block2, &cur.func.layout));
761 assert!(!dt.dominates(trap, jmp21, &cur.func.layout));
762
763 assert!(!dt.dominates(block2, block0, &cur.func.layout));
764 assert!(!dt.dominates(block2, jmp02, &cur.func.layout));
765 assert!(dt.dominates(block2, block1, &cur.func.layout));
766 assert!(dt.dominates(block2, trap, &cur.func.layout));
767 assert!(dt.dominates(block2, block2, &cur.func.layout));
768 assert!(dt.dominates(block2, jmp21, &cur.func.layout));
769
770 assert!(!dt.dominates(jmp21, block0, &cur.func.layout));
771 assert!(!dt.dominates(jmp21, jmp02, &cur.func.layout));
772 assert!(dt.dominates(jmp21, block1, &cur.func.layout));
773 assert!(dt.dominates(jmp21, trap, &cur.func.layout));
774 assert!(!dt.dominates(jmp21, block2, &cur.func.layout));
775 assert!(dt.dominates(jmp21, jmp21, &cur.func.layout));
776 }
777
778 #[test]
779 fn insts_same_block() {
780 let mut func = Function::new();
781 let block0 = func.dfg.make_block();
782
783 let mut cur = FuncCursor::new(&mut func);
784
785 cur.insert_block(block0);
786 let v1 = cur.ins().iconst(I32, 1);
787 let v2 = cur.ins().iadd(v1, v1);
788 let v3 = cur.ins().iadd(v2, v2);
789 cur.ins().return_(&[]);
790
791 let cfg = ControlFlowGraph::with_function(cur.func);
792 let dt = DominatorTree::with_function(cur.func, &cfg);
793
794 let v1_def = cur.func.dfg.value_def(v1).unwrap_inst();
795 let v2_def = cur.func.dfg.value_def(v2).unwrap_inst();
796 let v3_def = cur.func.dfg.value_def(v3).unwrap_inst();
797
798 assert!(dt.dominates(v1_def, v2_def, &cur.func.layout));
799 assert!(dt.dominates(v2_def, v3_def, &cur.func.layout));
800 assert!(dt.dominates(v1_def, v3_def, &cur.func.layout));
801
802 assert!(!dt.dominates(v2_def, v1_def, &cur.func.layout));
803 assert!(!dt.dominates(v3_def, v2_def, &cur.func.layout));
804 assert!(!dt.dominates(v3_def, v1_def, &cur.func.layout));
805
806 assert!(dt.dominates(v2_def, v2_def, &cur.func.layout));
807 assert!(dt.dominates(block0, block0, &cur.func.layout));
808
809 assert!(dt.dominates(block0, v1_def, &cur.func.layout));
810 assert!(dt.dominates(block0, v2_def, &cur.func.layout));
811 assert!(dt.dominates(block0, v3_def, &cur.func.layout));
812
813 assert!(!dt.dominates(v1_def, block0, &cur.func.layout));
814 assert!(!dt.dominates(v2_def, block0, &cur.func.layout));
815 assert!(!dt.dominates(v3_def, block0, &cur.func.layout));
816 }
817}