rapx/analysis/core/call_graph/
call_graph_helper.rs

1use rustc_hir::def_id::DefId;
2use std::collections::HashSet;
3use std::{collections::HashMap, hash::Hash};
4
5use crate::rap_info;
6
7#[derive(Debug, Clone, Eq, PartialEq, Hash)]
8pub struct Node {
9    def_id: DefId,
10    def_path: String,
11}
12
13impl Node {
14    pub fn new(def_id: DefId, def_path: &String) -> Self {
15        Self {
16            def_id: def_id,
17            def_path: def_path.clone(),
18        }
19    }
20
21    pub fn get_def_id(&self) -> DefId {
22        self.def_id
23    }
24
25    pub fn get_def_path(&self) -> String {
26        self.def_path.clone()
27    }
28}
29
30pub struct CallGraphInfo {
31    pub functions: HashMap<usize, Node>, // id -> node
32    // pub function_calls: Vec<(usize, usize)>,   // (id, id)
33    pub function_calls: HashMap<usize, HashSet<usize>>,
34    pub node_registry: HashMap<String, usize>, // path -> id
35}
36
37impl CallGraphInfo {
38    pub fn new() -> Self {
39        Self {
40            functions: HashMap::new(),
41            function_calls: HashMap::new(),
42            node_registry: HashMap::new(),
43        }
44    }
45
46    pub fn get_node_num(&self) -> usize {
47        self.functions.len()
48    }
49
50    pub fn get_callees_path(&self, caller_def_path: &String) -> Option<HashSet<String>> {
51        let mut callees_path: HashSet<String> = HashSet::new();
52        if let Some(caller_id) = self.node_registry.get(caller_def_path) {
53            if let Some(callee_ids) = self.function_calls.get(&caller_id) {
54                for id in callee_ids {
55                    if let Some(callee_node) = self.functions.get(id) {
56                        callees_path.insert(callee_node.get_def_path());
57                    }
58                }
59            }
60            Some(callees_path)
61        } else {
62            None
63        }
64    }
65
66    pub fn add_node(&mut self, def_id: DefId, def_path: &String) {
67        if let None = self.get_noed_by_path(def_path) {
68            let id = self.node_registry.len();
69            let node = Node::new(def_id, def_path);
70            self.node_registry.insert(def_path.clone(), id);
71            self.functions.insert(id, node);
72        }
73    }
74
75    pub fn add_funciton_call_edge(&mut self, caller_id: usize, callee_id: usize) {
76        if !self.function_calls.contains_key(&caller_id) {
77            self.function_calls.insert(caller_id, HashSet::new());
78        }
79        if let Some(callees) = self.function_calls.get_mut(&caller_id) {
80            callees.insert(callee_id);
81        }
82    }
83
84    pub fn get_noed_by_path(&self, def_path: &String) -> Option<usize> {
85        if let Some(&id) = self.node_registry.get(def_path) {
86            Some(id)
87        } else {
88            None
89        }
90    }
91
92    pub fn print_call_graph(&self) {
93        rap_info!("CallGraph Analysis:");
94        // println!("There are {} functions calls!", self.function_calls.len());
95        for (caller_id, callees) in self.function_calls.clone() {
96            if let Some(caller_node) = self.functions.get(&caller_id) {
97                for callee_id in callees {
98                    if let Some(callee_node) = self.functions.get(&callee_id) {
99                        let caller_def_path = caller_node.get_def_path();
100                        let callee_def_path = callee_node.get_def_path();
101                        rap_info!(
102                            "{}:{} -> {}:{}",
103                            caller_id,
104                            caller_def_path,
105                            callee_id,
106                            callee_def_path
107                        );
108                    }
109                }
110            }
111        }
112        // rap_info!("There are {} functions:", self.functions.len());
113        // for (id, node) in self.functions.clone() {
114        //     rap_info!("{}:{}", id, node.get_def_path());
115        // }
116    }
117}