use crate::edge::Edge;
use crate::error::Result;
use crate::hyperedge::{Hyperedge, HyperedgeId};
use crate::index::{AdjacencyIndex, EdgeTypeIndex, HyperedgeNodeIndex, LabelIndex, PropertyIndex};
use crate::node::Node;
#[cfg(feature = "storage")]
use crate::storage::GraphStorage;
use crate::types::{EdgeId, NodeId, PropertyValue};
use dashmap::DashMap;
#[cfg(feature = "storage")]
use std::path::Path;
use std::sync::Arc;
pub struct GraphDB {
nodes: Arc<DashMap<NodeId, Node>>,
edges: Arc<DashMap<EdgeId, Edge>>,
hyperedges: Arc<DashMap<HyperedgeId, Hyperedge>>,
label_index: LabelIndex,
property_index: PropertyIndex,
edge_type_index: EdgeTypeIndex,
adjacency_index: AdjacencyIndex,
hyperedge_node_index: HyperedgeNodeIndex,
#[cfg(feature = "storage")]
storage: Option<GraphStorage>,
}
impl GraphDB {
pub fn new() -> Self {
Self {
nodes: Arc::new(DashMap::new()),
edges: Arc::new(DashMap::new()),
hyperedges: Arc::new(DashMap::new()),
label_index: LabelIndex::new(),
property_index: PropertyIndex::new(),
edge_type_index: EdgeTypeIndex::new(),
adjacency_index: AdjacencyIndex::new(),
hyperedge_node_index: HyperedgeNodeIndex::new(),
#[cfg(feature = "storage")]
storage: None,
}
}
#[cfg(feature = "storage")]
pub fn with_storage<P: AsRef<Path>>(path: P) -> anyhow::Result<Self> {
let storage = GraphStorage::new(path)?;
let mut db = Self::new();
db.storage = Some(storage);
db.load_from_storage()?;
Ok(db)
}
#[cfg(feature = "storage")]
fn load_from_storage(&mut self) -> anyhow::Result<()> {
if let Some(storage) = &self.storage {
for node_id in storage.all_node_ids()? {
if let Some(node) = storage.get_node(&node_id)? {
self.nodes.insert(node_id.clone(), node.clone());
self.label_index.add_node(&node);
self.property_index.add_node(&node);
}
}
for edge_id in storage.all_edge_ids()? {
if let Some(edge) = storage.get_edge(&edge_id)? {
self.edges.insert(edge_id.clone(), edge.clone());
self.edge_type_index.add_edge(&edge);
self.adjacency_index.add_edge(&edge);
}
}
for hyperedge_id in storage.all_hyperedge_ids()? {
if let Some(hyperedge) = storage.get_hyperedge(&hyperedge_id)? {
self.hyperedges
.insert(hyperedge_id.clone(), hyperedge.clone());
self.hyperedge_node_index.add_hyperedge(&hyperedge);
}
}
}
Ok(())
}
pub fn create_node(&self, node: Node) -> Result<NodeId> {
let id = node.id.clone();
self.label_index.add_node(&node);
self.property_index.add_node(&node);
self.nodes.insert(id.clone(), node.clone());
#[cfg(feature = "storage")]
if let Some(storage) = &self.storage {
storage.insert_node(&node)?;
}
Ok(id)
}
pub fn get_node(&self, id: impl AsRef<str>) -> Option<Node> {
self.nodes.get(id.as_ref()).map(|entry| entry.clone())
}
pub fn with_node<R>(&self, id: &str, f: impl FnOnce(&Node) -> R) -> Option<R> {
self.nodes.get(id).map(|entry| f(entry.value()))
}
pub fn node_ids_by_label(&self, label: &str) -> Vec<NodeId> {
self.label_index.get_nodes_by_label(label)
}
pub fn delete_node(&self, id: impl AsRef<str>) -> Result<bool> {
if let Some((_, node)) = self.nodes.remove(id.as_ref()) {
self.label_index.remove_node(&node);
self.property_index.remove_node(&node);
#[cfg(feature = "storage")]
if let Some(storage) = &self.storage {
storage.delete_node(id.as_ref())?;
}
Ok(true)
} else {
Ok(false)
}
}
pub fn update_node<F>(&self, id: impl AsRef<str>, f: F) -> Result<bool>
where
F: FnOnce(&mut Node),
{
let id_ref = id.as_ref();
let Some(mut entry) = self.nodes.get_mut(id_ref) else {
return Ok(false);
};
let old_node = entry.value().clone();
let mut new_node = old_node.clone();
f(&mut new_node);
if new_node.id != old_node.id {
return Err(crate::error::GraphError::ConstraintViolation(
"A node's ID cannot be changed by update_node".to_string(),
));
}
#[cfg(feature = "storage")]
if let Some(storage) = &self.storage {
storage.insert_node(&new_node)?;
}
self.label_index.remove_node(&old_node);
self.property_index.remove_node(&old_node);
*entry.value_mut() = new_node.clone();
self.label_index.add_node(&new_node);
self.property_index.add_node(&new_node);
Ok(true)
}
pub fn keyword_search(
&self,
label: &str,
text_field: &str,
query: &str,
k: usize,
) -> Result<Vec<(NodeId, f32)>> {
let docs: Vec<(NodeId, String)> = self
.node_ids_by_label(label)
.into_iter()
.filter_map(|id| {
self.with_node(&id, |node| {
node.get_property(text_field).and_then(|value| match value {
PropertyValue::String(s) => Some(s.clone()),
_ => None,
})
})
.flatten()
.map(|text| (id, text))
})
.collect();
if docs.is_empty() {
return Ok(Vec::new());
}
let index = crate::bm25::Bm25Index::build(docs, crate::bm25::Bm25Params::default());
Ok(index.search(query, k))
}
pub fn get_nodes_by_label(&self, label: &str) -> Vec<Node> {
self.label_index
.get_nodes_by_label(label)
.into_iter()
.filter_map(|id| self.get_node(&id))
.collect()
}
pub fn get_nodes_by_property(&self, key: &str, value: &PropertyValue) -> Vec<Node> {
self.property_index
.get_nodes_by_property(key, value)
.into_iter()
.filter_map(|id| self.get_node(&id))
.collect()
}
pub fn create_edge(&self, edge: Edge) -> Result<EdgeId> {
let id = edge.id.clone();
if !self.nodes.contains_key(&edge.from) || !self.nodes.contains_key(&edge.to) {
return Err(crate::error::GraphError::NodeNotFound(
"Source or target node not found".to_string(),
));
}
self.edge_type_index.add_edge(&edge);
self.adjacency_index.add_edge(&edge);
self.edges.insert(id.clone(), edge.clone());
#[cfg(feature = "storage")]
if let Some(storage) = &self.storage {
storage.insert_edge(&edge)?;
}
Ok(id)
}
pub fn get_edge(&self, id: impl AsRef<str>) -> Option<Edge> {
self.edges.get(id.as_ref()).map(|entry| entry.clone())
}
pub fn delete_edge(&self, id: impl AsRef<str>) -> Result<bool> {
if let Some((_, edge)) = self.edges.remove(id.as_ref()) {
self.edge_type_index.remove_edge(&edge);
self.adjacency_index.remove_edge(&edge);
#[cfg(feature = "storage")]
if let Some(storage) = &self.storage {
storage.delete_edge(id.as_ref())?;
}
Ok(true)
} else {
Ok(false)
}
}
pub fn delete_edges_batch(&self, ids: &[impl AsRef<str>]) -> Result<usize> {
let mut deleted = 0;
let mut edges_to_update = Vec::with_capacity(ids.len());
for id in ids {
let key: &str = id.as_ref();
if let Some((_, edge)) = self.edges.remove(key) {
edges_to_update.push(edge);
deleted += 1;
}
}
for edge in &edges_to_update {
self.edge_type_index.remove_edge(edge);
self.adjacency_index.remove_edge(edge);
}
#[cfg(feature = "storage")]
if let Some(storage) = &self.storage {
let str_ids = ids.iter().map(|id| id.as_ref()).collect::<Vec<_>>();
storage.delete_edges_batch(&str_ids)?;
}
Ok(deleted)
}
pub fn get_edges_by_type(&self, edge_type: &str) -> Vec<Edge> {
self.edge_type_index
.get_edges_by_type(edge_type)
.into_iter()
.filter_map(|id| self.get_edge(&id))
.collect()
}
pub fn get_outgoing_edges(&self, node_id: &NodeId) -> Vec<Edge> {
self.adjacency_index
.get_outgoing_edges(node_id)
.into_iter()
.filter_map(|id| self.get_edge(&id))
.collect()
}
pub fn get_incoming_edges(&self, node_id: &NodeId) -> Vec<Edge> {
self.adjacency_index
.get_incoming_edges(node_id)
.into_iter()
.filter_map(|id| self.get_edge(&id))
.collect()
}
pub fn has_edge(&self, from: &NodeId, to: &NodeId, edge_type: &str) -> bool {
self.adjacency_index
.get_outgoing_edges(from)
.into_iter()
.any(|id| {
self.edges
.get(&id)
.is_some_and(|e| e.to == *to && e.edge_type == edge_type)
})
}
pub fn get_edges_for_nodes(&self, node_ids: &[NodeId]) -> Vec<Edge> {
let mut result = Vec::with_capacity(node_ids.len() * 4);
self.adjacency_index
.for_each_outgoing_edge(node_ids, |edge_id| {
if let Some(edge) = self.edges.get(edge_id.as_str()) {
result.push(edge.clone());
}
});
result
}
pub fn create_hyperedge(&self, hyperedge: Hyperedge) -> Result<HyperedgeId> {
let id = hyperedge.id.clone();
for node_id in &hyperedge.nodes {
if !self.nodes.contains_key(node_id) {
return Err(crate::error::GraphError::NodeNotFound(format!(
"Node {} not found",
node_id
)));
}
}
self.hyperedge_node_index.add_hyperedge(&hyperedge);
self.hyperedges.insert(id.clone(), hyperedge.clone());
#[cfg(feature = "storage")]
if let Some(storage) = &self.storage {
storage.insert_hyperedge(&hyperedge)?;
}
Ok(id)
}
pub fn get_hyperedge(&self, id: &HyperedgeId) -> Option<Hyperedge> {
self.hyperedges.get(id).map(|entry| entry.clone())
}
pub fn get_hyperedges_by_node(&self, node_id: &NodeId) -> Vec<Hyperedge> {
self.hyperedge_node_index
.get_hyperedges_by_node(node_id)
.into_iter()
.filter_map(|id| self.get_hyperedge(&id))
.collect()
}
pub fn delete_hyperedge(&self, id: &HyperedgeId) -> Result<bool> {
if let Some((_, hyperedge)) = self.hyperedges.remove(id) {
self.hyperedge_node_index.remove_hyperedge(&hyperedge);
#[cfg(feature = "storage")]
if let Some(storage) = &self.storage {
storage.delete_hyperedge(id)?;
}
Ok(true)
} else {
Ok(false)
}
}
pub fn delete_hyperedges_by_node(&self, node_id: &NodeId) -> Result<usize> {
let ids: Vec<HyperedgeId> = self.hyperedge_node_index.get_hyperedges_by_node(node_id);
let mut deleted = 0;
for id in &ids {
if self.delete_hyperedge(id)? {
deleted += 1;
}
}
Ok(deleted)
}
pub fn node_count(&self) -> usize {
self.nodes.len()
}
pub fn edge_count(&self) -> usize {
self.edges.len()
}
pub fn hyperedge_count(&self) -> usize {
self.hyperedges.len()
}
}
impl Default for GraphDB {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::edge::EdgeBuilder;
use crate::hyperedge::HyperedgeBuilder;
use crate::node::NodeBuilder;
use std::sync::{Arc, Barrier};
#[test]
fn test_graph_creation() {
let db = GraphDB::new();
assert_eq!(db.node_count(), 0);
assert_eq!(db.edge_count(), 0);
}
#[test]
fn test_node_operations() {
let db = GraphDB::new();
let node = NodeBuilder::new()
.label("Person")
.property("name", "Alice")
.build();
let id = db.create_node(node.clone()).unwrap();
assert_eq!(db.node_count(), 1);
let retrieved = db.get_node(&id);
assert!(retrieved.is_some());
let deleted = db.delete_node(&id).unwrap();
assert!(deleted);
assert_eq!(db.node_count(), 0);
}
#[test]
fn test_edge_operations() {
let db = GraphDB::new();
let node1 = NodeBuilder::new().build();
let node2 = NodeBuilder::new().build();
let id1 = db.create_node(node1.clone()).unwrap();
let id2 = db.create_node(node2.clone()).unwrap();
let edge = EdgeBuilder::new(id1.clone(), id2.clone(), "KNOWS")
.property("since", 2020i64)
.build();
let edge_id = db.create_edge(edge).unwrap();
assert_eq!(db.edge_count(), 1);
let retrieved = db.get_edge(&edge_id);
assert!(retrieved.is_some());
}
#[test]
fn test_label_index() {
let db = GraphDB::new();
let node1 = NodeBuilder::new().label("Person").build();
let node2 = NodeBuilder::new().label("Person").build();
let node3 = NodeBuilder::new().label("Organization").build();
db.create_node(node1).unwrap();
db.create_node(node2).unwrap();
db.create_node(node3).unwrap();
let people = db.get_nodes_by_label("Person");
assert_eq!(people.len(), 2);
let orgs = db.get_nodes_by_label("Organization");
assert_eq!(orgs.len(), 1);
}
#[test]
fn test_hyperedge_operations() {
let db = GraphDB::new();
let node1 = NodeBuilder::new().build();
let node2 = NodeBuilder::new().build();
let node3 = NodeBuilder::new().build();
let id1 = db.create_node(node1).unwrap();
let id2 = db.create_node(node2).unwrap();
let id3 = db.create_node(node3).unwrap();
let hyperedge =
HyperedgeBuilder::new(vec![id1.clone(), id2.clone(), id3.clone()], "MEETING")
.description("Team meeting")
.build();
let hedge_id = db.create_hyperedge(hyperedge).unwrap();
assert_eq!(db.hyperedge_count(), 1);
let hedges = db.get_hyperedges_by_node(&id1);
assert_eq!(hedges.len(), 1);
}
#[test]
fn test_update_node() {
let db = GraphDB::new();
let node = NodeBuilder::new()
.id("mem-001")
.label("Memory")
.property("content", "original content")
.property("status", "active")
.build();
db.create_node(node).unwrap();
let updated = db
.update_node("mem-001", |n| {
n.set_property("status", PropertyValue::from("deprecated"));
n.set_property("deprecated_at", PropertyValue::from("2026-07-12T12:00:00Z"));
})
.unwrap();
assert!(updated);
let retrieved = db.get_node("mem-001").unwrap();
assert_eq!(
retrieved.get_property("status").unwrap(),
&PropertyValue::from("deprecated")
);
assert!(retrieved.get_property("deprecated_at").is_some());
assert!(db
.get_nodes_by_property("status", &PropertyValue::from("active"))
.is_empty());
assert_eq!(
db.get_nodes_by_property("status", &PropertyValue::from("deprecated"))
.len(),
1
);
}
#[test]
fn test_update_node_refreshes_label_and_property_indexes() {
let db = GraphDB::new();
db.create_node(
NodeBuilder::new()
.id("indexed")
.label("OldLabel")
.property("state", "old")
.property("removed", true)
.build(),
)
.unwrap();
db.update_node("indexed", |node| {
node.remove_label("OldLabel");
node.add_label("NewLabel");
node.set_property("state", PropertyValue::from("new"));
node.properties.remove("removed");
})
.unwrap();
assert!(db.get_nodes_by_label("OldLabel").is_empty());
assert_eq!(db.get_nodes_by_label("NewLabel").len(), 1);
assert!(db
.get_nodes_by_property("state", &PropertyValue::from("old"))
.is_empty());
assert_eq!(
db.get_nodes_by_property("state", &PropertyValue::from("new"))
.len(),
1
);
assert!(db
.get_nodes_by_property("removed", &PropertyValue::from(true))
.is_empty());
}
#[test]
fn test_update_node_rejects_id_changes_without_side_effects() {
let db = GraphDB::new();
db.create_node(NodeBuilder::new().id("original").label("Old").build())
.unwrap();
let error = db
.update_node("original", |node| {
node.id = "replacement".to_string();
node.add_label("New");
})
.unwrap_err();
assert!(matches!(
error,
crate::error::GraphError::ConstraintViolation(_)
));
assert!(db.get_node("replacement").is_none());
assert!(db.get_node("original").unwrap().has_label("Old"));
assert!(db.get_nodes_by_label("New").is_empty());
}
#[test]
fn test_concurrent_updates_do_not_lose_writes() {
const THREADS: usize = 8;
const UPDATES_PER_THREAD: usize = 100;
let db = Arc::new(GraphDB::new());
db.create_node(
NodeBuilder::new()
.id("counter")
.property("value", 0_i64)
.build(),
)
.unwrap();
let barrier = Arc::new(Barrier::new(THREADS));
let mut handles = Vec::new();
for _ in 0..THREADS {
let db = Arc::clone(&db);
let barrier = Arc::clone(&barrier);
handles.push(std::thread::spawn(move || {
barrier.wait();
for _ in 0..UPDATES_PER_THREAD {
db.update_node("counter", |node| {
let value = match node.get_property("value") {
Some(PropertyValue::Integer(value)) => *value,
_ => panic!("counter property is missing"),
};
node.set_property("value", PropertyValue::Integer(value + 1));
})
.unwrap();
}
}));
}
for handle in handles {
handle.join().unwrap();
}
assert_eq!(
db.get_node("counter")
.unwrap()
.get_property("value")
.cloned(),
Some(PropertyValue::Integer(
(THREADS * UPDATES_PER_THREAD) as i64
))
);
}
#[test]
fn test_update_node_not_found() {
let db = GraphDB::new();
let result = db.update_node("nonexistent", |_| {}).unwrap();
assert!(!result);
}
#[test]
fn test_keyword_search() {
let db = GraphDB::new();
let docs = vec![
("mem-1", "the quick brown fox jumps over the lazy dog"),
("mem-2", "machine learning models for vector search"),
("mem-3", "vector databases enable semantic search at scale"),
("mem-4", "a recipe for italian pasta with tomato sauce"),
];
for (id, text) in docs {
let node = NodeBuilder::new()
.id(id)
.label("Memory")
.property("content", text)
.build();
db.create_node(node).unwrap();
}
let hits = db
.keyword_search("Memory", "content", "vector search", 4)
.unwrap();
assert!(!hits.is_empty());
assert!(hits[0].0 == "mem-2" || hits[0].0 == "mem-3");
assert!(hits.iter().all(|(id, _)| id != "mem-4") || hits.last().unwrap().0 == "mem-4");
}
#[test]
fn test_keyword_search_empty_label() {
let db = GraphDB::new();
let hits = db
.keyword_search("Nonexistent", "content", "anything", 5)
.unwrap();
assert!(hits.is_empty());
}
#[test]
fn test_keyword_search_ignores_other_labels_and_non_string_fields() {
let db = GraphDB::new();
for node in [
NodeBuilder::new()
.id("wanted")
.label("Memory")
.property("content", "unique needle")
.build(),
NodeBuilder::new()
.id("wrong-label")
.label("Other")
.property("content", "unique needle")
.build(),
NodeBuilder::new()
.id("wrong-type")
.label("Memory")
.property("content", 42_i64)
.build(),
] {
db.create_node(node).unwrap();
}
let hits = db
.keyword_search("Memory", "content", "needle", 10)
.unwrap();
assert_eq!(hits.len(), 1);
assert_eq!(hits[0].0, "wanted");
assert!(db
.keyword_search("Memory", "content", "needle", 0)
.unwrap()
.is_empty());
}
#[cfg(feature = "storage")]
#[test]
fn test_update_node_persists() {
let temp = tempfile::tempdir().unwrap();
let path = temp.path().join("graph.redb");
{
let db = GraphDB::with_storage(&path).unwrap();
db.create_node(
NodeBuilder::new()
.id("persistent")
.label("Old")
.property("state", "old")
.build(),
)
.unwrap();
db.update_node("persistent", |node| {
node.remove_label("Old");
node.add_label("New");
node.set_property("state", PropertyValue::from("new"));
})
.unwrap();
}
let reopened = GraphDB::with_storage(&path).unwrap();
let node = reopened.get_node("persistent").unwrap();
assert!(node.has_label("New"));
assert_eq!(
node.get_property("state"),
Some(&PropertyValue::from("new"))
);
assert_eq!(reopened.get_nodes_by_label("New").len(), 1);
assert_eq!(
reopened
.get_nodes_by_property("state", &PropertyValue::from("new"))
.len(),
1
);
}
}