rapx/analysis/callgraph/
visitor.rs

1use super::default::CallGraph;
2use rustc_hir::def_id::DefId;
3use rustc_middle::mir;
4#[cfg(rapx_ge_99)]
5use rustc_middle::ty::ShimKind;
6use rustc_middle::ty::{FnDef, Instance, InstanceKind, TyCtxt, TypingEnv};
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                    #[cfg(rapx_ge_99)]
110                    let callee_substs = callee_substs.skip_binder();
111                    if let Ok(Some(instance)) =
112                        Instance::try_resolve(self.tcx, ty_env, *callee_def_id, callee_substs)
113                    {
114                        let mut is_virtual = false;
115                        // Try to analysis the specific type of callee.
116                        #[cfg(rapx_ge_99)]
117                        let instance_def_id = match instance.def {
118                            InstanceKind::Item(def_id) => Some(def_id),
119                            InstanceKind::Intrinsic(def_id) => Some(def_id),
120                            InstanceKind::Virtual(def_id, _) => {
121                                is_virtual = true;
122                                Some(def_id)
123                            }
124                            InstanceKind::Shim(shim_kind) => match shim_kind {
125                                ShimKind::VTable(def_id) => Some(def_id),
126                                ShimKind::Reify(def_id, _) => Some(def_id),
127                                ShimKind::FnPtr(def_id, _) => Some(def_id),
128                                ShimKind::ClosureOnce { call_once, .. } => Some(call_once),
129                                ShimKind::ConstructCoroutineInClosure {
130                                    coroutine_closure_def_id,
131                                    ..
132                                } => Some(coroutine_closure_def_id),
133                                ShimKind::ThreadLocal(def_id) => Some(def_id),
134                                ShimKind::DropGlue(def_id, _) => Some(def_id),
135                                #[cfg(rapx_has_fnptr_asptr)]
136                                ShimKind::FnPtrAsPtr(def_id, _) => Some(def_id),
137                                #[cfg(not(rapx_has_fnptr_asptr))]
138                                ShimKind::FnPtrAddr(def_id, _) => Some(def_id),
139                                ShimKind::AsyncDropGlueCtor(def_id, _) => Some(def_id),
140                                ShimKind::Clone(def_id, _) => {
141                                    if !self.tcx.is_closure_like(def_id) {
142                                        Some(def_id)
143                                    } else {
144                                        None
145                                    }
146                                }
147                                _ => todo!(),
148                            },
149                            #[allow(unreachable_patterns)]
150                            _ => None,
151                        };
152
153                        #[cfg(not(rapx_ge_99))]
154                        let instance_def_id = match instance.def {
155                            InstanceKind::Item(def_id) => Some(def_id),
156                            InstanceKind::Intrinsic(def_id) => Some(def_id),
157                            InstanceKind::VTableShim(def_id) => Some(def_id),
158                            InstanceKind::ReifyShim(def_id, _) => Some(def_id),
159                            InstanceKind::FnPtrShim(def_id, _) => Some(def_id),
160                            InstanceKind::Virtual(def_id, _) => {
161                                is_virtual = true;
162                                Some(def_id)
163                            }
164                            InstanceKind::ClosureOnceShim { call_once, .. } => Some(call_once),
165                            InstanceKind::ConstructCoroutineInClosureShim {
166                                coroutine_closure_def_id,
167                                ..
168                            } => Some(coroutine_closure_def_id),
169                            InstanceKind::ThreadLocalShim(def_id) => Some(def_id),
170                            InstanceKind::DropGlue(def_id, _) => Some(def_id),
171                            InstanceKind::FnPtrAddrShim(def_id, _) => Some(def_id),
172                            InstanceKind::AsyncDropGlueCtorShim(def_id, _) => Some(def_id),
173                            InstanceKind::CloneShim(def_id, _) => {
174                                if !self.tcx.is_closure_like(def_id) {
175                                    Some(def_id)
176                                } else {
177                                    None
178                                }
179                            }
180                            _ => todo!(),
181                        };
182
183                        if let Some(instance_def_id) = instance_def_id {
184                            self.handle_fn_call(instance_def_id, is_virtual, terminator);
185                        }
186                    } else {
187                        // Although failing to get specific type, callee is still useful.
188                        self.handle_fn_call(*callee_def_id, false, terminator);
189                    }
190                }
191            }
192        }
193    }
194}
rapx/analysis/callgraph/visitor.rs

rapx/analysis/callgraph/
visitor.rs
