duskphantom_middle/analysis/
dominator_tree.rs

1// Copyright 2024 Duskphantom Authors
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14//
15// SPDX-License-Identifier: Apache-2.0
16
17use std::collections::{HashMap, HashSet};
18
19use crate::ir::{BBPtr, FunPtr};
20
21pub struct DominatorTree {
22    fun: FunPtr,
23    dominator_map: HashMap<BBPtr, HashSet<BBPtr>>,
24    dominatee_map: Option<HashMap<BBPtr, HashSet<BBPtr>>>,
25    idom_map: Option<HashMap<BBPtr, BBPtr>>,
26    postorder_map: Option<HashMap<BBPtr, i32>>,
27    df_map: Option<HashMap<BBPtr, HashSet<BBPtr>>>,
28}
29
30#[allow(unused)]
31impl DominatorTree {
32    pub fn new(fun: FunPtr) -> Self {
33        DominatorTree {
34            fun,
35            dominator_map: HashMap::new(),
36            dominatee_map: None,
37            idom_map: None,
38            postorder_map: None,
39            df_map: None,
40        }
41    }
42
43    pub fn is_dominate(&mut self, dominator: BBPtr, dominatee: BBPtr) -> bool {
44        self.get_dominator(dominatee).contains(&dominator)
45    }
46
47    pub fn get_dominator(&mut self, dominatee: BBPtr) -> HashSet<BBPtr> {
48        match self.dominator_map.get(&dominatee) {
49            Some(dom) => dom.clone(),
50            None => {
51                // Traverse up dominator tree
52                let mut dom = HashSet::new();
53                let mut cursor = dominatee;
54                loop {
55                    dom.insert(cursor);
56                    match self.get_idom(cursor) {
57                        Some(idom) => cursor = idom,
58                        None => break,
59                    }
60                }
61                dom
62            }
63        }
64    }
65
66    pub fn get_lca(&mut self, a: BBPtr, b: BBPtr) -> BBPtr {
67        self.get_idom_map();
68        let idoms = self.idom_map.as_ref().unwrap();
69        let depths = self.postorder_map.as_ref().unwrap();
70        intersect(a, b, depths, idoms)
71    }
72
73    pub fn get_dominatee(&mut self, dominator: BBPtr) -> HashSet<BBPtr> {
74        self.get_dominatee_map()
75            .get(&dominator)
76            .cloned()
77            .unwrap_or_default()
78    }
79
80    pub fn get_idom(&mut self, dominatee: BBPtr) -> Option<BBPtr> {
81        self.get_idom_map().get(&dominatee).copied()
82    }
83
84    pub fn get_df(&mut self, dominator: BBPtr) -> HashSet<BBPtr> {
85        self.get_df_map()
86            .get(&dominator)
87            .cloned()
88            .unwrap_or_default()
89    }
90
91    fn get_idom_map(&mut self) -> &HashMap<BBPtr, BBPtr> {
92        match self.idom_map {
93            Some(ref idoms) => idoms,
94            None => {
95                self.calculate_idom();
96                self.idom_map.as_ref().unwrap()
97            }
98        }
99    }
100
101    fn get_dominatee_map(&mut self) -> &HashMap<BBPtr, HashSet<BBPtr>> {
102        match self.dominatee_map {
103            Some(ref doms) => doms,
104            None => {
105                self.calculate_idom();
106                self.dominatee_map.as_ref().unwrap()
107            }
108        }
109    }
110
111    fn get_df_map(&mut self) -> &HashMap<BBPtr, HashSet<BBPtr>> {
112        match self.df_map {
113            Some(ref df) => df,
114            None => {
115                self.calculate_df();
116                self.df_map.as_ref().unwrap()
117            }
118        }
119    }
120
121    fn calculate_idom(&mut self) {
122        let entry = self.fun.entry.unwrap();
123        let mut idom_map = HashMap::new();
124        let mut dominatee_map = HashMap::new();
125        let mut postorder_map = HashMap::new();
126        self.fun.po_iter().enumerate().for_each(|(i, bb)| {
127            postorder_map.insert(bb, i as i32);
128        });
129
130        // Calculate idom with reverse postorder
131        for current_bb in self.fun.rpo_iter() {
132            if current_bb == entry {
133                continue;
134            }
135            let mut new_idom = None;
136            for pred in current_bb.get_pred_bb() {
137                if idom_map.contains_key(pred) {
138                    // Set idom as intersection of predecessors
139                    if let Some(idom) = new_idom {
140                        new_idom = Some(intersect(*pred, idom, &postorder_map, &idom_map));
141                    } else {
142                        new_idom = Some(*pred);
143                    }
144                } else if *pred == entry {
145                    // If one of predecessor is entry, intersection is entry
146                    new_idom = Some(entry);
147                }
148            }
149            let new_idom = new_idom.unwrap_or(entry);
150            idom_map.insert(current_bb, new_idom);
151            dominatee_map
152                .entry(new_idom)
153                .or_insert(HashSet::new())
154                .insert(current_bb);
155        }
156
157        // Assign idom map and dominatee map to self
158        self.idom_map = Some(idom_map);
159        self.postorder_map = Some(postorder_map);
160        self.dominatee_map = Some(dominatee_map);
161    }
162
163    fn calculate_df(&mut self) {
164        let fun = self.fun;
165        let idoms = self.get_idom_map();
166        let mut df_map = HashMap::new();
167        for bb in fun.dfs_iter() {
168            for pred in bb.get_pred_bb() {
169                let mut runner = *pred;
170
171                // Hop up from each predecessor until runner is a dominator of bb
172                // For non-entry block, the first hit must be it's immediate dominator
173                // For entry block, the first hit must be itself, as doms(entry) = { entry }
174                while runner != idoms.get(&bb).copied().unwrap_or(bb) {
175                    df_map.entry(runner).or_insert(HashSet::new()).insert(bb);
176
177                    // Only update runner if it's not dead block
178                    if let Some(new_runner) = idoms.get(&runner) {
179                        runner = *new_runner;
180                    } else {
181                        break;
182                    }
183                }
184            }
185        }
186
187        // Assign df map to self
188        self.df_map = Some(df_map);
189    }
190}
191
192/// Function to get lowest common ancestor of two basic blocks in the dominator tree
193fn intersect(
194    mut n: BBPtr,
195    mut m: BBPtr,
196    postorder_map: &HashMap<BBPtr, i32>,
197    idoms: &HashMap<BBPtr, BBPtr>,
198) -> BBPtr {
199    while n != m {
200        while postorder_map[&n] < postorder_map[&m] {
201            n = idoms[&n];
202        }
203        while postorder_map[&m] < postorder_map[&n] {
204            m = idoms[&m];
205        }
206    }
207    n
208}
209
210#[cfg(test)]
211pub mod tests_dominator_tree {
212    use super::*;
213    use crate::ir::IRBuilder;
214    use std::iter::zip;
215
216    struct TestContext {
217        bb_vec: Vec<BBPtr>,
218        dom_tree: DominatorTree,
219    }
220
221    impl TestContext {
222        fn new(pool: &mut IRBuilder, down_stream: Vec<[i32; 2]>) -> Self {
223            let bb_vec: Vec<BBPtr> = (0..down_stream.len())
224                .map(|_| pool.new_basicblock("no_name".to_string()))
225                .collect();
226            for (mut bb, down) in zip(bb_vec.clone(), down_stream) {
227                match down {
228                    [-1, -1] => {}
229                    [t, -1] => bb.set_true_bb(bb_vec[t as usize]),
230                    [t, f] => {
231                        bb.set_true_bb(bb_vec[t as usize]);
232                        bb.set_false_bb(bb_vec[f as usize]);
233                    }
234                }
235            }
236            let mut fun =
237                pool.new_function("no_name".to_string(), crate::ir::ValueType::Void);
238            fun.entry = Some(bb_vec.first().cloned().unwrap());
239            fun.exit = Some(bb_vec.last().cloned().unwrap());
240            let dom_tree = DominatorTree::new(fun);
241            Self { bb_vec, dom_tree }
242        }
243
244        /// Check if dom(i) == j.
245        fn check_dominator(&mut self, i: usize, j: Vec<usize>) {
246            let dom = self.dom_tree.get_dominator(self.bb_vec[i]);
247            let expected: HashSet<BBPtr> = j.into_iter().map(|index| self.bb_vec[index]).collect();
248            assert_eq!(dom, expected);
249        }
250
251        /// Check if idom(i) == j.
252        fn check_idom(&mut self, i: usize, j: Option<usize>) {
253            let idom = self.dom_tree.get_idom(self.bb_vec[i]);
254            let expected = j.map(|index| self.bb_vec[index]);
255            assert_eq!(idom, expected);
256        }
257
258        /// Check if df(i) == j.
259        fn check_df(&mut self, i: usize, j: Vec<usize>) {
260            let df = self.dom_tree.get_df(self.bb_vec[i]);
261            let expected: HashSet<BBPtr> = j.into_iter().map(|index| self.bb_vec[index]).collect();
262            assert_eq!(df, expected);
263        }
264    }
265
266    #[test]
267    fn basic_test() {
268        // 0 ──► 1
269        let mut pool = IRBuilder::new();
270        let mut ctx = TestContext::new(&mut pool, vec![[1, -1], [-1, -1]]);
271        ctx.check_dominator(0, vec![0]);
272        ctx.check_dominator(1, vec![0, 1]);
273        ctx.check_idom(0, None);
274        ctx.check_idom(1, Some(0));
275        ctx.check_df(0, vec![]);
276        ctx.check_df(1, vec![]);
277    }
278
279    #[test]
280    fn no_back_edge() {
281        //    ┌─► 1 ──► 3 ──► 4
282        //    │         ▲     │
283        // 0 ─┤         │     │
284        //    │         │     ▼
285        //    └─► 2 ────┴───► 5
286        let mut pool = IRBuilder::new();
287        let mut ctx = TestContext::new(
288            &mut pool,
289            vec![[1, 2], [3, -1], [3, 5], [4, -1], [5, -1], [-1, -1]],
290        );
291        ctx.check_dominator(0, vec![0]);
292        ctx.check_dominator(1, vec![0, 1]);
293        ctx.check_dominator(2, vec![0, 2]);
294        ctx.check_dominator(3, vec![0, 3]);
295        ctx.check_dominator(4, vec![0, 3, 4]);
296        ctx.check_dominator(5, vec![0, 5]);
297        ctx.check_idom(0, None);
298        ctx.check_idom(1, Some(0));
299        ctx.check_idom(2, Some(0));
300        ctx.check_idom(3, Some(0));
301        ctx.check_idom(4, Some(3));
302        ctx.check_idom(5, Some(0));
303        ctx.check_df(0, vec![]);
304        ctx.check_df(1, vec![3]);
305        ctx.check_df(2, vec![3, 5]);
306        ctx.check_df(3, vec![5]);
307        ctx.check_df(4, vec![5]);
308        ctx.check_df(5, vec![]);
309    }
310
311    #[test]
312    fn back_edge() {
313        //       ┌──────────┐
314        //       │          │
315        //       │  ┌─► 2 ──┤
316        //       ▼  │       │
317        // 0 ──► 1 ─┤       │
318        //          │       │
319        //          └─► 3 ◄─┘
320        let mut pool = IRBuilder::new();
321        let mut ctx = TestContext::new(&mut pool, vec![[1, -1], [2, 3], [1, 3], [-1, -1]]);
322        ctx.check_dominator(0, vec![0]);
323        ctx.check_dominator(1, vec![0, 1]);
324        ctx.check_dominator(2, vec![0, 1, 2]);
325        ctx.check_dominator(3, vec![0, 1, 3]);
326        ctx.check_idom(0, None);
327        ctx.check_idom(1, Some(0));
328        ctx.check_idom(2, Some(1));
329        ctx.check_idom(3, Some(1));
330        ctx.check_df(0, vec![]);
331        ctx.check_df(1, vec![1]);
332        ctx.check_df(2, vec![1, 3]);
333        ctx.check_df(3, vec![]);
334    }
335
336    #[test]
337    fn branch_nested_loop() {
338        //          ┌─► 2 ─┐
339        //          │      │
340        // 0 ──► 1 ─┤      ├─► 4 ─┐
341        // ▲     ▲  │      │      │
342        // │     │  └─► 3 ─┘      │
343        // │     │                │
344        // └─────┴────────────────┘
345        let mut pool = IRBuilder::new();
346        let mut ctx = TestContext::new(&mut pool, vec![[1, -1], [2, 3], [4, -1], [4, -1], [0, 1]]);
347        ctx.check_dominator(0, vec![0]);
348        ctx.check_dominator(1, vec![0, 1]);
349        ctx.check_dominator(2, vec![0, 1, 2]);
350        ctx.check_dominator(3, vec![0, 1, 3]);
351        ctx.check_dominator(4, vec![0, 1, 4]);
352        ctx.check_idom(0, None);
353        ctx.check_idom(1, Some(0));
354        ctx.check_idom(2, Some(1));
355        ctx.check_idom(3, Some(1));
356        ctx.check_idom(4, Some(1));
357        ctx.check_df(0, vec![]);
358        ctx.check_df(1, vec![0, 1]);
359        ctx.check_df(2, vec![4]);
360        ctx.check_df(3, vec![4]);
361        ctx.check_df(4, vec![0, 1]);
362    }
363}
duskphantom_middle/analysis/dominator_tree.rs

duskphantom_middle/analysis/
dominator_tree.rs
