rapx/analysis/core/range_analysis/
default.rs

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