codesynapse_core/
graph.rs

1use crate::error::{CodeSynapseError, Result};
2use crate::types::{Edge, Node, NodeId};
3use std::collections::HashMap;
4use std::path::Path;
5use std::sync::RwLock;
6
7pub trait GraphStore: Send + Sync {
8    fn add_node(&self, node: Node) -> Result<()>;
9    fn add_edge(&self, edge: Edge) -> Result<()>;
10    fn get_node(&self, id: &str) -> Result<Option<Node>>;
11    fn get_all_nodes(&self) -> Result<Vec<Node>>;
12    fn get_all_edges(&self) -> Result<Vec<Edge>>;
13    fn neighbors(&self, id: &str, relation_filter: Option<&str>) -> Result<Vec<(Node, Edge)>>;
14    fn search(&self, query: &str, top_k: usize) -> Result<Vec<(f64, Node)>>;
15    fn shortest_path(&self, src: &str, tgt: &str) -> Result<Option<Vec<Node>>>;
16    fn dijkstra_shortest_path(&self, src: &str, tgt: &str) -> Result<Option<Vec<Node>>>;
17    fn subgraph(&self, node_ids: &[&str]) -> Result<(Vec<Node>, Vec<Edge>)>;
18    fn node_count(&self) -> Result<usize>;
19    fn edge_count(&self) -> Result<usize>;
20    fn remove_node(&self, id: &str) -> Result<()>;
21    fn remove_edge(&self, source: &str, target: &str, relation: &str) -> Result<()>;
22    fn clear(&self) -> Result<()>;
23}
24
25pub enum StoreBackend {
26    Sled(SledGraphStore),
27    Memory(MemoryGraphStore),
28}
29
30impl GraphStore for StoreBackend {
31    fn add_node(&self, node: Node) -> Result<()> {
32        match self {
33            StoreBackend::Sled(s) => s.add_node(node),
34            StoreBackend::Memory(m) => m.add_node(node),
35        }
36    }
37
38    fn add_edge(&self, edge: Edge) -> Result<()> {
39        match self {
40            StoreBackend::Sled(s) => s.add_edge(edge),
41            StoreBackend::Memory(m) => m.add_edge(edge),
42        }
43    }
44
45    fn get_node(&self, id: &str) -> Result<Option<Node>> {
46        match self {
47            StoreBackend::Sled(s) => s.get_node(id),
48            StoreBackend::Memory(m) => m.get_node(id),
49        }
50    }
51
52    fn get_all_nodes(&self) -> Result<Vec<Node>> {
53        match self {
54            StoreBackend::Sled(s) => s.get_all_nodes(),
55            StoreBackend::Memory(m) => m.get_all_nodes(),
56        }
57    }
58
59    fn get_all_edges(&self) -> Result<Vec<Edge>> {
60        match self {
61            StoreBackend::Sled(s) => s.get_all_edges(),
62            StoreBackend::Memory(m) => m.get_all_edges(),
63        }
64    }
65
66    fn neighbors(&self, id: &str, relation_filter: Option<&str>) -> Result<Vec<(Node, Edge)>> {
67        match self {
68            StoreBackend::Sled(s) => s.neighbors(id, relation_filter),
69            StoreBackend::Memory(m) => m.neighbors(id, relation_filter),
70        }
71    }
72
73    fn search(&self, query: &str, top_k: usize) -> Result<Vec<(f64, Node)>> {
74        match self {
75            StoreBackend::Sled(s) => s.search(query, top_k),
76            StoreBackend::Memory(m) => m.search(query, top_k),
77        }
78    }
79
80    fn shortest_path(&self, src: &str, tgt: &str) -> Result<Option<Vec<Node>>> {
81        match self {
82            StoreBackend::Sled(s) => s.shortest_path(src, tgt),
83            StoreBackend::Memory(m) => m.shortest_path(src, tgt),
84        }
85    }
86
87    fn dijkstra_shortest_path(&self, src: &str, tgt: &str) -> Result<Option<Vec<Node>>> {
88        match self {
89            StoreBackend::Sled(s) => s.dijkstra_shortest_path(src, tgt),
90            StoreBackend::Memory(m) => m.dijkstra_shortest_path(src, tgt),
91        }
92    }
93
94    fn subgraph(&self, node_ids: &[&str]) -> Result<(Vec<Node>, Vec<Edge>)> {
95        match self {
96            StoreBackend::Sled(s) => s.subgraph(node_ids),
97            StoreBackend::Memory(m) => m.subgraph(node_ids),
98        }
99    }
100
101    fn node_count(&self) -> Result<usize> {
102        match self {
103            StoreBackend::Sled(s) => s.node_count(),
104            StoreBackend::Memory(m) => m.node_count(),
105        }
106    }
107
108    fn edge_count(&self) -> Result<usize> {
109        match self {
110            StoreBackend::Sled(s) => s.edge_count(),
111            StoreBackend::Memory(m) => m.edge_count(),
112        }
113    }
114
115    fn remove_node(&self, id: &str) -> Result<()> {
116        match self {
117            StoreBackend::Sled(s) => s.remove_node(id),
118            StoreBackend::Memory(m) => m.remove_node(id),
119        }
120    }
121
122    fn remove_edge(&self, source: &str, target: &str, relation: &str) -> Result<()> {
123        match self {
124            StoreBackend::Sled(s) => s.remove_edge(source, target, relation),
125            StoreBackend::Memory(m) => m.remove_edge(source, target, relation),
126        }
127    }
128
129    fn clear(&self) -> Result<()> {
130        match self {
131            StoreBackend::Sled(s) => s.clear(),
132            StoreBackend::Memory(m) => m.clear(),
133        }
134    }
135}
136
137pub struct SledGraphStore {
138    db: sled::Db,
139    inverted_index: RwLock<HashMap<String, Vec<String>>>,
140}
141
142impl SledGraphStore {
143    pub fn open(path: &Path) -> Result<Self> {
144        let db = sled::Config::new()
145            .path(path)
146            .flush_every_ms(if cfg!(test) { None } else { Some(500) })
147            .open()
148            .map_err(|e| CodeSynapseError::Database(e.to_string()))?;
149        let store = SledGraphStore {
150            db,
151            inverted_index: RwLock::new(HashMap::new()),
152        };
153        store.rebuild_index()?;
154        Ok(store)
155    }
156
157    pub fn temporary() -> Result<Self> {
158        let dir = std::env::temp_dir().join(format!("codesynapse-{}", uuid::Uuid::new_v4()));
159        std::fs::create_dir_all(&dir).ok();
160        Self::open(&dir)
161    }
162
163    fn node_key(id: &str) -> Vec<u8> {
164        format!("node:{}", id).into_bytes()
165    }
166
167    fn edge_key(id: &str) -> Vec<u8> {
168        format!("edge:{}", id).into_bytes()
169    }
170
171    fn adjacency_key(src: &str, target: &str) -> Vec<u8> {
172        format!("adj:{}:{}", src, target).into_bytes()
173    }
174
175    fn rebuild_index(&self) -> Result<()> {
176        let mut index = self.inverted_index.write().unwrap();
177        index.clear();
178        for result in self.db.iter() {
179            let (key, value) = result.map_err(|e| CodeSynapseError::Database(e.to_string()))?;
180            let key_str = String::from_utf8_lossy(&key);
181            if key_str.starts_with("node:") {
182                let node: Node = bincode::deserialize(&value).map_err(CodeSynapseError::Bincode)?;
183                for token in tokenize(&node.label) {
184                    index.entry(token).or_default().push(node.id.clone());
185                }
186                for token in tokenize(&node.id) {
187                    index.entry(token).or_default().push(node.id.clone());
188                }
189            }
190        }
191        Ok(())
192    }
193}
194
195impl GraphStore for SledGraphStore {
196    fn add_node(&self, node: Node) -> Result<()> {
197        let key = Self::node_key(&node.id);
198        let value = bincode::serialize(&node).map_err(CodeSynapseError::Bincode)?;
199        self.db
200            .insert(key, value)
201            .map_err(|e| CodeSynapseError::Database(e.to_string()))?;
202
203        // Update inverted index
204        let mut index = self.inverted_index.write().unwrap();
205        for token in tokenize(&node.label) {
206            index.entry(token).or_default().push(node.id.clone());
207        }
208        for token in tokenize(&node.id) {
209            index.entry(token).or_default().push(node.id.clone());
210        }
211
212        self.db
213            .flush()
214            .map_err(|e| CodeSynapseError::Database(e.to_string()))?;
215        Ok(())
216    }
217
218    fn add_edge(&self, edge: Edge) -> Result<()> {
219        let edge_id = format!("{}->{}:{}", edge.source, edge.target, edge.relation);
220        let key = Self::edge_key(&edge_id);
221        let value = bincode::serialize(&edge).map_err(CodeSynapseError::Bincode)?;
222        self.db
223            .insert(key, value)
224            .map_err(|e| CodeSynapseError::Database(e.to_string()))?;
225
226        // Adjacency index
227        let adj_key = Self::adjacency_key(&edge.source, &edge.target);
228        let adj_bytes = bincode::serialize(&edge).map_err(CodeSynapseError::Bincode)?;
229        self.db
230            .insert(adj_key, adj_bytes)
231            .map_err(|e| CodeSynapseError::Database(e.to_string()))?;
232
233        Ok(())
234    }
235
236    fn get_node(&self, id: &str) -> Result<Option<Node>> {
237        let key = Self::node_key(id);
238        match self.db.get(key) {
239            Ok(Some(value)) => {
240                let node: Node = bincode::deserialize(&value).map_err(CodeSynapseError::Bincode)?;
241                Ok(Some(node))
242            }
243            Ok(None) => Ok(None),
244            Err(e) => Err(CodeSynapseError::Database(e.to_string())),
245        }
246    }
247
248    fn get_all_nodes(&self) -> Result<Vec<Node>> {
249        let mut nodes = Vec::new();
250        for result in self.db.iter() {
251            let (key, value) = result.map_err(|e| CodeSynapseError::Database(e.to_string()))?;
252            let key_str = String::from_utf8_lossy(&key);
253            if key_str.starts_with("node:") {
254                let node: Node = bincode::deserialize(&value).map_err(CodeSynapseError::Bincode)?;
255                nodes.push(node);
256            }
257        }
258        Ok(nodes)
259    }
260
261    fn get_all_edges(&self) -> Result<Vec<Edge>> {
262        let mut edges = Vec::new();
263        for result in self.db.iter() {
264            let (key, value) = result.map_err(|e| CodeSynapseError::Database(e.to_string()))?;
265            let key_str = String::from_utf8_lossy(&key);
266            if key_str.starts_with("edge:") {
267                let edge: Edge = bincode::deserialize(&value).map_err(CodeSynapseError::Bincode)?;
268                edges.push(edge);
269            }
270        }
271        Ok(edges)
272    }
273
274    fn neighbors(&self, id: &str, _relation_filter: Option<&str>) -> Result<Vec<(Node, Edge)>> {
275        let mut result = Vec::new();
276        // Scan adjacency keys starting with adj:{id}:
277        let prefix = format!("adj:{}:", id);
278        for entry in self.db.scan_prefix(prefix.as_bytes()) {
279            let (_key, value) = entry.map_err(|e| CodeSynapseError::Database(e.to_string()))?;
280            let edge: Edge = bincode::deserialize(&value).map_err(CodeSynapseError::Bincode)?;
281            let neighbor_id = if edge.source == id {
282                &edge.target
283            } else {
284                &edge.source
285            };
286            if let Some(node) = self.get_node(neighbor_id)? {
287                result.push((node, edge));
288            }
289        }
290        Ok(result)
291    }
292
293    fn search(&self, query: &str, top_k: usize) -> Result<Vec<(f64, Node)>> {
294        let tokens = tokenize(query);
295        if tokens.is_empty() {
296            return Ok(vec![]);
297        }
298
299        let index = self.inverted_index.read().unwrap();
300        let mut scores: HashMap<String, f64> = HashMap::new();
301
302        // Use inverted index if available, fall back to full scan
303        if !index.is_empty() {
304            for token in &tokens {
305                if let Some(ids) = index.get(token) {
306                    for id in ids {
307                        *scores.entry(id.clone()).or_insert(0.0) += 1.0;
308                    }
309                }
310            }
311        } else {
312            // Full scan fallback
313            for result in self.db.iter() {
314                let (_key, value) =
315                    result.map_err(|e| CodeSynapseError::Database(e.to_string()))?;
316                let key_str = String::from_utf8_lossy(&_key);
317                if key_str.starts_with("node:") {
318                    let node: Node =
319                        bincode::deserialize(&value).map_err(CodeSynapseError::Bincode)?;
320                    let mut score = 0.0;
321                    for token in &tokens {
322                        if node.label.to_lowercase().contains(token) {
323                            score += 1.0;
324                        }
325                        if node.id.to_lowercase().contains(token) {
326                            score += 0.5;
327                        }
328                    }
329                    if score > 0.0 {
330                        scores.insert(node.id.clone(), score);
331                    }
332                }
333            }
334        }
335
336        let mut ranked: Vec<(f64, Node)> = scores
337            .into_iter()
338            .filter_map(|(id, score)| self.get_node(&id).ok().flatten().map(|n| (score, n)))
339            .collect();
340
341        ranked.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal));
342        ranked.truncate(top_k);
343        Ok(ranked)
344    }
345
346    fn shortest_path(&self, src: &str, tgt: &str) -> Result<Option<Vec<Node>>> {
347        if src == tgt {
348            return self.get_node(src).map(|n| n.map(|n| vec![n]));
349        }
350
351        let mut visited = std::collections::HashSet::new();
352        let mut queue = std::collections::VecDeque::new();
353        let mut parent: HashMap<String, String> = HashMap::new();
354
355        visited.insert(src.to_string());
356        queue.push_back(src.to_string());
357
358        while let Some(current) = queue.pop_front() {
359            if current == tgt {
360                let mut path = Vec::new();
361                let mut node_id = tgt.to_string();
362                while node_id != src {
363                    if let Some(n) = self.get_node(&node_id)? {
364                        path.push(n);
365                    }
366                    node_id = parent.get(&node_id).cloned().unwrap_or_default();
367                }
368                if let Some(n) = self.get_node(src)? {
369                    path.push(n);
370                }
371                path.reverse();
372                return Ok(Some(path));
373            }
374
375            if let Ok(neighbors) = self.neighbors(&current, None) {
376                for (neighbor, _) in neighbors {
377                    if visited.insert(neighbor.id.clone()) {
378                        parent.insert(neighbor.id.clone(), current.clone());
379                        queue.push_back(neighbor.id.clone());
380                    }
381                }
382            }
383        }
384
385        Ok(None)
386    }
387
388    fn dijkstra_shortest_path(&self, src: &str, tgt: &str) -> Result<Option<Vec<Node>>> {
389        use std::collections::BinaryHeap;
390
391        if src == tgt {
392            return self.get_node(src).map(|n| n.map(|n| vec![n]));
393        }
394
395        let mut distances: HashMap<String, f64> = HashMap::new();
396        let mut parents: HashMap<String, String> = HashMap::new();
397        let mut heap = BinaryHeap::new();
398
399        distances.insert(src.to_string(), 0.0);
400
401        #[derive(PartialEq)]
402        struct State(String, f64);
403        impl Eq for State {}
404        impl PartialOrd for State {
405            fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
406                Some(self.cmp(other))
407            }
408        }
409        impl Ord for State {
410            fn cmp(&self, other: &Self) -> std::cmp::Ordering {
411                other
412                    .1
413                    .partial_cmp(&self.1)
414                    .unwrap_or(std::cmp::Ordering::Equal)
415            }
416        }
417
418        heap.push(State(src.to_string(), 0.0));
419
420        while let Some(State(current, cost)) = heap.pop() {
421            if current == tgt {
422                let mut path = Vec::new();
423                let mut node_id = tgt.to_string();
424                while node_id != src {
425                    if let Some(n) = self.get_node(&node_id)? {
426                        path.push(n);
427                    }
428                    node_id = parents.get(&node_id).cloned().unwrap_or_default();
429                }
430                if let Some(n) = self.get_node(src)? {
431                    path.push(n);
432                }
433                path.reverse();
434                return Ok(Some(path));
435            }
436
437            if let Some(&best) = distances.get(&current) {
438                if cost > best {
439                    continue;
440                }
441            }
442
443            if let Ok(neighbors) = self.neighbors(&current, None) {
444                for (neighbor, edge) in neighbors {
445                    let next_cost = cost + edge.weight;
446                    let is_better = distances.get(&neighbor.id).is_none_or(|&d| next_cost < d);
447                    if is_better {
448                        distances.insert(neighbor.id.clone(), next_cost);
449                        parents.insert(neighbor.id.clone(), current.clone());
450                        heap.push(State(neighbor.id.clone(), next_cost));
451                    }
452                }
453            }
454        }
455
456        Ok(None)
457    }
458
459    fn subgraph(&self, node_ids: &[&str]) -> Result<(Vec<Node>, Vec<Edge>)> {
460        let mut nodes = Vec::new();
461        let mut edges = Vec::new();
462        let id_set: std::collections::HashSet<&str> = node_ids.iter().copied().collect();
463
464        for id in node_ids {
465            if let Some(node) = self.get_node(id)? {
466                nodes.push(node);
467            }
468        }
469
470        let all_edges = self.get_all_edges()?;
471        for edge in all_edges {
472            if id_set.contains(edge.source.as_str()) && id_set.contains(edge.target.as_str()) {
473                edges.push(edge);
474            }
475        }
476
477        Ok((nodes, edges))
478    }
479
480    fn node_count(&self) -> Result<usize> {
481        let mut count = 0;
482        for result in self.db.iter() {
483            let (key, _) = result.map_err(|e| CodeSynapseError::Database(e.to_string()))?;
484            if String::from_utf8_lossy(&key).starts_with("node:") {
485                count += 1;
486            }
487        }
488        Ok(count)
489    }
490
491    fn edge_count(&self) -> Result<usize> {
492        let mut count = 0;
493        for result in self.db.iter() {
494            let (key, _) = result.map_err(|e| CodeSynapseError::Database(e.to_string()))?;
495            if String::from_utf8_lossy(&key).starts_with("edge:") {
496                count += 1;
497            }
498        }
499        Ok(count)
500    }
501
502    fn remove_node(&self, id: &str) -> Result<()> {
503        let key = Self::node_key(id);
504        self.db
505            .remove(key)
506            .map_err(|e| CodeSynapseError::Database(e.to_string()))?;
507        Ok(())
508    }
509
510    fn remove_edge(&self, source: &str, target: &str, relation: &str) -> Result<()> {
511        let edge_id = format!("{}->{}:{}", source, target, relation);
512        let key = Self::edge_key(&edge_id);
513        self.db
514            .remove(key)
515            .map_err(|e| CodeSynapseError::Database(e.to_string()))?;
516        let adj_key = Self::adjacency_key(source, target);
517        self.db
518            .remove(adj_key)
519            .map_err(|e| CodeSynapseError::Database(e.to_string()))?;
520        Ok(())
521    }
522
523    fn clear(&self) -> Result<()> {
524        self.db
525            .clear()
526            .map_err(|e| CodeSynapseError::Database(e.to_string()))?;
527        self.inverted_index.write().unwrap().clear();
528        Ok(())
529    }
530}
531
532pub struct MemoryGraphStore {
533    nodes: RwLock<HashMap<NodeId, Node>>,
534    edges: RwLock<Vec<Edge>>,
535    inverted_index: RwLock<HashMap<String, Vec<String>>>,
536}
537
538impl MemoryGraphStore {
539    pub fn new() -> Self {
540        MemoryGraphStore {
541            nodes: RwLock::new(HashMap::new()),
542            edges: RwLock::new(Vec::new()),
543            inverted_index: RwLock::new(HashMap::new()),
544        }
545    }
546}
547
548impl Default for MemoryGraphStore {
549    fn default() -> Self {
550        Self::new()
551    }
552}
553
554impl GraphStore for MemoryGraphStore {
555    fn add_node(&self, node: Node) -> Result<()> {
556        let mut nodes = self.nodes.write().unwrap();
557        nodes.insert(node.id.clone(), node.clone());
558
559        let mut index = self.inverted_index.write().unwrap();
560        for token in tokenize(&node.label) {
561            index.entry(token).or_default().push(node.id.clone());
562        }
563        for token in tokenize(&node.id) {
564            index.entry(token).or_default().push(node.id.clone());
565        }
566
567        Ok(())
568    }
569
570    fn add_edge(&self, edge: Edge) -> Result<()> {
571        self.edges.write().unwrap().push(edge);
572        Ok(())
573    }
574
575    fn get_node(&self, id: &str) -> Result<Option<Node>> {
576        Ok(self.nodes.read().unwrap().get(id).cloned())
577    }
578
579    fn get_all_nodes(&self) -> Result<Vec<Node>> {
580        Ok(self.nodes.read().unwrap().values().cloned().collect())
581    }
582
583    fn get_all_edges(&self) -> Result<Vec<Edge>> {
584        Ok(self.edges.read().unwrap().clone())
585    }
586
587    fn neighbors(&self, id: &str, relation_filter: Option<&str>) -> Result<Vec<(Node, Edge)>> {
588        let edges = self.edges.read().unwrap();
589        let nodes = self.nodes.read().unwrap();
590        let mut result = Vec::new();
591        for edge in edges.iter() {
592            if let Some(filter) = relation_filter {
593                if edge.relation != filter {
594                    continue;
595                }
596            }
597            if edge.source == id {
598                if let Some(node) = nodes.get(&edge.target) {
599                    result.push((node.clone(), edge.clone()));
600                }
601            } else if edge.target == id {
602                if let Some(node) = nodes.get(&edge.source) {
603                    result.push((node.clone(), edge.clone()));
604                }
605            }
606        }
607        Ok(result)
608    }
609
610    fn search(&self, query: &str, top_k: usize) -> Result<Vec<(f64, Node)>> {
611        let tokens = tokenize(query);
612        if tokens.is_empty() {
613            return Ok(vec![]);
614        }
615
616        let nodes = self.nodes.read().unwrap();
617        let index = self.inverted_index.read().unwrap();
618
619        let mut scores: HashMap<String, f64> = HashMap::new();
620
621        if !index.is_empty() {
622            for token in &tokens {
623                if let Some(ids) = index.get(token) {
624                    for id in ids {
625                        *scores.entry(id.clone()).or_insert(0.0) += 1.0;
626                    }
627                }
628            }
629        } else {
630            for node in nodes.values() {
631                let mut score = 0.0;
632                for token in &tokens {
633                    if node.label.to_lowercase().contains(token) {
634                        score += 1.0;
635                    }
636                    if node.id.to_lowercase().contains(token) {
637                        score += 0.5;
638                    }
639                }
640                if score > 0.0 {
641                    scores.insert(node.id.clone(), score);
642                }
643            }
644        }
645
646        let mut ranked: Vec<(f64, Node)> = scores
647            .into_iter()
648            .filter_map(|(id, score)| nodes.get(&id).map(|n| (score, n.clone())))
649            .collect();
650
651        ranked.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal));
652        ranked.truncate(top_k);
653        Ok(ranked)
654    }
655
656    fn shortest_path(&self, src: &str, tgt: &str) -> Result<Option<Vec<Node>>> {
657        if src == tgt {
658            return Ok(self
659                .nodes
660                .read()
661                .unwrap()
662                .get(src)
663                .cloned()
664                .map(|n| vec![n]));
665        }
666
667        let nodes = self.nodes.read().unwrap();
668        let edges = self.edges.read().unwrap();
669
670        let mut visited = std::collections::HashSet::new();
671        let mut queue = std::collections::VecDeque::new();
672        let mut parent: HashMap<String, String> = HashMap::new();
673
674        visited.insert(src.to_string());
675        queue.push_back(src.to_string());
676
677        while let Some(current) = queue.pop_front() {
678            if current == tgt {
679                let mut path = Vec::new();
680                let mut node_id = tgt.to_string();
681                while node_id != src {
682                    if let Some(n) = nodes.get(&node_id) {
683                        path.push(n.clone());
684                    }
685                    node_id = parent.get(&node_id).cloned().unwrap_or_default();
686                }
687                if let Some(n) = nodes.get(src) {
688                    path.push(n.clone());
689                }
690                path.reverse();
691                return Ok(Some(path));
692            }
693
694            for edge in edges.iter() {
695                let neighbor = if edge.source == current {
696                    Some(&edge.target)
697                } else if edge.target == current {
698                    Some(&edge.source)
699                } else {
700                    None
701                };
702
703                if let Some(neighbor) = neighbor {
704                    if visited.insert(neighbor.clone()) {
705                        parent.insert(neighbor.clone(), current.clone());
706                        queue.push_back(neighbor.clone());
707                    }
708                }
709            }
710        }
711
712        Ok(None)
713    }
714
715    fn dijkstra_shortest_path(&self, src: &str, tgt: &str) -> Result<Option<Vec<Node>>> {
716        if src == tgt {
717            return Ok(self
718                .nodes
719                .read()
720                .unwrap()
721                .get(src)
722                .cloned()
723                .map(|n| vec![n]));
724        }
725
726        let nodes = self.nodes.read().unwrap();
727        let edges = self.edges.read().unwrap();
728        let mut distances: HashMap<String, f64> = HashMap::new();
729        let mut parents: HashMap<String, String> = HashMap::new();
730        let mut heap = std::collections::BinaryHeap::new();
731
732        distances.insert(src.to_string(), 0.0);
733
734        #[derive(PartialEq)]
735        struct State(String, f64);
736        impl Eq for State {}
737        impl PartialOrd for State {
738            fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
739                Some(self.cmp(other))
740            }
741        }
742        impl Ord for State {
743            fn cmp(&self, other: &Self) -> std::cmp::Ordering {
744                other
745                    .1
746                    .partial_cmp(&self.1)
747                    .unwrap_or(std::cmp::Ordering::Equal)
748            }
749        }
750
751        heap.push(State(src.to_string(), 0.0));
752
753        while let Some(State(current, cost)) = heap.pop() {
754            if current == tgt {
755                let mut path = Vec::new();
756                let mut node_id = tgt.to_string();
757                while node_id != src {
758                    if let Some(n) = nodes.get(&node_id) {
759                        path.push(n.clone());
760                    }
761                    node_id = parents.get(&node_id).cloned().unwrap_or_default();
762                }
763                if let Some(n) = nodes.get(src) {
764                    path.push(n.clone());
765                }
766                path.reverse();
767                return Ok(Some(path));
768            }
769
770            if let Some(&best) = distances.get(&current) {
771                if cost > best {
772                    continue;
773                }
774            }
775
776            for edge in edges.iter() {
777                let neighbor = if edge.source == current {
778                    Some(&edge.target)
779                } else if edge.target == current {
780                    Some(&edge.source)
781                } else {
782                    None
783                };
784                if let Some(neighbor) = neighbor {
785                    let next_cost = cost + edge.weight;
786                    let is_better = distances.get(neighbor).is_none_or(|&d| next_cost < d);
787                    if is_better {
788                        distances.insert(neighbor.clone(), next_cost);
789                        parents.insert(neighbor.clone(), current.clone());
790                        heap.push(State(neighbor.clone(), next_cost));
791                    }
792                }
793            }
794        }
795
796        Ok(None)
797    }
798
799    fn subgraph(&self, node_ids: &[&str]) -> Result<(Vec<Node>, Vec<Edge>)> {
800        let nodes = self.nodes.read().unwrap();
801        let edges = self.edges.read().unwrap();
802        let id_set: std::collections::HashSet<&str> = node_ids.iter().copied().collect();
803
804        let sg_nodes: Vec<Node> = nodes
805            .values()
806            .filter(|n| id_set.contains(n.id.as_str()))
807            .cloned()
808            .collect();
809        let sg_edges: Vec<Edge> = edges
810            .iter()
811            .filter(|e| id_set.contains(e.source.as_str()) && id_set.contains(e.target.as_str()))
812            .cloned()
813            .collect();
814
815        Ok((sg_nodes, sg_edges))
816    }
817
818    fn node_count(&self) -> Result<usize> {
819        Ok(self.nodes.read().unwrap().len())
820    }
821
822    fn edge_count(&self) -> Result<usize> {
823        Ok(self.edges.read().unwrap().len())
824    }
825
826    fn remove_node(&self, id: &str) -> Result<()> {
827        self.nodes.write().unwrap().remove(id);
828        Ok(())
829    }
830
831    fn remove_edge(&self, source: &str, target: &str, relation: &str) -> Result<()> {
832        let mut edges = self.edges.write().unwrap();
833        edges.retain(|e| !(e.source == source && e.target == target && e.relation == relation));
834        Ok(())
835    }
836
837    fn clear(&self) -> Result<()> {
838        self.nodes.write().unwrap().clear();
839        self.edges.write().unwrap().clear();
840        self.inverted_index.write().unwrap().clear();
841        Ok(())
842    }
843}
844
845fn tokenize(s: &str) -> Vec<String> {
846    s.to_lowercase()
847        .split(|c: char| !c.is_alphanumeric())
848        .filter(|t| !t.is_empty() && t.len() >= 2)
849        .map(|t| t.to_string())
850        .collect()
851}
852
853#[cfg(test)]
854mod tests {
855    use super::*;
856    use std::collections::HashMap;
857    use tempfile;
858
859    fn make_node(id: &str, label: &str) -> Node {
860        Node {
861            id: id.to_string(),
862            label: label.to_string(),
863            file_type: "code".to_string(),
864            source_file: "test.py".to_string(),
865            source_location: None,
866            community: None,
867            rationale: None,
868            docstring: None,
869            metadata: HashMap::new(),
870        }
871    }
872
873    fn make_edge(src: &str, tgt: &str, relation: &str) -> Edge {
874        Edge {
875            source: src.to_string(),
876            target: tgt.to_string(),
877            relation: relation.to_string(),
878            confidence: "EXTRACTED".to_string(),
879            source_file: Some("test.py".to_string()),
880            weight: 1.0,
881            context: None,
882        }
883    }
884
885    #[test]
886    fn test_memory_store_add_get_node() {
887        let store = MemoryGraphStore::new();
888        store.add_node(make_node("a", "Alpha")).unwrap();
889        let node = store.get_node("a").unwrap().unwrap();
890        assert_eq!(node.label, "Alpha");
891    }
892
893    #[test]
894    fn test_memory_store_search() {
895        let store = MemoryGraphStore::new();
896        store.add_node(make_node("auth", "AuthService")).unwrap();
897        store.add_node(make_node("user", "UserService")).unwrap();
898
899        let results = store.search("auth", 10).unwrap();
900        assert_eq!(results.len(), 1);
901        assert_eq!(results[0].1.label, "AuthService");
902    }
903
904    #[test]
905    fn test_memory_store_neighbors() {
906        let store = MemoryGraphStore::new();
907        store.add_node(make_node("a", "A")).unwrap();
908        store.add_node(make_node("b", "B")).unwrap();
909        store.add_edge(make_edge("a", "b", "connects")).unwrap();
910
911        let neighbors = store.neighbors("a", None).unwrap();
912        assert_eq!(neighbors.len(), 1);
913        assert_eq!(neighbors[0].0.label, "B");
914    }
915
916    #[test]
917    fn test_memory_store_shortest_path() {
918        let store = MemoryGraphStore::new();
919        store.add_node(make_node("a", "A")).unwrap();
920        store.add_node(make_node("b", "B")).unwrap();
921        store.add_node(make_node("c", "C")).unwrap();
922        store.add_edge(make_edge("a", "b", "connects")).unwrap();
923        store.add_edge(make_edge("b", "c", "connects")).unwrap();
924
925        let path = store.shortest_path("a", "c").unwrap().unwrap();
926        assert_eq!(path.len(), 3);
927        assert_eq!(path[0].label, "A");
928        assert_eq!(path[1].label, "B");
929        assert_eq!(path[2].label, "C");
930    }
931
932    #[test]
933    fn test_memory_store_shortest_path_same_node() {
934        let store = MemoryGraphStore::new();
935        store.add_node(make_node("a", "A")).unwrap();
936        let path = store.shortest_path("a", "a").unwrap().unwrap();
937        assert_eq!(path.len(), 1);
938    }
939
940    #[test]
941    fn test_memory_store_shortest_path_no_path() {
942        let store = MemoryGraphStore::new();
943        store.add_node(make_node("a", "A")).unwrap();
944        store.add_node(make_node("b", "B")).unwrap();
945        let path = store.shortest_path("a", "b").unwrap();
946        assert!(path.is_none());
947    }
948
949    #[test]
950    fn test_memory_store_subgraph() {
951        let store = MemoryGraphStore::new();
952        store.add_node(make_node("a", "A")).unwrap();
953        store.add_node(make_node("b", "B")).unwrap();
954        store.add_node(make_node("c", "C")).unwrap();
955        store.add_edge(make_edge("a", "b", "connects")).unwrap();
956        store.add_edge(make_edge("b", "c", "connects")).unwrap();
957
958        let (sg_nodes, sg_edges) = store.subgraph(&["a", "b"]).unwrap();
959        assert_eq!(sg_nodes.len(), 2);
960        assert_eq!(sg_edges.len(), 1);
961    }
962
963    #[test]
964    fn test_tokenize() {
965        let tokens = tokenize("AuthService");
966        assert!(tokens.contains(&"authservice".to_string()));
967
968        let tokens = tokenize("Hello World");
969        assert!(tokens.contains(&"hello".to_string()));
970        assert!(tokens.contains(&"world".to_string()));
971    }
972
973    #[test]
974    fn test_sled_store_basic_ops() {
975        let store = SledGraphStore::temporary().unwrap();
976        store.add_node(make_node("a", "Alpha")).unwrap();
977        store.add_node(make_node("b", "Beta")).unwrap();
978        store.add_edge(make_edge("a", "b", "connects")).unwrap();
979
980        assert_eq!(store.node_count().unwrap(), 2);
981        assert_eq!(store.edge_count().unwrap(), 1);
982
983        let node = store.get_node("a").unwrap().unwrap();
984        assert_eq!(node.label, "Alpha");
985    }
986
987    #[test]
988    fn test_sled_store_shortest_path() {
989        let store = SledGraphStore::temporary().unwrap();
990        store.add_node(make_node("a", "A")).unwrap();
991        store.add_node(make_node("b", "B")).unwrap();
992        store.add_node(make_node("c", "C")).unwrap();
993        store.add_edge(make_edge("a", "b", "connects")).unwrap();
994        store.add_edge(make_edge("b", "c", "connects")).unwrap();
995
996        let path = store.shortest_path("a", "c").unwrap().unwrap();
997        assert_eq!(path.len(), 3);
998        assert_eq!(path[0].label, "A");
999        assert_eq!(path[1].label, "B");
1000        assert_eq!(path[2].label, "C");
1001    }
1002
1003    #[test]
1004    fn test_sled_store_shortest_path_no_path() {
1005        let store = SledGraphStore::temporary().unwrap();
1006        store.add_node(make_node("a", "A")).unwrap();
1007        store.add_node(make_node("b", "B")).unwrap();
1008        // no edge between a and b
1009        let path = store.shortest_path("a", "b").unwrap();
1010        assert!(path.is_none());
1011    }
1012
1013    #[test]
1014    fn test_sled_store_rebuild_index() {
1015        let store = SledGraphStore::temporary().unwrap();
1016        store.add_node(make_node("auth", "AuthService")).unwrap();
1017        store.rebuild_index().unwrap();
1018
1019        let results = store.search("auth", 10).unwrap();
1020        assert_eq!(results.len(), 1);
1021    }
1022
1023    #[test]
1024    fn test_store_backend_memory() {
1025        let backend = StoreBackend::Memory(MemoryGraphStore::new());
1026        backend.add_node(make_node("a", "A")).unwrap();
1027        assert_eq!(backend.node_count().unwrap(), 1);
1028    }
1029
1030    // --- Edge case tests ---
1031
1032    #[test]
1033    fn test_memory_get_node_nonexistent() {
1034        let store = MemoryGraphStore::new();
1035        let node = store.get_node("nonexistent").unwrap();
1036        assert!(node.is_none());
1037    }
1038
1039    #[test]
1040    fn test_memory_neighbors_nonexistent() {
1041        let store = MemoryGraphStore::new();
1042        let nbrs = store.neighbors("ghost", None).unwrap();
1043        assert!(nbrs.is_empty());
1044    }
1045
1046    #[test]
1047    fn test_memory_neighbors_relation_filter() {
1048        let store = MemoryGraphStore::new();
1049        store.add_node(make_node("a", "A")).unwrap();
1050        store.add_node(make_node("b", "B")).unwrap();
1051        store.add_edge(make_edge("a", "b", "calls")).unwrap();
1052        store.add_edge(make_edge("a", "b", "imports")).unwrap();
1053
1054        let calls = store.neighbors("a", Some("calls")).unwrap();
1055        assert_eq!(calls.len(), 1);
1056        assert_eq!(calls[0].1.relation, "calls");
1057
1058        let imports = store.neighbors("a", Some("imports")).unwrap();
1059        assert_eq!(imports.len(), 1);
1060
1061        let none = store.neighbors("a", Some("inherits")).unwrap();
1062        assert!(none.is_empty());
1063    }
1064
1065    #[test]
1066    fn test_memory_search_empty_query() {
1067        let store = MemoryGraphStore::new();
1068        store.add_node(make_node("a", "Alpha")).unwrap();
1069        let results = store.search("", 10).unwrap();
1070        assert!(results.is_empty());
1071    }
1072
1073    #[test]
1074    fn test_memory_search_single_char() {
1075        let store = MemoryGraphStore::new();
1076        store.add_node(make_node("a", "Alpha")).unwrap();
1077        let results = store.search("A", 10).unwrap();
1078        assert!(results.is_empty(), "single char should be tokenized away");
1079    }
1080
1081    #[test]
1082    fn test_memory_search_no_match() {
1083        let store = MemoryGraphStore::new();
1084        store.add_node(make_node("a", "Alpha")).unwrap();
1085        let results = store.search("Zeta", 10).unwrap();
1086        assert!(results.is_empty());
1087    }
1088
1089    #[test]
1090    fn test_memory_remove_node_nonexistent() {
1091        let store = MemoryGraphStore::new();
1092        store.remove_node("ghost").unwrap();
1093        // no panic — should be a no-op
1094        assert_eq!(store.node_count().unwrap(), 0);
1095    }
1096
1097    #[test]
1098    fn test_memory_remove_edge_nonexistent() {
1099        let store = MemoryGraphStore::new();
1100        store.remove_edge("a", "b", "calls").unwrap();
1101        assert_eq!(store.edge_count().unwrap(), 0);
1102    }
1103
1104    #[test]
1105    fn test_memory_add_duplicate_edge() {
1106        let store = MemoryGraphStore::new();
1107        store.add_node(make_node("a", "A")).unwrap();
1108        store.add_node(make_node("b", "B")).unwrap();
1109        let e1 = make_edge("a", "b", "calls");
1110        let e2 = make_edge("a", "b", "calls");
1111        store.add_edge(e1).unwrap();
1112        store.add_edge(e2).unwrap();
1113        assert_eq!(
1114            store.edge_count().unwrap(),
1115            2,
1116            "duplicate edges are allowed"
1117        );
1118    }
1119
1120    #[test]
1121    fn test_memory_subgraph_nonexistent_ids() {
1122        let store = MemoryGraphStore::new();
1123        let (nodes, edges) = store.subgraph(&["ghost1", "ghost2"]).unwrap();
1124        assert!(nodes.is_empty());
1125        assert!(edges.is_empty());
1126    }
1127
1128    #[test]
1129    fn test_memory_shortest_path_missing_src() {
1130        let store = MemoryGraphStore::new();
1131        store.add_node(make_node("b", "B")).unwrap();
1132        let path = store.shortest_path("a", "b").unwrap();
1133        assert!(path.is_none(), "no path when src node missing");
1134    }
1135
1136    #[test]
1137    fn test_memory_clear_then_ops() {
1138        let store = MemoryGraphStore::new();
1139        store.add_node(make_node("a", "A")).unwrap();
1140        store.add_edge(make_edge("a", "b", "calls")).unwrap();
1141        store.clear().unwrap();
1142        assert_eq!(store.node_count().unwrap(), 0);
1143        assert_eq!(store.edge_count().unwrap(), 0);
1144        // After clear, adding should still work
1145        store.add_node(make_node("b", "B")).unwrap();
1146        assert_eq!(store.node_count().unwrap(), 1);
1147    }
1148
1149    #[test]
1150    fn test_sled_store_clear() {
1151        let store = SledGraphStore::temporary().unwrap();
1152        store.add_node(make_node("a", "A")).unwrap();
1153        store.clear().unwrap();
1154        assert_eq!(store.node_count().unwrap(), 0);
1155    }
1156
1157    #[test]
1158    fn test_dijkstra_shortest_path_weighted() {
1159        let store = MemoryGraphStore::new();
1160        store.add_node(make_node("a", "A")).unwrap();
1161        store.add_node(make_node("b", "B")).unwrap();
1162        store.add_node(make_node("c", "C")).unwrap();
1163        store.add_node(make_node("d", "D")).unwrap();
1164        // a -> b (1.0), a -> c (10.0), b -> d (1.0), c -> d (1.0)
1165        store
1166            .add_edge(make_edge_weighted("a", "b", "connects", 1.0))
1167            .unwrap();
1168        store
1169            .add_edge(make_edge_weighted("a", "c", "connects", 10.0))
1170            .unwrap();
1171        store
1172            .add_edge(make_edge_weighted("b", "d", "connects", 1.0))
1173            .unwrap();
1174        store
1175            .add_edge(make_edge_weighted("c", "d", "connects", 1.0))
1176            .unwrap();
1177
1178        let path = store.dijkstra_shortest_path("a", "d").unwrap().unwrap();
1179        assert_eq!(path.len(), 3, "shortest path should be a->b->d");
1180        assert_eq!(path[0].label, "A");
1181        assert_eq!(path[1].label, "B");
1182        assert_eq!(path[2].label, "D");
1183    }
1184
1185    #[test]
1186    fn test_dijkstra_no_path() {
1187        let store = MemoryGraphStore::new();
1188        store.add_node(make_node("a", "A")).unwrap();
1189        store.add_node(make_node("b", "B")).unwrap();
1190        let path = store.dijkstra_shortest_path("a", "b").unwrap();
1191        assert!(path.is_none());
1192    }
1193
1194    #[test]
1195    fn test_dijkstra_same_node() {
1196        let store = MemoryGraphStore::new();
1197        store.add_node(make_node("a", "A")).unwrap();
1198        let path = store.dijkstra_shortest_path("a", "a").unwrap().unwrap();
1199        assert_eq!(path.len(), 1);
1200    }
1201
1202    fn make_edge_weighted(src: &str, tgt: &str, relation: &str, weight: f64) -> Edge {
1203        Edge {
1204            source: src.to_string(),
1205            target: tgt.to_string(),
1206            relation: relation.to_string(),
1207            confidence: "EXTRACTED".to_string(),
1208            source_file: Some("test.py".to_string()),
1209            weight,
1210            context: None,
1211        }
1212    }
1213
1214    #[test]
1215    fn test_sled_store_persistence() {
1216        let tmpdir = tempfile::tempdir().unwrap();
1217        let db_path = tmpdir.path().join("sled.db");
1218
1219        {
1220            let store = SledGraphStore::open(&db_path).unwrap();
1221            store.add_node(make_node("persist_a", "PersistA")).unwrap();
1222            store.add_node(make_node("persist_b", "PersistB")).unwrap();
1223            store
1224                .add_edge(make_edge("persist_a", "persist_b", "persists"))
1225                .unwrap();
1226            assert_eq!(store.node_count().unwrap(), 2);
1227            assert_eq!(store.edge_count().unwrap(), 1);
1228            drop(store);
1229        }
1230
1231        {
1232            let store = SledGraphStore::open(&db_path).unwrap();
1233            assert_eq!(store.node_count().unwrap(), 2);
1234            assert_eq!(store.edge_count().unwrap(), 1);
1235            let node = store.get_node("persist_a").unwrap().unwrap();
1236            assert_eq!(node.label, "PersistA");
1237            let all_nodes = store.get_all_nodes().unwrap();
1238            let ids: Vec<&str> = all_nodes.iter().map(|n| n.id.as_str()).collect();
1239            assert!(ids.contains(&"persist_a"));
1240            assert!(ids.contains(&"persist_b"));
1241        }
1242    }
1243
1244    #[test]
1245    fn test_sled_store_reopen_empty() {
1246        let tmpdir = tempfile::tempdir().unwrap();
1247        let db_path = tmpdir.path().join("sled.db");
1248
1249        {
1250            let store = SledGraphStore::open(&db_path).unwrap();
1251            assert_eq!(store.node_count().unwrap(), 0);
1252            assert_eq!(store.edge_count().unwrap(), 0);
1253            drop(store);
1254        }
1255
1256        {
1257            let store = SledGraphStore::open(&db_path).unwrap();
1258            assert_eq!(store.node_count().unwrap(), 0);
1259            assert_eq!(store.edge_count().unwrap(), 0);
1260        }
1261    }
1262
1263    #[test]
1264    fn test_sled_store_search_after_reopen_with_new_add() {
1265        let tmpdir = tempfile::tempdir().unwrap();
1266        let db_path = tmpdir.path().join("sled.db");
1267
1268        {
1269            let store = SledGraphStore::open(&db_path).unwrap();
1270            store.add_node(make_node("auth", "AuthService")).unwrap();
1271            drop(store);
1272        }
1273
1274        {
1275            let store = SledGraphStore::open(&db_path).unwrap();
1276            // Adding a new node makes the in-memory index non-empty.
1277            // Without rebuild on open, search uses the index and misses "auth".
1278            store.add_node(make_node("new", "NewService")).unwrap();
1279            let results = store.search("auth", 10).unwrap();
1280            assert_eq!(
1281                results.len(),
1282                1,
1283                "persisted node must be found after reopen+add"
1284            );
1285        }
1286    }
1287}
codesynapse_core/graph.rs

codesynapse_core/
graph.rs
