rapx/analysis/core/ssa_transform/
mod.rs

1#![allow(unused_imports)]
2#![allow(unused_variables)]
3#![allow(dead_code)]
4#![allow(unused_assignments)]
5#![allow(unused_parens)]
6#![allow(non_snake_case)]
7
8pub mod Replacer;
9pub mod SSATransformer;
10
11use crate::{rap_info, rap_warn};
12use rustc_hir::def::DefKind;
13use rustc_hir::def_id::DefId;
14use rustc_hir::def_id::LocalDefId;
15use rustc_middle::mir::pretty::{write_mir_fn, PrettyPrintMirOptions};
16use rustc_middle::mir::*;
17use rustc_middle::ty::TyCtxt;
18use std::collections::HashMap;
19use std::collections::HashSet;
20use std::fs;
21use std::fs::File;
22use std::io;
23use std::io::Cursor;
24use std::io::Write;
25use std::path::PathBuf;
26pub struct SSATrans<'tcx> {
27    pub tcx: TyCtxt<'tcx>,
28    pub debug: bool,
29}
30
31impl<'tcx> SSATrans<'tcx> {
32    pub fn new(tcx: TyCtxt<'tcx>, debug: bool) -> Self {
33        Self { tcx: tcx, debug }
34    }
35
36    pub fn start(&mut self) {
37        for local_def_id in self.tcx.iter_local_def_id() {
38            if matches!(self.tcx.def_kind(local_def_id), DefKind::Fn) {
39                if self.tcx.hir_maybe_body_owned_by(local_def_id).is_some() {
40                    if let Some(def_id) = self
41                        .tcx
42                        .hir_body_owners()
43                        .find(|id| self.tcx.def_path_str(*id) == "main")
44                    {
45                        if let Some(ssa_def_id) =
46                            self.tcx.hir_crate_items(()).free_items().find(|id| {
47                                let hir_id = id.hir_id();
48                                if let Some(ident_name) = self.tcx.hir_opt_name(hir_id) {
49                                    ident_name.to_string() == "SSAstmt"
50                                } else {
51                                    false
52                                }
53                            })
54                        {
55                            let ssa_def_id = ssa_def_id.owner_id.to_def_id();
56                            if let Some(essa_def_id) =
57                                self.tcx.hir_crate_items(()).free_items().find(|id| {
58                                    let hir_id = id.hir_id();
59                                    if let Some(ident_name) = self.tcx.hir_opt_name(hir_id) {
60                                        ident_name.to_string() == "ESSAstmt"
61                                    } else {
62                                        false
63                                    }
64                                })
65                            {
66                                let essa_def_id = essa_def_id.owner_id.to_def_id();
67                                self.analyze_mir(self.tcx, def_id, ssa_def_id, essa_def_id);
68                            }
69                        }
70                    }
71                }
72            }
73        }
74    }
75    fn analyze_mir(
76        &mut self,
77        tcx: TyCtxt<'tcx>,
78        def_id: LocalDefId,
79        ssa_def_id: DefId,
80        essa_def_id: DefId,
81    ) {
82        let mut body = tcx.optimized_mir(def_id).clone();
83        {
84            let body_mut_ref: &mut Body<'tcx> = unsafe { &mut *(&mut body as *mut Body<'tcx>) };
85            let mut passrunner = PassRunner::new(tcx);
86            passrunner.run_pass(body_mut_ref, ssa_def_id, essa_def_id);
87            // passrunner.print_diff(body_mut_ref);
88            let essa_mir_string = passrunner.get_final_ssa_as_string(body_mut_ref);
89            // rap_info!("final SSA {:?}\n", &essa_mir_string);
90            rap_info!("ssa lvalue check {:?}", lvalue_check(&essa_mir_string));
91        }
92    }
93}
94pub struct PassRunner<'tcx> {
95    tcx: TyCtxt<'tcx>,
96    pub places_map: HashMap<Place<'tcx>, HashSet<Place<'tcx>>>,
97}
98pub fn lvalue_check(mir_string: &str) -> bool {
99    let re = regex::Regex::new(r"_(\d+)\s*=").unwrap();
100    let mut counts = HashMap::new();
101    let mut has_duplicate = false;
102
103    for cap in re.captures_iter(mir_string) {
104        let var = cap[1].parse::<u32>().unwrap();
105        let counter = counts.entry(var).or_insert(0);
106        *counter += 1;
107        if *counter > 1 {
108            has_duplicate = true;
109        }
110    }
111
112    for (var, count) in counts {
113        if count > 1 {
114            rap_warn!("Variable _ {} is used {} times", var, count);
115        }
116    }
117
118    !has_duplicate
119}
120pub fn print_diff<'tcx>(tcx: TyCtxt<'tcx>, body: &Body<'tcx>, def_id: DefId) {
121    let dir_path = "ssa_mir";
122    fs::create_dir_all(dir_path).unwrap();
123    // PassRunner::new(self.tcx);
124    let name = tcx.def_path_str(def_id);
125    let mir_file_path = format!("{}/origin_mir.txt", dir_path);
126    let phi_mir_file_path = format!("{}/{}_after_rename_mir.txt", dir_path, name);
127    let mut file = File::create(&mir_file_path).unwrap();
128    let mut w = io::BufWriter::new(&mut file);
129    write_mir_pretty(tcx, None, &mut w).unwrap();
130    let mut file2 = File::create(&phi_mir_file_path).unwrap();
131    let mut w2 = io::BufWriter::new(&mut file2);
132    let options = PrettyPrintMirOptions::from_cli(tcx);
133    write_mir_fn(tcx, body, &mut |_, _| Ok(()), &mut w2, options).unwrap();
134}
135pub fn print_mir_graph<'tcx>(tcx: TyCtxt<'tcx>, body: &Body<'tcx>, def_id: DefId) {
136    let dir_path = PathBuf::from("passrunner_mir_dot");
137    fs::create_dir_all(dir_path.clone()).unwrap();
138
139    let dot_graph = mir_to_dot(tcx, body);
140    let function_name = tcx.def_path_str(def_id);
141    let safe_filename = format!("{}_after_rename_mir.dot", function_name);
142    let output_path = dir_path.join(format!("{}", safe_filename));
143
144    let mut file = File::create(&output_path).expect("cannot create file");
145    let _ = file.write_all(dot_graph.as_bytes());
146}
147//for f in *.dot; do dot -Tpng "$f" -o "${f%.dot}.png"; done
148fn mir_to_dot<'tcx>(tcx: TyCtxt<'tcx>, body: &Body<'tcx>) -> String {
149    let mut dot = String::new();
150    dot.push_str("digraph MIR {\n");
151    dot.push_str("  node [shape=box];\n");
152
153    for (bb, bb_data) in body.basic_blocks.iter_enumerated() {
154        let statements_str = bb_data
155            .statements
156            .iter()
157            .filter(|stmt| {
158                !matches!(
159                    stmt.kind,
160                    StatementKind::StorageLive(_) | StatementKind::StorageDead(_)
161                )
162            })
163            .map(|stmt| format!("{:?}", stmt).replace('\n', " "))
164            .collect::<Vec<_>>()
165            .join("\\l");
166
167        let terminator_str = match &bb_data.terminator {
168            Some(term) => match term.kind {
169                TerminatorKind::Assert { .. } => String::new(),
170                _ => format!("{:?}", term.kind).replace("->", "-->"),
171            },
172            None => "NoTerminator".to_string(),
173        };
174
175        let label = format!("{}\\l{}\\l", statements_str, terminator_str);
176
177        dot.push_str(&format!(
178            "  {} [label=\"{:?}:\\l{}\"];\n",
179            bb.index(),
180            bb,
181            label
182        ));
183
184        if let Some(terminator) = &bb_data.terminator {
185            for successor in terminator.successors() {
186                dot.push_str(&format!("  {} -> {};\n", bb.index(), successor.index()));
187            }
188        }
189    }
190
191    dot.push_str("}\n");
192    dot
193}
194
195impl<'tcx> PassRunner<'tcx> {
196    pub fn new(tcx: TyCtxt<'tcx>) -> Self {
197        Self {
198            tcx,
199            places_map: HashMap::default(),
200        }
201    }
202
203    pub fn get_final_ssa_as_string(&self, body: &Body<'tcx>) -> String {
204        let mut buffer2 = Cursor::new(Vec::new());
205        let options = PrettyPrintMirOptions::from_cli(self.tcx);
206        write_mir_fn(self.tcx, body, &mut |_, _| Ok(()), &mut buffer2, options).unwrap();
207        let after_mir = String::from_utf8(buffer2.into_inner()).unwrap();
208        after_mir
209    }
210
211    pub fn run_pass(&mut self, body: &mut Body<'tcx>, ssa_def_id: DefId, essa_def_id: DefId) {
212        let arg_count = body.arg_count;
213        let ssatransformer =
214            SSATransformer::SSATransformer::new(self.tcx, body, ssa_def_id, essa_def_id, arg_count);
215        let mut replacer = Replacer::Replacer {
216            tcx: self.tcx,
217            ssatransformer,
218            new_local_collection: HashSet::default(),
219        };
220        replacer.insert_phi_statment(body);
221        replacer.insert_essa_statement(body);
222        replacer.rename_variables(body);
223        self.places_map = replacer.ssatransformer.places_map.clone();
224    }
225}
rapx/analysis/core/ssa_transform/mod.rs

rapx/analysis/core/ssa_transform/
mod.rs
