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 post_number: u32,
36}
37
38const NOT_VISITED: u32 = 0;
40
41#[derive(Clone, Default)]
48struct SpanningTree {
49 nodes: Vec<SpanningTreeNode>,
50}
51
52impl SpanningTree {
53 fn new() -> Self {
54 Self {
56 nodes: vec![Default::default()],
57 }
58 }
59
60 fn with_capacity(capacity: usize) -> Self {
61 let mut nodes = Vec::with_capacity(capacity + 1);
63 nodes.push(Default::default());
64 Self { nodes }
65 }
66
67 fn len(&self) -> usize {
68 self.nodes.len()
69 }
70
71 fn reserve(&mut self, capacity: usize) {
72 self.nodes.reserve(capacity);
74 }
75
76 fn clear(&mut self) {
77 self.nodes.resize(1, Default::default());
78 }
79
80 fn push(&mut self, ancestor: u32, block: Block) -> u32 {
82 debug_assert!(!self.nodes.is_empty());
84
85 let pre_number = self.nodes.len() as u32;
86
87 self.nodes.push(SpanningTreeNode {
88 block: block.into(),
89 ancestor,
90 label: pre_number,
91 semi: pre_number,
92 idom: ancestor,
93 post_number: 0,
94 });
95
96 pre_number
97 }
98}
99
100impl core::ops::Index<u32> for SpanningTree {
101 type Output = SpanningTreeNode;
102
103 fn index(&self, idx: u32) -> &Self::Output {
104 &self.nodes[idx as usize]
105 }
106}
107
108impl core::ops::IndexMut<u32> for SpanningTree {
109 fn index_mut(&mut self, idx: u32) -> &mut Self::Output {
110 &mut self.nodes[idx as usize]
111 }
112}
113
114enum TraversalEvent {
118 Enter(u32, Block),
119 Exit(Block),
120}
121
122#[derive(Clone, Default)]
124struct DominatorTreeNode {
125 idom: PackedOption<Block>,
127 pre_number: u32,
129
130 child: PackedOption<Block>,
132
133 sibling: PackedOption<Block>,
136
137 dom_pre_number: u32,
140
141 dom_pre_max: u32,
144
145 post_number: u32,
147}
148
149pub struct DominatorTree {
152 stree: SpanningTree,
154 postorder: Vec<Block>,
156 nodes: SecondaryMap<Block, DominatorTreeNode>,
158
159 dfs_worklist: Vec<TraversalEvent>,
161 eval_worklist: Vec<u32>,
164
165 valid: bool,
166}
167
168impl core::fmt::Debug for DominatorTree {
169 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
170 if !self.is_valid() {
171 return f.write_str("DominatorTree { <invalid> }");
172 }
173
174 let mut s = f.debug_tuple("DominatorTree");
175
176 if let Some((mut root, _)) = self.nodes.iter().find(|n| n.1.pre_number != NOT_VISITED) {
177 loop {
178 if let Some(b) = self.idom(root)
179 && b != root
180 {
181 root = b;
182 } else {
183 break;
184 }
185 }
186
187 fn fmt_block(domtree: &DominatorTree, block: Block) -> impl core::fmt::Debug {
188 core::fmt::from_fn(move |f| {
189 let children = domtree
190 .children(block)
191 .map(|c| fmt_block(domtree, c))
192 .collect::<Vec<_>>();
193 let mut s = f.debug_tuple(&format!("{block}"));
194 if !children.is_empty() {
195 s.field(&children);
196 }
197 s.finish()
198 })
199 }
200
201 s.field(&fmt_block(self, root));
202 }
203
204 s.finish()
205 }
206}
207
208impl DominatorTree {
210 pub fn is_reachable(&self, block: Block) -> bool {
212 self.nodes[block].pre_number != NOT_VISITED
213 }
214
215 pub fn cfg_postorder(&self) -> &[Block] {
220 debug_assert!(self.is_valid());
221 &self.postorder
222 }
223
224 pub fn cfg_rpo(&self) -> impl Iterator<Item = &Block> {
229 debug_assert!(self.is_valid());
230 self.postorder.iter().rev()
231 }
232
233 pub fn idom(&self, block: Block) -> Option<Block> {
244 self.nodes[block].idom.into()
245 }
246
247 pub fn dominates<A, B>(&self, a: A, b: B, layout: &Layout) -> bool
258 where
259 A: Into<ProgramPoint>,
260 B: Into<ProgramPoint>,
261 {
262 let a = a.into();
263 let b = b.into();
264 match a {
265 ProgramPoint::Block(block_a) => match b {
266 ProgramPoint::Block(block_b) => self.block_dominates(block_a, block_b),
267 ProgramPoint::Inst(inst_b) => {
268 let block_b = layout
269 .inst_block(inst_b)
270 .expect("Instruction not in layout.");
271 self.block_dominates(block_a, block_b)
272 }
273 },
274 ProgramPoint::Inst(inst_a) => {
275 let block_a: Block = layout
276 .inst_block(inst_a)
277 .expect("Instruction not in layout.");
278 match b {
279 ProgramPoint::Block(block_b) => {
280 block_a != block_b && self.block_dominates(block_a, block_b)
281 }
282 ProgramPoint::Inst(inst_b) => {
283 let block_b = layout
284 .inst_block(inst_b)
285 .expect("Instruction not in layout.");
286 if block_a == block_b {
287 layout.pp_cmp(a, b) != Ordering::Greater
288 } else {
289 self.block_dominates(block_a, block_b)
290 }
291 }
292 }
293 }
294 }
295 }
296
297 pub fn block_dominates(&self, block_a: Block, block_b: Block) -> bool {
302 let na = &self.nodes[block_a];
303 let nb = &self.nodes[block_b];
304 na.dom_pre_number <= nb.dom_pre_number && na.dom_pre_max >= nb.dom_pre_max
305 }
306
307 pub fn post_number(&self, block: Block) -> u32 {
309 self.nodes[block].post_number
310 }
311
312 pub fn children(&self, block: Block) -> ChildIter<'_> {
317 ChildIter {
318 domtree: self,
319 next: self.nodes[block].child,
320 }
321 }
322}
323
324impl DominatorTree {
325 pub fn new() -> Self {
328 Self {
329 stree: SpanningTree::new(),
330 nodes: SecondaryMap::new(),
331 postorder: Vec::new(),
332 dfs_worklist: Vec::new(),
333 eval_worklist: Vec::new(),
334 valid: false,
335 }
336 }
337
338 pub fn with_function(func: &Function, cfg: &ControlFlowGraph) -> Self {
340 let block_capacity = func.layout.block_capacity();
341 let mut domtree = Self {
342 stree: SpanningTree::with_capacity(block_capacity),
343 nodes: SecondaryMap::with_capacity(block_capacity),
344 postorder: Vec::with_capacity(block_capacity),
345 dfs_worklist: Vec::new(),
346 eval_worklist: Vec::new(),
347 valid: false,
348 };
349 domtree.compute(func, cfg);
350 domtree
351 }
352
353 pub fn compute(&mut self, func: &Function, cfg: &ControlFlowGraph) {
362 let _tt = timing::domtree();
363 debug_assert!(cfg.is_valid());
364
365 self.clear();
366 self.compute_spanning_tree(func);
367 self.compute_domtree(cfg);
368 self.compute_domtree_preorder();
369
370 self.valid = true;
371 }
372
373 pub fn clear(&mut self) {
376 self.stree.clear();
377 self.nodes.clear();
378 self.postorder.clear();
379 self.valid = false;
380 }
381
382 pub fn is_valid(&self) -> bool {
388 self.valid
389 }
390
391 fn compute_spanning_tree(&mut self, func: &Function) {
394 self.nodes.resize(func.dfg.num_blocks());
395 self.stree.reserve(func.dfg.num_blocks());
396
397 if let Some(block) = func.layout.entry_block() {
398 self.dfs_worklist.push(TraversalEvent::Enter(0, block));
399 }
400
401 loop {
402 match self.dfs_worklist.pop() {
403 Some(TraversalEvent::Enter(parent, block)) => {
404 let node = &mut self.nodes[block];
405 if node.pre_number != NOT_VISITED {
406 continue;
407 }
408
409 self.dfs_worklist.push(TraversalEvent::Exit(block));
410
411 let pre_number = self.stree.push(parent, block);
412 node.pre_number = pre_number;
413
414 self.dfs_worklist.extend(
416 func.block_successors(block)
417 .rev()
427 .filter(|successor| self.nodes[*successor].pre_number == NOT_VISITED)
429 .map(|successor| TraversalEvent::Enter(pre_number, successor)),
430 );
431 }
432 Some(TraversalEvent::Exit(block)) => {
433 let pre_number = self.nodes[block].pre_number;
434 self.stree[pre_number].post_number = self.postorder.len() as u32;
435 self.postorder.push(block);
436 }
437 None => break,
438 }
439 }
440 }
441
442 fn eval(&mut self, v: u32, last_linked: u32) -> u32 {
447 if self.stree[v].ancestor < last_linked {
448 return self.stree[v].label;
449 }
450
451 let mut root = v;
453 loop {
454 self.eval_worklist.push(root);
455 root = self.stree[root].ancestor;
456
457 if self.stree[root].ancestor < last_linked {
458 break;
459 }
460 }
461
462 let mut prev = root;
463 let root = self.stree[prev].ancestor;
464
465 while let Some(curr) = self.eval_worklist.pop() {
468 if self.stree[prev].label < self.stree[curr].label {
469 self.stree[curr].label = self.stree[prev].label;
470 }
471
472 self.stree[curr].ancestor = root;
473 prev = curr;
474 }
475
476 self.stree[v].label
477 }
478
479 fn compute_domtree(&mut self, cfg: &ControlFlowGraph) {
480 for w in (1..self.stree.len() as u32).rev() {
482 let w_node = &mut self.stree[w];
483 let block = w_node.block.expect("Virtual root must have been excluded");
484 let mut semi = w_node.ancestor;
485
486 let last_linked = w + 1;
487
488 for pred in cfg
489 .pred_iter(block)
490 .map(|pred: BlockPredecessor| pred.block)
491 {
492 if self.nodes[pred].pre_number == NOT_VISITED {
494 continue;
495 }
496
497 let semi_candidate = self.eval(self.nodes[pred].pre_number, last_linked);
498 semi = core::cmp::min(semi, semi_candidate);
499 }
500
501 let w_node = &mut self.stree[w];
502 w_node.label = semi;
503 w_node.semi = semi;
504 }
505
506 for v in 1..self.stree.len() as u32 {
508 let semi = self.stree[v].semi;
509 let block = self.stree[v]
510 .block
511 .expect("Virtual root must have been excluded");
512 let mut idom = self.stree[v].idom;
513
514 while idom > semi {
515 idom = self.stree[idom].idom;
516 }
517
518 self.stree[v].idom = idom;
519
520 self.nodes[block].idom = self.stree[idom].block;
521 }
522 }
523
524 fn compute_domtree_preorder(&mut self) {
528 for (i, &block) in self.postorder.iter().enumerate() {
530 self.nodes[block].post_number = i as u32;
531 }
532
533 for &block in &self.postorder {
539 if let Some(idom) = self.idom(block) {
540 let sib = mem::replace(&mut self.nodes[idom].child, block.into());
541 self.nodes[block].sibling = sib;
542 } else {
543 self.dfs_worklist.push(TraversalEvent::Enter(0, block));
545 }
546 }
547
548 debug_assert!(self.dfs_worklist.len() <= 1);
550 let mut n = 0;
551 while let Some(event) = self.dfs_worklist.pop() {
552 if let TraversalEvent::Enter(_, block) = event {
553 n += 1;
554 let node = &mut self.nodes[block];
555 node.dom_pre_number = n;
556 node.dom_pre_max = n;
557 if let Some(sibling) = node.sibling.expand() {
558 self.dfs_worklist.push(TraversalEvent::Enter(0, sibling));
559 }
560 if let Some(child) = node.child.expand() {
561 self.dfs_worklist.push(TraversalEvent::Enter(0, child));
562 }
563 }
564 }
565
566 for &block in &self.postorder {
570 if let Some(idom) = self.idom(block) {
571 let pre_max = cmp::max(self.nodes[block].dom_pre_max, self.nodes[idom].dom_pre_max);
572 self.nodes[idom].dom_pre_max = pre_max;
573 }
574 }
575 }
576}
577
578pub struct ChildIter<'a> {
580 domtree: &'a DominatorTree,
581 next: PackedOption<Block>,
582}
583
584impl<'a> Iterator for ChildIter<'a> {
585 type Item = Block;
586
587 fn next(&mut self) -> Option<Block> {
588 let n = self.next.expand();
589 if let Some(block) = n {
590 self.next = self.domtree.nodes[block].sibling;
591 }
592 n
593 }
594}
595
596#[cfg(test)]
597mod tests {
598 use super::*;
599 use crate::cursor::{Cursor, FuncCursor};
600 use crate::ir::types::*;
601 use crate::ir::{InstBuilder, TrapCode};
602
603 #[test]
604 fn empty() {
605 let func = Function::new();
606 let cfg = ControlFlowGraph::with_function(&func);
607 debug_assert!(cfg.is_valid());
608 let dtree = DominatorTree::with_function(&func, &cfg);
609 assert_eq!(0, dtree.nodes.keys().count());
610 assert_eq!(dtree.cfg_postorder(), &[]);
611 }
612
613 #[test]
614 fn unreachable_node() {
615 let mut func = Function::new();
616 let block0 = func.dfg.make_block();
617 let v0 = func.dfg.append_block_param(block0, I32);
618 let block1 = func.dfg.make_block();
619 let block2 = func.dfg.make_block();
620 let trap_block = func.dfg.make_block();
621
622 let mut cur = FuncCursor::new(&mut func);
623
624 cur.insert_block(block0);
625 cur.ins().brif(v0, block2, &[], trap_block, &[]);
626
627 cur.insert_block(trap_block);
628 cur.ins().trap(TrapCode::unwrap_user(1));
629
630 cur.insert_block(block1);
631 let v1 = cur.ins().iconst(I32, 1);
632 let v2 = cur.ins().iadd(v0, v1);
633 cur.ins().jump(block0, &[v2.into()]);
634
635 cur.insert_block(block2);
636 cur.ins().return_(&[v0]);
637
638 let cfg = ControlFlowGraph::with_function(cur.func);
639 let dt = DominatorTree::with_function(cur.func, &cfg);
640
641 assert_eq!(dt.cfg_postorder(), &[block2, trap_block, block0]);
651
652 let v2_def = cur.func.dfg.value_def(v2).unwrap_inst();
653 assert!(!dt.dominates(v2_def, block0, &cur.func.layout));
654 assert!(!dt.dominates(block0, v2_def, &cur.func.layout));
655
656 assert!(dt.block_dominates(block0, block0));
657 assert!(!dt.block_dominates(block0, block1));
658 assert!(dt.block_dominates(block0, block2));
659 assert!(!dt.block_dominates(block1, block0));
660 assert!(dt.block_dominates(block1, block1));
661 assert!(!dt.block_dominates(block1, block2));
662 assert!(!dt.block_dominates(block2, block0));
663 assert!(!dt.block_dominates(block2, block1));
664 assert!(dt.block_dominates(block2, block2));
665 }
666
667 #[test]
668 fn non_zero_entry_block() {
669 let mut func = Function::new();
670 let block0 = func.dfg.make_block();
671 let block1 = func.dfg.make_block();
672 let block2 = func.dfg.make_block();
673 let block3 = func.dfg.make_block();
674 let cond = func.dfg.append_block_param(block3, I32);
675
676 let mut cur = FuncCursor::new(&mut func);
677
678 cur.insert_block(block3);
679 let jmp_block3_block1 = cur.ins().jump(block1, &[]);
680
681 cur.insert_block(block1);
682 let br_block1_block0_block2 = cur.ins().brif(cond, block0, &[], block2, &[]);
683
684 cur.insert_block(block2);
685 cur.ins().jump(block0, &[]);
686
687 cur.insert_block(block0);
688
689 let cfg = ControlFlowGraph::with_function(cur.func);
690 let dt = DominatorTree::with_function(cur.func, &cfg);
691
692 assert_eq!(dt.cfg_postorder(), &[block0, block2, block1, block3]);
711
712 assert_eq!(cur.func.layout.entry_block().unwrap(), block3);
713 assert_eq!(dt.idom(block3), None);
714 assert_eq!(dt.idom(block1).unwrap(), block3);
715 assert_eq!(dt.idom(block2).unwrap(), block1);
716 assert_eq!(dt.idom(block0).unwrap(), block1);
717
718 assert!(dt.dominates(
719 br_block1_block0_block2,
720 br_block1_block0_block2,
721 &cur.func.layout
722 ));
723 assert!(!dt.dominates(br_block1_block0_block2, jmp_block3_block1, &cur.func.layout));
724 assert!(dt.dominates(jmp_block3_block1, br_block1_block0_block2, &cur.func.layout));
725 }
726
727 #[test]
728 fn backwards_layout() {
729 let mut func = Function::new();
730 let block0 = func.dfg.make_block();
731 let block1 = func.dfg.make_block();
732 let block2 = func.dfg.make_block();
733
734 let mut cur = FuncCursor::new(&mut func);
735
736 cur.insert_block(block0);
737 let jmp02 = cur.ins().jump(block2, &[]);
738
739 cur.insert_block(block1);
740 let trap = cur.ins().trap(TrapCode::unwrap_user(5));
741
742 cur.insert_block(block2);
743 let jmp21 = cur.ins().jump(block1, &[]);
744
745 let cfg = ControlFlowGraph::with_function(cur.func);
746 let dt = DominatorTree::with_function(cur.func, &cfg);
747
748 assert_eq!(cur.func.layout.entry_block(), Some(block0));
749 assert_eq!(dt.idom(block0), None);
750 assert_eq!(dt.idom(block1), Some(block2));
751 assert_eq!(dt.idom(block2), Some(block0));
752
753 assert!(dt.dominates(block0, block0, &cur.func.layout));
754 assert!(dt.dominates(block0, jmp02, &cur.func.layout));
755 assert!(dt.dominates(block0, block1, &cur.func.layout));
756 assert!(dt.dominates(block0, trap, &cur.func.layout));
757 assert!(dt.dominates(block0, block2, &cur.func.layout));
758 assert!(dt.dominates(block0, jmp21, &cur.func.layout));
759
760 assert!(!dt.dominates(jmp02, block0, &cur.func.layout));
761 assert!(dt.dominates(jmp02, jmp02, &cur.func.layout));
762 assert!(dt.dominates(jmp02, block1, &cur.func.layout));
763 assert!(dt.dominates(jmp02, trap, &cur.func.layout));
764 assert!(dt.dominates(jmp02, block2, &cur.func.layout));
765 assert!(dt.dominates(jmp02, jmp21, &cur.func.layout));
766
767 assert!(!dt.dominates(block1, block0, &cur.func.layout));
768 assert!(!dt.dominates(block1, jmp02, &cur.func.layout));
769 assert!(dt.dominates(block1, block1, &cur.func.layout));
770 assert!(dt.dominates(block1, trap, &cur.func.layout));
771 assert!(!dt.dominates(block1, block2, &cur.func.layout));
772 assert!(!dt.dominates(block1, jmp21, &cur.func.layout));
773
774 assert!(!dt.dominates(trap, block0, &cur.func.layout));
775 assert!(!dt.dominates(trap, jmp02, &cur.func.layout));
776 assert!(!dt.dominates(trap, block1, &cur.func.layout));
777 assert!(dt.dominates(trap, trap, &cur.func.layout));
778 assert!(!dt.dominates(trap, block2, &cur.func.layout));
779 assert!(!dt.dominates(trap, jmp21, &cur.func.layout));
780
781 assert!(!dt.dominates(block2, block0, &cur.func.layout));
782 assert!(!dt.dominates(block2, jmp02, &cur.func.layout));
783 assert!(dt.dominates(block2, block1, &cur.func.layout));
784 assert!(dt.dominates(block2, trap, &cur.func.layout));
785 assert!(dt.dominates(block2, block2, &cur.func.layout));
786 assert!(dt.dominates(block2, jmp21, &cur.func.layout));
787
788 assert!(!dt.dominates(jmp21, block0, &cur.func.layout));
789 assert!(!dt.dominates(jmp21, jmp02, &cur.func.layout));
790 assert!(dt.dominates(jmp21, block1, &cur.func.layout));
791 assert!(dt.dominates(jmp21, trap, &cur.func.layout));
792 assert!(!dt.dominates(jmp21, block2, &cur.func.layout));
793 assert!(dt.dominates(jmp21, jmp21, &cur.func.layout));
794 }
795
796 #[test]
797 fn insts_same_block() {
798 let mut func = Function::new();
799 let block0 = func.dfg.make_block();
800
801 let mut cur = FuncCursor::new(&mut func);
802
803 cur.insert_block(block0);
804 let v1 = cur.ins().iconst(I32, 1);
805 let v2 = cur.ins().iadd(v1, v1);
806 let v3 = cur.ins().iadd(v2, v2);
807 cur.ins().return_(&[]);
808
809 let cfg = ControlFlowGraph::with_function(cur.func);
810 let dt = DominatorTree::with_function(cur.func, &cfg);
811
812 let v1_def = cur.func.dfg.value_def(v1).unwrap_inst();
813 let v2_def = cur.func.dfg.value_def(v2).unwrap_inst();
814 let v3_def = cur.func.dfg.value_def(v3).unwrap_inst();
815
816 assert!(dt.dominates(v1_def, v2_def, &cur.func.layout));
817 assert!(dt.dominates(v2_def, v3_def, &cur.func.layout));
818 assert!(dt.dominates(v1_def, v3_def, &cur.func.layout));
819
820 assert!(!dt.dominates(v2_def, v1_def, &cur.func.layout));
821 assert!(!dt.dominates(v3_def, v2_def, &cur.func.layout));
822 assert!(!dt.dominates(v3_def, v1_def, &cur.func.layout));
823
824 assert!(dt.dominates(v2_def, v2_def, &cur.func.layout));
825 assert!(dt.dominates(block0, block0, &cur.func.layout));
826
827 assert!(dt.dominates(block0, v1_def, &cur.func.layout));
828 assert!(dt.dominates(block0, v2_def, &cur.func.layout));
829 assert!(dt.dominates(block0, v3_def, &cur.func.layout));
830
831 assert!(!dt.dominates(v1_def, block0, &cur.func.layout));
832 assert!(!dt.dominates(v2_def, block0, &cur.func.layout));
833 assert!(!dt.dominates(v3_def, block0, &cur.func.layout));
834 }
835}