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::*;
8use rustc_hir::def::DefKind;
9use rustc_hir::def_id::DefId;
10use rustc_middle::mir::{Body, Local};
11use rustc_middle::ty::TyCtxt;
12
13pub struct DataflowAnalyzer<'tcx> {
14    pub tcx: TyCtxt<'tcx>,
15    pub graphs: HashMap<DefId, DataflowGraph>,
16    pub debug: bool,
17    pub draw: bool,
18}
19
20impl<'tcx> DataflowAnalysis for DataflowAnalyzer<'tcx> {
21    fn get_fn_dataflow(&self, def_id: DefId) -> Option<DataFlowGraph> {
22        self.graphs.get(&def_id).cloned().map(Into::into)
23    }
24
25    fn get_all_dataflow(&self) -> DataFlowGraphMap {
26        self.graphs
27            .iter()
28            .map(|(&def_id, graph)| (def_id, graph.clone().into()))
29            .collect()
30    }
31
32    fn has_flow_between(&self, def_id: DefId, local1: Local, local2: Local) -> bool {
33        let graph = self.graphs.get(&def_id).unwrap();
34        graph.is_connected(local1, local2)
35    }
36
37    fn collect_equivalent_locals(&self, def_id: DefId, local: Local) -> HashSet<Local> {
38        let graph = self.graphs.get(&def_id).unwrap();
39        graph.collect_equivalent_locals(local, true)
40    }
41
42    fn get_fn_arg2ret(&self, def_id: DefId) -> Arg2Ret {
43        let graph = self.graphs.get(&def_id).unwrap();
44        graph.param_return_deps()
45    }
46
47    fn get_all_arg2ret(&self) -> Arg2RetMap {
48        let mut result = HashMap::new();
49        for (def_id, graph) in &self.graphs {
50            let deps = graph.param_return_deps();
51            result.insert(*def_id, deps);
52        }
53        result
54    }
55}
56
57impl<'tcx> Analysis for DataflowAnalyzer<'tcx> {
58    fn name(&self) -> &'static str {
59        "DataFlow Analysis"
60    }
61
62    fn run(&mut self) {
63        self.build_graphs();
64        if self.draw {
65            self.draw_graphs();
66        }
67    }
68
69    fn reset(&mut self) {
70        self.graphs.clear();
71    }
72}
73
74impl<'tcx> DataflowAnalyzer<'tcx> {
75    pub fn new(tcx: TyCtxt<'tcx>, debug: bool) -> Self {
76        Self {
77            tcx: tcx,
78            graphs: HashMap::new(),
79            debug,
80            draw: false,
81        }
82    }
83
84    pub fn with_draw(mut self, draw: bool) -> Self {
85        self.draw = draw;
86        self
87    }
88
89    pub fn start(&mut self) {
90        self.build_graphs();
91        if self.draw {
92            self.draw_graphs();
93        }
94    }
95
96    pub fn build_graphs(&mut self) {
97        for local_def_id in self.tcx.iter_local_def_id() {
98            let def_kind = self.tcx.def_kind(local_def_id);
99            if matches!(def_kind, DefKind::Fn) || matches!(def_kind, DefKind::AssocFn) {
100                if self.tcx.hir_maybe_body_owned_by(local_def_id).is_some() {
101                    let def_id = local_def_id.to_def_id();
102                    self.build_graph(def_id);
103                }
104            }
105        }
106    }
107
108    pub fn build_graph(&mut self, def_id: DefId) {
109        if self.graphs.contains_key(&def_id) {
110            return;
111        }
112        let body: &Body = self.tcx.optimized_mir(def_id);
113        let mut graph =
114            DataflowGraph::new(def_id, body.span, body.arg_count, body.local_decls.len());
115        let basic_blocks = &body.basic_blocks;
116        for (block_idx, basic_block_data) in basic_blocks.iter().enumerate() {
117            graph.block = block_idx;
118            for (stmt_idx, statement) in basic_block_data.statements.iter().enumerate() {
119                graph.statement_index = stmt_idx;
120                graph.add_statm_to_graph(&statement);
121            }
122            if let Some(terminator) = &basic_block_data.terminator {
123                graph.statement_index = basic_block_data.statements.len();
124                graph.add_terminator_to_graph(&terminator);
125            }
126        }
127        for closure_id in graph.closures.iter() {
128            self.build_graph(*closure_id);
129        }
130        self.graphs.insert(def_id, graph);
131    }
132
133    pub fn draw_graphs(&self) {
134        let dir_name = "DataflowGraph";
135
136        Command::new("rm")
137            .args(&["-rf", dir_name])
138            .output()
139            .expect("Failed to remove directory.");
140
141        Command::new("mkdir")
142            .args(&[dir_name])
143            .output()
144            .expect("Failed to create directory.");
145
146        for (def_id, graph) in self.graphs.iter() {
147            let name = self.tcx.def_path_str(*def_id);
148            let dot_file_name = format!("DataflowGraph/{}.dot", &name);
149            let png_file_name = format!("DataflowGraph/{}.png", &name);
150            let mut file = File::create(&dot_file_name).expect("Unable to create file.");
151            let dot = graph.to_dot_graph(&self.tcx);
152            file.write_all(dot.as_bytes())
153                .expect("Unable to write data.");
154
155            Command::new("dot")
156                .args(&["-Tpng", &dot_file_name, "-o", &png_file_name])
157                .output()
158                .expect("Failed to execute Graphviz dot command.");
159        }
160    }
161}
rapx/analysis/dataflow/default.rs

rapx/analysis/dataflow/
default.rs
