rapx/analysis/callgraph/
default.rs

1use rustc_hir::{def::DefKind, def_id::DefId};
2use rustc_middle::{
3    mir::{self, Body},
4    ty::TyCtxt,
5};
6use std::collections::HashMap;
7use std::collections::HashSet;
8
9use super::visitor::CallGraphVisitor;
10use crate::{
11    Analysis,
12    analysis::callgraph::{CallGraphAnalysis, FnCallMap},
13};
14
15pub struct CallGraphAnalyzer<'tcx> {
16    pub tcx: TyCtxt<'tcx>,
17    pub graph: CallGraph<'tcx>,
18}
19
20impl<'tcx> Analysis for CallGraphAnalyzer<'tcx> {
21    fn run(&mut self) {
22        self.start();
23    }
24
25}
26
27impl<'tcx> CallGraphAnalysis for CallGraphAnalyzer<'tcx> {
28    fn get_fn_calls(&self) -> FnCallMap {
29        let fn_calls: HashMap<DefId, Vec<DefId>> = self
30            .graph
31            .fn_calls
32            .clone()
33            .into_iter()
34            .map(|(caller, callees)| {
35                let callee_ids = callees.into_iter().map(|(did, _)| did).collect::<Vec<_>>();
36                (caller, callee_ids)
37            })
38            .collect();
39        fn_calls
40    }
41}
42
43impl<'tcx> CallGraphAnalyzer<'tcx> {
44    pub fn new(tcx: TyCtxt<'tcx>) -> Self {
45        Self {
46            tcx: tcx,
47            graph: CallGraph::new(tcx),
48        }
49    }
50
51    pub fn start(&mut self) {
52        for local_def_id in self.tcx.mir_keys(()) {
53            let def_id = local_def_id.to_def_id();
54            if self.tcx.is_mir_available(def_id) {
55                let def_kind = self.tcx.def_kind(def_id);
56
57                let body: &Body<'_> = match def_kind {
58                    DefKind::Fn | DefKind::AssocFn | DefKind::Closure => {
59                        &self.tcx.optimized_mir(def_id)
60                    }
61                    #[cfg(rapx_ge_99)]
62                    DefKind::Const { .. }
63                    | DefKind::Static { .. }
64                    | DefKind::AssocConst { .. }
65                    | DefKind::AnonConst => {
66                        // NOTE: safer fallback for constants
67                        &self.tcx.mir_for_ctfe(def_id)
68                    }
69                    #[cfg(not(rapx_ge_99))]
70                    DefKind::Const
71                    | DefKind::Static { .. }
72                    | DefKind::AssocConst
73                    | DefKind::AnonConst => {
74                        // NOTE: safer fallback for constants
75                        &self.tcx.mir_for_ctfe(def_id)
76                    }
77                    #[cfg(not(rapx_ge_99))]
78                    DefKind::InlineConst => &self.tcx.mir_for_ctfe(def_id),
79                    // These don't have MIR or shouldn't be visited
80                    _ => {
81                        rap_debug!("Skipping def_id {:?} with kind {:?}", def_id, def_kind);
82                        continue;
83                    }
84                };
85
86                let mut call_graph_visitor =
87                    CallGraphVisitor::new(self.tcx, def_id.into(), body, &mut self.graph);
88                call_graph_visitor.visit();
89            }
90        }
91    }
92}
93
94pub type CallMap<'tcx> = HashMap<DefId, Vec<(DefId, Option<&'tcx mir::Terminator<'tcx>>)>>;
95
96pub struct CallGraph<'tcx> {
97    pub tcx: TyCtxt<'tcx>,
98    pub functions: HashSet<DefId>, // Function-like, including closures
99    pub fn_calls: CallMap<'tcx>,   // caller -> Vec<(callee, terminator)>
100}
101
102/// Internal apis for constructing a call graph
103impl<'tcx> CallGraph<'tcx> {
104    pub fn new(tcx: TyCtxt<'tcx>) -> Self {
105        Self {
106            tcx,
107            functions: HashSet::new(),
108            fn_calls: HashMap::new(),
109        }
110    }
111
112    /// Register a function to the call graph. Return true on insert, false if that DefId already exists.
113    pub fn register_fn(&mut self, def_id: DefId) -> bool {
114        if let Some(_) = self.functions.iter().find(|func_id| **func_id == def_id) {
115            false
116        } else {
117            self.functions.insert(def_id);
118            true
119        }
120    }
121
122    /// Add a function call to the call graph.
123    pub fn add_funciton_call(
124        &mut self,
125        caller_id: DefId,
126        callee_id: DefId,
127        terminator_stmt: Option<&'tcx mir::Terminator<'tcx>>,
128    ) {
129        let entry = self.fn_calls.entry(caller_id).or_insert_with(Vec::new);
130        entry.push((callee_id, terminator_stmt));
131    }
132}
133
134/// Public apis to get information from the call graph
135impl<'tcx> CallGraph<'tcx> {
136    pub fn get_reverse_post_order(&self) -> Vec<DefId> {
137        let mut result = self.get_post_order();
138        result.reverse();
139        result
140    }
141
142    pub fn get_post_order(&self) -> Vec<DefId> {
143        let mut visited = HashSet::new();
144        let mut post_order_ids = Vec::new(); // Will store the post-order traversal of `usize` IDs
145
146        // Iterate over all functions defined in the graph to handle disconnected components
147        for &func_def_id in self.functions.iter() {
148            if !visited.contains(&func_def_id) {
149                self.dfs_post_order(func_def_id, &mut visited, &mut post_order_ids);
150            }
151        }
152
153        post_order_ids
154    }
155
156    /// Helper function to perform a recursive depth-first search.
157    fn dfs_post_order(
158        &self,
159        func_def_id: DefId,
160        visited: &mut HashSet<DefId>,
161        post_order_ids: &mut Vec<DefId>,
162    ) {
163        // Mark the current node as visited
164        visited.insert(func_def_id);
165
166        // Visit all callees (children) of the current node
167        if let Some(callees) = self.fn_calls.get(&func_def_id) {
168            for (callee_id, _terminator) in callees {
169                if !visited.contains(callee_id) {
170                    self.dfs_post_order(*callee_id, visited, post_order_ids);
171                }
172            }
173        }
174
175        // After visiting all children, add the current node to the post-order list
176        post_order_ids.push(func_def_id);
177    }
178
179    /// Get a reversed (callee -> Vec<Caller>) call map.
180    pub fn get_callers_map(&self) -> CallMap<'tcx> {
181        let mut callers_map: CallMap<'tcx> = HashMap::new();
182
183        for (&caller_id, calls_vec) in &self.fn_calls {
184            for (callee_id, terminator) in calls_vec {
185                callers_map
186                    .entry(*callee_id)
187                    .or_insert_with(Vec::new)
188                    .push((caller_id, *terminator));
189            }
190        }
191        callers_map
192    }
193
194    /// Get all direct callees' DefId of the caller function
195    pub fn get_callees(&self, caller_def_id: DefId) -> Vec<DefId> {
196        if let Some(callees) = self.fn_calls.get(&caller_def_id) {
197            callees
198                .clone()
199                .into_iter()
200                .map(|(did, _)| did)
201                .collect::<Vec<_>>()
202        } else {
203            vec![]
204        }
205    }
206
207    /// Get all recursively reachable callee's DefId
208    pub fn get_callees_recursive(&self, caller_def_id: DefId) -> Vec<DefId> {
209        let mut visited = HashSet::new();
210        let mut result = Vec::new();
211        self.dfs_post_order(caller_def_id, &mut visited, &mut result);
212        result
213    }
214
215    /// Get all direct callers' DefId of the callee function
216    pub fn get_callers(&self, callee_def_id: DefId) -> Vec<DefId> {
217        let callers_map = self.get_callers_map();
218        if let Some(callers) = callers_map.get(&callee_def_id) {
219            callers
220                .clone()
221                .into_iter()
222                .map(|(did, _)| did)
223                .collect::<Vec<_>>()
224        } else {
225            vec![]
226        }
227    }
228}
rapx/analysis/callgraph/default.rs

rapx/analysis/callgraph/
default.rs
