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

rapx/analysis/core/range_analysis/SSA/
SSATransformer.rs
