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

rapx/analysis/ssa_transform/
mod.rs
