rapx/analysis/core/range_analysis/
default.rs

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

rapx/analysis/core/range_analysis/
default.rs
