vicis_core/pass/analysis/
dom_tree.rs

1use crate::traits::basic_block::{BasicBlock, BasicBlockData, BasicBlockLayout};
2use id_arena::Id;
3use rustc_hash::{FxHashMap, FxHashSet};
4
5#[derive(Debug)]
6pub struct DominatorTree<BB: BasicBlock> {
7    dom: DomTree<BB>,
8    frontier: FxHashMap<Id<BB>, FxHashSet<Id<BB>>>,
9    level: FxHashMap<Id<BB>, usize>,
10    root: Id<BB>,
11}
12
13type DomTree<BB> = FxHashMap<Id<BB>, FxHashSet<Id<BB>>>;
14
15type Map<T> = FxHashMap<T, T>;
16
17struct Context<'a, BB: BasicBlock, F: BasicBlockData<BB> + BasicBlockLayout<BB>> {
18    f: &'a F,
19    dfnum: FxHashMap<Id<BB>, usize>,
20    vertex: Vec<Id<BB>>,
21    semi: Map<Id<BB>>,
22    ancestor: Map<Id<BB>>,
23    idom: Map<Id<BB>>,
24    samedom: Map<Id<BB>>,
25    parent: Map<Id<BB>>,
26    best: Map<Id<BB>>,
27}
28
29impl<BB: BasicBlock> DominatorTree<BB> {
30    pub fn new<F: BasicBlockData<BB> + BasicBlockLayout<BB>>(f: &F) -> Self {
31        let mut dom = FxHashMap::default();
32        let ctx = Context::new(f).compute();
33
34        // x dominates y
35        for (&y, &x) in &ctx.idom {
36            dom.entry(x).or_insert_with(FxHashSet::default).insert(y);
37        }
38
39        fn leveling<BB: BasicBlock>(
40            level: &mut FxHashMap<Id<BB>, usize>,
41            dom: &FxHashMap<Id<BB>, FxHashSet<Id<BB>>>,
42            cur: Id<BB>,
43            cur_level: usize,
44        ) {
45            level.insert(cur, cur_level);
46            if dom.get(&cur).is_none() {
47                return;
48            }
49            for &child in dom.get(&cur).unwrap() {
50                leveling(level, dom, child, cur_level + 1);
51            }
52        }
53
54        let entry = f.order().next().unwrap();
55        let frontier = Self::compute_dom_frontier(&ctx, &dom, entry);
56        let mut level = FxHashMap::default();
57        leveling(&mut level, &dom, entry, 0);
58
59        Self {
60            dom,
61            frontier,
62            level,
63            root: entry,
64        }
65    }
66
67    pub fn dominates(&self, x: Id<BB>, y: Id<BB>) -> bool {
68        x == y
69            || self.dom.get(&x).map_or(false, |children| {
70                children.contains(&y) || children.iter().any(|&child| self.dominates(child, y))
71            })
72    }
73
74    pub fn dominance_frontier_of(&self, x: Id<BB>) -> Option<&FxHashSet<Id<BB>>> {
75        self.frontier.get(&x)
76    }
77
78    pub fn level_of(&self, x: Id<BB>) -> Option<usize> {
79        self.level.get(&x).copied()
80    }
81
82    pub fn root(&self) -> &Id<BB> {
83        &self.root
84    }
85
86    pub fn children_of(&self, x: Id<BB>) -> Option<&FxHashSet<Id<BB>>> {
87        self.dom.get(&x)
88    }
89
90    fn compute_dom_frontier_of<F: BasicBlockData<BB> + BasicBlockLayout<BB>>(
91        ctx: &Context<BB, F>,
92        dom: &FxHashMap<Id<BB>, FxHashSet<Id<BB>>>,
93        x: Id<BB>,
94        frontier: &mut FxHashMap<Id<BB>, FxHashSet<Id<BB>>>,
95    ) {
96        if frontier.contains_key(&x) {
97            // dominance frontier for x is already computed
98            return;
99        }
100
101        frontier.insert(x, FxHashSet::default());
102
103        for succ in ctx.f.get(x).succs() {
104            if ctx.idom.get(succ).map_or(true, |&x_| x != x_) {
105                frontier.get_mut(&x).unwrap().insert(*succ);
106            }
107            for child in dom.get(&x).unwrap_or(&FxHashSet::default()) {
108                Self::compute_dom_frontier_of(ctx, dom, *child, frontier);
109                for y in frontier.get(child).unwrap().clone() {
110                    if ctx.idom.get(&y).map_or(true, |&x_| x_ != x) {
111                        frontier.get_mut(&x).unwrap().insert(y);
112                    }
113                }
114            }
115        }
116    }
117
118    fn compute_dom_frontier<F: BasicBlockData<BB> + BasicBlockLayout<BB>>(
119        ctx: &Context<BB, F>,
120        dom: &FxHashMap<Id<BB>, FxHashSet<Id<BB>>>,
121        start: Id<BB>,
122    ) -> FxHashMap<Id<BB>, FxHashSet<Id<BB>>> {
123        let mut frontier = FxHashMap::default();
124        for &child in dom.get(&start).unwrap_or(&FxHashSet::default()) {
125            Self::compute_dom_frontier_of(ctx, dom, child, &mut frontier);
126        }
127        Self::compute_dom_frontier_of(ctx, dom, start, &mut frontier);
128        frontier
129    }
130}
131
132impl<'a, BB: BasicBlock, F: BasicBlockData<BB> + BasicBlockLayout<BB>> Context<'a, BB, F> {
133    fn new(f: &'a F) -> Self {
134        Self {
135            f,
136            dfnum: FxHashMap::default(),
137            semi: FxHashMap::default(),
138            ancestor: FxHashMap::default(),
139            idom: FxHashMap::default(),
140            samedom: FxHashMap::default(),
141            vertex: Vec::new(),
142            parent: FxHashMap::default(),
143            best: FxHashMap::default(),
144        }
145    }
146
147    fn compute(mut self) -> Self {
148        let entry = self.f.order().next().unwrap();
149        let mut bucket = FxHashMap::default();
150        let mut num = 0;
151
152        self.number_by_dfs(None, entry, &mut num);
153
154        for i in (1..num).rev() {
155            let node = self.vertex[i];
156            let pred = *self.parent.get(&node).unwrap();
157            let mut s = pred;
158
159            for v in self.f.get(node).preds() {
160                if !self.dfnum.contains_key(v) {
161                    // `v` is unreachable. skip it.
162                    continue;
163                }
164
165                let s_ = if self.dfnum[v] <= self.dfnum[&node] {
166                    *v
167                } else {
168                    let n = self.ancestor_with_lowest_semi(*v);
169                    *self.semi.get(&n).unwrap()
170                };
171                if self.dfnum[&s_] < self.dfnum[&s] {
172                    s = s_;
173                }
174            }
175
176            self.semi.insert(node, s);
177            bucket
178                .entry(s)
179                .or_insert_with(FxHashSet::default)
180                .insert(node);
181            self.link(pred, node);
182
183            if let Some(set) = bucket.get_mut(&pred) {
184                for v in &*set {
185                    let y = self.ancestor_with_lowest_semi(*v);
186                    if self.semi[&y] == self.semi[v] {
187                        self.idom.insert(*v, pred);
188                    } else {
189                        self.samedom.insert(*v, y);
190                    }
191                }
192                set.clear();
193            }
194        }
195
196        for &n in &self.vertex[1..] {
197            if let Some(s) = self.samedom.get(&n) {
198                self.idom.insert(n, *s);
199            }
200        }
201
202        self
203    }
204
205    fn number_by_dfs(&mut self, pred: Option<Id<BB>>, node: Id<BB>, num: &mut usize) {
206        if self.dfnum.contains_key(&node) {
207            return;
208        }
209
210        self.dfnum.insert(node, *num);
211        self.vertex.insert(*num, node);
212        if let Some(pred) = pred {
213            self.parent.insert(node, pred);
214        }
215        *num += 1;
216
217        for succ in self.f.get(node).succs() {
218            self.number_by_dfs(Some(node), *succ, num);
219        }
220    }
221
222    fn ancestor_with_lowest_semi(&mut self, node: Id<BB>) -> Id<BB> {
223        let a = *self.ancestor.get(&node).unwrap();
224        if self.ancestor.contains_key(&a) {
225            let b = self.ancestor_with_lowest_semi(a);
226            let aa = *self.ancestor.get(&a).unwrap();
227            self.ancestor.insert(node, aa);
228            if self.dfnum[&self.semi[&b]] < self.dfnum[&self.semi[&self.best[&node]]] {
229                self.best.insert(node, b);
230            }
231        }
232        *self.best.get(&node).unwrap()
233    }
234
235    fn link(&mut self, pred: Id<BB>, node: Id<BB>) {
236        self.ancestor.insert(node, pred);
237        self.best.insert(node, node);
238    }
239}
vicis_core/pass/analysis/dom_tree.rs

vicis_core/pass/analysis/
dom_tree.rs
