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
61impl<'tcx> DataflowAnalyzer<'tcx> {
62    pub fn new(tcx: TyCtxt<'tcx>, debug: bool) -> Self {
63        Self {
64            tcx,
65            graphs: HashMap::new(),
66            debug,
67            draw: false,
68        }
69    }
70
71    pub fn with_draw(mut self, draw: bool) -> Self {
72        self.draw = draw;
73        self
74    }
75
76    pub fn start(&mut self) {
77        self.build_graphs();
78        if self.draw {
79            self.draw_graphs();
80        }
81    }
82
83    pub fn build_graphs(&mut self) {
84        for local_def_id in self.tcx.iter_local_def_id() {
85            let def_kind = self.tcx.def_kind(local_def_id);
86            if matches!(def_kind, DefKind::Fn) || matches!(def_kind, DefKind::AssocFn) {
87                if self.tcx.hir_maybe_body_owned_by(local_def_id).is_some() {
88                    let def_id = local_def_id.to_def_id();
89                    self.build_graph(def_id);
90                }
91            }
92        }
93    }
94
95    pub fn build_graph(&mut self, def_id: DefId) {
96        if self.graphs.contains_key(&def_id) {
97            return;
98        }
99        let body: &Body = self.tcx.optimized_mir(def_id);
100        let graph = build_dataflow_graph_from_body(def_id, body);
101        for closure_id in graph.closures.iter() {
102            self.build_graph(*closure_id);
103        }
104        self.graphs.insert(def_id, graph);
105    }
106
107    pub fn draw_graphs(&self) {
108        let dir_name = "DataflowGraph";
109
110        Command::new("rm")
111            .args(&["-rf", dir_name])
112            .output()
113            .expect("Failed to remove directory.");
114
115        Command::new("mkdir")
116            .args(&[dir_name])
117            .output()
118            .expect("Failed to create directory.");
119
120        for (def_id, graph) in self.graphs.iter() {
121            let name = self.tcx.def_path_str(*def_id);
122            let dot_file_name = format!("DataflowGraph/{}.dot", &name);
123            let png_file_name = format!("DataflowGraph/{}.png", &name);
124            let mut file = File::create(&dot_file_name).expect("Unable to create file.");
125            let dot = graph.to_dot_graph(&self.tcx);
126            file.write_all(dot.as_bytes())
127                .expect("Unable to write data.");
128
129            Command::new("dot")
130                .args(&["-Tpng", &dot_file_name, "-o", &png_file_name])
131                .output()
132                .expect("Failed to execute Graphviz dot command.");
133        }
134    }
135}