use std::fs;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use anyhow::Result;
pub struct GraphStore {
root: PathBuf,
mem: Option<Arc<dashmap::DashMap<(String, String), std::collections::HashSet<String>>>>,
}
impl GraphStore {
pub fn new(db_root: &Path) -> Result<Self> {
let root = db_root.join("graph");
fs::create_dir_all(&root)?;
Ok(Self { root, mem: None })
}
pub fn in_memory() -> Self {
Self {
root: PathBuf::from(":memory:"),
mem: Some(Arc::new(dashmap::DashMap::new())),
}
}
fn edge_path(&self, from: &str, edge_type: &str, to: &str) -> PathBuf {
self.root.join(from).join(edge_type).join(to)
}
pub fn add_edge(&self, from: &str, edge_type: &str, to: &str) -> Result<()> {
if let Some(ref mem) = self.mem {
mem.entry((from.to_string(), edge_type.to_string()))
.or_default()
.insert(to.to_string());
return Ok(());
}
let path = self.edge_path(from, edge_type, to);
fs::create_dir_all(path.parent().unwrap())?;
if !path.exists() {
fs::write(&path, b"")?;
}
Ok(())
}
pub fn outgoing(&self, from: &str, edge_type: &str) -> Vec<String> {
if let Some(ref mem) = self.mem {
return mem.get(&(from.to_string(), edge_type.to_string()))
.map(|s| s.iter().cloned().collect())
.unwrap_or_default();
}
let dir = self.root.join(from).join(edge_type);
fs::read_dir(&dir)
.into_iter()
.flatten()
.filter_map(|e| e.ok())
.map(|e| e.file_name().to_string_lossy().to_string())
.collect()
}
pub fn incoming(&self, to: &str, reverse_edge_type: &str) -> Vec<String> {
self.outgoing(to, reverse_edge_type)
}
pub fn trace(
&self,
start: &str,
edge_type: &str,
reverse: bool,
limit: usize,
) -> Vec<String> {
let mut result = Vec::new();
let mut queue = vec![start.to_string()];
let mut seen = std::collections::HashSet::new();
seen.insert(start.to_string());
while !queue.is_empty() && result.len() < limit {
let current = queue.remove(0);
result.push(current.clone());
let next_hashes = if reverse {
let rev_type = format!("{}_rev", edge_type);
self.outgoing(¤t, &rev_type)
} else {
self.outgoing(¤t, edge_type)
};
for next in next_hashes {
if !seen.contains(&next) {
seen.insert(next.clone());
queue.push(next);
}
}
}
result
}
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::tempdir;
#[test]
fn add_and_traverse_edge() {
let dir = tempdir().unwrap();
let g = GraphStore::new(dir.path()).unwrap();
g.add_edge("hash_c", "caused_by", "hash_b").unwrap();
g.add_edge("hash_b", "caused_by", "hash_a").unwrap();
let trace = g.trace("hash_c", "caused_by", false, 10);
assert_eq!(trace, vec!["hash_c", "hash_b", "hash_a"]);
}
#[test]
fn idempotent_edge() {
let dir = tempdir().unwrap();
let g = GraphStore::new(dir.path()).unwrap();
g.add_edge("a", "caused_by", "b").unwrap();
g.add_edge("a", "caused_by", "b").unwrap(); assert_eq!(g.outgoing("a", "caused_by").len(), 1);
}
}