rapx/analysis/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<Location, usize>,
31    pub phi_def_id: DefId,
32    pub essa_def_id: DefId,
33    pub ref_local_map: HashMap<Local, Local>,
34    pub places_map: HashMap<Place<'tcx>, HashSet<Place<'tcx>>>,
35    pub ssa_locals_map: HashMap<Place<'tcx>, HashSet<Place<'tcx>>>,
36}
37
38impl<'tcx> SSATransformer<'tcx> {
39    fn find_phi_placeholder(tcx: TyCtxt<'_>, crate_name: &str) -> Option<DefId> {
40        let sym_crate = Symbol::intern(crate_name);
41        let krate = tcx
42            .crates(())
43            .iter()
44            .find(|&&c| tcx.crate_name(c) == sym_crate)?;
45        let root_def_id = DefId {
46            krate: *krate,
47            index: CRATE_DEF_INDEX,
48        };
49        // print!("Phid\n");
50
51        for item in tcx.module_children(root_def_id) {
52            // println!("Module child: {:?}", item.ident.name.as_str());
53
54            if item.ident.name.as_str() == "PhiPlaceholder" {
55                if let Some(def_id) = item.res.opt_def_id() {
56                    return Some(def_id);
57                }
58            }
59        }
60        // print!("Phid\n");
61        return Some(root_def_id);
62    }
63    pub fn new(
64        tcx: TyCtxt<'tcx>,
65        body: &Body<'tcx>,
66        ssa_def_id: DefId,
67        essa_def_id: DefId,
68        arg_count: usize,
69    ) -> Self {
70        let cfg: HashMap<BasicBlock, Vec<BasicBlock>> = Self::extract_cfg_from_predecessors(&body);
71
72        let dominators: Dominators<BasicBlock> = body.basic_blocks.dominators().clone();
73
74        let dom_tree: HashMap<BasicBlock, Vec<BasicBlock>> = Self::construct_dominance_tree(&body);
75
76        let df: HashMap<BasicBlock, HashSet<BasicBlock>> =
77            Self::compute_dominance_frontier(&body, &dom_tree);
78
79        let local_assign_blocks: HashMap<Local, HashSet<BasicBlock>> =
80            Self::map_locals_to_assign_blocks(&body);
81        let local_defination_block: HashMap<Local, BasicBlock> =
82            Self::map_locals_to_definition_block(&body);
83        let len = body.local_decls.len() as usize;
84        let mut skipped = HashSet::new();
85        if len > 0 {
86            skipped.extend(arg_count + 1..len + 1);
87            // skipped.insert(0); // Skip the return place
88        }
89
90        SSATransformer {
91            tcx,
92            body: body.clone(),
93            cfg,
94            dominators,
95            dom_tree,
96            df,
97            local_assign_blocks,
98            reaching_def: HashMap::default(),
99            local_index: len,
100            local_defination_block: local_defination_block,
101            skipped: skipped,
102            phi_index: HashMap::default(),
103            phi_def_id: ssa_def_id,
104            essa_def_id: essa_def_id,
105            ref_local_map: HashMap::default(),
106            places_map: HashMap::default(),
107            ssa_locals_map: HashMap::default(),
108        }
109    }
110
111    pub fn return_body_ref(&self) -> &Body<'tcx> {
112        &self.body
113    }
114
115    fn map_locals_to_definition_block(body: &Body) -> HashMap<Local, BasicBlock> {
116        let mut local_to_block_map: HashMap<Local, BasicBlock> = HashMap::new();
117
118        for (bb, block_data) in body.basic_blocks.iter_enumerated() {
119            for statement in &block_data.statements {
120                match &statement.kind {
121                    StatementKind::Assign(assign) => {
122                        let (place, _) = &**assign;
123                        if let Some(local) = place.as_local() {
124                            if local.as_u32() == 0 {
125                                continue; // Skip the return place
126                            }
127                            local_to_block_map.entry(local).or_insert(bb);
128                        }
129                    }
130                    _ => {}
131                }
132            }
133            if let Some(terminator) = &block_data.terminator {
134                match &terminator.kind {
135                    TerminatorKind::Call { destination, .. } => {
136                        if let Some(local) = destination.as_local() {
137                            if local.as_u32() == 0 {
138                                continue; // Skip the return place
139                            }
140                            local_to_block_map.entry(local).or_insert(bb);
141                        }
142                    }
143                    _ => {}
144                }
145            }
146        }
147
148        local_to_block_map
149    }
150    pub fn depth_first_search_preorder(
151        dom_tree: &HashMap<BasicBlock, Vec<BasicBlock>>,
152        root: BasicBlock,
153    ) -> Vec<BasicBlock> {
154        let mut visited: HashSet<BasicBlock> = HashSet::new();
155        let mut preorder = Vec::new();
156
157        fn dfs(
158            node: BasicBlock,
159            dom_tree: &HashMap<BasicBlock, Vec<BasicBlock>>,
160            visited: &mut HashSet<BasicBlock>,
161            preorder: &mut Vec<BasicBlock>,
162        ) {
163            if visited.insert(node) {
164                preorder.push(node);
165
166                if let Some(children) = dom_tree.get(&node) {
167                    for &child in children {
168                        dfs(child, dom_tree, visited, preorder);
169                    }
170                }
171            }
172        }
173
174        dfs(root, dom_tree, &mut visited, &mut preorder);
175        preorder
176    }
177    pub fn depth_first_search_postorder(
178        dom_tree: &HashMap<BasicBlock, Vec<BasicBlock>>,
179        root: &BasicBlock,
180    ) -> Vec<BasicBlock> {
181        let mut visited: HashSet<BasicBlock> = HashSet::new();
182        let mut postorder = Vec::new();
183
184        fn dfs(
185            node: BasicBlock,
186            dom_tree: &HashMap<BasicBlock, Vec<BasicBlock>>,
187            visited: &mut HashSet<BasicBlock>,
188            postorder: &mut Vec<BasicBlock>,
189        ) {
190            if visited.insert(node) {
191                if let Some(children) = dom_tree.get(&node) {
192                    for &child in children {
193                        dfs(child, dom_tree, visited, postorder);
194                    }
195                }
196                postorder.push(node);
197            }
198        }
199
200        dfs(*root, dom_tree, &mut visited, &mut postorder);
201        postorder
202    }
203
204    fn map_locals_to_assign_blocks(body: &Body) -> HashMap<Local, HashSet<BasicBlock>> {
205        let mut local_to_blocks: HashMap<Local, HashSet<BasicBlock>> = HashMap::new();
206
207        for (bb, data) in body.basic_blocks.iter_enumerated() {
208            for stmt in &data.statements {
209                if let StatementKind::Assign(assign) = &stmt.kind {
210                    let (place, _) = &**assign;
211                    let local = place.local;
212                    if local.as_u32() == 0 {
213                        continue; // Skip the return place
214                    }
215                    local_to_blocks
216                        .entry(local)
217                        .or_insert_with(HashSet::new)
218                        .insert(bb);
219                }
220            }
221        }
222        for arg in body.args_iter() {
223            local_to_blocks
224                .entry(arg)
225                .or_insert_with(HashSet::new)
226                .insert(BasicBlock::from_u32(0)); // Assuming arg block is 0
227        }
228        local_to_blocks
229    }
230    fn construct_dominance_tree(body: &Body<'_>) -> HashMap<BasicBlock, Vec<BasicBlock>> {
231        let mut dom_tree: HashMap<BasicBlock, Vec<BasicBlock>> = HashMap::new();
232        let dominators = body.basic_blocks.dominators();
233        for (block, _) in body.basic_blocks.iter_enumerated() {
234            if let Some(idom) = dominators.immediate_dominator(block) {
235                dom_tree.entry(idom).or_default().push(block);
236            }
237        }
238
239        dom_tree
240    }
241    fn compute_dominance_frontier(
242        body: &Body<'_>,
243        dom_tree: &HashMap<BasicBlock, Vec<BasicBlock>>,
244    ) -> HashMap<BasicBlock, HashSet<BasicBlock>> {
245        let mut dominance_frontier: HashMap<BasicBlock, HashSet<BasicBlock>> = HashMap::new();
246        let dominators = body.basic_blocks.dominators();
247        let predecessors = body.basic_blocks.predecessors();
248        for (block, _) in body.basic_blocks.iter_enumerated() {
249            dominance_frontier.entry(block).or_default();
250        }
251
252        for (block, _) in body.basic_blocks.iter_enumerated() {
253            if predecessors[block].len() > 1 {
254                let preds = body.basic_blocks.predecessors()[block].clone();
255
256                for &pred in &preds {
257                    let mut runner = pred;
258                    while runner != dominators.immediate_dominator(block).unwrap() {
259                        dominance_frontier.entry(runner).or_default().insert(block);
260                        runner = dominators.immediate_dominator(runner).unwrap();
261                    }
262                }
263            }
264        }
265
266        dominance_frontier
267    }
268    fn extract_cfg_from_predecessors(body: &Body<'_>) -> HashMap<BasicBlock, Vec<BasicBlock>> {
269        let mut cfg: HashMap<BasicBlock, Vec<BasicBlock>> = HashMap::new();
270
271        for (block, _) in body.basic_blocks.iter_enumerated() {
272            for &predecessor in body.basic_blocks.predecessors()[block].iter() {
273                cfg.entry(predecessor).or_default().push(block);
274            }
275        }
276
277        cfg
278    }
279    fn print_dominance_tree(
280        dom_tree: &HashMap<BasicBlock, Vec<BasicBlock>>,
281        current: BasicBlock,
282        depth: usize,
283    ) {
284        if let Some(children) = dom_tree.get(&current) {
285            for &child in children {
286                Self::print_dominance_tree(dom_tree, child, depth + 1);
287            }
288        }
289    }
290
291    pub fn is_phi_statement(&self, statement: &Statement<'tcx>) -> bool {
292        if let StatementKind::Assign(assign) = &statement.kind {
293            let (_, rvalue) = &**assign;
294            if let Rvalue::Aggregate(k_box, _) = rvalue {
295                let aggregate_kind = &**k_box;
296                if let AggregateKind::Adt(def_id, ..) = aggregate_kind {
297                    return *def_id == self.phi_def_id;
298                }
299            }
300        }
301        false
302    }
303
304    pub fn is_essa_statement(&self, statement: &Statement<'tcx>) -> bool {
305        if let StatementKind::Assign(assign) = &statement.kind {
306            let (_, rvalue) = &**assign;
307            if let Rvalue::Aggregate(k_box, _) = rvalue {
308                let aggregate_kind = &**k_box;
309                if let AggregateKind::Adt(def_id, ..) = aggregate_kind {
310                    return *def_id == self.essa_def_id;
311                }
312            }
313        }
314        false
315    }
316    pub fn get_essa_source_block(&self, statement: &Statement<'tcx>) -> Option<BasicBlock> {
317        if !self.is_essa_statement(statement) {
318            return None;
319        }
320
321        if let StatementKind::Assign(assign) = &statement.kind {
322            let (_, rvalue) = &**assign;
323            if let Rvalue::Aggregate(_, operands) = rvalue {
324                if let Some(last_op) = operands.into_iter().last() {
325                    if let Operand::Constant(c_box) = last_op {
326                        let ConstOperand { const_: c, .. } = &**c_box;
327                        if let Some(val) = self.try_const_to_usize(c) {
328                            return Some(BasicBlock::from_usize(val as usize));
329                        }
330                    }
331                }
332            }
333        }
334        None
335    }
336
337    fn try_const_to_usize(&self, c: &Const<'tcx>) -> Option<u64> {
338        if let Some(scalar_int) = c.try_to_scalar_int() {
339            let size = scalar_int.size();
340            if let Ok(bits) = scalar_int.try_to_bits(size) {
341                return Some(bits as u64);
342            }
343        }
344        None
345    }
346}
rapx/analysis/ssa_transform/SSATransformer.rs

rapx/analysis/ssa_transform/
SSATransformer.rs
