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

rapx/analysis/core/ssa_transform/
SSATransformer.rs
