rapx/analysis/core/range_analysis/
default.rs

1#![allow(unused_imports)]
2
3use crate::{
4    analysis::{
5        Analysis,
6        core::{
7            alias_analysis::default::graph::MopGraph,
8            callgraph::{default::CallGraphInfo, visitor::CallGraphVisitor},
9            range_analysis::{
10                Range, RangeAnalysis,
11                domain::{
12                    ConstraintGraph::ConstraintGraph,
13                    domain::{ConstConvert, IntervalArithmetic, VarNodes},
14                },
15            },
16            ssa_transform::*,
17        },
18        graphs::scc::Scc,
19    },
20    rap_debug, rap_info,
21};
22
23use rustc_data_structures::fx::FxHashMap;
24use rustc_hir::{def::DefKind, def_id::DefId};
25use rustc_middle::{
26    mir::{Body, Place},
27    ty::TyCtxt,
28};
29use std::{
30    cell::RefCell,
31    collections::{HashMap, HashSet},
32    fmt::Debug,
33    fs::{self, File},
34    io::Write,
35    path::PathBuf,
36    rc::Rc,
37};
38
39use super::{PathConstraint, PathConstraintMap, RAResult, RAResultMap, RAVecResultMap};
40pub struct RangeAnalyzer<'tcx, T: IntervalArithmetic + ConstConvert + Debug> {
41    pub tcx: TyCtxt<'tcx>,
42    pub debug: bool,
43    pub ssa_def_id: Option<DefId>,
44    pub essa_def_id: Option<DefId>,
45    pub final_vars: RAResultMap<'tcx, T>,
46    pub ssa_places_mapping: FxHashMap<DefId, HashMap<Place<'tcx>, HashSet<Place<'tcx>>>>,
47    pub fn_constraintgraph_mapping: FxHashMap<DefId, ConstraintGraph<'tcx, T>>,
48    pub callgraph: CallGraphInfo<'tcx>,
49    pub body_map: FxHashMap<DefId, Body<'tcx>>,
50    pub cg_map: FxHashMap<DefId, Rc<RefCell<ConstraintGraph<'tcx, T>>>>,
51    pub vars_map: FxHashMap<DefId, Vec<RefCell<VarNodes<'tcx, T>>>>,
52    pub final_vars_vec: RAVecResultMap<'tcx, T>,
53    pub path_constraints: PathConstraintMap<'tcx>,
54}
55impl<'tcx, T: IntervalArithmetic + ConstConvert + Debug> Analysis for RangeAnalyzer<'tcx, T>
56where
57    T: IntervalArithmetic + ConstConvert + Debug,
58{
59    fn name(&self) -> &'static str {
60        "Range Analysis"
61    }
62
63    fn run(&mut self) {
64        // self.start();
65        self.only_caller_range_analysis();
66        self.start_path_constraints_analysis();
67    }
68
69    fn reset(&mut self) {
70        self.final_vars.clear();
71        self.ssa_places_mapping.clear();
72    }
73}
74
75impl<'tcx, T: IntervalArithmetic + ConstConvert + Debug> RangeAnalysis<'tcx, T>
76    for RangeAnalyzer<'tcx, T>
77where
78    T: IntervalArithmetic + ConstConvert + Debug,
79{
80    fn get_fn_range(&self, def_id: DefId) -> Option<RAResult<'tcx, T>> {
81        self.final_vars.get(&def_id).cloned()
82    }
83    fn get_fn_ranges_percall(&self, def_id: DefId) -> Option<Vec<RAResult<'tcx, T>>> {
84        self.final_vars_vec.get(&def_id).cloned()
85    }
86    fn get_all_fn_ranges(&self) -> RAResultMap<'tcx, T> {
87        // REFACTOR: Using `.clone()` is more explicit that a copy is being returned.
88        self.final_vars.clone()
89    }
90    fn get_all_fn_ranges_percall(&self) -> RAVecResultMap<'tcx, T> {
91        self.final_vars_vec.clone()
92    }
93
94    // REFACTOR: This lookup is now much more efficient.
95    fn get_fn_local_range(&self, def_id: DefId, place: Place<'tcx>) -> Option<Range<T>> {
96        self.final_vars
97            .get(&def_id)
98            .and_then(|vars| vars.get(&place).cloned())
99    }
100
101    fn get_fn_path_constraints(&self, def_id: DefId) -> Option<PathConstraint<'tcx>> {
102        self.path_constraints.get(&def_id).cloned()
103    }
104
105    fn get_all_path_constraints(&self) -> PathConstraintMap<'tcx> {
106        self.path_constraints.clone()
107    }
108}
109
110impl<'tcx, T> RangeAnalyzer<'tcx, T>
111where
112    T: IntervalArithmetic + ConstConvert + Debug,
113{
114    pub fn new(tcx: TyCtxt<'tcx>, debug: bool) -> Self {
115        let mut ssa_id = None;
116        let mut essa_id = None;
117
118        if let Some(ssa_def_id) = tcx.hir_crate_items(()).free_items().find(|id| {
119            let hir_id = id.hir_id();
120            if let Some(ident_name) = tcx.hir_opt_name(hir_id) {
121                ident_name.to_string() == "SSAstmt"
122            } else {
123                false
124            }
125        }) {
126            ssa_id = Some(ssa_def_id.owner_id.to_def_id());
127            if let Some(essa_def_id) = tcx.hir_crate_items(()).free_items().find(|id| {
128                let hir_id = id.hir_id();
129                if let Some(ident_name) = tcx.hir_opt_name(hir_id) {
130                    ident_name.to_string() == "ESSAstmt"
131                } else {
132                    false
133                }
134            }) {
135                essa_id = Some(essa_def_id.owner_id.to_def_id());
136            }
137        }
138        Self {
139            tcx: tcx,
140            debug,
141            ssa_def_id: ssa_id,
142            essa_def_id: essa_id,
143            final_vars: FxHashMap::default(),
144            ssa_places_mapping: FxHashMap::default(),
145            fn_constraintgraph_mapping: FxHashMap::default(),
146            callgraph: CallGraphInfo::new(),
147            body_map: FxHashMap::default(),
148            cg_map: FxHashMap::default(),
149            vars_map: FxHashMap::default(),
150            final_vars_vec: FxHashMap::default(),
151            path_constraints: FxHashMap::default(),
152        }
153    }
154
155    fn build_constraintgraph(&mut self, body_mut_ref: &'tcx Body<'tcx>, def_id: DefId) {
156        let ssa_def_id = self.ssa_def_id.expect("SSA definition ID is not set");
157        let essa_def_id = self.essa_def_id.expect("ESSA definition ID is not set");
158        let mut cg: ConstraintGraph<'tcx, T> =
159            ConstraintGraph::new(body_mut_ref, self.tcx, def_id, essa_def_id, ssa_def_id);
160        cg.build_graph(body_mut_ref);
161        cg.build_nuutila(false);
162        // cg.rap_print_vars();
163        // cg.rap_print_final_vars();
164        let dot_output = cg.to_dot();
165        let vars_map = cg.get_vars().clone();
166
167        self.cg_map.insert(def_id, Rc::new(RefCell::new(cg)));
168        let mut vec = Vec::new();
169        vec.push(RefCell::new(vars_map));
170        self.vars_map.insert(def_id, vec);
171        let function_name = self.tcx.def_path_str(def_id);
172
173        let dir_path = PathBuf::from("cg_dot");
174        fs::create_dir_all(dir_path.clone()).unwrap();
175        let safe_filename = format!("{}_cg.dot", function_name);
176        let output_path = dir_path.join(format!("{}", safe_filename));
177
178        let mut file = File::create(&output_path).expect("cannot create file");
179        file.write_all(dot_output.as_bytes())
180            .expect("Could not write to file");
181
182        println!("Successfully generated graph.dot");
183        // println!("Run 'dot -Tpng -o graph.png graph.dot' to generate the image.");
184    }
185
186    fn only_caller_range_analysis(&mut self) {
187        let ssa_def_id = self.ssa_def_id.expect("SSA definition ID is not set");
188        let essa_def_id = self.essa_def_id.expect("ESSA definition ID is not set");
189        // ====================================================================
190        // PHASE 1: Build all ConstraintGraphs and the complete CallGraph first.
191        // ====================================================================
192        rap_debug!("PHASE 1: Building all ConstraintGraphs and the CallGraph...");
193        for local_def_id in self.tcx.iter_local_def_id() {
194            if matches!(self.tcx.def_kind(local_def_id), DefKind::Fn) {
195                let def_id = local_def_id.to_def_id();
196
197                if self.tcx.is_mir_available(def_id) {
198                    let mut body = self.tcx.optimized_mir(def_id).clone();
199                    let body_mut_ref = unsafe { &mut *(&mut body as *mut Body<'tcx>) };
200                    // Run SSA/ESSA passes
201                    let mut passrunner = PassRunner::new(self.tcx);
202                    passrunner.run_pass(body_mut_ref, ssa_def_id, essa_def_id);
203                    self.body_map.insert(def_id, body);
204                    // Print the MIR after SSA/ESSA passes
205                    rap_debug!("{:#?}", body_mut_ref.local_decls);
206                    if self.debug {
207                        print_diff(self.tcx, body_mut_ref, def_id.into());
208                        print_mir_graph(self.tcx, body_mut_ref, def_id.into());
209                    }
210
211                    self.ssa_places_mapping
212                        .insert(def_id, passrunner.places_map.clone());
213                    // rap_debug!("ssa_places_mapping: {:?}", self.ssa_places_mapping);
214                    // Build and store the constraint graph
215                    self.build_constraintgraph(body_mut_ref, def_id);
216                    // Visit for call graph construction
217                    let mut call_graph_visitor =
218                        CallGraphVisitor::new(self.tcx, def_id, body_mut_ref, &mut self.callgraph);
219                    call_graph_visitor.visit();
220                }
221            }
222        }
223        rap_debug!("PHASE 1 Complete. CallGraph built.");
224        // self.callgraph.print_call_graph(); // Optional: for debugging
225
226        // ====================================================================
227        // PHASE 2: Analyze only the call chain start functions.
228        // ====================================================================
229        rap_debug!("PHASE 2: Finding and analyzing call chain start functions...");
230
231        let mut call_chain_starts: Vec<DefId> = Vec::new();
232
233        let callers_by_callee_id = self.callgraph.get_callers_map();
234
235        for (&node_id_usize, node) in &self.callgraph.functions {
236            let def_id = node.get_def_id();
237
238            if !callers_by_callee_id.contains_key(&node_id_usize)
239                && self.cg_map.contains_key(&def_id)
240            {
241                call_chain_starts.push(def_id);
242            }
243        }
244
245        call_chain_starts.sort_by_key(|d| self.tcx.def_path_str(*d));
246
247        rap_debug!(
248            "Found call chain starts ({} functions): {:?}",
249            call_chain_starts.len(),
250            call_chain_starts
251                .iter()
252                .map(|d| self.tcx.def_path_str(*d))
253                .collect::<Vec<_>>()
254        );
255
256        for def_id in call_chain_starts {
257            rap_debug!(
258                "Analyzing function (call chain start): {}",
259                self.tcx.def_path_str(def_id)
260            );
261            if let Some(cg_cell) = self.cg_map.get(&def_id) {
262                let mut cg = cg_cell.borrow_mut();
263                cg.find_intervals(&self.cg_map, &mut self.vars_map);
264            } else {
265                rap_debug!(
266                    "Warning: No ConstraintGraph found for DefId {:?} during analysis of call chain starts.",
267                    def_id
268                );
269            }
270        }
271
272        let analysis_order = self.callgraph.get_reverse_post_order();
273        for def_id in analysis_order {
274            if let Some(cg_cell) = self.cg_map.get(&def_id) {
275                let mut cg = cg_cell.borrow_mut();
276                let (final_vars_for_fn, _) = cg.build_final_vars(&self.ssa_places_mapping[&def_id]);
277                let mut ranges_for_fn = HashMap::new();
278                for (&place, varnode) in final_vars_for_fn {
279                    ranges_for_fn.insert(place, varnode.get_range().clone());
280                }
281                let Some(varnodes_vec) = self.vars_map.get_mut(&def_id) else {
282                    rap_debug!(
283                        "Warning: No VarNodes found for DefId {:?} during analysis of call chain starts.",
284                        def_id
285                    );
286                    continue;
287                };
288                for varnodes in varnodes_vec.iter_mut() {
289                    let ranges_for_fn_recursive = ConstraintGraph::filter_final_vars(
290                        &varnodes.borrow(),
291                        &self.ssa_places_mapping[&def_id],
292                    );
293                    self.final_vars_vec
294                        .entry(def_id)
295                        .or_default()
296                        .push(ranges_for_fn_recursive);
297                }
298
299                self.final_vars.insert(def_id, ranges_for_fn);
300            }
301        }
302
303        rap_debug!("PHASE 2 Complete. Interval analysis finished for call chain start functions.");
304    }
305    pub fn start_path_constraints_analysis_for_defid(
306        &mut self,
307        def_id: DefId,
308    ) -> Option<PathConstraint<'tcx>> {
309        if self.tcx.is_mir_available(def_id) {
310            let mut body = self.tcx.optimized_mir(def_id).clone();
311            let body_mut_ref = unsafe { &mut *(&mut body as *mut Body<'tcx>) };
312
313            let mut cg: ConstraintGraph<'tcx, T> =
314                ConstraintGraph::new_without_ssa(body_mut_ref, self.tcx, def_id);
315            let mut graph = MopGraph::new(self.tcx, def_id);
316            graph.find_scc();
317            let paths: Vec<Vec<usize>> = graph.get_all_branch_sub_blocks_paths();
318            let result = cg.start_analyze_path_constraints(body_mut_ref, &paths);
319            rap_debug!(
320                "Paths for function {}: {:?}",
321                self.tcx.def_path_str(def_id),
322                paths
323            );
324            let switchbbs = cg.switchbbs.clone();
325            rap_debug!(
326                "Switch basicblocks for function {}: {:?}",
327                self.tcx.def_path_str(def_id),
328                switchbbs
329            );
330            rap_debug!(
331                "Path Constraints Analysis Result for function {}: {:?}",
332                self.tcx.def_path_str(def_id),
333                result
334            );
335            self.path_constraints.insert(def_id, result.clone());
336            Some(result)
337        } else {
338            None
339        }
340    }
341    pub fn start_path_constraints_analysis(&mut self) {
342        for local_def_id in self.tcx.iter_local_def_id() {
343            if matches!(self.tcx.def_kind(local_def_id), DefKind::Fn) {
344                let def_id = local_def_id.to_def_id();
345
346                if self.tcx.is_mir_available(def_id) {
347                    let mut body = self.tcx.optimized_mir(def_id).clone();
348                    let body_mut_ref = unsafe { &mut *(&mut body as *mut Body<'tcx>) };
349
350                    let mut cg: ConstraintGraph<'tcx, T> =
351                        ConstraintGraph::new_without_ssa(body_mut_ref, self.tcx, def_id);
352                    let mut graph = MopGraph::new(self.tcx, def_id);
353                    graph.find_scc();
354                    // rap_info!("child_scc: {:?}\n", graph.child_scc);
355                    // rap_info!("scc_indices: {:?}\n", graph.scc_indices);
356                    // rap_info!("blocks: {:?}\n", graph.blocks);
357                    let paths: Vec<Vec<usize>> = graph.get_all_branch_sub_blocks_paths();
358                    let result = cg.start_analyze_path_constraints(body_mut_ref, &paths);
359                    rap_debug!(
360                        "Paths for function {}: {:?}",
361                        self.tcx.def_path_str(def_id),
362                        paths
363                    );
364                    let switchbbs = cg.switchbbs.clone();
365                    rap_debug!(
366                        "Switch basicblocks for function {}: {:?}",
367                        self.tcx.def_path_str(def_id),
368                        switchbbs
369                    );
370                    rap_debug!(
371                        "Path Constraints Analysis Result for function {}: {:?}",
372                        self.tcx.def_path_str(def_id),
373                        result
374                    );
375                    self.path_constraints.insert(def_id, result);
376                }
377            }
378        }
379    }
380}
rapx/analysis/core/range_analysis/default.rs

rapx/analysis/core/range_analysis/
default.rs
