go_brrr/callgraph/
types.rs1use serde::{Deserialize, Serialize};
4use std::cell::OnceCell;
5use std::collections::{HashMap, HashSet, VecDeque};
6
7#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
9pub struct FunctionRef {
10 pub file: String,
12 pub name: String,
14 pub qualified_name: Option<String>,
16}
17
18impl FunctionRef {
19 pub fn matches(&self, other: &FunctionRef) -> bool {
21 if !self.file.is_empty() && !other.file.is_empty() && self.file != other.file {
22 return false;
23 }
24 self.name == other.name
25 }
26}
27
28#[derive(Debug, Clone, Serialize, Deserialize)]
30pub struct CallEdge {
31 pub caller: FunctionRef,
33 pub callee: FunctionRef,
35 pub call_line: usize,
37}
38
39#[derive(Debug, Default, Serialize, Deserialize)]
41pub struct CallGraph {
42 pub edges: Vec<CallEdge>,
44
45 #[serde(skip)]
47 pub callers: HashMap<FunctionRef, HashSet<FunctionRef>>,
48
49 #[serde(skip)]
51 pub callees: HashMap<FunctionRef, HashSet<FunctionRef>>,
52
53 #[serde(skip)]
55 all_functions_cache: OnceCell<HashSet<FunctionRef>>,
56}
57
58impl Clone for CallGraph {
59 fn clone(&self) -> Self {
60 Self {
61 edges: self.edges.clone(),
62 callers: self.callers.clone(),
63 callees: self.callees.clone(),
64 all_functions_cache: self
66 .all_functions_cache
67 .get()
68 .cloned()
69 .map(|v| {
70 let cell = OnceCell::new();
71 let _ = cell.set(v);
72 cell
73 })
74 .unwrap_or_default(),
75 }
76 }
77}
78
79impl CallGraph {
80 #[must_use]
84 pub fn from_edges(edges: Vec<CallEdge>) -> Self {
85 Self {
86 edges,
87 callers: HashMap::new(),
88 callees: HashMap::new(),
89 all_functions_cache: OnceCell::new(),
90 }
91 }
92
93 pub fn build_indexes(&mut self) {
95 self.invalidate_caches();
96 self.callers.clear();
97 self.callees.clear();
98
99 for edge in &self.edges {
100 self.callers
101 .entry(edge.callee.clone())
102 .or_default()
103 .insert(edge.caller.clone());
104 self.callees
105 .entry(edge.caller.clone())
106 .or_default()
107 .insert(edge.callee.clone());
108 }
109 }
110
111 pub fn invalidate_caches(&mut self) {
113 self.all_functions_cache.take();
114 }
115
116 pub fn get_callees(&self, source: &FunctionRef, depth: usize) -> HashSet<FunctionRef> {
121 if depth == 0 {
122 return HashSet::new();
123 }
124
125 let mut visited: HashSet<FunctionRef> = HashSet::new();
126 let mut queue: VecDeque<(FunctionRef, usize)> = VecDeque::new();
127
128 for (k, callees) in &self.callees {
130 if k.matches(source) {
131 for callee in callees {
132 if !visited.contains(callee) {
133 queue.push_back((callee.clone(), 1));
134 }
135 }
136 }
137 }
138
139 while let Some((current, current_depth)) = queue.pop_front() {
140 if !visited.insert(current.clone()) {
141 continue;
142 }
143
144 if current_depth < depth {
146 if let Some(callees) = self.callees.get(¤t) {
147 for callee in callees {
148 if !visited.contains(callee) {
149 queue.push_back((callee.clone(), current_depth + 1));
150 }
151 }
152 }
153 }
154 }
155
156 visited
157 }
158
159 pub fn all_functions(&self) -> &HashSet<FunctionRef> {
165 self.all_functions_cache.get_or_init(|| {
166 self.edges
167 .iter()
168 .flat_map(|e| [e.caller.clone(), e.callee.clone()])
169 .collect()
170 })
171 }
172}