code_repo_wiki/search/
callgraph.rs

1use std::collections::HashMap;
2use petgraph::visit::{EdgeRef, IntoEdgeReferences};
3use crate::model::{EdgeKind, KnowledgeGraph, NodeId};
4
5/// 调用图查询 — 在 KnowledgeGraph 上提供调用者/被调用者等查询
6pub struct CallGraph<'a> {
7    graph: &'a KnowledgeGraph,
8}
9
10/// 符号名 → (调用者列表, 被调用者列表) 的预计算表。
11/// 独立类型别名：该表可被 serde_json 序列化落盘做磁盘缓存（见 lib.rs
12/// load_call_index_cache），是调用索引在进程外的唯一持久形态。
13pub type CallIndex = HashMap<String, (Vec<String>, Vec<String>)>;
14
15impl<'a> CallGraph<'a> {
16    pub fn new(graph: &'a KnowledgeGraph) -> Self {
17        Self { graph }
18    }
19
20    /// 返回指定符号调用的所有符号（被调用者）
21    pub fn callee_of(&self, name: &str) -> Vec<NodeId> {
22        let mut callees = Vec::new();
23        for n in self.graph.graph.node_indices() {
24            if let Some(w) = self.graph.graph.node_weight(n) && w.name == name {
25                for e in self.graph.graph.edges(n) {
26                    if e.weight().kind == EdgeKind::Calls {
27                        callees.push(e.target());
28                    }
29                }
30            }
31        }
32        callees
33    }
34
35    /// 返回调用指定符号的所有符号（调用者）
36    pub fn caller_of(&self, name: &str) -> Vec<NodeId> {
37        let mut callers = Vec::new();
38        for n in self.graph.graph.node_indices() {
39            if let Some(w) = self.graph.graph.node_weight(n) && w.name == name {
40                for e in self.graph.graph.edges_directed(n, petgraph::Direction::Incoming) {
41                    if e.weight().kind == EdgeKind::Calls {
42                        callers.push(e.source());
43                    }
44                }
45            }
46        }
47        callers
48    }
49
50    /// 返回所有 Calls 边的列表
51    pub fn all_call_edges(&self) -> Vec<(NodeId, NodeId)> {
52        let mut edges = Vec::new();
53        for e in self.graph.graph.edge_references() {
54            if e.weight().kind == EdgeKind::Calls {
55                edges.push((e.source(), e.target()));
56            }
57        }
58        edges
59    }
60
61    /// 构建符号名 → (调用者列表, 被调用者列表) 预计算表。
62    /// 一次性遍历所有 Calls 边，避免查询时重复扫描全图。
63    pub fn build_call_index(&self) -> CallIndex {
64        let mut index: HashMap<String, (Vec<String>, Vec<String>)> = HashMap::new();
65        for (src, dst) in self.all_call_edges() {
66            if let (Some(s), Some(d)) = (self.graph.graph.node_weight(src), self.graph.graph.node_weight(dst)) {
67                index.entry(d.name.clone()).or_default().0.push(s.name.clone());
68                index.entry(s.name.clone()).or_default().1.push(d.name.clone());
69            }
70        }
71        index
72    }
73}
74
75#[cfg(test)]
76mod tests {
77    use super::*;
78    use crate::model::{CodeNode, CodeEdge, NodeKind};
79    use petgraph::stable_graph::StableDiGraph;
80
81    fn make_test_graph() -> KnowledgeGraph {
82        let mut g = StableDiGraph::<CodeNode, CodeEdge>::new();
83        let caller = g.add_node(CodeNode {
84            id: NodeId::new(0), kind: NodeKind::Function,
85            name: "caller".into(), file_path: None,
86            line_range: None, doc_comment: None, signature: None, visibility: None,
87            module_path: vec!["test".into()],
88        });
89        let callee = g.add_node(CodeNode {
90            id: NodeId::new(1), kind: NodeKind::Function,
91            name: "callee".into(), file_path: None,
92            line_range: None, doc_comment: None, signature: None, visibility: None,
93            module_path: vec!["test".into()],
94        });
95        g.add_edge(caller, callee, CodeEdge {
96            id: petgraph::stable_graph::EdgeIndex::new(0),
97            kind: EdgeKind::Calls, source: caller, target: callee,
98            weight: 1.0, location: None,
99        });
100        KnowledgeGraph { graph: g, modules: vec![], features: Vec::new() }
101    }
102
103    #[test]
104    fn test_callee_of() {
105        let kg = make_test_graph();
106        let cg = CallGraph::new(&kg);
107        let callees = cg.callee_of("caller");
108        assert_eq!(callees.len(), 1);
109    }
110
111    #[test]
112    fn test_caller_of() {
113        let kg = make_test_graph();
114        let cg = CallGraph::new(&kg);
115        let callers = cg.caller_of("callee");
116        assert_eq!(callers.len(), 1);
117    }
118
119    #[test]
120    fn test_call_index_serde_round_trip() {
121        // 磁盘缓存契约：CallIndex 必须可 JSON 序列化往返且内容不变
122        //（lib.rs load_call_index_cache 依赖此格式）
123        let kg = make_test_graph();
124        let cg = CallGraph::new(&kg);
125        let index = cg.build_call_index();
126        let json = serde_json::to_string(&index).unwrap();
127        let back: CallIndex = serde_json::from_str(&json).unwrap();
128        assert_eq!(back, index);
129        // caller 侧：callee 的调用者列表含 caller
130        assert_eq!(back.get("callee").unwrap().0, vec!["caller".to_string()]);
131        // callee 侧：caller 的被调用者列表含 callee
132        assert_eq!(back.get("caller").unwrap().1, vec!["callee".to_string()]);
133    }
134}
code_repo_wiki/search/callgraph.rs

code_repo_wiki/search/
callgraph.rs
