Skip to main content

icb_graph/
query.rs

1use crate::graph::{CodePropertyGraph, Edge, Node};
2use icb_common::NodeKind;
3use petgraph::visit::{EdgeRef, IntoEdgeReferences};
4use std::collections::{HashSet, VecDeque};
5
6/// Returns all nodes of the given `kind`.
7///
8/// Traverses the entire graph and collects every node whose
9/// [`NodeKind`] matches the requested kind.
10///
11/// # Examples
12///
13/// ```rust
14/// use icb_graph::graph::{CodePropertyGraph, Node};
15/// use icb_graph::query::find_by_kind;
16/// use icb_common::NodeKind;
17///
18/// let mut cpg = CodePropertyGraph::new();
19/// cpg.graph.add_node(Node {
20///     kind: NodeKind::Function,
21///     name: Some("main".into()),
22///     usr: None,
23///     start_line: 1,
24///     end_line: 5,
25/// });
26///
27/// let functions = find_by_kind(&cpg, NodeKind::Function);
28/// assert_eq!(functions.len(), 1);
29/// ```
30pub fn find_by_kind(cpg: &CodePropertyGraph, kind: NodeKind) -> Vec<&Node> {
31    cpg.graph
32        .node_weights()
33        .filter(|n| n.kind == kind)
34        .collect()
35}
36
37/// Returns all call sites that target a function with the given `func_name`.
38///
39/// The comparison is done only on the node's name; no overload resolution
40/// is performed.
41///
42/// # Examples
43///
44/// ```rust
45/// use icb_graph::graph::{CodePropertyGraph, Node};
46/// use icb_graph::query::find_calls_to;
47/// use icb_common::NodeKind;
48///
49/// let mut cpg = CodePropertyGraph::new();
50/// cpg.graph.add_node(Node {
51///     kind: NodeKind::CallSite,
52///     name: Some("foo".into()),
53///     usr: None,
54///     start_line: 1,
55///     end_line: 1,
56/// });
57///
58/// let calls = find_calls_to(&cpg, "foo");
59/// assert_eq!(calls.len(), 1);
60/// ```
61pub fn find_calls_to<'a>(cpg: &'a CodePropertyGraph, func_name: &str) -> Vec<&'a Node> {
62    cpg.graph
63        .node_weights()
64        .filter(|n| n.kind == NodeKind::CallSite && n.name.as_deref() == Some(func_name))
65        .collect()
66}
67
68/// For a given function name, returns all functions that directly call it.
69///
70/// Each returned tuple contains the caller function node and the callee
71/// definition node. The enclosing function is found by walking up
72/// [`Edge::AstChild`] edges from the call site.
73///
74/// # Examples
75///
76/// ```rust
77/// use icb_graph::graph::{CodePropertyGraph, Node, Edge};
78/// use icb_graph::query::callers_of;
79/// use icb_common::NodeKind;
80///
81/// let mut cpg = CodePropertyGraph::new();
82/// let foo = cpg.graph.add_node(Node {
83///     kind: NodeKind::Function,
84///     name: Some("foo".into()),
85///     usr: Some("foo".into()),
86///     start_line: 1,
87///     end_line: 2,
88/// });
89/// let bar = cpg.graph.add_node(Node {
90///     kind: NodeKind::Function,
91///     name: Some("bar".into()),
92///     usr: Some("bar".into()),
93///     start_line: 5,
94///     end_line: 6,
95/// });
96/// let call = cpg.graph.add_node(Node {
97///     kind: NodeKind::CallSite,
98///     name: Some("bar".into()),
99///     usr: None,
100///     start_line: 2,
101///     end_line: 2,
102/// });
103/// cpg.graph.add_edge(foo, call, Edge::AstChild);
104/// cpg.graph.add_edge(call, bar, Edge::Call);
105///
106/// let callers = callers_of(&cpg, "bar");
107/// assert_eq!(callers.len(), 1);
108/// assert_eq!(callers[0].0.name.as_deref(), Some("foo"));
109/// ```
110pub fn callers_of<'a>(cpg: &'a CodePropertyGraph, func_name: &str) -> Vec<(&'a Node, &'a Node)> {
111    let target_defs: HashSet<_> = cpg
112        .graph
113        .node_indices()
114        .filter(|&idx| {
115            let node = &cpg.graph[idx];
116            (node.kind == NodeKind::Function || node.kind == NodeKind::Class)
117                && node.name.as_deref() == Some(func_name)
118        })
119        .collect();
120
121    if target_defs.is_empty() {
122        return vec![];
123    }
124
125    let mut results = Vec::new();
126    for edge_ref in cpg.graph.edge_references() {
127        if *edge_ref.weight() != Edge::Call {
128            continue;
129        }
130        let call_idx = edge_ref.source();
131        let def_idx = edge_ref.target();
132        if target_defs.contains(&def_idx) {
133            if let Some(enclosing) = get_enclosing_function(cpg, call_idx) {
134                results.push((enclosing, &cpg.graph[def_idx]));
135            }
136        }
137    }
138    results
139}
140
141/// For a given function name, returns all functions it directly calls.
142///
143/// Each returned tuple contains the callee node and the call‑site node.
144/// The function first locates the caller by name, collects all
145/// [`NodeKind::CallSite`] nodes inside its AST subtree, and follows
146/// outgoing [`Edge::Call`] edges.
147///
148/// # Examples
149///
150/// ```rust
151/// use icb_graph::graph::{CodePropertyGraph, Node, Edge};
152/// use icb_graph::query::callees_of;
153/// use icb_common::NodeKind;
154///
155/// let mut cpg = CodePropertyGraph::new();
156/// let foo = cpg.graph.add_node(Node {
157///     kind: NodeKind::Function,
158///     name: Some("foo".into()),
159///     usr: Some("foo".into()),
160///     start_line: 1,
161///     end_line: 2,
162/// });
163/// let bar = cpg.graph.add_node(Node {
164///     kind: NodeKind::Function,
165///     name: Some("bar".into()),
166///     usr: Some("bar".into()),
167///     start_line: 5,
168///     end_line: 6,
169/// });
170/// let call = cpg.graph.add_node(Node {
171///     kind: NodeKind::CallSite,
172///     name: Some("bar".into()),
173///     usr: None,
174///     start_line: 2,
175///     end_line: 2,
176/// });
177/// cpg.graph.add_edge(foo, call, Edge::AstChild);
178/// cpg.graph.add_edge(call, bar, Edge::Call);
179///
180/// let callees = callees_of(&cpg, "foo");
181/// assert_eq!(callees.len(), 1);
182/// assert_eq!(callees[0].0.name.as_deref(), Some("bar"));
183/// ```
184pub fn callees_of<'a>(cpg: &'a CodePropertyGraph, func_name: &str) -> Vec<(&'a Node, &'a Node)> {
185    let caller_idx = cpg.graph.node_indices().find(|&idx| {
186        let node = &cpg.graph[idx];
187        node.kind == NodeKind::Function && node.name.as_deref() == Some(func_name)
188    });
189
190    let caller_idx = match caller_idx {
191        Some(idx) => idx,
192        None => return vec![],
193    };
194
195    let mut results = Vec::new();
196    let call_sites = collect_call_sites_in_subtree(cpg, caller_idx);
197
198    for call_idx in call_sites {
199        for edge_ref in cpg.graph.edges(call_idx) {
200            if *edge_ref.weight() == Edge::Call {
201                results.push((&cpg.graph[edge_ref.target()], &cpg.graph[call_idx]));
202            }
203        }
204    }
205    results
206}
207
208/// Returns all functions that are never the target of a direct call.
209///
210/// A function is considered unused if no [`Edge::Call`] edge points to its
211/// node. This is a simple direct‑call analysis; indirect calls through
212/// function pointers or references are not detected.
213///
214/// # Examples
215///
216/// ```rust
217/// use icb_graph::graph::{CodePropertyGraph, Node, Edge};
218/// use icb_graph::query::unused_functions;
219/// use icb_common::NodeKind;
220///
221/// let mut cpg = CodePropertyGraph::new();
222/// let foo = cpg.graph.add_node(Node {
223///     kind: NodeKind::Function,
224///     name: Some("foo".into()),
225///     usr: Some("foo".into()),
226///     start_line: 1,
227///     end_line: 1,
228/// });
229/// let bar = cpg.graph.add_node(Node {
230///     kind: NodeKind::Function,
231///     name: Some("bar".into()),
232///     usr: Some("bar".into()),
233///     start_line: 2,
234///     end_line: 2,
235/// });
236/// // foo calls bar
237/// cpg.graph.add_edge(foo, bar, Edge::Call);
238///
239/// let unused = unused_functions(&cpg);
240/// // foo is never called directly => unused
241/// assert!(unused.iter().any(|n| n.name.as_deref() == Some("foo")));
242/// assert!(!unused.iter().any(|n| n.name.as_deref() == Some("bar")));
243/// ```
244pub fn unused_functions(cpg: &CodePropertyGraph) -> Vec<&Node> {
245    let called_defs: HashSet<_> = cpg
246        .graph
247        .edge_references()
248        .filter(|e| *e.weight() == Edge::Call)
249        .map(|e| e.target())
250        .collect();
251
252    cpg.graph
253        .node_indices()
254        .filter_map(|idx| {
255            let node = &cpg.graph[idx];
256            if node.kind == NodeKind::Function && !called_defs.contains(&idx) {
257                Some(node)
258            } else {
259                None
260            }
261        })
262        .collect()
263}
264
265/// Walks up [`Edge::AstChild`] edges from `start` until a
266/// [`NodeKind::Function`] node is found.
267fn get_enclosing_function(
268    cpg: &CodePropertyGraph,
269    start: petgraph::stable_graph::NodeIndex,
270) -> Option<&Node> {
271    let mut current = start;
272    loop {
273        let node = &cpg.graph[current];
274        if node.kind == NodeKind::Function {
275            return Some(node);
276        }
277        let parent = cpg
278            .graph
279            .edges_directed(current, petgraph::Direction::Incoming)
280            .find(|e| *e.weight() == Edge::AstChild)
281            .map(|e| e.source());
282        match parent {
283            Some(p) => current = p,
284            None => return None,
285        }
286    }
287}
288
289/// Collects all [`NodeKind::CallSite`] nodes in the AST subtree of `root`.
290///
291/// The traversal follows [`Edge::AstChild`] edges.
292fn collect_call_sites_in_subtree(
293    cpg: &CodePropertyGraph,
294    root: petgraph::stable_graph::NodeIndex,
295) -> Vec<petgraph::stable_graph::NodeIndex> {
296    let mut result = Vec::new();
297    let mut queue = VecDeque::new();
298    queue.push_back(root);
299    while let Some(current) = queue.pop_front() {
300        if cpg.graph[current].kind == NodeKind::CallSite {
301            result.push(current);
302        }
303        for edge_ref in cpg.graph.edges(current) {
304            if *edge_ref.weight() == Edge::AstChild {
305                queue.push_back(edge_ref.target());
306            }
307        }
308    }
309    result
310}
311
312#[cfg(test)]
313mod tests {
314    use super::*;
315    use crate::graph::Edge;
316    use icb_common::NodeKind;
317
318    fn make_test_graph() -> CodePropertyGraph {
319        let mut cpg = CodePropertyGraph::new();
320        let foo = cpg.graph.add_node(Node {
321            kind: NodeKind::Function,
322            name: Some("foo".into()),
323            usr: Some("foo".into()),
324            start_line: 1,
325            end_line: 3,
326        });
327        let bar = cpg.graph.add_node(Node {
328            kind: NodeKind::Function,
329            name: Some("bar".into()),
330            usr: Some("bar".into()),
331            start_line: 10,
332            end_line: 12,
333        });
334
335        let call_in_foo = cpg.graph.add_node(Node {
336            kind: NodeKind::CallSite,
337            name: Some("bar".into()),
338            usr: None,
339            start_line: 2,
340            end_line: 2,
341        });
342        cpg.graph.add_edge(foo, call_in_foo, Edge::AstChild);
343        cpg.graph.add_edge(call_in_foo, bar, Edge::Call);
344
345        let call_in_bar = cpg.graph.add_node(Node {
346            kind: NodeKind::CallSite,
347            name: Some("foo".into()),
348            usr: None,
349            start_line: 11,
350            end_line: 11,
351        });
352        cpg.graph.add_edge(bar, call_in_bar, Edge::AstChild);
353        cpg.graph.add_edge(call_in_bar, foo, Edge::Call);
354
355        cpg.graph.add_node(Node {
356            kind: NodeKind::Function,
357            name: Some("baz".into()),
358            usr: Some("baz".into()),
359            start_line: 20,
360            end_line: 22,
361        });
362
363        cpg
364    }
365
366    #[test]
367    fn test_callers_of() {
368        let cpg = make_test_graph();
369        let callers_bar = callers_of(&cpg, "bar");
370        assert_eq!(callers_bar.len(), 1);
371        assert_eq!(callers_bar[0].0.name.as_deref(), Some("foo"));
372
373        let callers_foo = callers_of(&cpg, "foo");
374        assert_eq!(callers_foo.len(), 1);
375        assert_eq!(callers_foo[0].0.name.as_deref(), Some("bar"));
376    }
377
378    #[test]
379    fn test_callees_of() {
380        let cpg = make_test_graph();
381        let callees_foo = callees_of(&cpg, "foo");
382        assert_eq!(callees_foo.len(), 1);
383        assert_eq!(callees_foo[0].0.name.as_deref(), Some("bar"));
384
385        let callees_bar = callees_of(&cpg, "bar");
386        assert_eq!(callees_bar.len(), 1);
387        assert_eq!(callees_bar[0].0.name.as_deref(), Some("foo"));
388    }
389
390    #[test]
391    fn test_unused_functions() {
392        let cpg = make_test_graph();
393        let unused = unused_functions(&cpg);
394        assert_eq!(unused.len(), 1);
395        assert_eq!(unused[0].name.as_deref(), Some("baz"));
396    }
397}