1use std::collections::{HashMap, HashSet, VecDeque};
2
3use crate::ir::{
4 api_endpoint_key, external_api_key, module_key, ApiEndpointIr, BehaviourIr, CallbackIr,
5 ClassIr, EdgeIr, EdgeKind, ExternalApiIr, FileIr, FunctionIr, ModuleIr, PropertyIr, ProjectIr,
6};
7use crate::schema::NodeLabel;
8
9#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
11pub struct SymbolRef {
12 pub label: String,
13 pub key: String,
14 pub name: Option<String>,
15 pub path: Option<String>,
16}
17
18#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
20pub struct FileSymbols {
21 pub path: String,
22 pub classes: Vec<SymbolRef>,
23 pub functions: Vec<SymbolRef>,
24 pub modules: Vec<SymbolRef>,
25 pub api_endpoints: Vec<SymbolRef>,
26}
27
28#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
30pub struct ImpactReport {
31 pub symbol: String,
32 pub depth: u32,
33 pub callers: Vec<String>,
34 pub affected_files: Vec<String>,
35 pub truncated: bool,
36}
37
38#[derive(Debug, Clone, Copy)]
39pub struct QueryLimits {
40 pub max_depth: u32,
41 pub max_results: usize,
42}
43
44impl Default for QueryLimits {
45 fn default() -> Self {
46 Self {
47 max_depth: 2,
48 max_results: 50,
49 }
50 }
51}
52
53#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
54pub struct RefreshReport {
55 pub cleanup_targets: usize,
56 pub parse_targets: usize,
57 pub nodes_merged: usize,
58 pub edges_merged: usize,
59}
60
61#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
62enum NodeKind {
63 File,
64 Module,
65 Class,
66 Property,
67 Function,
68 Behaviour,
69 Callback,
70 ApiEndpoint,
71 ExternalApi,
72}
73
74#[derive(Debug, Clone)]
75struct NodeRecord {
76 kind: NodeKind,
77 key: String,
78 label: String,
79 name: Option<String>,
80 path: Option<String>,
81}
82
83type AdjKey = (usize, EdgeKind);
84
85#[derive(Debug, Default)]
87pub struct InMemoryGraph {
88 nodes: Vec<NodeRecord>,
89 nodes_by_key: HashMap<(NodeKind, String), usize>,
90 symbols_by_path: HashMap<String, HashSet<usize>>,
91 by_simple_name: HashMap<String, Vec<usize>>,
92 outgoing: HashMap<AdjKey, Vec<usize>>,
93 incoming: HashMap<AdjKey, Vec<usize>>,
94}
95
96pub trait GraphStore {
97 fn callers(&self, fqn: &str) -> Vec<SymbolRef>;
98 fn callees(&self, fqn: &str) -> Vec<SymbolRef>;
99 fn file_dependencies(&self, path: &str) -> Vec<String>;
100 fn impact(&self, fqn: &str, limits: QueryLimits) -> ImpactReport;
101 fn symbols_in_file(&self, path: &str) -> FileSymbols;
102 fn find_symbol(&self, query: &str) -> Vec<SymbolRef>;
103 fn node_count(&self) -> usize;
104 fn edge_count(&self) -> usize;
105}
106
107impl InMemoryGraph {
108 pub fn from_ir(ir: ProjectIr) -> Self {
109 let mut g = Self::default();
110 g.merge_ir(ir);
111 g
112 }
113
114 pub fn merge_ir(&mut self, delta: ProjectIr) {
115 for f in delta.files {
116 self.upsert_file(f);
117 }
118 for m in delta.modules {
119 self.upsert_module(m);
120 }
121 for c in delta.classes {
122 self.upsert_class(c);
123 }
124 for p in delta.properties {
125 self.upsert_property(p);
126 }
127 for f in delta.functions {
128 self.upsert_function(f);
129 }
130 for b in delta.behaviours {
131 self.upsert_behaviour(b);
132 }
133 for c in delta.callbacks {
134 self.upsert_callback(c);
135 }
136 for a in delta.api_endpoints {
137 self.upsert_api_endpoint(a);
138 }
139 for e in delta.external_apis {
140 self.upsert_external_api(e);
141 }
142 for edge in delta.edges {
143 self.add_edge(edge);
144 }
145 }
146
147 pub fn remove_file(&mut self, path: &str) {
148 let normalized = path.replace('\\', "/");
149 let to_remove: Vec<usize> = self
150 .nodes
151 .iter()
152 .enumerate()
153 .filter(|(_, n)| n.path.as_deref() == Some(normalized.as_str()))
154 .map(|(i, _)| i)
155 .collect();
156
157 if to_remove.is_empty() {
158 return;
159 }
160
161 let remove_set: HashSet<usize> = to_remove.into_iter().collect();
162 self.remove_nodes(&remove_set);
163 }
164
165 fn remove_nodes(&mut self, remove_set: &HashSet<usize>) {
166 let mut remap: HashMap<usize, Option<usize>> = HashMap::new();
167 let mut new_nodes: Vec<NodeRecord> = Vec::new();
168
169 for (old_id, node) in self.nodes.iter().enumerate() {
170 if remove_set.contains(&old_id) {
171 remap.insert(old_id, None);
172 self.nodes_by_key.remove(&(node.kind, node.key.clone()));
173 if let Some(p) = &node.path {
174 if let Some(set) = self.symbols_by_path.get_mut(p) {
175 set.remove(&old_id);
176 }
177 }
178 if let Some(name) = &node.name {
179 if let Some(ids) = self.by_simple_name.get_mut(name) {
180 ids.retain(|id| *id != old_id);
181 }
182 }
183 } else {
184 let new_id = new_nodes.len();
185 remap.insert(old_id, Some(new_id));
186 new_nodes.push(node.clone());
187 }
188 }
189
190 self.nodes = new_nodes;
191 self.rebuild_adjacency(&remap);
192 self.rebuild_indexes();
193 }
194
195 fn rebuild_adjacency(&mut self, remap: &HashMap<usize, Option<usize>>) {
196 let mut new_out: HashMap<AdjKey, Vec<usize>> = HashMap::new();
197 let mut new_in: HashMap<AdjKey, Vec<usize>> = HashMap::new();
198
199 for ((from, kind), targets) in &self.outgoing {
200 let Some(&Some(new_from)) = remap.get(from) else {
201 continue;
202 };
203 for to in targets {
204 if let Some(Some(new_to)) = remap.get(to) {
205 new_out.entry((new_from, *kind)).or_default().push(*new_to);
206 new_in.entry((*new_to, *kind)).or_default().push(new_from);
207 }
208 }
209 }
210 self.outgoing = new_out;
211 self.incoming = new_in;
212 }
213
214 fn rebuild_indexes(&mut self) {
215 self.nodes_by_key.clear();
216 self.symbols_by_path.clear();
217 self.by_simple_name.clear();
218 for (id, node) in self.nodes.iter().enumerate() {
219 self.nodes_by_key.insert((node.kind, node.key.clone()), id);
220 if let Some(p) = &node.path {
221 self.symbols_by_path.entry(p.clone()).or_default().insert(id);
222 }
223 if let Some(name) = &node.name {
224 self.by_simple_name.entry(name.clone()).or_default().push(id);
225 }
226 }
227 }
228
229 fn upsert_file(&mut self, f: FileIr) {
230 let _id = self.ensure_node(
231 NodeKind::File,
232 f.path.clone(),
233 NodeLabel::File.to_string(),
234 None,
235 Some(f.path),
236 );
237 }
238
239 fn upsert_module(&mut self, m: ModuleIr) {
240 let key = module_key(&m.name, &m.path);
241 let _id = self.ensure_node(
242 NodeKind::Module,
243 key,
244 NodeLabel::Module.to_string(),
245 Some(m.name),
246 Some(m.path),
247 );
248 }
249
250 fn upsert_class(&mut self, c: ClassIr) {
251 let _id = self.ensure_node(
252 NodeKind::Class,
253 c.fqn.clone(),
254 "Class".to_string(),
255 Some(c.name),
256 Some(c.path),
257 );
258 }
259
260 fn upsert_property(&mut self, p: PropertyIr) {
261 let _id = self.ensure_node(
262 NodeKind::Property,
263 p.fqn.clone(),
264 "Property".to_string(),
265 Some(p.name),
266 Some(p.path),
267 );
268 }
269
270 fn upsert_function(&mut self, f: FunctionIr) {
271 let _id = self.ensure_node(
272 NodeKind::Function,
273 f.fqn.clone(),
274 NodeLabel::Function.to_string(),
275 Some(f.name),
276 Some(f.path),
277 );
278 }
279
280 fn upsert_behaviour(&mut self, b: BehaviourIr) {
281 let _id = self.ensure_node(
282 NodeKind::Behaviour,
283 b.name.clone(),
284 NodeLabel::Behaviour.to_string(),
285 Some(b.name.clone()),
286 b.path,
287 );
288 }
289
290 fn upsert_callback(&mut self, c: CallbackIr) {
291 let _id = self.ensure_node(
292 NodeKind::Callback,
293 c.fqn.clone(),
294 NodeLabel::Callback.to_string(),
295 Some(c.name),
296 None,
297 );
298 }
299
300 fn upsert_api_endpoint(&mut self, a: ApiEndpointIr) {
301 let key = api_endpoint_key(&a.methods, &a.path);
302 let _id = self.ensure_node(
303 NodeKind::ApiEndpoint,
304 key,
305 NodeLabel::ApiEndpoint.to_string(),
306 Some(a.path.clone()),
307 None,
308 );
309 }
310
311 fn upsert_external_api(&mut self, e: ExternalApiIr) {
312 let key = if let (Some(base), Some(norm)) = (&e.base_url, &e.norm_path) {
313 external_api_key(base, norm)
314 } else {
315 e.name.clone()
316 };
317 let _id = self.ensure_node(
318 NodeKind::ExternalApi,
319 key,
320 NodeLabel::ExternalApi.to_string(),
321 Some(e.name),
322 None,
323 );
324 }
325
326 fn ensure_node(
327 &mut self,
328 kind: NodeKind,
329 key: String,
330 label: String,
331 name: Option<String>,
332 path: Option<String>,
333 ) -> usize {
334 if let Some(&id) = self.nodes_by_key.get(&(kind, key.clone())) {
335 if let Some(n) = self.nodes.get_mut(id) {
336 if name.is_some() {
337 n.name = name.clone();
338 }
339 if path.is_some() {
340 n.path = path.clone();
341 }
342 }
343 return id;
344 }
345 let id = self.nodes.len();
346 let record = NodeRecord {
347 kind,
348 key: key.clone(),
349 label,
350 name: name.clone(),
351 path: path.clone(),
352 };
353 self.nodes.push(record);
354 self.nodes_by_key.insert((kind, key), id);
355 if let Some(p) = path {
356 self.symbols_by_path.entry(p).or_default().insert(id);
357 }
358 if let Some(n) = name {
359 self.by_simple_name.entry(n).or_default().push(id);
360 }
361 id
362 }
363
364 fn add_edge(&mut self, edge: EdgeIr) {
365 let Some(from_id) = self.lookup_id(&edge.from_label, &edge.from_key) else {
366 return;
367 };
368 let Some(to_id) = self.lookup_id(&edge.to_label, &edge.to_key) else {
369 return;
370 };
371 let out_list = self.outgoing.entry((from_id, edge.kind)).or_default();
372 if !out_list.contains(&to_id) {
373 out_list.push(to_id);
374 }
375 let in_list = self.incoming.entry((to_id, edge.kind)).or_default();
376 if !in_list.contains(&from_id) {
377 in_list.push(from_id);
378 }
379 }
380
381 fn lookup_id(&self, label: &str, key: &str) -> Option<usize> {
382 let kind = node_kind_from_label(label)?;
383 self.nodes_by_key.get(&(kind, key.to_string())).copied()
384 }
385
386 fn function_id(&self, fqn: &str) -> Option<usize> {
387 self.nodes_by_key
388 .get(&(NodeKind::Function, fqn.to_string()))
389 .copied()
390 }
391
392 fn node_to_symbol(&self, id: usize) -> Option<SymbolRef> {
393 self.nodes.get(id).map(|n| SymbolRef {
394 label: n.label.clone(),
395 key: n.key.clone(),
396 name: n.name.clone(),
397 path: n.path.clone(),
398 })
399 }
400}
401
402impl GraphStore for InMemoryGraph {
403 fn callers(&self, fqn: &str) -> Vec<SymbolRef> {
404 let Some(fn_id) = self.function_id(fqn) else {
405 return Vec::new();
406 };
407 self.incoming
408 .get(&(fn_id, EdgeKind::CallsFunction))
409 .map(|ids| {
410 ids.iter()
411 .filter_map(|id| self.node_to_symbol(*id))
412 .collect()
413 })
414 .unwrap_or_default()
415 }
416
417 fn callees(&self, fqn: &str) -> Vec<SymbolRef> {
418 let Some(fn_id) = self.function_id(fqn) else {
419 return Vec::new();
420 };
421 self.outgoing
422 .get(&(fn_id, EdgeKind::CallsFunction))
423 .map(|ids| {
424 ids.iter()
425 .filter_map(|id| self.node_to_symbol(*id))
426 .collect()
427 })
428 .unwrap_or_default()
429 }
430
431 fn file_dependencies(&self, path: &str) -> Vec<String> {
432 let normalized = path.replace('\\', "/");
433 let Some(&file_id) = self.nodes_by_key.get(&(NodeKind::File, normalized.clone())) else {
434 return Vec::new();
435 };
436 self.outgoing
437 .get(&(file_id, EdgeKind::DependsOnFile))
438 .map(|ids| {
439 ids.iter()
440 .filter_map(|id| self.nodes.get(*id).and_then(|n| n.path.clone()))
441 .collect()
442 })
443 .unwrap_or_default()
444 }
445
446 fn impact(&self, fqn: &str, limits: QueryLimits) -> ImpactReport {
447 let mut report = ImpactReport {
448 symbol: fqn.to_string(),
449 depth: limits.max_depth,
450 ..Default::default()
451 };
452 let Some(start) = self.function_id(fqn) else {
453 return report;
454 };
455
456 let mut visited: HashSet<usize> = HashSet::new();
457 let mut queue: VecDeque<(usize, u32)> = VecDeque::new();
458 queue.push_back((start, 0));
459 visited.insert(start);
460
461 while let Some((node_id, depth)) = queue.pop_front() {
462 if depth >= limits.max_depth {
463 continue;
464 }
465 if let Some(callers) = self.incoming.get(&(node_id, EdgeKind::CallsFunction)) {
466 for caller_id in callers {
467 if visited.contains(caller_id) {
468 continue;
469 }
470 if report.callers.len() >= limits.max_results {
471 report.truncated = true;
472 return report;
473 }
474 visited.insert(*caller_id);
475 if let Some(sym) = self.node_to_symbol(*caller_id) {
476 if sym.label == NodeLabel::Function.to_string() {
477 report.callers.push(sym.key.clone());
478 }
479 if let Some(p) = sym.path {
480 if !report.affected_files.contains(&p) {
481 report.affected_files.push(p);
482 }
483 }
484 }
485 queue.push_back((*caller_id, depth + 1));
486 }
487 }
488 }
489 report
490 }
491
492 fn symbols_in_file(&self, path: &str) -> FileSymbols {
493 let normalized = path.replace('\\', "/");
494 let mut out = FileSymbols {
495 path: normalized.clone(),
496 ..Default::default()
497 };
498 let Some(ids) = self.symbols_by_path.get(&normalized) else {
499 return out;
500 };
501 for id in ids {
502 let Some(n) = self.nodes.get(*id) else {
503 continue;
504 };
505 let sym = SymbolRef {
506 label: n.label.clone(),
507 key: n.key.clone(),
508 name: n.name.clone(),
509 path: n.path.clone(),
510 };
511 match n.kind {
512 NodeKind::Class => out.classes.push(sym),
513 NodeKind::Function => out.functions.push(sym),
514 NodeKind::Module => out.modules.push(sym),
515 NodeKind::ApiEndpoint => out.api_endpoints.push(sym),
516 _ => {}
517 }
518 }
519 out
520 }
521
522 fn find_symbol(&self, query: &str) -> Vec<SymbolRef> {
523 let q = query.trim();
524 if q.is_empty() {
525 return Vec::new();
526 }
527
528 if let Some(&id) = self.nodes_by_key.get(&(NodeKind::Function, q.to_string())) {
529 return vec![self.node_to_symbol(id).unwrap()];
530 }
531 if let Some(&id) = self.nodes_by_key.get(&(NodeKind::Class, q.to_string())) {
532 return vec![self.node_to_symbol(id).unwrap()];
533 }
534
535 let q_lower = q.to_lowercase();
536 let mut results = Vec::new();
537 for (id, node) in self.nodes.iter().enumerate() {
538 if node.key.to_lowercase().contains(&q_lower)
539 || node
540 .name
541 .as_ref()
542 .is_some_and(|n| n.to_lowercase().contains(&q_lower))
543 {
544 if let Some(sym) = self.node_to_symbol(id) {
545 results.push(sym);
546 }
547 }
548 }
549 results.sort_by(|a, b| a.key.cmp(&b.key));
550 results.truncate(50);
551 results
552 }
553
554 fn node_count(&self) -> usize {
555 self.nodes.len()
556 }
557
558 fn edge_count(&self) -> usize {
559 self.outgoing.values().map(|v| v.len()).sum()
560 }
561}
562
563fn node_kind_from_label(label: &str) -> Option<NodeKind> {
564 match label {
565 "File" => Some(NodeKind::File),
566 "Module" => Some(NodeKind::Module),
567 "Class" => Some(NodeKind::Class),
568 "Property" => Some(NodeKind::Property),
569 "Function" => Some(NodeKind::Function),
570 "Behaviour" => Some(NodeKind::Behaviour),
571 "Callback" => Some(NodeKind::Callback),
572 "ApiEndpoint" => Some(NodeKind::ApiEndpoint),
573 "ExternalApi" => Some(NodeKind::ExternalApi),
574 _ => None,
575 }
576}
577
578#[cfg(test)]
579mod tests {
580 use super::*;
581 use crate::ir::{EdgeKind, ProjectIr};
582
583 fn sample_ir() -> ProjectIr {
584 let mut ir = ProjectIr::empty();
585 ir.files.push(FileIr {
586 path: "src/a.rs".into(),
587 language: "rust".into(),
588 framework: None,
589 project_name: None,
590 });
591 ir.functions.push(FunctionIr {
592 name: "main".into(),
593 fqn: "src/a.rs::main".into(),
594 path: "src/a.rs".into(),
595 language: "rust".into(),
596 framework: None,
597 project_name: None,
598 arity: None,
599 return_type: None,
600 param_count: None,
601 param_types: vec![],
602 });
603 ir.functions.push(FunctionIr {
604 name: "helper".into(),
605 fqn: "src/a.rs::helper".into(),
606 path: "src/a.rs".into(),
607 language: "rust".into(),
608 framework: None,
609 project_name: None,
610 arity: None,
611 return_type: None,
612 param_count: None,
613 param_types: vec![],
614 });
615 ir.edges.push(EdgeIr {
616 kind: EdgeKind::DeclaresFunction,
617 from_label: "File".into(),
618 from_key: "src/a.rs".into(),
619 to_label: "Function".into(),
620 to_key: "src/a.rs::main".into(),
621 });
622 ir.edges.push(EdgeIr {
623 kind: EdgeKind::CallsFunction,
624 from_label: "Function".into(),
625 from_key: "src/a.rs::main".into(),
626 to_label: "Function".into(),
627 to_key: "src/a.rs::helper".into(),
628 });
629 ir
630 }
631
632 #[test]
633 fn in_memory_graph_queries_callers_and_callees() {
634 let graph = InMemoryGraph::from_ir(sample_ir());
635 let callees = graph.callees("src/a.rs::main");
636 assert_eq!(callees.len(), 1);
637 assert_eq!(callees[0].key, "src/a.rs::helper");
638 let callers = graph.callers("src/a.rs::helper");
639 assert_eq!(callers.len(), 1);
640 assert_eq!(callers[0].key, "src/a.rs::main");
641 }
642
643 #[test]
644 fn remove_file_strips_symbols() {
645 let mut graph = InMemoryGraph::from_ir(sample_ir());
646 graph.remove_file("src/a.rs");
647 assert!(graph.find_symbol("main").is_empty());
648 }
649}