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

rapx/analysis/ssa_transform/
mod.rs
