go_brrr/callgraph/
types.rs

1//! Call graph type definitions.
2
3use serde::{Deserialize, Serialize};
4use std::cell::OnceCell;
5use std::collections::{HashMap, HashSet, VecDeque};
6
7/// Reference to a function.
8#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
9pub struct FunctionRef {
10    /// File containing the function
11    pub file: String,
12    /// Function name
13    pub name: String,
14    /// Fully qualified name (e.g., "module.Class.method")
15    pub qualified_name: Option<String>,
16}
17
18impl FunctionRef {
19    /// Check if this ref matches another (by name, ignoring file if empty).
20    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/// A call edge in the graph.
29#[derive(Debug, Clone, Serialize, Deserialize)]
30pub struct CallEdge {
31    /// Calling function
32    pub caller: FunctionRef,
33    /// Called function
34    pub callee: FunctionRef,
35    /// Line number of the call
36    pub call_line: usize,
37}
38
39/// Complete call graph for a project.
40#[derive(Debug, Default, Serialize, Deserialize)]
41pub struct CallGraph {
42    /// All call edges
43    pub edges: Vec<CallEdge>,
44
45    /// Index: callee -> callers (who calls this?)
46    #[serde(skip)]
47    pub callers: HashMap<FunctionRef, HashSet<FunctionRef>>,
48
49    /// Index: caller -> callees (what does this call?)
50    #[serde(skip)]
51    pub callees: HashMap<FunctionRef, HashSet<FunctionRef>>,
52
53    /// Cached set of all functions (lazily computed on first access)
54    #[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            // Clone the cache content if present, otherwise leave empty
65            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    /// Create a new CallGraph from edges.
81    ///
82    /// Indexes are not built automatically. Call `build_indexes()` after creation.
83    #[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    /// Build indexes from edges (call after deserialization).
94    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    /// Invalidate all caches (call after modifying edges).
112    pub fn invalidate_caches(&mut self) {
113        self.all_functions_cache.take();
114    }
115
116    /// Forward call graph: what does function X call?
117    ///
118    /// Uses partial matching via `FunctionRef::matches()` to find the source function,
119    /// then performs BFS traversal up to the specified depth.
120    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        // Find all matching source functions and seed the queue with their direct callees
129        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            // Only expand if we haven't reached max depth
145            if current_depth < depth {
146                if let Some(callees) = self.callees.get(&current) {
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    /// Get all unique functions in the graph (cached).
160    ///
161    /// The result is lazily computed on first access and cached for subsequent calls.
162    /// The cache is automatically invalidated when `build_indexes()` or
163    /// `invalidate_caches()` is called.
164    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}