arbor_graph/
store.rs

1use crate::builder::GraphBuilder;
2use crate::graph::ArborGraph;
3use arbor_core::CodeNode;
4use sled::{Batch, Db};
5use std::path::Path;
6use thiserror::Error;
7
8#[derive(Error, Debug)]
9pub enum StoreError {
10    #[error("Database error: {0}")]
11    Sled(#[from] sled::Error),
12    #[error("Serialization error: {0}")]
13    Bincode(#[from] bincode::Error),
14    #[error("Corrupted data: {0}")]
15    Corrupted(String),
16}
17
18pub struct GraphStore {
19    db: Db,
20}
21
22impl GraphStore {
23    /// Opens or creates a graph store at the specified path.
24    pub fn open<P: AsRef<Path>>(path: P) -> Result<Self, StoreError> {
25        let db = sled::open(path)?;
26        Ok(Self { db })
27    }
28
29    /// Updates the nodes for a specific file.
30    ///
31    /// This operation is atomic: it removes old nodes associated with the file
32    /// and inserts the new ones.
33    pub fn update_file(&self, file_path: &str, nodes: &[CodeNode]) -> Result<(), StoreError> {
34        let file_key = format!("f:{}", file_path);
35        let mut batch = Batch::default();
36
37        // 1. Get old nodes for this file
38        if let Some(old_bytes) = self.db.get(&file_key)? {
39            let old_ids: Vec<String> = bincode::deserialize(&old_bytes)?;
40            for id in old_ids {
41                batch.remove(format!("n:{}", id).as_bytes());
42            }
43        }
44
45        // 2. Insert new nodes
46        let mut new_ids = Vec::with_capacity(nodes.len());
47        for node in nodes {
48            let node_key = format!("n:{}", node.id);
49            let bytes = bincode::serialize(node)?;
50            batch.insert(node_key.as_bytes(), bytes);
51            new_ids.push(node.id.clone());
52        }
53
54        // 3. Update file index
55        let index_bytes = bincode::serialize(&new_ids)?;
56        batch.insert(file_key.as_bytes(), index_bytes);
57
58        // 4. Commit batch
59        self.db.apply_batch(batch)?;
60        self.db.flush()?; // flushing optional for perf, but good for safety
61        Ok(())
62    }
63
64    /// Loads the entire graph from the store.
65    ///
66    /// This iterates over all stored nodes and reconstructs the ArborGraph
67    /// using the GraphBuilder (which re-links edges).
68    pub fn load_graph(&self) -> Result<ArborGraph, StoreError> {
69        let mut builder = GraphBuilder::new();
70        let mut nodes = Vec::new();
71
72        // Iterate over all keys starting with "n:"
73        let prefix = b"n:";
74        for item in self.db.scan_prefix(prefix) {
75            let (_key, value) = item?;
76            let node: CodeNode = bincode::deserialize(&value)?;
77            nodes.push(node);
78        }
79
80        if nodes.is_empty() {
81            // Return empty graph
82            return Ok(ArborGraph::new());
83        }
84
85        // Reconstruct graph
86        builder.add_nodes(nodes);
87        // resolve_edges() is called by build()
88        let graph = builder.build();
89
90        Ok(graph)
91    }
92
93    /// Clears the stored graph.
94    pub fn clear(&self) -> Result<(), StoreError> {
95        self.db.clear()?;
96        self.db.flush()?;
97        Ok(())
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104    use arbor_core::NodeKind;
105    use tempfile::tempdir;
106
107    #[test]
108    fn test_incremental_updates() {
109        let dir = tempdir().unwrap();
110        let store = GraphStore::open(dir.path()).unwrap();
111
112        let node1 = CodeNode::new("foo", "foo", NodeKind::Function, "test.rs");
113        let node2 = CodeNode::new("bar", "bar", NodeKind::Function, "test.rs");
114
115        // Initial update
116        store
117            .update_file("test.rs", &[node1.clone(), node2.clone()])
118            .unwrap();
119
120        // Verify load
121        let graph = store.load_graph().unwrap();
122        assert_eq!(graph.node_count(), 2);
123
124        // Update with one node removed
125        store.update_file("test.rs", &[node1.clone()]).unwrap();
126        let graph2 = store.load_graph().unwrap();
127        assert_eq!(graph2.node_count(), 1);
128        assert!(graph2.find_by_name("foo").len() > 0);
129        assert!(graph2.find_by_name("bar").is_empty());
130    }
131}
arbor_graph/store.rs

arbor_graph/
store.rs
