Skip to main content

rapx/analysis/range/
default.rs

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