circomspect_program_structure/static_single_assignment/
dominator_tree.rs

1use log::trace;
2use std::collections::HashSet;
3use std::marker::PhantomData;
4
5use super::traits::DirectedGraphNode;
6
7type Index = usize;
8type DominatorInfo = Vec<HashSet<Index>>;
9type ImmediateDominatorInfo = Vec<Option<Index>>;
10
11// A structure which encapsulates the dominance relation on a CFG.
12pub struct DominatorTree<T: DirectedGraphNode> {
13    dominators: DominatorInfo,
14    immediate_dominators: ImmediateDominatorInfo,
15    dominator_successors: DominatorInfo,
16    dominance_frontier: DominatorInfo,
17    marker: PhantomData<T>,
18}
19
20impl<T: DirectedGraphNode> DominatorTree<T> {
21    pub fn new(basic_blocks: &[T]) -> DominatorTree<T> {
22        let dominators = compute_dominators(basic_blocks);
23        let (immediate_dominators, dominator_successors) =
24            compute_immediate_dominators(basic_blocks, &dominators);
25        let dominance_frontier = compute_dominance_frontier(basic_blocks, &immediate_dominators);
26        // We assume that the first block (with index 0) represents the entry block.
27        assert!(immediate_dominators[0].is_none());
28        DominatorTree {
29            dominators,
30            immediate_dominators,
31            dominator_successors,
32            dominance_frontier,
33            marker: PhantomData,
34        }
35    }
36
37    pub fn entry_block(&self) -> Index {
38        Index::default()
39    }
40
41    pub fn get_dominators(&self, i: Index) -> HashSet<Index> {
42        self.dominators[i].clone()
43    }
44
45    pub fn get_immediate_dominator(&self, i: Index) -> Option<Index> {
46        self.immediate_dominators[i]
47    }
48
49    pub fn get_dominator_successors(&self, i: Index) -> HashSet<Index> {
50        self.dominator_successors[i].clone()
51    }
52
53    pub fn get_dominance_frontier(&self, i: Index) -> HashSet<Index> {
54        self.dominance_frontier[i].clone()
55    }
56}
57
58// This is a stupid simple (quadratic) algorithm based on an iterative data-flow analysis.
59fn compute_dominators<T: DirectedGraphNode>(basic_blocks: &[T]) -> DominatorInfo {
60    let mut dominators = Vec::new();
61    let nof_blocks = basic_blocks.len();
62    dominators.push(HashSet::from([0]));
63    for _ in 1..basic_blocks.len() {
64        dominators.push((0..nof_blocks).collect());
65    }
66
67    let mut done = false;
68    while !done {
69        done = true;
70        for i in 1..nof_blocks {
71            let mut new_dominators: HashSet<usize> = (0..nof_blocks).collect();
72            for &j in basic_blocks[i].predecessors() {
73                new_dominators = new_dominators.intersection(&dominators[j]).copied().collect();
74            }
75            new_dominators.insert(i);
76            if new_dominators != dominators[i] {
77                dominators[i] = new_dominators;
78                done = false;
79            }
80        }
81    }
82    dominators
83}
84
85// Compute immediate dominators (a `Vec<Option<usize>>`) and the dominator tree relation (a
86// `Vec<HashSet<usize>>`). (Note that the entry block of the CFG has no immediate dominator.)
87fn compute_immediate_dominators<T: DirectedGraphNode>(
88    basic_blocks: &[T],
89    dominators: &DominatorInfo,
90) -> (ImmediateDominatorInfo, DominatorInfo) {
91    let nof_blocks = basic_blocks.len();
92    let mut immediate_dominators = vec![None; nof_blocks];
93    let mut dominator_successors = vec![HashSet::new(); nof_blocks];
94
95    for i in 0..nof_blocks {
96        trace!("the dominator set of block {i} is {:?}", dominators[i]);
97        let mut idom_candidates: HashSet<usize> = dominators[i].clone();
98        idom_candidates.remove(&i);
99
100        if idom_candidates.len() > 1 {
101            // The set `all_dominators` is the strict up set of the nodes dominators. I.e.
102            //
103            //     `all_dominators(i) = U {Dom(j) - {j}; j strictly dominates i}`.
104            //
105            // The immediate dominator of the node will be the unique element in the set
106            // `idom_candidates - all_dominators` when this set is non-empty.
107            let mut all_dominators: HashSet<usize> = HashSet::new();
108            for j in &idom_candidates {
109                // 'all_dominators' is upwards closed.
110                if all_dominators.contains(j) {
111                    continue;
112                }
113                // Set `all_dominators = all_dominators U (Dom(i) \ {i}`.
114                all_dominators = dominators[*j]
115                    .clone()
116                    .into_iter()
117                    .filter(|&k| k != *j) // Remove i.
118                    .collect::<HashSet<usize>>()
119                    .union(&all_dominators)
120                    .copied()
121                    .collect();
122            }
123            idom_candidates = &idom_candidates - &all_dominators;
124            assert!(idom_candidates.len() <= 1);
125        }
126        if let Some(&j) = idom_candidates.iter().next() {
127            trace!("the immediate dominator of {i} is {j}");
128            immediate_dominators[i] = Some(j);
129            dominator_successors[j].insert(i);
130        }
131    }
132    (immediate_dominators, dominator_successors)
133}
134
135// Compute dominance frontiers (a `Vec<HashSet<usize>>`) of all nodes. The node
136// `i` is in the _dominance frontier_ of the node `j` if `j` dominates an
137// immediate predecessor of `i`, but `j` does not strictly dominate `i`.
138fn compute_dominance_frontier<T: DirectedGraphNode>(
139    basic_blocks: &[T],
140    immediate_dominators: &ImmediateDominatorInfo,
141) -> DominatorInfo {
142    let nof_blocks = basic_blocks.len();
143    let mut dominance_frontier = vec![HashSet::new(); nof_blocks];
144    for i in 0..nof_blocks {
145        if basic_blocks[i].predecessors().len() > 1 {
146            for &j in basic_blocks[i].predecessors() {
147                let mut k = j;
148                while Some(k) != immediate_dominators[i] {
149                    dominance_frontier[k].insert(i);
150                    k = match immediate_dominators[k] {
151                        Some(idom) => idom,
152                        None => break,
153                    };
154                }
155            }
156        }
157    }
158    dominance_frontier
159}