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

rapx/analysis/callgraph/
default.rs
