use std::io::{Read, Write};
use std::path::{Path, PathBuf};
use flate2::read::GzDecoder;
use flate2::write::GzEncoder;
use flate2::Compression;
use crate::error::VasariError;
use crate::schema::{Attribution, AttributionTarget, Node, NodeId};
#[derive(Debug)]
pub struct ObjectStore {
root: PathBuf,
}
impl ObjectStore {
pub const FORMAT_VERSION: &'static str = "2";
pub fn open(repo_root: &Path) -> Result<Self, VasariError> {
let root = repo_root.join(".vasari");
let preexisting = root.join("objects").exists();
let format_path = root.join("format");
std::fs::create_dir_all(root.join("objects"))?;
std::fs::create_dir_all(root.join("index").join("targets"))?;
std::fs::create_dir_all(root.join("refs"))?;
if !root.join("HEAD").exists() {
std::fs::write(root.join("HEAD"), "")?;
}
if preexisting {
let found = std::fs::read_to_string(&format_path)
.ok()
.map(|s| s.trim().to_string())
.filter(|s| !s.is_empty())
.unwrap_or_else(|| "1 (pre-v2)".to_string());
if found != Self::FORMAT_VERSION {
return Err(VasariError::IncompatibleStore {
found,
expected: Self::FORMAT_VERSION.to_string(),
});
}
} else {
std::fs::write(&format_path, Self::FORMAT_VERSION)?;
}
Ok(Self { root })
}
pub fn put(&self, node: &Node) -> Result<NodeId, VasariError> {
let id = node.id();
let (dir, file) = self.object_path(id)?;
if file.exists() {
return Ok(id.clone());
}
std::fs::create_dir_all(&dir)?;
let json = serde_json::to_vec(node)?;
let mut encoder = GzEncoder::new(Vec::new(), Compression::fast());
encoder.write_all(&json)?;
let compressed = encoder.finish()?;
std::fs::write(&file, compressed)?;
if let Node::Attribution(attr) = node {
self.index_attribution(attr)?;
}
Ok(id.clone())
}
pub fn get(&self, id: &NodeId) -> Result<Option<Node>, VasariError> {
let (_, file) = self.object_path(id)?;
if !file.exists() {
return Ok(None);
}
let compressed = std::fs::read(&file)?;
let mut decoder = GzDecoder::new(&compressed[..]);
let mut json = Vec::new();
decoder.read_to_end(&mut json)?;
let node: Node = serde_json::from_slice(&json)?;
Ok(Some(node))
}
pub fn lookup_attributions(&self, path: &str, line: u32) -> Result<Vec<NodeId>, VasariError> {
let path_dir = self
.root
.join("index")
.join("targets")
.join(encode_path(path));
if !path_dir.exists() {
return Ok(vec![]);
}
let mut ids = Vec::new();
for entry in std::fs::read_dir(&path_dir)? {
let entry = entry?;
let name = entry.file_name();
let name_str = name.to_string_lossy();
let covers_line = if name_str == "whole" {
true
} else if let Some((start_s, end_s)) = name_str.split_once('-') {
match (start_s.parse::<u32>(), end_s.parse::<u32>()) {
(Ok(start), Ok(end)) => line >= start && line <= end,
_ => false,
}
} else {
false
};
if covers_line {
let content = std::fs::read_to_string(entry.path())?;
for id_str in content.lines() {
if !id_str.is_empty()
&& id_str.len() >= 4
&& id_str.chars().all(|c| c.is_ascii_hexdigit())
{
ids.push(NodeId(id_str.to_string()));
}
}
}
}
ids.sort_unstable_by(|a, b| a.0.cmp(&b.0));
ids.dedup();
Ok(ids)
}
pub fn iter_all(&self) -> Result<Vec<Node>, VasariError> {
let objects_dir = self.root.join("objects");
if !objects_dir.exists() {
return Ok(vec![]);
}
let mut nodes = Vec::new();
for prefix_entry in std::fs::read_dir(&objects_dir)? {
let prefix_entry = prefix_entry?;
let prefix = prefix_entry.file_name().to_string_lossy().to_string();
for obj_entry in std::fs::read_dir(prefix_entry.path())? {
let obj_entry = obj_entry?;
let suffix = obj_entry.file_name().to_string_lossy().to_string();
let id = NodeId(format!("{prefix}{suffix}"));
match self.get(&id) {
Ok(Some(node)) => nodes.push(node),
Ok(None) | Err(VasariError::InvalidNodeId(_)) => continue,
Err(e) => return Err(e),
}
}
}
Ok(nodes)
}
pub fn resolve_prefix(&self, prefix: &str) -> Result<NodeId, VasariError> {
if prefix.len() < 4 || !prefix.chars().all(|c| c.is_ascii_hexdigit()) {
return Err(VasariError::InvalidNodeId(prefix.to_string()));
}
let lower = prefix.to_ascii_lowercase();
let (shard, rest) = lower.split_at(2);
let shard_dir = self.root.join("objects").join(shard);
if !shard_dir.exists() {
return Err(VasariError::NodeNotFound(prefix.to_string()));
}
let mut matches: Vec<NodeId> = Vec::new();
for entry in std::fs::read_dir(&shard_dir)? {
let entry = entry?;
let fname = entry.file_name().to_string_lossy().to_string();
if fname.starts_with(rest) {
matches.push(NodeId(format!("{shard}{fname}")));
}
}
match matches.len() {
0 => Err(VasariError::NodeNotFound(prefix.to_string())),
1 => Ok(matches.pop().expect("len checked == 1")),
count => Err(VasariError::AmbiguousPrefix {
prefix: prefix.to_string(),
count,
}),
}
}
pub fn rebuild_index(&self) -> Result<usize, VasariError> {
let objects_dir = self.root.join("objects");
let index_targets = self.root.join("index").join("targets");
if index_targets.exists() {
std::fs::remove_dir_all(&index_targets)?;
}
std::fs::create_dir_all(&index_targets)?;
let mut count = 0;
for prefix_entry in std::fs::read_dir(&objects_dir)? {
let prefix_entry = prefix_entry?;
for obj_entry in std::fs::read_dir(prefix_entry.path())? {
let obj_entry = obj_entry?;
let prefix = prefix_entry.file_name();
let suffix = obj_entry.file_name();
let id = NodeId(format!(
"{}{}",
prefix.to_string_lossy(),
suffix.to_string_lossy()
));
match self.get(&id) {
Ok(Some(Node::Attribution(attr))) => {
self.index_attribution(&attr)?;
count += 1;
}
Ok(_) | Err(VasariError::InvalidNodeId(_)) => continue,
Err(e) => return Err(e),
}
}
}
Ok(count)
}
fn object_path(&self, id: &NodeId) -> Result<(PathBuf, PathBuf), VasariError> {
let s = id.as_str();
if s.len() < 4 || !s.chars().all(|c| c.is_ascii_hexdigit()) {
return Err(VasariError::InvalidNodeId(s.to_string()));
}
let dir = self.root.join("objects").join(&s[..2]);
let file = dir.join(&s[2..]);
Ok((dir, file))
}
fn index_attribution(&self, attr: &Attribution) -> Result<(), VasariError> {
let (path, index_name) = match &attr.target {
AttributionTarget::LineRange { path, start, end } => (path, format!("{start}-{end}")),
AttributionTarget::WholeFile { path } => (path, "whole".to_string()),
AttributionTarget::CommitSha { .. } => return Ok(()),
};
let path_dir = self
.root
.join("index")
.join("targets")
.join(encode_path(path));
std::fs::create_dir_all(&path_dir)?;
let index_file = path_dir.join(index_name);
let mut f = std::fs::OpenOptions::new()
.create(true)
.append(true)
.open(&index_file)?;
writeln!(f, "{}", attr.id.as_str())?;
Ok(())
}
}
fn encode_path(path: &str) -> String {
path.split('/')
.map(|component| match component {
".." => "%2E%2E".to_string(),
"." => "%2E".to_string(),
other => other.replace('%', "%25"),
})
.collect::<Vec<_>>()
.join("%2F")
}
#[cfg(test)]
mod tests {
use super::*;
use crate::schema::{Attribution, AttributionTarget, Intent, Node};
#[test]
fn round_trip_intent() {
let dir = tempfile::tempdir().unwrap();
let store = ObjectStore::open(dir.path()).unwrap();
let intent = Intent::new("ACME-411".into(), "Add JWT verification".into(), vec![]);
let id = store.put(&Node::Intent(intent.clone())).unwrap();
let got = store.get(&id).unwrap().unwrap();
match got {
Node::Intent(i) => assert_eq!(i.id, intent.id),
_ => panic!("wrong node type"),
}
}
#[test]
fn put_is_idempotent() {
let dir = tempfile::tempdir().unwrap();
let store = ObjectStore::open(dir.path()).unwrap();
let intent = Intent::new("src".into(), "test".into(), vec![]);
let node = Node::Intent(intent);
let id1 = store.put(&node).unwrap();
let id2 = store.put(&node).unwrap();
assert_eq!(id1, id2);
}
#[test]
fn attribution_index_lookup() {
let dir = tempfile::tempdir().unwrap();
let store = ObjectStore::open(dir.path()).unwrap();
let action_id = NodeId("deadbeef".repeat(8));
let attr = Attribution::new(
action_id,
AttributionTarget::LineRange {
path: "src/auth.ts".into(),
start: 40,
end: 55,
},
1.0,
vec![],
vec![],
);
let attr_id = attr.id.clone();
store.put(&Node::Attribution(attr)).unwrap();
let found = store.lookup_attributions("src/auth.ts", 47).unwrap();
assert!(found.contains(&attr_id));
let not_found = store.lookup_attributions("src/auth.ts", 60).unwrap();
assert!(not_found.is_empty());
}
#[test]
fn fresh_store_is_stamped_and_reopens() {
let dir = tempfile::tempdir().unwrap();
let store = ObjectStore::open(dir.path()).unwrap();
let version = std::fs::read_to_string(dir.path().join(".vasari").join("format")).unwrap();
assert_eq!(version.trim(), ObjectStore::FORMAT_VERSION);
drop(store);
assert!(ObjectStore::open(dir.path()).is_ok());
}
#[test]
fn pre_v2_store_without_format_marker_is_refused() {
let dir = tempfile::tempdir().unwrap();
std::fs::create_dir_all(dir.path().join(".vasari").join("objects")).unwrap();
let err = ObjectStore::open(dir.path()).unwrap_err();
match err {
VasariError::IncompatibleStore { expected, .. } => {
assert_eq!(expected, ObjectStore::FORMAT_VERSION);
}
other => panic!("expected IncompatibleStore, got {other:?}"),
}
}
#[test]
fn whole_file_attribution_matches_any_line() {
let dir = tempfile::tempdir().unwrap();
let store = ObjectStore::open(dir.path()).unwrap();
let attr = Attribution::new(
NodeId("deadbeef".repeat(8)),
AttributionTarget::WholeFile {
path: "src/auth.rs".into(),
},
0.7,
vec![],
vec![],
);
let attr_id = attr.id.clone();
store.put(&Node::Attribution(attr)).unwrap();
for line in [1u32, 47, 9999] {
let found = store.lookup_attributions("src/auth.rs", line).unwrap();
assert!(
found.contains(&attr_id),
"whole-file should match line {line}"
);
}
assert!(store
.lookup_attributions("src/other.rs", 1)
.unwrap()
.is_empty());
}
#[test]
fn resolve_prefix_unique_match() {
let dir = tempfile::tempdir().unwrap();
let store = ObjectStore::open(dir.path()).unwrap();
let intent = Intent::new("s1".into(), "unique intent".into(), vec![]);
let full = store.put(&Node::Intent(intent)).unwrap();
let resolved = store.resolve_prefix(&full.as_str()[..10]).unwrap();
assert_eq!(resolved, full);
assert_eq!(store.resolve_prefix(full.as_str()).unwrap(), full);
}
#[test]
fn resolve_prefix_not_found() {
let dir = tempfile::tempdir().unwrap();
let store = ObjectStore::open(dir.path()).unwrap();
store
.put(&Node::Intent(Intent::new("s".into(), "x".into(), vec![])))
.unwrap();
let err = store.resolve_prefix(&"f".repeat(12)).unwrap_err();
assert!(matches!(err, VasariError::NodeNotFound(_)));
}
#[test]
fn resolve_prefix_rejects_short_and_non_hex_without_panic() {
let dir = tempfile::tempdir().unwrap();
let store = ObjectStore::open(dir.path()).unwrap();
assert!(matches!(
store.resolve_prefix("ab").unwrap_err(),
VasariError::InvalidNodeId(_)
));
assert!(matches!(
store.resolve_prefix("").unwrap_err(),
VasariError::InvalidNodeId(_)
));
assert!(matches!(
store.resolve_prefix("zzzz").unwrap_err(),
VasariError::InvalidNodeId(_)
));
}
#[test]
fn resolve_prefix_ambiguous_reports_count() {
let dir = tempfile::tempdir().unwrap();
let store = ObjectStore::open(dir.path()).unwrap();
for suffix in ["aaaa", "bbbb"] {
let id = NodeId(format!("abcdef{}", suffix.repeat(14)));
let (dir_p, file_p) = store.object_path(&id).unwrap();
std::fs::create_dir_all(&dir_p).unwrap();
std::fs::write(&file_p, b"x").unwrap();
}
let err = store.resolve_prefix("abcdef").unwrap_err();
match err {
VasariError::AmbiguousPrefix { count, .. } => assert_eq!(count, 2),
other => panic!("expected AmbiguousPrefix, got {other:?}"),
}
}
#[test]
fn iter_all_on_empty_store_returns_empty() {
let dir = tempfile::tempdir().unwrap();
let store = ObjectStore::open(dir.path()).unwrap();
let nodes = store.iter_all().unwrap();
assert!(nodes.is_empty());
}
#[test]
fn iter_all_returns_all_stored_nodes() {
let dir = tempfile::tempdir().unwrap();
let store = ObjectStore::open(dir.path()).unwrap();
let i1 = Intent::new("s1".into(), "first intent".into(), vec![]);
let i2 = Intent::new("s2".into(), "second intent".into(), vec![]);
store.put(&Node::Intent(i1.clone())).unwrap();
store.put(&Node::Intent(i2.clone())).unwrap();
let nodes = store.iter_all().unwrap();
assert_eq!(nodes.len(), 2);
}
#[test]
fn rebuild_index_restores_attribution_lookup() {
let dir = tempfile::tempdir().unwrap();
let store = ObjectStore::open(dir.path()).unwrap();
let action_id = NodeId("deadbeef".repeat(8));
let attr = Attribution::new(
action_id,
AttributionTarget::LineRange {
path: "src/lib.rs".into(),
start: 1,
end: u32::MAX,
},
0.9,
vec![],
vec![],
);
let attr_id = attr.id.clone();
store.put(&Node::Attribution(attr)).unwrap();
let index_dir = dir.path().join(".vasari").join("index").join("targets");
std::fs::remove_dir_all(&index_dir).unwrap();
std::fs::create_dir_all(&index_dir).unwrap();
let before = store.lookup_attributions("src/lib.rs", 42).unwrap();
assert!(before.is_empty());
let count = store.rebuild_index().unwrap();
assert_eq!(count, 1);
let after = store.lookup_attributions("src/lib.rs", 42).unwrap();
assert!(after.contains(&attr_id));
}
#[test]
fn lookup_attributions_deduplicates_on_re_ingest() {
let dir = tempfile::tempdir().unwrap();
let store = ObjectStore::open(dir.path()).unwrap();
let action_id = NodeId("deadbeef".repeat(8));
let attr = Attribution::new(
action_id,
AttributionTarget::LineRange {
path: "src/dup.rs".into(),
start: 1,
end: 100,
},
1.0,
vec![],
vec![],
);
let attr_id = attr.id.clone();
store.put(&Node::Attribution(attr.clone())).unwrap();
let encoded = encode_path("src/dup.rs");
let index_file = dir
.path()
.join(".vasari")
.join("index")
.join("targets")
.join(&encoded)
.join("1-100");
let mut f = std::fs::OpenOptions::new()
.append(true)
.open(&index_file)
.unwrap();
use std::io::Write;
writeln!(f, "{}", attr_id.as_str()).unwrap();
let found = store.lookup_attributions("src/dup.rs", 50).unwrap();
assert_eq!(found.len(), 1, "dedup-on-read should remove the duplicate");
}
#[test]
fn encode_path_encodes_slashes() {
let encoded = encode_path("src/auth/mod.rs");
assert!(!encoded.contains('/'));
assert!(encoded.contains("%2F"));
}
#[test]
fn encode_path_encodes_dotdot() {
let encoded = encode_path("../escape/path.rs");
assert!(!encoded.contains(".."));
assert!(encoded.contains("%2E%2E"));
}
#[test]
fn encode_path_encodes_single_dot() {
let encoded = encode_path("./relative.rs");
assert!(
encoded.contains("%2E"),
"single dot should be percent-encoded"
);
assert!(
!encoded.contains(".."),
"single dot should not be mistaken for dotdot"
);
}
#[test]
fn encode_path_encodes_percent() {
let encoded = encode_path("src/100%done.rs");
assert!(encoded.contains("%25"));
}
#[test]
fn lookup_attributions_returns_empty_for_unknown_path() {
let dir = tempfile::tempdir().unwrap();
let store = ObjectStore::open(dir.path()).unwrap();
let result = store.lookup_attributions("nonexistent/file.rs", 1).unwrap();
assert!(result.is_empty());
}
}