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