rapx/analysis/callgraph/
visitor.rs

1use super::default::CallGraph;
2use rustc_hir::def_id::DefId;
3use rustc_middle::mir;
4use rustc_middle::ty::{FnDef, Instance, InstanceKind, TyCtxt, TypingEnv};
5#[cfg(rapx_rustc_ge_198)]
6use rustc_middle::ty::ShimKind;
7use std::collections::HashSet;
8
9pub struct CallGraphVisitor<'b, 'tcx> {
10    tcx: TyCtxt<'tcx>,
11    def_id: DefId,
12    body: &'tcx mir::Body<'tcx>,
13    call_graph_info: &'b mut CallGraph<'tcx>,
14}
15
16impl<'b, 'tcx> CallGraphVisitor<'b, 'tcx> {
17    pub fn new(
18        tcx: TyCtxt<'tcx>,
19        def_id: DefId,
20        body: &'tcx mir::Body<'tcx>,
21        call_graph_info: &'b mut CallGraph<'tcx>,
22    ) -> Self {
23        Self {
24            tcx: tcx,
25            def_id: def_id,
26            body: body,
27            call_graph_info: call_graph_info,
28        }
29    }
30
31    fn add_fn_call(&mut self, callee_def_id: DefId, terminator: &'tcx mir::Terminator<'tcx>) {
32        self.call_graph_info.register_fn(callee_def_id);
33        self.call_graph_info.add_funciton_call(
34            self.def_id.clone(),
35            callee_def_id,
36            Some(terminator),
37        );
38    }
39
40    fn handle_fn_call(
41        &mut self,
42        callee_def_id: DefId,
43        is_virtual: bool,
44        terminator: &'tcx mir::Terminator<'tcx>,
45    ) {
46        if is_virtual {
47            // Handle dynamic dispatch for trait objects
48            self.handle_virtual_call(callee_def_id, terminator);
49        } else {
50            self.add_fn_call(callee_def_id, terminator);
51        }
52    }
53
54    fn handle_virtual_call(
55        &mut self,
56        stub_def_id: DefId, // Callee is the dynamic call stub, i.e. the fn definition in trait
57        terminator: &'tcx mir::Terminator<'tcx>,
58    ) {
59        // Step 1: Add an edge from caller to the virtual function (stub);
60        // If the DefId exists, we assume that stub has been analyzed.
61        let visited = !self.call_graph_info.register_fn(stub_def_id);
62        self.add_fn_call(stub_def_id, terminator);
63
64        // If this function has already been analyzed, return;
65        if visited {
66            return;
67        }
68
69        // Step 2: Find all impls of the virtual function;
70        let mut candidates: HashSet<DefId> = HashSet::new();
71        if let Some(trait_def_id) = self.tcx.trait_of_assoc(stub_def_id) {
72            rap_debug!(
73                "[Callgraph] Virtual fn {:?} belongs to trait {:?}",
74                stub_def_id,
75                trait_def_id
76            );
77            for impl_id in self.tcx.all_impls(trait_def_id) {
78                let impl_map = self.tcx.impl_item_implementor_ids(impl_id);
79                if let Some(candidate_def_id) = impl_map.get(&stub_def_id) {
80                    candidates.insert(*candidate_def_id);
81                }
82            }
83        }
84        rap_debug!(
85            "[Callgraph] Implementors of {:?}: {:?}",
86            stub_def_id,
87            candidates
88        );
89
90        // Step 3: For each implementor, add an edge from the stub to it.
91        for candidate_def_id in candidates {
92            self.add_fn_call(candidate_def_id, terminator);
93        }
94    }
95
96    pub fn visit(&mut self) {
97        self.call_graph_info.register_fn(self.def_id);
98        for (_, data) in self.body.basic_blocks.iter().enumerate() {
99            let terminator = data.terminator();
100            self.visit_terminator(&terminator);
101        }
102    }
103
104    fn visit_terminator(&mut self, terminator: &'tcx mir::Terminator<'tcx>) {
105        if let mir::TerminatorKind::Call { func, .. } = &terminator.kind {
106            if let mir::Operand::Constant(constant) = func {
107                if let FnDef(callee_def_id, callee_substs) = constant.const_.ty().kind() {
108                    let ty_env = TypingEnv::post_analysis(self.tcx, self.def_id);
109                    if let Ok(Some(instance)) =
110                        Instance::try_resolve(self.tcx, ty_env, *callee_def_id, callee_substs)
111                    {
112                        let mut is_virtual = false;
113                        // Try to analysis the specific type of callee.
114                        #[cfg(rapx_rustc_ge_198)]
115                        let instance_def_id = match instance.def {
116                            InstanceKind::Item(def_id) => Some(def_id),
117                            InstanceKind::Intrinsic(def_id) => Some(def_id),
118                            InstanceKind::Virtual(def_id, _) => {
119                                is_virtual = true;
120                                Some(def_id)
121                            }
122                            InstanceKind::Shim(shim_kind) => match shim_kind {
123                                ShimKind::VTable(def_id) => Some(def_id),
124                                ShimKind::Reify(def_id, _) => Some(def_id),
125                                ShimKind::FnPtr(def_id, _) => Some(def_id),
126                                ShimKind::ClosureOnce { call_once, .. } => Some(call_once),
127                                ShimKind::ConstructCoroutineInClosure {
128                                    coroutine_closure_def_id,
129                                    ..
130                                } => Some(coroutine_closure_def_id),
131                                ShimKind::ThreadLocal(def_id) => Some(def_id),
132                                ShimKind::DropGlue(def_id, _) => Some(def_id),
133                                ShimKind::FnPtrAddr(def_id, _) => Some(def_id),
134                                ShimKind::AsyncDropGlueCtor(def_id, _) => Some(def_id),
135                                ShimKind::Clone(def_id, _) => {
136                                    if !self.tcx.is_closure_like(def_id) {
137                                        Some(def_id)
138                                    } else {
139                                        None
140                                    }
141                                }
142                                _ => todo!(),
143                            },
144                        };
145
146                        #[cfg(not(rapx_rustc_ge_198))]
147                        let instance_def_id = match instance.def {
148                            InstanceKind::Item(def_id) => Some(def_id),
149                            InstanceKind::Intrinsic(def_id) => Some(def_id),
150                            InstanceKind::VTableShim(def_id) => Some(def_id),
151                            InstanceKind::ReifyShim(def_id, _) => Some(def_id),
152                            InstanceKind::FnPtrShim(def_id, _) => Some(def_id),
153                            InstanceKind::Virtual(def_id, _) => {
154                                is_virtual = true;
155                                Some(def_id)
156                            }
157                            InstanceKind::ClosureOnceShim { call_once, .. } => Some(call_once),
158                            InstanceKind::ConstructCoroutineInClosureShim {
159                                coroutine_closure_def_id,
160                                ..
161                            } => Some(coroutine_closure_def_id),
162                            InstanceKind::ThreadLocalShim(def_id) => Some(def_id),
163                            InstanceKind::DropGlue(def_id, _) => Some(def_id),
164                            InstanceKind::FnPtrAddrShim(def_id, _) => Some(def_id),
165                            InstanceKind::AsyncDropGlueCtorShim(def_id, _) => Some(def_id),
166                            InstanceKind::CloneShim(def_id, _) => {
167                                if !self.tcx.is_closure_like(def_id) {
168                                    Some(def_id)
169                                } else {
170                                    None
171                                }
172                            }
173                            _ => todo!(),
174                        };
175
176                        if let Some(instance_def_id) = instance_def_id {
177                            self.handle_fn_call(instance_def_id, is_virtual, terminator);
178                        }
179                    } else {
180                        // Although failing to get specific type, callee is still useful.
181                        self.handle_fn_call(*callee_def_id, false, terminator);
182                    }
183                }
184            }
185        }
186    }
187}
rapx/analysis/callgraph/visitor.rs

rapx/analysis/callgraph/
visitor.rs
