1use crate::edge::Edge;
4use crate::error::Result;
5use crate::hyperedge::{Hyperedge, HyperedgeId};
6use crate::index::{AdjacencyIndex, EdgeTypeIndex, HyperedgeNodeIndex, LabelIndex, PropertyIndex};
7use crate::node::Node;
8#[cfg(feature = "storage")]
9use crate::storage::GraphStorage;
10use crate::types::{EdgeId, NodeId, PropertyValue};
11use dashmap::DashMap;
12#[cfg(feature = "storage")]
13use std::path::Path;
14use std::sync::Arc;
15
16pub struct GraphDB {
18 nodes: Arc<DashMap<NodeId, Node>>,
20 edges: Arc<DashMap<EdgeId, Edge>>,
22 hyperedges: Arc<DashMap<HyperedgeId, Hyperedge>>,
24 label_index: LabelIndex,
26 property_index: PropertyIndex,
28 edge_type_index: EdgeTypeIndex,
30 adjacency_index: AdjacencyIndex,
32 hyperedge_node_index: HyperedgeNodeIndex,
34 #[cfg(feature = "storage")]
36 storage: Option<GraphStorage>,
37}
38
39impl GraphDB {
40 pub fn new() -> Self {
42 Self {
43 nodes: Arc::new(DashMap::new()),
44 edges: Arc::new(DashMap::new()),
45 hyperedges: Arc::new(DashMap::new()),
46 label_index: LabelIndex::new(),
47 property_index: PropertyIndex::new(),
48 edge_type_index: EdgeTypeIndex::new(),
49 adjacency_index: AdjacencyIndex::new(),
50 hyperedge_node_index: HyperedgeNodeIndex::new(),
51 #[cfg(feature = "storage")]
52 storage: None,
53 }
54 }
55
56 #[cfg(feature = "storage")]
58 pub fn with_storage<P: AsRef<Path>>(path: P) -> anyhow::Result<Self> {
59 let storage = GraphStorage::new(path)?;
60
61 let mut db = Self::new();
62 db.storage = Some(storage);
63
64 db.load_from_storage()?;
66
67 Ok(db)
68 }
69
70 #[cfg(feature = "storage")]
72 fn load_from_storage(&mut self) -> anyhow::Result<()> {
73 if let Some(storage) = &self.storage {
74 for node_id in storage.all_node_ids()? {
76 if let Some(node) = storage.get_node(&node_id)? {
77 self.nodes.insert(node_id.clone(), node.clone());
78 self.label_index.add_node(&node);
79 self.property_index.add_node(&node);
80 }
81 }
82
83 for edge_id in storage.all_edge_ids()? {
85 if let Some(edge) = storage.get_edge(&edge_id)? {
86 self.edges.insert(edge_id.clone(), edge.clone());
87 self.edge_type_index.add_edge(&edge);
88 self.adjacency_index.add_edge(&edge);
89 }
90 }
91
92 for hyperedge_id in storage.all_hyperedge_ids()? {
94 if let Some(hyperedge) = storage.get_hyperedge(&hyperedge_id)? {
95 self.hyperedges
96 .insert(hyperedge_id.clone(), hyperedge.clone());
97 self.hyperedge_node_index.add_hyperedge(&hyperedge);
98 }
99 }
100 }
101 Ok(())
102 }
103
104 pub fn create_node(&self, node: Node) -> Result<NodeId> {
108 let id = node.id.clone();
109
110 self.label_index.add_node(&node);
112 self.property_index.add_node(&node);
113
114 self.nodes.insert(id.clone(), node.clone());
116
117 #[cfg(feature = "storage")]
119 if let Some(storage) = &self.storage {
120 storage.insert_node(&node)?;
121 }
122
123 Ok(id)
124 }
125
126 pub fn get_node(&self, id: impl AsRef<str>) -> Option<Node> {
128 self.nodes.get(id.as_ref()).map(|entry| entry.clone())
129 }
130
131 pub fn with_node<R>(&self, id: &str, f: impl FnOnce(&Node) -> R) -> Option<R> {
137 self.nodes.get(id).map(|entry| f(entry.value()))
138 }
139
140 pub fn node_ids_by_label(&self, label: &str) -> Vec<NodeId> {
142 self.label_index.get_nodes_by_label(label)
143 }
144
145 pub fn delete_node(&self, id: impl AsRef<str>) -> Result<bool> {
147 if let Some((_, node)) = self.nodes.remove(id.as_ref()) {
148 self.label_index.remove_node(&node);
150 self.property_index.remove_node(&node);
151
152 #[cfg(feature = "storage")]
154 if let Some(storage) = &self.storage {
155 storage.delete_node(id.as_ref())?;
156 }
157
158 Ok(true)
159 } else {
160 Ok(false)
161 }
162 }
163
164 pub fn update_node<F>(&self, id: impl AsRef<str>, f: F) -> Result<bool>
188 where
189 F: FnOnce(&mut Node),
190 {
191 let id_ref = id.as_ref();
192 let Some(mut entry) = self.nodes.get_mut(id_ref) else {
193 return Ok(false);
194 };
195 let old_node = entry.value().clone();
196 let mut new_node = old_node.clone();
197 f(&mut new_node);
198
199 if new_node.id != old_node.id {
200 return Err(crate::error::GraphError::ConstraintViolation(
201 "A node's ID cannot be changed by update_node".to_string(),
202 ));
203 }
204
205 #[cfg(feature = "storage")]
208 if let Some(storage) = &self.storage {
209 storage.insert_node(&new_node)?;
210 }
211
212 self.label_index.remove_node(&old_node);
213 self.property_index.remove_node(&old_node);
214 *entry.value_mut() = new_node.clone();
215 self.label_index.add_node(&new_node);
216 self.property_index.add_node(&new_node);
217
218 Ok(true)
219 }
220
221 pub fn keyword_search(
239 &self,
240 label: &str,
241 text_field: &str,
242 query: &str,
243 k: usize,
244 ) -> Result<Vec<(NodeId, f32)>> {
245 let docs: Vec<(NodeId, String)> = self
246 .node_ids_by_label(label)
247 .into_iter()
248 .filter_map(|id| {
249 self.with_node(&id, |node| {
250 node.get_property(text_field).and_then(|value| match value {
251 PropertyValue::String(s) => Some(s.clone()),
252 _ => None,
253 })
254 })
255 .flatten()
256 .map(|text| (id, text))
257 })
258 .collect();
259
260 if docs.is_empty() {
261 return Ok(Vec::new());
262 }
263
264 let index = crate::bm25::Bm25Index::build(docs, crate::bm25::Bm25Params::default());
265 Ok(index.search(query, k))
266 }
267
268 pub fn get_nodes_by_label(&self, label: &str) -> Vec<Node> {
270 self.label_index
271 .get_nodes_by_label(label)
272 .into_iter()
273 .filter_map(|id| self.get_node(&id))
274 .collect()
275 }
276
277 pub fn get_nodes_by_property(&self, key: &str, value: &PropertyValue) -> Vec<Node> {
279 self.property_index
280 .get_nodes_by_property(key, value)
281 .into_iter()
282 .filter_map(|id| self.get_node(&id))
283 .collect()
284 }
285
286 pub fn create_edge(&self, edge: Edge) -> Result<EdgeId> {
290 let id = edge.id.clone();
291
292 if !self.nodes.contains_key(&edge.from) || !self.nodes.contains_key(&edge.to) {
294 return Err(crate::error::GraphError::NodeNotFound(
295 "Source or target node not found".to_string(),
296 ));
297 }
298
299 self.edge_type_index.add_edge(&edge);
301 self.adjacency_index.add_edge(&edge);
302
303 self.edges.insert(id.clone(), edge.clone());
305
306 #[cfg(feature = "storage")]
308 if let Some(storage) = &self.storage {
309 storage.insert_edge(&edge)?;
310 }
311
312 Ok(id)
313 }
314
315 pub fn get_edge(&self, id: impl AsRef<str>) -> Option<Edge> {
317 self.edges.get(id.as_ref()).map(|entry| entry.clone())
318 }
319
320 pub fn delete_edge(&self, id: impl AsRef<str>) -> Result<bool> {
322 if let Some((_, edge)) = self.edges.remove(id.as_ref()) {
323 self.edge_type_index.remove_edge(&edge);
325 self.adjacency_index.remove_edge(&edge);
326
327 #[cfg(feature = "storage")]
329 if let Some(storage) = &self.storage {
330 storage.delete_edge(id.as_ref())?;
331 }
332
333 Ok(true)
334 } else {
335 Ok(false)
336 }
337 }
338
339 pub fn delete_edges_batch(&self, ids: &[impl AsRef<str>]) -> Result<usize> {
341 let mut deleted = 0;
342 let mut edges_to_update = Vec::with_capacity(ids.len());
343
344 for id in ids {
345 let key: &str = id.as_ref();
346 if let Some((_, edge)) = self.edges.remove(key) {
347 edges_to_update.push(edge);
348 deleted += 1;
349 }
350 }
351
352 for edge in &edges_to_update {
353 self.edge_type_index.remove_edge(edge);
354 self.adjacency_index.remove_edge(edge);
355 }
356
357 #[cfg(feature = "storage")]
358 if let Some(storage) = &self.storage {
359 let str_ids = ids.iter().map(|id| id.as_ref()).collect::<Vec<_>>();
360 storage.delete_edges_batch(&str_ids)?;
361 }
362
363 Ok(deleted)
364 }
365
366 pub fn get_edges_by_type(&self, edge_type: &str) -> Vec<Edge> {
368 self.edge_type_index
369 .get_edges_by_type(edge_type)
370 .into_iter()
371 .filter_map(|id| self.get_edge(&id))
372 .collect()
373 }
374
375 pub fn get_outgoing_edges(&self, node_id: &NodeId) -> Vec<Edge> {
377 self.adjacency_index
378 .get_outgoing_edges(node_id)
379 .into_iter()
380 .filter_map(|id| self.get_edge(&id))
381 .collect()
382 }
383
384 pub fn get_incoming_edges(&self, node_id: &NodeId) -> Vec<Edge> {
386 self.adjacency_index
387 .get_incoming_edges(node_id)
388 .into_iter()
389 .filter_map(|id| self.get_edge(&id))
390 .collect()
391 }
392
393 pub fn has_edge(&self, from: &NodeId, to: &NodeId, edge_type: &str) -> bool {
399 self.adjacency_index
400 .get_outgoing_edges(from)
401 .into_iter()
402 .any(|id| {
403 self.edges
404 .get(&id)
405 .is_some_and(|e| e.to == *to && e.edge_type == edge_type)
406 })
407 }
408
409 pub fn get_edges_for_nodes(&self, node_ids: &[NodeId]) -> Vec<Edge> {
411 let mut result = Vec::with_capacity(node_ids.len() * 4);
412 self.adjacency_index
413 .for_each_outgoing_edge(node_ids, |edge_id| {
414 if let Some(edge) = self.edges.get(edge_id.as_str()) {
415 result.push(edge.clone());
416 }
417 });
418
419 result
420 }
421
422 pub fn create_hyperedge(&self, hyperedge: Hyperedge) -> Result<HyperedgeId> {
426 let id = hyperedge.id.clone();
427
428 for node_id in &hyperedge.nodes {
430 if !self.nodes.contains_key(node_id) {
431 return Err(crate::error::GraphError::NodeNotFound(format!(
432 "Node {} not found",
433 node_id
434 )));
435 }
436 }
437
438 self.hyperedge_node_index.add_hyperedge(&hyperedge);
440
441 self.hyperedges.insert(id.clone(), hyperedge.clone());
443
444 #[cfg(feature = "storage")]
446 if let Some(storage) = &self.storage {
447 storage.insert_hyperedge(&hyperedge)?;
448 }
449
450 Ok(id)
451 }
452
453 pub fn get_hyperedge(&self, id: &HyperedgeId) -> Option<Hyperedge> {
455 self.hyperedges.get(id).map(|entry| entry.clone())
456 }
457
458 pub fn get_hyperedges_by_node(&self, node_id: &NodeId) -> Vec<Hyperedge> {
460 self.hyperedge_node_index
461 .get_hyperedges_by_node(node_id)
462 .into_iter()
463 .filter_map(|id| self.get_hyperedge(&id))
464 .collect()
465 }
466
467 pub fn delete_hyperedge(&self, id: &HyperedgeId) -> Result<bool> {
469 if let Some((_, hyperedge)) = self.hyperedges.remove(id) {
470 self.hyperedge_node_index.remove_hyperedge(&hyperedge);
471
472 #[cfg(feature = "storage")]
473 if let Some(storage) = &self.storage {
474 storage.delete_hyperedge(id)?;
475 }
476
477 Ok(true)
478 } else {
479 Ok(false)
480 }
481 }
482
483 pub fn delete_hyperedges_by_node(&self, node_id: &NodeId) -> Result<usize> {
485 let ids: Vec<HyperedgeId> = self.hyperedge_node_index.get_hyperedges_by_node(node_id);
486 let mut deleted = 0;
487 for id in &ids {
488 if self.delete_hyperedge(id)? {
489 deleted += 1;
490 }
491 }
492 Ok(deleted)
493 }
494
495 pub fn node_count(&self) -> usize {
499 self.nodes.len()
500 }
501
502 pub fn edge_count(&self) -> usize {
504 self.edges.len()
505 }
506
507 pub fn hyperedge_count(&self) -> usize {
509 self.hyperedges.len()
510 }
511}
512
513impl Default for GraphDB {
514 fn default() -> Self {
515 Self::new()
516 }
517}
518
519#[cfg(test)]
520mod tests {
521 use super::*;
522 use crate::edge::EdgeBuilder;
523 use crate::hyperedge::HyperedgeBuilder;
524 use crate::node::NodeBuilder;
525 use std::sync::{Arc, Barrier};
526
527 #[test]
528 fn test_graph_creation() {
529 let db = GraphDB::new();
530 assert_eq!(db.node_count(), 0);
531 assert_eq!(db.edge_count(), 0);
532 }
533
534 #[test]
535 fn test_node_operations() {
536 let db = GraphDB::new();
537
538 let node = NodeBuilder::new()
539 .label("Person")
540 .property("name", "Alice")
541 .build();
542
543 let id = db.create_node(node.clone()).unwrap();
544 assert_eq!(db.node_count(), 1);
545
546 let retrieved = db.get_node(&id);
547 assert!(retrieved.is_some());
548
549 let deleted = db.delete_node(&id).unwrap();
550 assert!(deleted);
551 assert_eq!(db.node_count(), 0);
552 }
553
554 #[test]
555 fn test_edge_operations() {
556 let db = GraphDB::new();
557
558 let node1 = NodeBuilder::new().build();
559 let node2 = NodeBuilder::new().build();
560
561 let id1 = db.create_node(node1.clone()).unwrap();
562 let id2 = db.create_node(node2.clone()).unwrap();
563
564 let edge = EdgeBuilder::new(id1.clone(), id2.clone(), "KNOWS")
565 .property("since", 2020i64)
566 .build();
567
568 let edge_id = db.create_edge(edge).unwrap();
569 assert_eq!(db.edge_count(), 1);
570
571 let retrieved = db.get_edge(&edge_id);
572 assert!(retrieved.is_some());
573 }
574
575 #[test]
576 fn test_label_index() {
577 let db = GraphDB::new();
578
579 let node1 = NodeBuilder::new().label("Person").build();
580 let node2 = NodeBuilder::new().label("Person").build();
581 let node3 = NodeBuilder::new().label("Organization").build();
582
583 db.create_node(node1).unwrap();
584 db.create_node(node2).unwrap();
585 db.create_node(node3).unwrap();
586
587 let people = db.get_nodes_by_label("Person");
588 assert_eq!(people.len(), 2);
589
590 let orgs = db.get_nodes_by_label("Organization");
591 assert_eq!(orgs.len(), 1);
592 }
593
594 #[test]
595 fn test_hyperedge_operations() {
596 let db = GraphDB::new();
597
598 let node1 = NodeBuilder::new().build();
599 let node2 = NodeBuilder::new().build();
600 let node3 = NodeBuilder::new().build();
601
602 let id1 = db.create_node(node1).unwrap();
603 let id2 = db.create_node(node2).unwrap();
604 let id3 = db.create_node(node3).unwrap();
605
606 let hyperedge =
607 HyperedgeBuilder::new(vec![id1.clone(), id2.clone(), id3.clone()], "MEETING")
608 .description("Team meeting")
609 .build();
610
611 let hedge_id = db.create_hyperedge(hyperedge).unwrap();
612 assert_eq!(db.hyperedge_count(), 1);
613
614 let hedges = db.get_hyperedges_by_node(&id1);
615 assert_eq!(hedges.len(), 1);
616 }
617
618 #[test]
619 fn test_update_node() {
620 let db = GraphDB::new();
621
622 let node = NodeBuilder::new()
623 .id("mem-001")
624 .label("Memory")
625 .property("content", "original content")
626 .property("status", "active")
627 .build();
628
629 db.create_node(node).unwrap();
630
631 let updated = db
633 .update_node("mem-001", |n| {
634 n.set_property("status", PropertyValue::from("deprecated"));
635 n.set_property("deprecated_at", PropertyValue::from("2026-07-12T12:00:00Z"));
636 })
637 .unwrap();
638 assert!(updated);
639
640 let retrieved = db.get_node("mem-001").unwrap();
641 assert_eq!(
642 retrieved.get_property("status").unwrap(),
643 &PropertyValue::from("deprecated")
644 );
645 assert!(retrieved.get_property("deprecated_at").is_some());
646
647 assert!(db
648 .get_nodes_by_property("status", &PropertyValue::from("active"))
649 .is_empty());
650 assert_eq!(
651 db.get_nodes_by_property("status", &PropertyValue::from("deprecated"))
652 .len(),
653 1
654 );
655 }
656
657 #[test]
658 fn test_update_node_refreshes_label_and_property_indexes() {
659 let db = GraphDB::new();
660 db.create_node(
661 NodeBuilder::new()
662 .id("indexed")
663 .label("OldLabel")
664 .property("state", "old")
665 .property("removed", true)
666 .build(),
667 )
668 .unwrap();
669
670 db.update_node("indexed", |node| {
671 node.remove_label("OldLabel");
672 node.add_label("NewLabel");
673 node.set_property("state", PropertyValue::from("new"));
674 node.properties.remove("removed");
675 })
676 .unwrap();
677
678 assert!(db.get_nodes_by_label("OldLabel").is_empty());
679 assert_eq!(db.get_nodes_by_label("NewLabel").len(), 1);
680 assert!(db
681 .get_nodes_by_property("state", &PropertyValue::from("old"))
682 .is_empty());
683 assert_eq!(
684 db.get_nodes_by_property("state", &PropertyValue::from("new"))
685 .len(),
686 1
687 );
688 assert!(db
689 .get_nodes_by_property("removed", &PropertyValue::from(true))
690 .is_empty());
691 }
692
693 #[test]
694 fn test_update_node_rejects_id_changes_without_side_effects() {
695 let db = GraphDB::new();
696 db.create_node(NodeBuilder::new().id("original").label("Old").build())
697 .unwrap();
698
699 let error = db
700 .update_node("original", |node| {
701 node.id = "replacement".to_string();
702 node.add_label("New");
703 })
704 .unwrap_err();
705
706 assert!(matches!(
707 error,
708 crate::error::GraphError::ConstraintViolation(_)
709 ));
710 assert!(db.get_node("replacement").is_none());
711 assert!(db.get_node("original").unwrap().has_label("Old"));
712 assert!(db.get_nodes_by_label("New").is_empty());
713 }
714
715 #[test]
716 fn test_concurrent_updates_do_not_lose_writes() {
717 const THREADS: usize = 8;
718 const UPDATES_PER_THREAD: usize = 100;
719
720 let db = Arc::new(GraphDB::new());
721 db.create_node(
722 NodeBuilder::new()
723 .id("counter")
724 .property("value", 0_i64)
725 .build(),
726 )
727 .unwrap();
728 let barrier = Arc::new(Barrier::new(THREADS));
729 let mut handles = Vec::new();
730
731 for _ in 0..THREADS {
732 let db = Arc::clone(&db);
733 let barrier = Arc::clone(&barrier);
734 handles.push(std::thread::spawn(move || {
735 barrier.wait();
736 for _ in 0..UPDATES_PER_THREAD {
737 db.update_node("counter", |node| {
738 let value = match node.get_property("value") {
739 Some(PropertyValue::Integer(value)) => *value,
740 _ => panic!("counter property is missing"),
741 };
742 node.set_property("value", PropertyValue::Integer(value + 1));
743 })
744 .unwrap();
745 }
746 }));
747 }
748
749 for handle in handles {
750 handle.join().unwrap();
751 }
752 assert_eq!(
753 db.get_node("counter")
754 .unwrap()
755 .get_property("value")
756 .cloned(),
757 Some(PropertyValue::Integer(
758 (THREADS * UPDATES_PER_THREAD) as i64
759 ))
760 );
761 }
762
763 #[test]
764 fn test_update_node_not_found() {
765 let db = GraphDB::new();
766 let result = db.update_node("nonexistent", |_| {}).unwrap();
767 assert!(!result);
768 }
769
770 #[test]
771 fn test_keyword_search() {
772 let db = GraphDB::new();
773
774 let docs = vec![
775 ("mem-1", "the quick brown fox jumps over the lazy dog"),
776 ("mem-2", "machine learning models for vector search"),
777 ("mem-3", "vector databases enable semantic search at scale"),
778 ("mem-4", "a recipe for italian pasta with tomato sauce"),
779 ];
780
781 for (id, text) in docs {
782 let node = NodeBuilder::new()
783 .id(id)
784 .label("Memory")
785 .property("content", text)
786 .build();
787 db.create_node(node).unwrap();
788 }
789
790 let hits = db
791 .keyword_search("Memory", "content", "vector search", 4)
792 .unwrap();
793
794 assert!(!hits.is_empty());
795 assert!(hits[0].0 == "mem-2" || hits[0].0 == "mem-3");
797 assert!(hits.iter().all(|(id, _)| id != "mem-4") || hits.last().unwrap().0 == "mem-4");
798 }
799
800 #[test]
801 fn test_keyword_search_empty_label() {
802 let db = GraphDB::new();
803 let hits = db
804 .keyword_search("Nonexistent", "content", "anything", 5)
805 .unwrap();
806 assert!(hits.is_empty());
807 }
808
809 #[test]
810 fn test_keyword_search_ignores_other_labels_and_non_string_fields() {
811 let db = GraphDB::new();
812 for node in [
813 NodeBuilder::new()
814 .id("wanted")
815 .label("Memory")
816 .property("content", "unique needle")
817 .build(),
818 NodeBuilder::new()
819 .id("wrong-label")
820 .label("Other")
821 .property("content", "unique needle")
822 .build(),
823 NodeBuilder::new()
824 .id("wrong-type")
825 .label("Memory")
826 .property("content", 42_i64)
827 .build(),
828 ] {
829 db.create_node(node).unwrap();
830 }
831
832 let hits = db
833 .keyword_search("Memory", "content", "needle", 10)
834 .unwrap();
835 assert_eq!(hits.len(), 1);
836 assert_eq!(hits[0].0, "wanted");
837 assert!(db
838 .keyword_search("Memory", "content", "needle", 0)
839 .unwrap()
840 .is_empty());
841 }
842
843 #[cfg(feature = "storage")]
844 #[test]
845 fn test_update_node_persists() {
846 let temp = tempfile::tempdir().unwrap();
847 let path = temp.path().join("graph.redb");
848
849 {
850 let db = GraphDB::with_storage(&path).unwrap();
851 db.create_node(
852 NodeBuilder::new()
853 .id("persistent")
854 .label("Old")
855 .property("state", "old")
856 .build(),
857 )
858 .unwrap();
859 db.update_node("persistent", |node| {
860 node.remove_label("Old");
861 node.add_label("New");
862 node.set_property("state", PropertyValue::from("new"));
863 })
864 .unwrap();
865 }
866
867 let reopened = GraphDB::with_storage(&path).unwrap();
868 let node = reopened.get_node("persistent").unwrap();
869 assert!(node.has_label("New"));
870 assert_eq!(
871 node.get_property("state"),
872 Some(&PropertyValue::from("new"))
873 );
874 assert_eq!(reopened.get_nodes_by_label("New").len(), 1);
875 assert_eq!(
876 reopened
877 .get_nodes_by_property("state", &PropertyValue::from("new"))
878 .len(),
879 1
880 );
881 }
882}