rapx/analysis/core/call_graph/
call_graph_visitor.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
use super::call_graph_helper::CallGraphInfo;
use regex::Regex;
use rustc_hir::def_id::DefId;
use rustc_middle::mir;
use rustc_middle::ty::{FnDef, Instance, InstanceKind, TyCtxt};

pub struct CallGraphVisitor<'b, 'tcx> {
    tcx: TyCtxt<'tcx>,
    def_id: DefId,
    body: &'tcx mir::Body<'tcx>,
    call_graph_info: &'b mut CallGraphInfo,
}

impl<'b, 'tcx> CallGraphVisitor<'b, 'tcx> {
    pub fn new(
        tcx: TyCtxt<'tcx>,
        def_id: DefId,
        body: &'tcx mir::Body<'tcx>,
        call_graph_info: &'b mut CallGraphInfo,
    ) -> Self {
        Self {
            tcx: tcx,
            def_id: def_id,
            body: body,
            call_graph_info: call_graph_info,
        }
    }

    pub fn add_in_call_graph(
        &mut self,
        caller_def_path: &String,
        callee_def_id: DefId,
        callee_def_path: &String,
    ) {
        if let Some(caller_id) = self.call_graph_info.get_noed_by_path(caller_def_path) {
            if let Some(callee_id) = self.call_graph_info.get_noed_by_path(callee_def_path) {
                self.call_graph_info
                    .add_funciton_call_edge(caller_id, callee_id);
            } else {
                self.call_graph_info
                    .add_node(callee_def_id, callee_def_path);
                if let Some(callee_id) = self.call_graph_info.get_noed_by_path(callee_def_path) {
                    self.call_graph_info
                        .add_funciton_call_edge(caller_id, callee_id);
                }
            }
        }
    }

    pub fn visit(&mut self) {
        let caller_path_str = self.tcx.def_path_str(self.def_id);
        self.call_graph_info.add_node(self.def_id, &caller_path_str);
        for (_, data) in self.body.basic_blocks.iter().enumerate() {
            let terminator = data.terminator();
            self.visit_terminator(&terminator);
        }
    }

    fn add_to_call_graph(&mut self, callee_def_id: DefId, is_virtual: Option<bool>) {
        let caller_def_path = self.tcx.def_path_str(self.def_id);
        let mut callee_def_path = self.tcx.def_path_str(callee_def_id);
        if let Some(judge) = is_virtual {
            if judge {
                let re = Regex::new(r"(?<dyn>\w+)::(?<func>\w+)").unwrap();
                let Some(caps) = re.captures(&callee_def_path) else {
                    return;
                };
                callee_def_path = format!("(dyn trait) <* as {}>::{}", &caps["dyn"], &caps["func"]);
            }
        }

        // let callee_location = self.tcx.def_span(callee_def_id);
        if callee_def_id == self.def_id {
            // Recursion
            println!("Warning! Find a recursion function which may cause stackoverflow!")
        }
        self.add_in_call_graph(&caller_def_path, callee_def_id, &callee_def_path);
    }

    fn visit_terminator(&mut self, terminator: &mir::Terminator<'tcx>) {
        if let mir::TerminatorKind::Call { func, .. } = &terminator.kind {
            if let mir::Operand::Constant(constant) = func {
                if let FnDef(callee_def_id, callee_substs) = constant.const_.ty().kind() {
                    let param_env = self.tcx.param_env(self.def_id);
                    if let Ok(Some(instance)) =
                        Instance::try_resolve(self.tcx, param_env, *callee_def_id, callee_substs)
                    {
                        let mut is_virtual = false;
                        // Try to analysis the specific type of callee.
                        let instance_def_id = match instance.def {
                            InstanceKind::Item(def_id) => Some(def_id),
                            InstanceKind::Intrinsic(def_id) => Some(def_id),
                            InstanceKind::VTableShim(def_id) => Some(def_id),
                            InstanceKind::ReifyShim(def_id, _) => Some(def_id),
                            InstanceKind::FnPtrShim(def_id, _) => Some(def_id),
                            InstanceKind::Virtual(def_id, _) => {
                                is_virtual = true;
                                Some(def_id)
                            }
                            InstanceKind::ClosureOnceShim { call_once, .. } => Some(call_once),
                            InstanceKind::ConstructCoroutineInClosureShim {
                                coroutine_closure_def_id,
                                ..
                            } => Some(coroutine_closure_def_id),
                            InstanceKind::ThreadLocalShim(def_id) => Some(def_id),
                            InstanceKind::DropGlue(def_id, _) => Some(def_id),
                            InstanceKind::FnPtrAddrShim(def_id, _) => Some(def_id),
                            InstanceKind::AsyncDropGlueCtorShim(def_id, _) => Some(def_id),
                            InstanceKind::CloneShim(def_id, _) => {
                                if !self.tcx.is_closure_like(def_id) {
                                    // Not a closure
                                    Some(def_id)
                                } else {
                                    None
                                }
                            }
                        };
                        if let Some(instance_def_id) = instance_def_id {
                            self.add_to_call_graph(instance_def_id, Some(is_virtual));
                        }
                    } else {
                        // Although failing to get specific type, callee is still useful.
                        self.add_to_call_graph(*callee_def_id, None);
                    }
                }
            }
        }
    }
}