1use std::collections::{HashMap, VecDeque};
2
3#[derive(Clone, Debug, PartialEq)]
4pub struct GraphNode {
5 pub id: String,
6 pub label: String,
7 pub source_file: String,
8 pub source_location: String,
9 pub community: i64,
10}
11
12#[derive(Clone, Debug, PartialEq)]
13pub struct GraphEdge {
14 pub source: String,
15 pub target: String,
16 pub relation: String,
17 pub confidence: String,
18}
19
20#[derive(Default)]
21pub struct GraphState {
22 nodes: HashMap<String, GraphNode>,
23 edges: Vec<GraphEdge>,
24 adj: HashMap<String, Vec<String>>,
25}
26
27impl GraphState {
28 pub fn new() -> Self {
29 Self::default()
30 }
31
32 pub fn add_node(&mut self, node: GraphNode) {
33 self.adj.entry(node.id.clone()).or_default();
34 self.nodes.insert(node.id.clone(), node);
35 }
36
37 pub fn add_edge(&mut self, edge: GraphEdge) {
38 self.adj
39 .entry(edge.source.clone())
40 .or_default()
41 .push(edge.target.clone());
42 self.adj.entry(edge.target.clone()).or_default();
43 self.edges.push(edge);
44 }
45
46 pub fn get_nodes(&self) -> Vec<GraphNode> {
47 self.nodes.values().cloned().collect()
48 }
49
50 pub fn get_edges(&self) -> Vec<GraphEdge> {
51 self.edges.clone()
52 }
53
54 pub fn get_node(&self, id: &str) -> Option<&GraphNode> {
55 self.nodes.get(id)
56 }
57
58 pub fn search_nodes(&self, query: &str, limit: usize) -> Vec<GraphNode> {
59 let q = query.to_lowercase();
60 self.nodes
61 .values()
62 .filter(|n| n.label.to_lowercase().contains(&q))
63 .take(limit)
64 .cloned()
65 .collect()
66 }
67
68 pub fn shortest_path(&self, source: &str, target: &str) -> Option<Vec<String>> {
69 if source == target {
70 return Some(vec![source.to_string()]);
71 }
72 let mut visited: HashMap<String, Option<String>> = HashMap::new();
73 let mut queue = VecDeque::new();
74 queue.push_back(source.to_string());
75 visited.insert(source.to_string(), None);
76
77 while let Some(current) = queue.pop_front() {
78 if current == target {
79 let mut path = vec![target.to_string()];
80 let mut node = target.to_string();
81 while let Some(Some(prev)) = visited.get(&node) {
82 path.push(prev.clone());
83 node = prev.clone();
84 }
85 path.reverse();
86 return Some(path);
87 }
88 if let Some(neighbors) = self.adj.get(¤t) {
89 for next in neighbors {
90 if !visited.contains_key(next) {
91 visited.insert(next.clone(), Some(current.clone()));
92 queue.push_back(next.clone());
93 }
94 }
95 }
96 }
97 None
98 }
99
100 pub fn node_count(&self) -> usize {
101 self.nodes.len()
102 }
103
104 pub fn edge_count(&self) -> usize {
105 self.edges.len()
106 }
107}
108
109#[cfg(test)]
110mod tests {
111 use super::*;
112
113 fn node(id: &str, label: &str) -> GraphNode {
114 GraphNode {
115 id: id.to_string(),
116 label: label.to_string(),
117 source_file: "test.rs".to_string(),
118 source_location: "1:1".to_string(),
119 community: 0,
120 }
121 }
122
123 fn edge(src: &str, tgt: &str) -> GraphEdge {
124 GraphEdge {
125 source: src.to_string(),
126 target: tgt.to_string(),
127 relation: "calls".to_string(),
128 confidence: "1.0".to_string(),
129 }
130 }
131
132 #[test]
133 fn test_empty_graph() {
134 let s = GraphState::new();
135 assert_eq!(s.node_count(), 0);
136 assert_eq!(s.edge_count(), 0);
137 }
138
139 #[test]
140 fn test_add_node() {
141 let mut s = GraphState::new();
142 s.add_node(node("n1", "Foo"));
143 assert_eq!(s.node_count(), 1);
144 assert!(s.get_node("n1").is_some());
145 }
146
147 #[test]
148 fn test_add_duplicate_node_overwrites() {
149 let mut s = GraphState::new();
150 s.add_node(node("n1", "Foo"));
151 s.add_node(node("n1", "Bar"));
152 assert_eq!(s.node_count(), 1);
153 assert_eq!(s.get_node("n1").unwrap().label, "Bar");
154 }
155
156 #[test]
157 fn test_add_edge() {
158 let mut s = GraphState::new();
159 s.add_node(node("a", "A"));
160 s.add_node(node("b", "B"));
161 s.add_edge(edge("a", "b"));
162 assert_eq!(s.edge_count(), 1);
163 }
164
165 #[test]
166 fn test_get_nodes_returns_all() {
167 let mut s = GraphState::new();
168 s.add_node(node("a", "A"));
169 s.add_node(node("b", "B"));
170 assert_eq!(s.get_nodes().len(), 2);
171 }
172
173 #[test]
174 fn test_get_edges_returns_all() {
175 let mut s = GraphState::new();
176 s.add_node(node("a", "A"));
177 s.add_node(node("b", "B"));
178 s.add_node(node("c", "C"));
179 s.add_edge(edge("a", "b"));
180 s.add_edge(edge("b", "c"));
181 assert_eq!(s.get_edges().len(), 2);
182 }
183
184 #[test]
185 fn test_get_node_found() {
186 let mut s = GraphState::new();
187 s.add_node(node("x", "X"));
188 assert_eq!(s.get_node("x").unwrap().label, "X");
189 }
190
191 #[test]
192 fn test_get_node_not_found() {
193 let s = GraphState::new();
194 assert!(s.get_node("missing").is_none());
195 }
196
197 #[test]
198 fn test_search_nodes_exact() {
199 let mut s = GraphState::new();
200 s.add_node(node("1", "AuthService"));
201 s.add_node(node("2", "UserService"));
202 let results = s.search_nodes("AuthService", 10);
203 assert_eq!(results.len(), 1);
204 assert_eq!(results[0].id, "1");
205 }
206
207 #[test]
208 fn test_search_nodes_case_insensitive() {
209 let mut s = GraphState::new();
210 s.add_node(node("1", "AuthService"));
211 s.add_node(node("2", "AuthController"));
212 s.add_node(node("3", "UserService"));
213 let results = s.search_nodes("auth", 10);
214 assert_eq!(results.len(), 2);
215 }
216
217 #[test]
218 fn test_search_nodes_empty_query_matches_all() {
219 let mut s = GraphState::new();
220 s.add_node(node("1", "A"));
221 s.add_node(node("2", "B"));
222 let results = s.search_nodes("", 10);
223 assert_eq!(results.len(), 2);
224 }
225
226 #[test]
227 fn test_search_nodes_no_match() {
228 let mut s = GraphState::new();
229 s.add_node(node("1", "Foo"));
230 let results = s.search_nodes("zzz", 10);
231 assert!(results.is_empty());
232 }
233
234 #[test]
235 fn test_search_nodes_limit_applied() {
236 let mut s = GraphState::new();
237 for i in 0..5 {
238 s.add_node(node(&i.to_string(), &format!("Service{}", i)));
239 }
240 let results = s.search_nodes("service", 3);
241 assert_eq!(results.len(), 3);
242 }
243
244 #[test]
245 fn test_shortest_path_direct() {
246 let mut s = GraphState::new();
247 s.add_node(node("a", "A"));
248 s.add_node(node("b", "B"));
249 s.add_edge(edge("a", "b"));
250 let path = s.shortest_path("a", "b").unwrap();
251 assert_eq!(path, vec!["a", "b"]);
252 }
253
254 #[test]
255 fn test_shortest_path_multi_hop() {
256 let mut s = GraphState::new();
257 for id in ["a", "b", "c"] {
258 s.add_node(node(id, id));
259 }
260 s.add_edge(edge("a", "b"));
261 s.add_edge(edge("b", "c"));
262 let path = s.shortest_path("a", "c").unwrap();
263 assert_eq!(path, vec!["a", "b", "c"]);
264 }
265
266 #[test]
267 fn test_shortest_path_self() {
268 let mut s = GraphState::new();
269 s.add_node(node("a", "A"));
270 let path = s.shortest_path("a", "a").unwrap();
271 assert_eq!(path, vec!["a"]);
272 }
273
274 #[test]
275 fn test_shortest_path_not_found() {
276 let mut s = GraphState::new();
277 s.add_node(node("a", "A"));
278 s.add_node(node("b", "B"));
279 assert!(s.shortest_path("a", "b").is_none());
280 }
281}