rapx/analysis/ssa_transform/
ssa_transformer.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};
17
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
40    pub fn new(
41        tcx: TyCtxt<'tcx>,
42        body: &Body<'tcx>,
43        ssa_def_id: DefId,
44        essa_def_id: DefId,
45        arg_count: usize,
46    ) -> Self {
47        let cfg: HashMap<BasicBlock, Vec<BasicBlock>> = Self::extract_cfg_from_predecessors(&body);
48
49        let dominators: Dominators<BasicBlock> = body.basic_blocks.dominators().clone();
50
51        let dom_tree: HashMap<BasicBlock, Vec<BasicBlock>> = Self::construct_dominance_tree(&body);
52
53        let df: HashMap<BasicBlock, HashSet<BasicBlock>> =
54            Self::compute_dominance_frontier(&body, &dom_tree);
55
56        let local_assign_blocks: HashMap<Local, HashSet<BasicBlock>> =
57            Self::map_locals_to_assign_blocks(&body);
58        let local_defination_block: HashMap<Local, BasicBlock> =
59            Self::map_locals_to_definition_block(&body);
60        let len = body.local_decls.len() as usize;
61        let mut skipped = HashSet::new();
62        if len > 0 {
63            skipped.extend(arg_count + 1..len + 1);
64            // skipped.insert(0); // Skip the return place
65        }
66
67        SSATransformer {
68            tcx,
69            body: body.clone(),
70            cfg,
71            dominators,
72            dom_tree,
73            df,
74            local_assign_blocks,
75            reaching_def: HashMap::default(),
76            local_index: len,
77            local_defination_block: local_defination_block,
78            skipped: skipped,
79            phi_index: HashMap::default(),
80            phi_def_id: ssa_def_id,
81            essa_def_id: essa_def_id,
82            ref_local_map: HashMap::default(),
83            places_map: HashMap::default(),
84            ssa_locals_map: HashMap::default(),
85        }
86    }
87
88
89
90    fn map_locals_to_definition_block(body: &Body) -> HashMap<Local, BasicBlock> {
91        let mut local_to_block_map: HashMap<Local, BasicBlock> = HashMap::new();
92
93        for (bb, block_data) in body.basic_blocks.iter_enumerated() {
94            for statement in &block_data.statements {
95                match &statement.kind {
96                    StatementKind::Assign(assign) => {
97                        let (place, _) = &**assign;
98                        if let Some(local) = place.as_local() {
99                            if local.as_u32() == 0 {
100                                continue; // Skip the return place
101                            }
102                            local_to_block_map.entry(local).or_insert(bb);
103                        }
104                    }
105                    _ => {}
106                }
107            }
108            if let Some(terminator) = &block_data.terminator {
109                match &terminator.kind {
110                    TerminatorKind::Call { destination, .. } => {
111                        if let Some(local) = destination.as_local() {
112                            if local.as_u32() == 0 {
113                                continue; // Skip the return place
114                            }
115                            local_to_block_map.entry(local).or_insert(bb);
116                        }
117                    }
118                    _ => {}
119                }
120            }
121        }
122
123        local_to_block_map
124    }
125    pub fn depth_first_search_preorder(
126        dom_tree: &HashMap<BasicBlock, Vec<BasicBlock>>,
127        root: BasicBlock,
128    ) -> Vec<BasicBlock> {
129        let mut visited: HashSet<BasicBlock> = HashSet::new();
130        let mut preorder = Vec::new();
131
132        fn dfs(
133            node: BasicBlock,
134            dom_tree: &HashMap<BasicBlock, Vec<BasicBlock>>,
135            visited: &mut HashSet<BasicBlock>,
136            preorder: &mut Vec<BasicBlock>,
137        ) {
138            if visited.insert(node) {
139                preorder.push(node);
140
141                if let Some(children) = dom_tree.get(&node) {
142                    for &child in children {
143                        dfs(child, dom_tree, visited, preorder);
144                    }
145                }
146            }
147        }
148
149        dfs(root, dom_tree, &mut visited, &mut preorder);
150        preorder
151    }
152
153
154    fn map_locals_to_assign_blocks(body: &Body) -> HashMap<Local, HashSet<BasicBlock>> {
155        let mut local_to_blocks: HashMap<Local, HashSet<BasicBlock>> = HashMap::new();
156
157        for (bb, data) in body.basic_blocks.iter_enumerated() {
158            for stmt in &data.statements {
159                if let StatementKind::Assign(assign) = &stmt.kind {
160                    let (place, _) = &**assign;
161                    let local = place.local;
162                    if local.as_u32() == 0 {
163                        continue; // Skip the return place
164                    }
165                    local_to_blocks
166                        .entry(local)
167                        .or_insert_with(HashSet::new)
168                        .insert(bb);
169                }
170            }
171        }
172        for arg in body.args_iter() {
173            local_to_blocks
174                .entry(arg)
175                .or_insert_with(HashSet::new)
176                .insert(BasicBlock::from_u32(0)); // Assuming arg block is 0
177        }
178        local_to_blocks
179    }
180    fn construct_dominance_tree(body: &Body<'_>) -> HashMap<BasicBlock, Vec<BasicBlock>> {
181        let mut dom_tree: HashMap<BasicBlock, Vec<BasicBlock>> = HashMap::new();
182        let dominators = body.basic_blocks.dominators();
183        for (block, _) in body.basic_blocks.iter_enumerated() {
184            if let Some(idom) = dominators.immediate_dominator(block) {
185                dom_tree.entry(idom).or_default().push(block);
186            }
187        }
188
189        dom_tree
190    }
191    fn compute_dominance_frontier(
192        body: &Body<'_>,
193        dom_tree: &HashMap<BasicBlock, Vec<BasicBlock>>,
194    ) -> HashMap<BasicBlock, HashSet<BasicBlock>> {
195        let mut dominance_frontier: HashMap<BasicBlock, HashSet<BasicBlock>> = HashMap::new();
196        let dominators = body.basic_blocks.dominators();
197        let predecessors = body.basic_blocks.predecessors();
198        for (block, _) in body.basic_blocks.iter_enumerated() {
199            dominance_frontier.entry(block).or_default();
200        }
201
202        for (block, _) in body.basic_blocks.iter_enumerated() {
203            if predecessors[block].len() > 1 {
204                let preds = body.basic_blocks.predecessors()[block].clone();
205
206                for &pred in &preds {
207                    let mut runner = pred;
208                    while runner != dominators.immediate_dominator(block).unwrap() {
209                        dominance_frontier.entry(runner).or_default().insert(block);
210                        runner = dominators.immediate_dominator(runner).unwrap();
211                    }
212                }
213            }
214        }
215
216        dominance_frontier
217    }
218    fn extract_cfg_from_predecessors(body: &Body<'_>) -> HashMap<BasicBlock, Vec<BasicBlock>> {
219        let mut cfg: HashMap<BasicBlock, Vec<BasicBlock>> = HashMap::new();
220
221        for (block, _) in body.basic_blocks.iter_enumerated() {
222            for &predecessor in body.basic_blocks.predecessors()[block].iter() {
223                cfg.entry(predecessor).or_default().push(block);
224            }
225        }
226
227        cfg
228    }
229
230
231    pub fn is_phi_statement(&self, statement: &Statement<'tcx>) -> bool {
232        if let StatementKind::Assign(assign) = &statement.kind {
233            let (_, rvalue) = &**assign;
234            if let Rvalue::Aggregate(k_box, _) = rvalue {
235                let aggregate_kind = &**k_box;
236                if let AggregateKind::Adt(def_id, ..) = aggregate_kind {
237                    return *def_id == self.phi_def_id;
238                }
239            }
240        }
241        false
242    }
243
244    pub fn is_essa_statement(&self, statement: &Statement<'tcx>) -> bool {
245        if let StatementKind::Assign(assign) = &statement.kind {
246            let (_, rvalue) = &**assign;
247            if let Rvalue::Aggregate(k_box, _) = rvalue {
248                let aggregate_kind = &**k_box;
249                if let AggregateKind::Adt(def_id, ..) = aggregate_kind {
250                    return *def_id == self.essa_def_id;
251                }
252            }
253        }
254        false
255    }
256    pub fn get_essa_source_block(&self, statement: &Statement<'tcx>) -> Option<BasicBlock> {
257        if !self.is_essa_statement(statement) {
258            return None;
259        }
260
261        if let StatementKind::Assign(assign) = &statement.kind {
262            let (_, rvalue) = &**assign;
263            if let Rvalue::Aggregate(_, operands) = rvalue {
264                if let Some(last_op) = operands.into_iter().last() {
265                    if let Operand::Constant(c_box) = last_op {
266                        let ConstOperand { const_: c, .. } = &**c_box;
267                        if let Some(val) = self.try_const_to_usize(c) {
268                            return Some(BasicBlock::from_usize(val as usize));
269                        }
270                    }
271                }
272            }
273        }
274        None
275    }
276
277    fn try_const_to_usize(&self, c: &Const<'tcx>) -> Option<u64> {
278        if let Some(scalar_int) = c.try_to_scalar_int() {
279            let size = scalar_int.size();
280            let bits = scalar_int.to_bits(size);
281            return Some(bits as u64);
282        }
283        None
284    }
285}
rapx/analysis/ssa_transform/ssa_transformer.rs

rapx/analysis/ssa_transform/
ssa_transformer.rs
