rapx/analysis/core/ssa_transform/
SSATransformer.rs

1#![allow(unused_imports)]
2#![allow(unused_variables)]
3#![allow(dead_code)]
4
5use rustc_data_structures::graph::dominators::Dominators;
6use rustc_data_structures::graph::{Predecessors, dominators};
7use rustc_driver::args;
8use rustc_hir::def_id::DefId;
9use rustc_hir::def_id::{CRATE_DEF_INDEX, CrateNum, DefIndex, LOCAL_CRATE, LocalDefId};
10use rustc_middle::mir::*;
11use rustc_middle::{
12    mir::{Body, Local, Location, visit::Visitor},
13    ty::TyCtxt,
14};
15use rustc_span::symbol::Symbol;
16use std::collections::{HashMap, HashSet};
17pub struct PhiPlaceholder;
18pub struct SSATransformer<'tcx> {
19    pub tcx: TyCtxt<'tcx>,
20    pub body: Body<'tcx>,
21    pub cfg: HashMap<BasicBlock, Vec<BasicBlock>>,
22    pub dominators: Dominators<BasicBlock>,
23    pub dom_tree: HashMap<BasicBlock, Vec<BasicBlock>>,
24    pub df: HashMap<BasicBlock, HashSet<BasicBlock>>,
25    pub local_assign_blocks: HashMap<Local, HashSet<BasicBlock>>,
26    pub reaching_def: HashMap<Local, Option<Local>>,
27    pub local_index: usize,
28    pub local_defination_block: HashMap<Local, BasicBlock>,
29    pub skipped: HashSet<usize>,
30    pub phi_index: HashMap<*const Statement<'tcx>, usize>,
31    pub phi_statements: HashMap<*const Statement<'tcx>, bool>,
32    pub essa_statements: HashMap<*const Statement<'tcx>, bool>,
33    pub phi_def_id: DefId,
34    pub essa_def_id: DefId,
35    pub ref_local_map: HashMap<Local, Local>,
36    pub places_map: HashMap<Place<'tcx>, HashSet<Place<'tcx>>>,
37    pub ssa_locals_map: HashMap<Place<'tcx>, HashSet<Place<'tcx>>>,
38}
39
40impl<'tcx> SSATransformer<'tcx> {
41    fn find_phi_placeholder(tcx: TyCtxt<'_>, crate_name: &str) -> Option<DefId> {
42        let sym_crate = Symbol::intern(crate_name);
43        let krate = tcx
44            .crates(())
45            .iter()
46            .find(|&&c| tcx.crate_name(c) == sym_crate)?;
47        let root_def_id = DefId {
48            krate: *krate,
49            index: CRATE_DEF_INDEX,
50        };
51        // print!("Phid\n");
52
53        for item in tcx.module_children(root_def_id) {
54            // println!("Module child: {:?}", item.ident.name.as_str());
55
56            if item.ident.name.as_str() == "PhiPlaceholder" {
57                if let Some(def_id) = item.res.opt_def_id() {
58                    return Some(def_id);
59                }
60            }
61        }
62        // print!("Phid\n");
63        return Some(root_def_id);
64    }
65    pub fn new(
66        tcx: TyCtxt<'tcx>,
67        body: &Body<'tcx>,
68        ssa_def_id: DefId,
69        essa_def_id: DefId,
70        arg_count: usize,
71    ) -> Self {
72        let cfg: HashMap<BasicBlock, Vec<BasicBlock>> = Self::extract_cfg_from_predecessors(&body);
73
74        let dominators: Dominators<BasicBlock> = body.basic_blocks.dominators().clone();
75
76        let dom_tree: HashMap<BasicBlock, Vec<BasicBlock>> = Self::construct_dominance_tree(&body);
77
78        let df: HashMap<BasicBlock, HashSet<BasicBlock>> =
79            Self::compute_dominance_frontier(&body, &dom_tree);
80
81        let local_assign_blocks: HashMap<Local, HashSet<BasicBlock>> =
82            Self::map_locals_to_assign_blocks(&body);
83        let local_defination_block: HashMap<Local, BasicBlock> =
84            Self::map_locals_to_definition_block(&body);
85        let len = body.local_decls.len() as usize;
86        let mut skipped = HashSet::new();
87        if len > 0 {
88            skipped.extend(arg_count + 1..len + 1);
89            // skipped.insert(0); // Skip the return place
90        }
91
92        SSATransformer {
93            tcx,
94            body: body.clone(),
95            cfg,
96            dominators,
97            dom_tree,
98            df,
99            local_assign_blocks,
100            reaching_def: HashMap::default(),
101            local_index: len,
102            local_defination_block: local_defination_block,
103            skipped: skipped,
104            phi_index: HashMap::default(),
105            phi_statements: HashMap::default(),
106            essa_statements: HashMap::default(),
107            phi_def_id: ssa_def_id,
108            essa_def_id: essa_def_id,
109            ref_local_map: HashMap::default(),
110            places_map: HashMap::default(),
111            ssa_locals_map: HashMap::default(),
112        }
113    }
114
115    pub fn return_body_ref(&self) -> &Body<'tcx> {
116        &self.body
117    }
118
119    fn map_locals_to_definition_block(body: &Body) -> HashMap<Local, BasicBlock> {
120        let mut local_to_block_map: HashMap<Local, BasicBlock> = HashMap::new();
121
122        for (bb, block_data) in body.basic_blocks.iter_enumerated() {
123            for statement in &block_data.statements {
124                match &statement.kind {
125                    StatementKind::Assign(box (place, _)) => {
126                        if let Some(local) = place.as_local() {
127                            if local.as_u32() == 0 {
128                                continue; // Skip the return place
129                            }
130                            local_to_block_map.entry(local).or_insert(bb);
131                        }
132                    }
133                    _ => {}
134                }
135            }
136            if let Some(terminator) = &block_data.terminator {
137                match &terminator.kind {
138                    TerminatorKind::Call { destination, .. } => {
139                        if let Some(local) = destination.as_local() {
140                            if local.as_u32() == 0 {
141                                continue; // Skip the return place
142                            }
143                            local_to_block_map.entry(local).or_insert(bb);
144                        }
145                    }
146                    _ => {}
147                }
148            }
149        }
150
151        local_to_block_map
152    }
153    pub fn depth_first_search_preorder(
154        dom_tree: &HashMap<BasicBlock, Vec<BasicBlock>>,
155        root: BasicBlock,
156    ) -> Vec<BasicBlock> {
157        let mut visited: HashSet<BasicBlock> = HashSet::new();
158        let mut preorder = Vec::new();
159
160        fn dfs(
161            node: BasicBlock,
162            dom_tree: &HashMap<BasicBlock, Vec<BasicBlock>>,
163            visited: &mut HashSet<BasicBlock>,
164            preorder: &mut Vec<BasicBlock>,
165        ) {
166            if visited.insert(node) {
167                preorder.push(node);
168
169                if let Some(children) = dom_tree.get(&node) {
170                    for &child in children {
171                        dfs(child, dom_tree, visited, preorder);
172                    }
173                }
174            }
175        }
176
177        dfs(root, dom_tree, &mut visited, &mut preorder);
178        preorder
179    }
180    pub fn depth_first_search_postorder(
181        dom_tree: &HashMap<BasicBlock, Vec<BasicBlock>>,
182        root: &BasicBlock,
183    ) -> Vec<BasicBlock> {
184        let mut visited: HashSet<BasicBlock> = HashSet::new();
185        let mut postorder = Vec::new();
186
187        fn dfs(
188            node: BasicBlock,
189            dom_tree: &HashMap<BasicBlock, Vec<BasicBlock>>,
190            visited: &mut HashSet<BasicBlock>,
191            postorder: &mut Vec<BasicBlock>,
192        ) {
193            if visited.insert(node) {
194                if let Some(children) = dom_tree.get(&node) {
195                    for &child in children {
196                        dfs(child, dom_tree, visited, postorder);
197                    }
198                }
199                postorder.push(node);
200            }
201        }
202
203        dfs(*root, dom_tree, &mut visited, &mut postorder);
204        postorder
205    }
206
207    fn map_locals_to_assign_blocks(body: &Body) -> HashMap<Local, HashSet<BasicBlock>> {
208        let mut local_to_blocks: HashMap<Local, HashSet<BasicBlock>> = HashMap::new();
209
210        for (bb, data) in body.basic_blocks.iter_enumerated() {
211            for stmt in &data.statements {
212                if let StatementKind::Assign(box (place, _)) = &stmt.kind {
213                    let local = place.local;
214                    if local.as_u32() == 0 {
215                        continue; // Skip the return place
216                    }
217                    local_to_blocks
218                        .entry(local)
219                        .or_insert_with(HashSet::new)
220                        .insert(bb);
221                }
222            }
223        }
224        for arg in body.args_iter() {
225            local_to_blocks
226                .entry(arg)
227                .or_insert_with(HashSet::new)
228                .insert(BasicBlock::from_u32(0)); // Assuming arg block is 0
229        }
230        local_to_blocks
231    }
232    fn construct_dominance_tree(body: &Body<'_>) -> HashMap<BasicBlock, Vec<BasicBlock>> {
233        let mut dom_tree: HashMap<BasicBlock, Vec<BasicBlock>> = HashMap::new();
234        let dominators = body.basic_blocks.dominators();
235        for (block, _) in body.basic_blocks.iter_enumerated() {
236            if let Some(idom) = dominators.immediate_dominator(block) {
237                dom_tree.entry(idom).or_default().push(block);
238            }
239        }
240
241        dom_tree
242    }
243    fn compute_dominance_frontier(
244        body: &Body<'_>,
245        dom_tree: &HashMap<BasicBlock, Vec<BasicBlock>>,
246    ) -> HashMap<BasicBlock, HashSet<BasicBlock>> {
247        let mut dominance_frontier: HashMap<BasicBlock, HashSet<BasicBlock>> = HashMap::new();
248        let dominators = body.basic_blocks.dominators();
249        let predecessors = body.basic_blocks.predecessors();
250        for (block, _) in body.basic_blocks.iter_enumerated() {
251            dominance_frontier.entry(block).or_default();
252        }
253
254        for (block, _) in body.basic_blocks.iter_enumerated() {
255            if predecessors[block].len() > 1 {
256                let preds = body.basic_blocks.predecessors()[block].clone();
257
258                for &pred in &preds {
259                    let mut runner = pred;
260                    while runner != dominators.immediate_dominator(block).unwrap() {
261                        dominance_frontier.entry(runner).or_default().insert(block);
262                        runner = dominators.immediate_dominator(runner).unwrap();
263                    }
264                }
265            }
266        }
267
268        dominance_frontier
269    }
270    fn extract_cfg_from_predecessors(body: &Body<'_>) -> HashMap<BasicBlock, Vec<BasicBlock>> {
271        let mut cfg: HashMap<BasicBlock, Vec<BasicBlock>> = HashMap::new();
272
273        for (block, _) in body.basic_blocks.iter_enumerated() {
274            for &predecessor in body.basic_blocks.predecessors()[block].iter() {
275                cfg.entry(predecessor).or_default().push(block);
276            }
277        }
278
279        cfg
280    }
281    fn print_dominance_tree(
282        dom_tree: &HashMap<BasicBlock, Vec<BasicBlock>>,
283        current: BasicBlock,
284        depth: usize,
285    ) {
286        if let Some(children) = dom_tree.get(&current) {
287            for &child in children {
288                Self::print_dominance_tree(dom_tree, child, depth + 1);
289            }
290        }
291    }
292
293    pub fn is_phi_statement(&self, statement: &Statement<'tcx>) -> bool {
294        let phi_stmt = statement as *const Statement<'tcx>;
295        if self.phi_statements.contains_key(&phi_stmt) {
296            return true;
297        } else {
298            return false;
299        }
300    }
301    pub fn is_essa_statement(&self, statement: &Statement<'tcx>) -> bool {
302        let essa_stmt = statement as *const Statement<'tcx>;
303        if self.essa_statements.contains_key(&essa_stmt) {
304            return true;
305        } else {
306            return false;
307        }
308    }
309}
rapx/analysis/core/ssa_transform/SSATransformer.rs

rapx/analysis/core/ssa_transform/
SSATransformer.rs
