Skip to main content

impactsense_parser/store/
in_memory.rs

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/// Reference to a symbol in the graph.
10#[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/// Symbols declared in a single file.
19#[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/// Bounded impact analysis result.
29#[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/// In-memory dependency graph with indexed lookups for IDE/MCP queries.
86#[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}