Skip to main content

rapx/analysis/dataflow/
default.rs

1use std::collections::{HashMap, HashSet};
2use std::fs::File;
3use std::io::Write;
4use std::process::Command;
5
6use crate::analysis::Analysis;
7use crate::analysis::dataflow::graph::build_dataflow_graph_from_body;
8use crate::analysis::dataflow::*;
9use rustc_hir::def::DefKind;
10use rustc_hir::def_id::DefId;
11use rustc_middle::mir::{Body, Local};
12use rustc_middle::ty::TyCtxt;
13
14pub struct DataflowAnalyzer<'tcx> {
15    pub tcx: TyCtxt<'tcx>,
16    pub graphs: HashMap<DefId, DataflowGraph>,
17    pub debug: bool,
18    pub draw: bool,
19}
20
21impl<'tcx> DataflowAnalysis for DataflowAnalyzer<'tcx> {
22    fn get_fn_dataflow(&self, def_id: DefId) -> Option<DataflowGraph> {
23        self.graphs.get(&def_id).cloned()
24    }
25
26    fn get_all_dataflow(&self) -> DataflowGraphMap {
27        self.graphs.clone()
28    }
29
30    fn has_flow_between(&self, def_id: DefId, local1: Local, local2: Local) -> bool {
31        let graph = self.graphs.get(&def_id).unwrap();
32        graph.is_connected(local1, local2)
33    }
34
35    fn collect_equivalent_locals(&self, def_id: DefId, local: Local) -> HashSet<Local> {
36        let graph = self.graphs.get(&def_id).unwrap();
37        graph.collect_equivalent_locals(local, true)
38    }
39
40    fn get_fn_arg2ret(&self, def_id: DefId) -> Arg2Ret {
41        let graph = self.graphs.get(&def_id).unwrap();
42        graph.param_return_deps()
43    }
44
45    fn get_all_arg2ret(&self) -> Arg2RetMap {
46        let mut result = HashMap::new();
47        for (def_id, graph) in &self.graphs {
48            let deps = graph.param_return_deps();
49            result.insert(*def_id, deps);
50        }
51        result
52    }
53}
54
55impl<'tcx> Analysis for DataflowAnalyzer<'tcx> {
56    fn run(&mut self) {
57        self.start();
58    }
59
60}
61
62impl<'tcx> DataflowAnalyzer<'tcx> {
63    pub fn new(tcx: TyCtxt<'tcx>, debug: bool) -> Self {
64        Self {
65            tcx: tcx,
66            graphs: HashMap::new(),
67            debug,
68            draw: false,
69        }
70    }
71
72    pub fn with_draw(mut self, draw: bool) -> Self {
73        self.draw = draw;
74        self
75    }
76
77    pub fn start(&mut self) {
78        self.build_graphs();
79        if self.draw {
80            self.draw_graphs();
81        }
82    }
83
84    pub fn build_graphs(&mut self) {
85        for local_def_id in self.tcx.iter_local_def_id() {
86            let def_kind = self.tcx.def_kind(local_def_id);
87            if matches!(def_kind, DefKind::Fn) || matches!(def_kind, DefKind::AssocFn) {
88                if self.tcx.hir_maybe_body_owned_by(local_def_id).is_some() {
89                    let def_id = local_def_id.to_def_id();
90                    self.build_graph(def_id);
91                }
92            }
93        }
94    }
95
96    pub fn build_graph(&mut self, def_id: DefId) {
97        if self.graphs.contains_key(&def_id) {
98            return;
99        }
100        let body: &Body = self.tcx.optimized_mir(def_id);
101        let graph = build_dataflow_graph_from_body(def_id, body);
102        for closure_id in graph.closures.iter() {
103            self.build_graph(*closure_id);
104        }
105        self.graphs.insert(def_id, graph);
106    }
107
108    pub fn draw_graphs(&self) {
109        let dir_name = "DataflowGraph";
110
111        Command::new("rm")
112            .args(&["-rf", dir_name])
113            .output()
114            .expect("Failed to remove directory.");
115
116        Command::new("mkdir")
117            .args(&[dir_name])
118            .output()
119            .expect("Failed to create directory.");
120
121        for (def_id, graph) in self.graphs.iter() {
122            let name = self.tcx.def_path_str(*def_id);
123            let dot_file_name = format!("DataflowGraph/{}.dot", &name);
124            let png_file_name = format!("DataflowGraph/{}.png", &name);
125            let mut file = File::create(&dot_file_name).expect("Unable to create file.");
126            let dot = graph.to_dot_graph(&self.tcx);
127            file.write_all(dot.as_bytes())
128                .expect("Unable to write data.");
129
130            Command::new("dot")
131                .args(&["-Tpng", &dot_file_name, "-o", &png_file_name])
132                .output()
133                .expect("Failed to execute Graphviz dot command.");
134        }
135    }
136}