1use crate::dominator_tree::{ChildIter, DomTreeGraph, DominatorTree};
25use crate::flowgraph::{BlockPredecessor, ControlFlowGraph};
26use crate::ir::{Block, Layout, ProgramPoint};
27use core::cmp::Ordering;
28
29struct ReverseGraph<'a> {
33 cfg: &'a ControlFlowGraph,
34}
35
36impl DomTreeGraph for ReverseGraph<'_> {
37 fn num_blocks(&self) -> usize {
38 self.cfg.num_blocks()
39 }
40
41 fn roots(&self) -> impl Iterator<Item = Block> {
42 self.cfg
53 .blocks()
54 .filter(|&block| self.cfg.succ_iter(block).next().is_none())
55 }
56
57 fn successors(&self, block: Block) -> impl Iterator<Item = Block> {
58 self.cfg
61 .pred_iter(block)
62 .map(|pred: BlockPredecessor| pred.block)
63 }
64
65 fn predecessors(&self, block: Block) -> impl Iterator<Item = Block> {
66 self.cfg.succ_iter(block)
69 }
70}
71
72pub struct PostDominatorTree {
74 dom_tree: DominatorTree,
77}
78
79impl Default for PostDominatorTree {
80 fn default() -> Self {
81 Self::new()
82 }
83}
84
85impl PostDominatorTree {
86 pub fn new() -> Self {
90 Self {
91 dom_tree: DominatorTree::new(),
92 }
93 }
94
95 pub fn with_cfg(cfg: &ControlFlowGraph) -> Self {
97 let mut post_domtree = Self::new();
98 post_domtree.compute(cfg);
99 post_domtree
100 }
101
102 pub fn compute(&mut self, cfg: &ControlFlowGraph) {
105 debug_assert!(cfg.is_valid());
106 self.dom_tree.compute_from_graph(&ReverseGraph { cfg });
107 }
108
109 pub fn clear(&mut self) {
114 self.dom_tree.clear();
115 }
116
117 pub fn is_valid(&self) -> bool {
124 self.dom_tree.is_valid()
125 }
126
127 pub fn immediate_post_dominator(&self, block: Block) -> Option<Block> {
140 self.dom_tree.idom(block)
141 }
142
143 pub fn post_dominates<A, B>(&self, a: A, b: B, layout: &Layout) -> bool
146 where
147 A: Into<ProgramPoint>,
148 B: Into<ProgramPoint>,
149 {
150 let a = a.into();
151 let b = b.into();
152 match a {
153 ProgramPoint::Block(block_a) => match b {
154 ProgramPoint::Block(block_b) => self.block_post_dominates(block_a, block_b),
155 ProgramPoint::Inst(inst_b) => {
156 let block_b = layout
157 .inst_block(inst_b)
158 .expect("instruction not in layout");
159 block_a != block_b && self.block_post_dominates(block_a, block_b)
163 }
164 },
165 ProgramPoint::Inst(inst_a) => {
166 let block_a: Block = layout
167 .inst_block(inst_a)
168 .expect("Instruction not in layout.");
169 match b {
170 ProgramPoint::Block(block_b) => {
171 self.block_post_dominates(block_a, block_b)
175 }
176 ProgramPoint::Inst(inst_b) => {
177 let block_b = layout
178 .inst_block(inst_b)
179 .expect("instruction not in layout");
180 if block_a == block_b {
181 layout.pp_cmp(a, b) != Ordering::Less
184 } else {
185 self.block_post_dominates(block_a, block_b)
186 }
187 }
188 }
189 }
190 }
191 }
192
193 pub fn block_post_dominates(&self, block_a: Block, block_b: Block) -> bool {
196 self.dom_tree.block_dominates(block_a, block_b)
197 }
198
199 pub fn children(&self, block: Block) -> ChildIter<'_> {
204 self.dom_tree.children(block)
205 }
206
207 pub fn diverges(&self, block: Block) -> bool {
209 !self.dom_tree.is_reachable(block)
213 }
214}
215
216#[cfg(test)]
217mod tests {
218 use super::*;
219 use crate::cursor::{Cursor, FuncCursor};
220 use crate::ir::types::*;
221 use crate::ir::{Function, InstBuilder, TrapCode};
222 use alloc::string::String;
223 use alloc::vec::Vec;
224 use mutatis::{Mutate, check::Check, mutators as m};
225
226 #[test]
227 fn empty() {
228 let func = Function::new();
229 let cfg = ControlFlowGraph::with_function(&func);
230 let pdt = PostDominatorTree::with_cfg(&cfg);
231 assert!(pdt.is_valid());
232 }
233
234 #[test]
235 fn lifecycle() {
236 let mut func = Function::new();
237 let block0 = func.dfg.make_block();
238 let mut cur = FuncCursor::new(&mut func);
239 cur.insert_block(block0);
240 cur.ins().return_(&[]);
241 let cfg = ControlFlowGraph::with_function(cur.func);
242
243 let mut pdt = PostDominatorTree::new();
244 assert!(!pdt.is_valid());
245 pdt.compute(&cfg);
246 assert!(pdt.is_valid());
247 pdt.clear();
248 assert!(!pdt.is_valid());
249 pdt.compute(&cfg);
251 assert!(pdt.is_valid());
252 }
253
254 #[test]
255 fn straight_line() {
256 let mut func = Function::new();
257 let block0 = func.dfg.make_block();
258 let mut cur = FuncCursor::new(&mut func);
259 cur.insert_block(block0);
260 let v0 = cur.ins().iconst(I32, 1);
261 let v1 = cur.ins().iadd(v0, v0);
262 cur.ins().return_(&[]);
263
264 let cfg = ControlFlowGraph::with_function(cur.func);
265 let pdt = PostDominatorTree::with_cfg(&cfg);
266
267 assert_eq!(pdt.immediate_post_dominator(block0), None);
270 assert!(pdt.block_post_dominates(block0, block0));
271 assert!(!pdt.diverges(block0));
272
273 let v0_def = cur.func.dfg.value_def(v0).unwrap_inst();
275 let v1_def = cur.func.dfg.value_def(v1).unwrap_inst();
276 assert!(pdt.post_dominates(v1_def, v0_def, &cur.func.layout));
277 assert!(!pdt.post_dominates(v0_def, v1_def, &cur.func.layout));
278 assert!(pdt.post_dominates(v0_def, v0_def, &cur.func.layout));
279 }
280
281 #[test]
282 fn if_else_diamond() {
283 let mut func = Function::new();
284 let block0 = func.dfg.make_block();
285 let block1 = func.dfg.make_block();
286 let block2 = func.dfg.make_block();
287 let join = func.dfg.make_block();
288
289 let mut cur = FuncCursor::new(&mut func);
290
291 cur.insert_block(block0);
292 let v0 = cur.ins().iconst(I32, 0);
293 cur.ins().brif(v0, block1, &[], block2, &[]);
294
295 cur.insert_block(block1);
296 cur.ins().jump(join, &[]);
297
298 cur.insert_block(block2);
299 cur.ins().jump(join, &[]);
300
301 cur.insert_block(join);
302 cur.ins().return_(&[]);
303
304 let cfg = ControlFlowGraph::with_function(cur.func);
305 let pdt = PostDominatorTree::with_cfg(&cfg);
306
307 assert_eq!(pdt.immediate_post_dominator(block0), Some(join));
309 assert_eq!(pdt.immediate_post_dominator(block1), Some(join));
310 assert_eq!(pdt.immediate_post_dominator(block2), Some(join));
311 assert_eq!(pdt.immediate_post_dominator(join), None);
312
313 assert!(pdt.block_post_dominates(join, block0));
314 assert!(pdt.block_post_dominates(join, block1));
315 assert!(!pdt.block_post_dominates(block1, block0));
317 assert!(!pdt.block_post_dominates(block2, block0));
318 assert!(!pdt.block_post_dominates(block0, join));
320
321 for block in [block0, block1, block2, join] {
322 assert!(!pdt.diverges(block));
323 }
324
325 let layout = &cur.func.layout;
327 let entry_term = layout.last_inst(block0).unwrap();
328 let join_term = layout.last_inst(join).unwrap();
329 assert!(pdt.post_dominates(join_term, entry_term, layout));
331 assert!(!pdt.post_dominates(entry_term, join_term, layout));
333 assert!(pdt.post_dominates(join, entry_term, layout));
335 assert!(!pdt.post_dominates(entry_term, join, layout));
336
337 let mut kids = pdt.children(join).collect::<alloc::vec::Vec<_>>();
339 kids.sort();
340 assert_eq!(kids, [block0, block1, block2]);
341 }
342
343 #[test]
344 fn terminating_loop() {
345 let mut func = Function::new();
346 let entry = func.dfg.make_block();
347 let header = func.dfg.make_block();
348 let body = func.dfg.make_block();
349 let exit = func.dfg.make_block();
350
351 let mut cur = FuncCursor::new(&mut func);
352
353 cur.insert_block(entry);
354 cur.ins().jump(header, &[]);
355 cur.insert_block(header);
356 let v0 = cur.ins().iconst(I32, 0);
357 cur.ins().brif(v0, body, &[], exit, &[]);
358
359 cur.insert_block(body);
360 cur.ins().jump(header, &[]);
361 cur.insert_block(exit);
362 cur.ins().return_(&[]);
363
364 let cfg = ControlFlowGraph::with_function(cur.func);
365 let pdt = PostDominatorTree::with_cfg(&cfg);
366
367 assert_eq!(pdt.immediate_post_dominator(entry), Some(header));
368 assert_eq!(pdt.immediate_post_dominator(header), Some(exit));
369 assert_eq!(pdt.immediate_post_dominator(body), Some(header));
370 assert_eq!(pdt.immediate_post_dominator(exit), None);
371
372 assert!(pdt.block_post_dominates(exit, entry));
373 assert!(pdt.block_post_dominates(exit, body));
374 assert!(pdt.block_post_dominates(header, body));
375 assert!(!pdt.block_post_dominates(body, header));
376
377 for block in [entry, header, body, exit] {
379 assert!(!pdt.diverges(block));
380 }
381 }
382
383 #[test]
384 fn infinite_loop() {
385 let mut func = Function::new();
386 let block0 = func.dfg.make_block();
387 let mut cur = FuncCursor::new(&mut func);
388 cur.insert_block(block0);
389 cur.ins().jump(block0, &[]);
390
391 let cfg = ControlFlowGraph::with_function(cur.func);
392 let pdt = PostDominatorTree::with_cfg(&cfg);
393
394 assert!(pdt.is_valid());
397 assert!(pdt.diverges(block0));
398 assert_eq!(pdt.immediate_post_dominator(block0), None);
399 }
400
401 #[test]
402 fn infinite_loop_with_side_exit() {
403 let mut func = Function::new();
404 let entry = func.dfg.make_block();
405 let header = func.dfg.make_block();
406 let body = func.dfg.make_block();
407 let exit = func.dfg.make_block();
408
409 let mut cur = FuncCursor::new(&mut func);
410
411 cur.insert_block(entry);
412 cur.ins().jump(header, &[]);
413
414 cur.insert_block(header);
415 let v0 = cur.ins().iconst(I32, 0);
416 cur.ins().brif(v0, exit, &[], body, &[]);
417
418 cur.insert_block(body);
420 cur.ins().jump(body, &[]);
421
422 cur.insert_block(exit);
423 cur.ins().return_(&[]);
424
425 let cfg = ControlFlowGraph::with_function(cur.func);
426 let pdt = PostDominatorTree::with_cfg(&cfg);
427
428 assert!(pdt.diverges(body));
430 assert_eq!(pdt.immediate_post_dominator(body), None);
431
432 assert!(!pdt.diverges(entry));
433 assert!(!pdt.diverges(header));
434 assert!(!pdt.diverges(exit));
435 assert_eq!(pdt.immediate_post_dominator(header), Some(exit));
436 assert_eq!(pdt.immediate_post_dominator(entry), Some(header));
437 assert_eq!(pdt.immediate_post_dominator(exit), None);
438 }
439
440 #[test]
441 fn multiple_returns() {
442 let mut func = Function::new();
443 let entry = func.dfg.make_block();
444 let block1 = func.dfg.make_block();
445 let block2 = func.dfg.make_block();
446
447 let mut cur = FuncCursor::new(&mut func);
448
449 cur.insert_block(entry);
450 let v0 = cur.ins().iconst(I32, 0);
451 cur.ins().brif(v0, block1, &[], block2, &[]);
452
453 cur.insert_block(block1);
454 cur.ins().return_(&[]);
455
456 cur.insert_block(block2);
457 cur.ins().return_(&[]);
458
459 let cfg = ControlFlowGraph::with_function(cur.func);
460 let pdt = PostDominatorTree::with_cfg(&cfg);
461
462 assert_eq!(pdt.immediate_post_dominator(block1), None);
465 assert_eq!(pdt.immediate_post_dominator(block2), None);
466 assert_eq!(pdt.immediate_post_dominator(entry), None);
467
468 assert!(!pdt.block_post_dominates(block1, entry));
469 assert!(!pdt.block_post_dominates(block2, entry));
470
471 assert!(!pdt.block_post_dominates(block1, block2));
473 assert!(!pdt.block_post_dominates(block2, block1));
474
475 for block in [entry, block1, block2] {
476 assert!(!pdt.diverges(block));
477 }
478 }
479
480 #[test]
481 fn trap_as_exit() {
482 let mut func = Function::new();
483 let entry = func.dfg.make_block();
484 let ret_block = func.dfg.make_block();
485 let trap_block = func.dfg.make_block();
486
487 let mut cur = FuncCursor::new(&mut func);
488
489 cur.insert_block(entry);
490 let v0 = cur.ins().iconst(I32, 0);
491 cur.ins().brif(v0, ret_block, &[], trap_block, &[]);
492
493 cur.insert_block(ret_block);
494 cur.ins().return_(&[]);
495
496 cur.insert_block(trap_block);
497 cur.ins().trap(TrapCode::unwrap_user(1));
498
499 let cfg = ControlFlowGraph::with_function(cur.func);
500 let pdt = PostDominatorTree::with_cfg(&cfg);
501
502 assert_eq!(pdt.immediate_post_dominator(trap_block), None);
505 assert_eq!(pdt.immediate_post_dominator(ret_block), None);
506 assert_eq!(pdt.immediate_post_dominator(entry), None);
507
508 assert!(!pdt.diverges(trap_block));
509 assert!(!pdt.diverges(ret_block));
510 assert!(!pdt.diverges(entry));
511 }
512
513 #[test]
514 fn insts_post_dominate_same_block() {
515 let mut func = Function::new();
516 let block0 = func.dfg.make_block();
517
518 let mut cur = FuncCursor::new(&mut func);
519 cur.insert_block(block0);
520 let v1 = cur.ins().iconst(I32, 1);
521 let v2 = cur.ins().iadd(v1, v1);
522 let v3 = cur.ins().iadd(v2, v2);
523 cur.ins().return_(&[]);
524
525 let cfg = ControlFlowGraph::with_function(cur.func);
526 let pdt = PostDominatorTree::with_cfg(&cfg);
527
528 let v1_def = cur.func.dfg.value_def(v1).unwrap_inst();
529 let v2_def = cur.func.dfg.value_def(v2).unwrap_inst();
530 let v3_def = cur.func.dfg.value_def(v3).unwrap_inst();
531 let layout = &cur.func.layout;
532
533 assert!(pdt.post_dominates(v2_def, v1_def, layout));
535 assert!(pdt.post_dominates(v3_def, v1_def, layout));
536 assert!(pdt.post_dominates(v3_def, v2_def, layout));
537
538 assert!(!pdt.post_dominates(v1_def, v2_def, layout));
540 assert!(!pdt.post_dominates(v1_def, v3_def, layout));
541
542 assert!(pdt.post_dominates(v2_def, v2_def, layout));
544
545 assert!(pdt.post_dominates(v1_def, block0, layout));
547 assert!(!pdt.post_dominates(block0, v1_def, layout));
550
551 assert!(pdt.post_dominates(block0, block0, layout));
553 }
554
555 #[test]
562 fn post_dominators_match_oracle() -> mutatis::check::CheckResult<GraphSpec> {
563 use Terminator::*;
564
565 let corpus = [
566 GraphSpec {
568 blocks: alloc::vec![Return],
569 },
570 GraphSpec {
572 blocks: alloc::vec![Jump(1), Jump(2), Return],
573 },
574 GraphSpec {
576 blocks: alloc::vec![Brif(1, 2), Jump(3), Jump(3), Return],
577 },
578 GraphSpec {
580 blocks: alloc::vec![Jump(1), Brif(1, 2), Return],
581 },
582 GraphSpec {
584 blocks: alloc::vec![Jump(1), Jump(0)],
585 },
586 ];
587
588 Check::new()
589 .iters(10_000)
590 .run_with(m::default::<GraphSpec>(), corpus, check_post_dominance)
591 }
592
593 const MAX_BLOCKS: usize = 12;
596
597 #[derive(Clone, Debug, Default, Mutate)]
599 struct GraphSpec {
600 blocks: Vec<Terminator>,
601 }
602
603 impl GraphSpec {
604 fn fixup(&self) -> Option<Self> {
605 let n = self.blocks.len().min(MAX_BLOCKS);
606 if n == 0 {
607 return None;
608 }
609 let mut graph = GraphSpec {
610 blocks: self.blocks[..n].to_vec(),
611 };
612 for terminator in &mut graph.blocks {
613 terminator.fixup(n);
614 }
615 Some(graph)
616 }
617
618 fn build(&self) -> (Function, Vec<Block>) {
620 let mut func = Function::new();
621
622 let blocks: Vec<Block> = (0..self.blocks.len())
623 .map(|_| func.dfg.make_block())
624 .collect();
625
626 let mut cur = FuncCursor::new(&mut func);
627 for (i, terminator) in self.blocks.iter().enumerate() {
628 cur.insert_block(blocks[i]);
629 match *terminator {
630 Terminator::Return => {
631 cur.ins().return_(&[]);
632 }
633 Terminator::Jump(t) => {
634 cur.ins().jump(blocks[t], &[]);
635 }
636 Terminator::Brif(t1, t2) => {
637 let c = cur.ins().iconst(I32, 0);
638 cur.ins().brif(c, blocks[t1], &[], blocks[t2], &[]);
639 }
640 }
641 }
642 (func, blocks)
643 }
644 }
645
646 #[derive(Clone, Copy, Debug, Default, Mutate)]
648 enum Terminator {
649 #[default]
650 Return,
651 Jump(usize),
652 Brif(usize, usize),
653 }
654
655 impl Terminator {
656 fn fixup(&mut self, n: usize) {
657 match self {
658 Self::Return => {}
659 Self::Jump(a) => {
660 *a %= n;
661 }
662 Self::Brif(a, b) => {
663 *a %= n;
664 *b %= n;
665 }
666 }
667 }
668 }
669
670 fn check_post_dominance(graph: &GraphSpec) -> Result<(), String> {
673 let Some(graph) = graph.fixup() else {
674 return Ok(());
675 };
676
677 let (func, blocks) = graph.build();
678 let cfg = ControlFlowGraph::with_function(&func);
679 let pdt = PostDominatorTree::with_cfg(&cfg);
680
681 let sink = graph.blocks.len();
683 let succ: Vec<Vec<usize>> = graph
684 .blocks
685 .iter()
686 .map(|t| match *t {
687 Terminator::Return => alloc::vec![sink],
688 Terminator::Jump(x) => alloc::vec![x],
689 Terminator::Brif(x, y) => alloc::vec![x, y],
690 })
691 .collect();
692
693 let mut reaches = alloc::vec![false; graph.blocks.len() + 1];
696 reaches[sink] = true;
697 loop {
698 let mut changed = false;
699 for i in 0..graph.blocks.len() {
700 if !reaches[i] && succ[i].iter().any(|&s| reaches[s]) {
701 reaches[i] = true;
702 changed = true;
703 }
704 }
705 if !changed {
706 break;
707 }
708 }
709
710 let bit = |x: usize| 1u64 << x;
713 let universe = bit(graph.blocks.len() + 1) - 1;
714 let mut pdom = alloc::vec![universe; graph.blocks.len() + 1];
715 pdom[sink] = bit(sink);
716 loop {
717 let mut changed = false;
718 for i in 0..graph.blocks.len() {
719 let mut inter = u64::MAX;
720 for &s in &succ[i] {
721 inter &= pdom[s];
722 }
723 let next = bit(i) | inter;
724 if next != pdom[i] {
725 pdom[i] = next;
726 changed = true;
727 }
728 }
729 if !changed {
730 break;
731 }
732 }
733
734 for i in 0..graph.blocks.len() {
735 let expect_diverges = !reaches[i];
736 if pdt.diverges(blocks[i]) != expect_diverges {
737 return Err(format!(
738 "diverges({i}) = {}, expected {expect_diverges}",
739 pdt.diverges(blocks[i]),
740 ));
741 }
742
743 if expect_diverges {
744 if pdt.immediate_post_dominator(blocks[i]).is_some() {
745 return Err(format!(
746 "immediate_post_dominator({i}) should be None for a diverging block;",
747 ));
748 }
749 continue;
750 }
751
752 for a in 0..graph.blocks.len() {
754 let expect = pdom[i] & bit(a) != 0;
755 if pdt.block_post_dominates(blocks[a], blocks[i]) != expect {
756 return Err(format!(
757 "block_post_dominates({a}, {i}) = {}, expected {expect}",
758 pdt.block_post_dominates(blocks[a], blocks[i]),
759 ));
760 }
761 }
762
763 let strict = pdom[i] & !bit(i);
768 let mut best: Option<(u32, usize)> = None;
769 for x in 0..=graph.blocks.len() {
770 if strict & bit(x) != 0 {
771 let size = pdom[x].count_ones();
772 if best.map_or(true, |(best_size, _)| size > best_size) {
773 best = Some((size, x));
774 }
775 }
776 }
777 let expect_ipdom = match best {
778 Some((_, x)) if x != sink => Some(blocks[x]),
779 _ => None,
780 };
781 if pdt.immediate_post_dominator(blocks[i]) != expect_ipdom {
782 return Err(format!(
783 "immediate_post_dominator({i}) = {:?}, expected {expect_ipdom:?}",
784 pdt.immediate_post_dominator(blocks[i]),
785 ));
786 }
787 }
788
789 Ok(())
790 }
791}