use ipfrs_core::{Error, Result};
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};
use std::io::Write;
use std::path::Path;
use tracing::{debug, info, warn};
const SNAPSHOT_VERSION: u32 = 1;
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct IndexEntry {
pub id: u32,
pub cid: String,
pub vector: Vec<f32>,
pub max_layer: usize,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IndexSnapshot {
pub version: u32,
pub dimension: usize,
pub ef_construction: usize,
pub m: usize,
pub entries: Vec<IndexEntry>,
pub layer_connections: Vec<Vec<Vec<u32>>>,
pub metadata_map: HashMap<String, String>,
pub created_at: u64,
pub entry_point: Option<u32>,
}
impl IndexSnapshot {
pub fn validate(&self) -> std::result::Result<(), String> {
if self.version != SNAPSHOT_VERSION {
return Err(format!(
"unsupported snapshot version {} (expected {})",
self.version, SNAPSHOT_VERSION
));
}
if self.dimension == 0 {
return Err("dimension must be > 0".into());
}
for (idx, entry) in self.entries.iter().enumerate() {
if entry.vector.len() != self.dimension {
return Err(format!(
"entry {} has vector length {} but dimension is {}",
idx,
entry.vector.len(),
self.dimension
));
}
}
let n = self.entries.len() as u32;
for (layer_idx, layer) in self.layer_connections.iter().enumerate() {
for (node_idx, neighbors) in layer.iter().enumerate() {
for &nb in neighbors {
if nb >= n {
return Err(format!(
"layer {} node {} has neighbor {} which is out of range (n={})",
layer_idx, node_idx, nb, n
));
}
}
}
}
if let Some(ep) = self.entry_point {
if ep >= n {
return Err(format!("entry_point {} is out of range (n={})", ep, n));
}
}
Ok(())
}
}
pub struct IndexPersistence {
path: std::path::PathBuf,
}
impl IndexPersistence {
pub fn new(path: impl Into<std::path::PathBuf>) -> Self {
Self { path: path.into() }
}
pub fn save(&self, snapshot: &IndexSnapshot) -> Result<()> {
if let Err(msg) = snapshot.validate() {
return Err(Error::InvalidInput(format!(
"snapshot validation failed: {}",
msg
)));
}
let bytes = oxicode::serde::encode_to_vec(snapshot, oxicode::config::standard())
.map_err(|e| Error::Serialization(format!("oxicode encode failed: {}", e)))?;
if let Some(parent) = self.path.parent() {
std::fs::create_dir_all(parent).map_err(Error::Io)?;
}
let tmp_path = self.path.with_extension("snap.tmp");
{
let mut f = std::fs::File::create(&tmp_path).map_err(Error::Io)?;
f.write_all(&bytes).map_err(Error::Io)?;
f.flush().map_err(Error::Io)?;
}
std::fs::rename(&tmp_path, &self.path).map_err(Error::Io)?;
info!(
path = %self.path.display(),
entries = snapshot.entries.len(),
bytes = bytes.len(),
"HNSW index snapshot saved"
);
Ok(())
}
pub fn load(&self) -> Result<IndexSnapshot> {
if !self.path.exists() {
return Err(Error::NotFound(format!(
"snapshot file not found: {}",
self.path.display()
)));
}
let bytes = std::fs::read(&self.path).map_err(Error::Io)?;
let (snapshot, _consumed): (IndexSnapshot, _) =
oxicode::serde::decode_from_slice(&bytes, oxicode::config::standard())
.map_err(|e| Error::Deserialization(format!("oxicode decode failed: {}", e)))?;
snapshot
.validate()
.map_err(|msg| Error::InvalidData(format!("loaded snapshot is corrupt: {}", msg)))?;
debug!(
path = %self.path.display(),
entries = snapshot.entries.len(),
dimension = snapshot.dimension,
"HNSW index snapshot loaded"
);
Ok(snapshot)
}
pub fn exists(&self) -> bool {
self.path.exists()
}
pub fn delete(&self) -> Result<()> {
if self.path.exists() {
std::fs::remove_file(&self.path).map_err(Error::Io)?;
warn!(path = %self.path.display(), "HNSW snapshot deleted");
}
Ok(())
}
pub fn path(&self) -> &Path {
&self.path
}
pub fn save_smart(&self, index: &crate::hnsw::VectorIndex) -> Result<()> {
let total = index.len();
let dirty = index.dirty_count();
let use_incremental = total > 0
&& dirty < total / 10 && self.exists();
if use_incremental {
let base_version = index.tracker_version();
let delta = index.snapshot_incremental(base_version)?;
self.save_incremental(&delta)?;
index.mark_tracker_clean();
info!(
dirty,
total,
base_version,
delta_version = delta.delta_version,
"HNSW index saved as incremental delta"
);
} else {
let snap = index.snapshot()?;
self.save(&snap)?;
index.record_full_snapshot();
debug!(entries = total, "HNSW index saved as full snapshot");
}
Ok(())
}
pub fn incremental_path(&self) -> std::path::PathBuf {
let mut p = self.path.clone();
let ext = p
.extension()
.map(|e| {
let mut s = e.to_os_string();
s.push(".delta");
s
})
.unwrap_or_else(|| std::ffi::OsString::from("snap.delta"));
p.set_extension(ext);
p
}
pub fn save_incremental(&self, snapshot: &IncrementalSnapshot) -> Result<()> {
let delta_path = self.incremental_path();
let bytes =
oxicode::serde::encode_to_vec(snapshot, oxicode::config::standard()).map_err(|e| {
Error::Serialization(format!("oxicode encode incremental failed: {}", e))
})?;
if let Some(parent) = delta_path.parent() {
std::fs::create_dir_all(parent).map_err(Error::Io)?;
}
let tmp_path = delta_path.with_extension("delta.tmp");
{
let mut f = std::fs::File::create(&tmp_path).map_err(Error::Io)?;
f.write_all(&bytes).map_err(Error::Io)?;
f.flush().map_err(Error::Io)?;
}
std::fs::rename(&tmp_path, &delta_path).map_err(Error::Io)?;
info!(
path = %delta_path.display(),
changed = snapshot.changed_entries.len(),
deleted = snapshot.deleted_ids.len(),
base_version = snapshot.base_version,
delta_version = snapshot.delta_version,
"Incremental HNSW snapshot saved"
);
Ok(())
}
pub fn load_incremental(&self) -> Result<IncrementalSnapshot> {
let delta_path = self.incremental_path();
if !delta_path.exists() {
return Err(Error::NotFound(format!(
"incremental snapshot file not found: {}",
delta_path.display()
)));
}
let bytes = std::fs::read(&delta_path).map_err(Error::Io)?;
let (snapshot, _consumed): (IncrementalSnapshot, _) =
oxicode::serde::decode_from_slice(&bytes, oxicode::config::standard()).map_err(
|e| Error::Deserialization(format!("oxicode decode incremental failed: {}", e)),
)?;
debug!(
path = %delta_path.display(),
changed = snapshot.changed_entries.len(),
deleted = snapshot.deleted_ids.len(),
"Incremental HNSW snapshot loaded"
);
Ok(snapshot)
}
pub fn apply_incremental(base: &mut IndexSnapshot, delta: &IncrementalSnapshot) -> Result<()> {
if !delta.deleted_ids.is_empty() {
let deleted_set: HashSet<u32> = delta.deleted_ids.iter().copied().collect();
base.entries.retain(|e| !deleted_set.contains(&e.id));
}
for changed in &delta.changed_entries {
if let Some(existing) = base.entries.iter_mut().find(|e| e.id == changed.id) {
*existing = changed.clone();
} else {
base.entries.push(changed.clone());
}
}
debug!(
base_version = delta.base_version,
delta_version = delta.delta_version,
entries_after = base.entries.len(),
"Applied incremental snapshot delta"
);
Ok(())
}
}
pub struct IncrementalTracker {
dirty_ids: HashSet<u32>,
version: u64,
last_full_snapshot: Option<std::time::SystemTime>,
}
impl IncrementalTracker {
pub fn new() -> Self {
Self {
dirty_ids: HashSet::new(),
version: 0,
last_full_snapshot: None,
}
}
pub fn mark_dirty(&mut self, id: u32) {
self.dirty_ids.insert(id);
}
pub fn mark_clean(&mut self) {
self.dirty_ids.clear();
self.version = self.version.saturating_add(1);
}
pub fn record_full_snapshot(&mut self, time: std::time::SystemTime) {
self.last_full_snapshot = Some(time);
self.mark_clean();
}
pub fn dirty_ids(&self) -> &HashSet<u32> {
&self.dirty_ids
}
pub fn is_dirty(&self) -> bool {
!self.dirty_ids.is_empty()
}
pub fn dirty_count(&self) -> usize {
self.dirty_ids.len()
}
pub fn version(&self) -> u64 {
self.version
}
pub fn last_full_snapshot(&self) -> Option<std::time::SystemTime> {
self.last_full_snapshot
}
}
impl Default for IncrementalTracker {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IncrementalSnapshot {
pub base_version: u64,
pub delta_version: u64,
pub changed_entries: Vec<IndexEntry>,
pub deleted_ids: Vec<u32>,
pub created_at: u64,
}
#[cfg(test)]
mod tests {
use super::*;
fn make_snapshot() -> IndexSnapshot {
IndexSnapshot {
version: 1,
dimension: 4,
ef_construction: 200,
m: 16,
entries: vec![
IndexEntry {
id: 0,
cid: "test_cid_0".to_string(),
vector: vec![1.0, 0.0, 0.0, 0.0],
max_layer: 0,
},
IndexEntry {
id: 1,
cid: "test_cid_1".to_string(),
vector: vec![0.0, 1.0, 0.0, 0.0],
max_layer: 0,
},
],
layer_connections: vec![vec![vec![1], vec![0]]],
metadata_map: HashMap::new(),
created_at: 12345,
entry_point: Some(0),
}
}
fn temp_snap_path(tag: &str) -> std::path::PathBuf {
let nanos = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.subsec_nanos())
.unwrap_or(0);
std::env::temp_dir().join(format!("ipfrs_hnsw_test_{}_{}", tag, nanos))
}
#[test]
fn test_snapshot_save_load_roundtrip() {
let dir = temp_snap_path("roundtrip");
std::fs::create_dir_all(&dir).expect("create temp dir");
let persistence = IndexPersistence::new(dir.join("index.snap"));
let snapshot = make_snapshot();
persistence.save(&snapshot).expect("save");
assert!(persistence.exists(), "file must exist after save");
let loaded = persistence.load().expect("load");
assert_eq!(loaded.version, 1);
assert_eq!(loaded.dimension, 4);
assert_eq!(loaded.ef_construction, 200);
assert_eq!(loaded.m, 16);
assert_eq!(loaded.entries.len(), 2);
assert_eq!(loaded.entries[0].cid, "test_cid_0");
assert_eq!(loaded.entries[0].vector, vec![1.0f32, 0.0, 0.0, 0.0]);
assert_eq!(loaded.entries[1].cid, "test_cid_1");
assert_eq!(loaded.entries[1].vector, vec![0.0f32, 1.0, 0.0, 0.0]);
assert_eq!(loaded.entry_point, Some(0));
assert_eq!(loaded.created_at, 12345);
persistence.delete().expect("delete");
assert!(!persistence.exists(), "file must be gone after delete");
}
#[test]
fn test_persistence_not_found() {
let persistence = IndexPersistence::new("/nonexistent/path/index.snap");
assert!(!persistence.exists());
assert!(persistence.load().is_err());
}
#[test]
fn test_save_creates_parent_dirs() {
let base = temp_snap_path("mkdir");
let nested = base.join("a").join("b").join("c").join("index.snap");
let persistence = IndexPersistence::new(&nested);
let snapshot = make_snapshot();
persistence
.save(&snapshot)
.expect("save should create parent dirs");
assert!(persistence.exists());
let _ = persistence.delete();
let _ = std::fs::remove_dir_all(&base);
}
#[test]
fn test_overwrite_is_atomic() {
let dir = temp_snap_path("atomic");
std::fs::create_dir_all(&dir).expect("create temp dir");
let persistence = IndexPersistence::new(dir.join("index.snap"));
let snap_a = make_snapshot();
let mut snap_b = make_snapshot();
snap_b.created_at = 99999;
snap_b.entries.push(IndexEntry {
id: 2,
cid: "test_cid_2".to_string(),
vector: vec![0.0, 0.0, 1.0, 0.0],
max_layer: 1,
});
persistence.save(&snap_a).expect("first save");
persistence.save(&snap_b).expect("second save");
let loaded = persistence.load().expect("load");
assert_eq!(loaded.created_at, 99999);
assert_eq!(loaded.entries.len(), 3);
let _ = persistence.delete();
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn test_snapshot_validate_version_mismatch() {
let mut snap = make_snapshot();
snap.version = 99;
assert!(snap.validate().is_err());
}
#[test]
fn test_snapshot_validate_dimension_mismatch() {
let mut snap = make_snapshot();
snap.entries[0].vector = vec![1.0, 2.0]; assert!(snap.validate().is_err());
}
#[test]
fn test_snapshot_validate_out_of_range_neighbor() {
let mut snap = make_snapshot();
snap.layer_connections[0][0] = vec![999]; assert!(snap.validate().is_err());
}
}
#[cfg(test)]
mod incremental_tests {
use super::*;
fn temp_dir(tag: &str) -> std::path::PathBuf {
let nanos = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.subsec_nanos())
.unwrap_or(0);
std::env::temp_dir().join(format!("ipfrs_incr_test_{}_{}", tag, nanos))
}
fn make_base_snapshot() -> IndexSnapshot {
IndexSnapshot {
version: 1,
dimension: 3,
ef_construction: 100,
m: 8,
entries: vec![
IndexEntry {
id: 0,
cid: "cid0".to_string(),
vector: vec![1.0, 0.0, 0.0],
max_layer: 0,
},
IndexEntry {
id: 1,
cid: "cid1".to_string(),
vector: vec![0.0, 1.0, 0.0],
max_layer: 0,
},
],
layer_connections: vec![vec![vec![1], vec![0]]],
metadata_map: HashMap::new(),
created_at: 1000,
entry_point: Some(0),
}
}
#[test]
fn test_incremental_tracker_dirty_tracking() {
let mut tracker = IncrementalTracker::new();
assert!(!tracker.is_dirty());
assert_eq!(tracker.dirty_count(), 0);
assert_eq!(tracker.version(), 0);
tracker.mark_dirty(5);
tracker.mark_dirty(10);
tracker.mark_dirty(5); assert!(tracker.is_dirty());
assert_eq!(tracker.dirty_count(), 2);
assert!(tracker.dirty_ids().contains(&5));
assert!(tracker.dirty_ids().contains(&10));
tracker.mark_clean();
assert!(!tracker.is_dirty());
assert_eq!(tracker.dirty_count(), 0);
assert_eq!(tracker.version(), 1);
tracker.mark_dirty(99);
assert_eq!(tracker.dirty_count(), 1);
tracker.mark_clean();
assert_eq!(tracker.version(), 2);
}
#[test]
fn test_incremental_snapshot_save_load() {
let dir = temp_dir("save_load");
std::fs::create_dir_all(&dir).expect("create temp dir");
let persistence = IndexPersistence::new(dir.join("index.snap"));
let delta = IncrementalSnapshot {
base_version: 3,
delta_version: 4,
changed_entries: vec![IndexEntry {
id: 7,
cid: "cid7".to_string(),
vector: vec![0.5, 0.5, 0.0],
max_layer: 1,
}],
deleted_ids: vec![2, 3],
created_at: 9999,
};
persistence
.save_incremental(&delta)
.expect("save incremental");
let loaded = persistence.load_incremental().expect("load incremental");
assert_eq!(loaded.base_version, 3);
assert_eq!(loaded.delta_version, 4);
assert_eq!(loaded.changed_entries.len(), 1);
assert_eq!(loaded.changed_entries[0].id, 7);
assert_eq!(loaded.deleted_ids, vec![2, 3]);
assert_eq!(loaded.created_at, 9999);
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn test_apply_incremental_to_base() {
let mut base = make_base_snapshot();
let delta = IncrementalSnapshot {
base_version: 0,
delta_version: 1,
changed_entries: vec![
IndexEntry {
id: 0,
cid: "cid0_v2".to_string(),
vector: vec![0.9, 0.1, 0.0],
max_layer: 0,
},
IndexEntry {
id: 2,
cid: "cid2".to_string(),
vector: vec![0.0, 0.0, 1.0],
max_layer: 0,
},
],
deleted_ids: vec![1], created_at: 2000,
};
IndexPersistence::apply_incremental(&mut base, &delta).expect("apply incremental");
assert!(base.entries.iter().all(|e| e.id != 1));
let e0 = base.entries.iter().find(|e| e.id == 0).expect("entry 0");
assert_eq!(e0.cid, "cid0_v2");
assert_eq!(e0.vector, vec![0.9f32, 0.1, 0.0]);
assert!(base.entries.iter().any(|e| e.id == 2));
assert_eq!(base.entries.len(), 2);
}
#[test]
fn test_incremental_path_naming() {
let p = IndexPersistence::new("/some/dir/index.snap");
let delta = p.incremental_path();
assert!(
delta.to_string_lossy().ends_with(".snap.delta"),
"unexpected delta path: {}",
delta.display()
);
assert_eq!(delta.parent(), std::path::Path::new("/some/dir").into());
}
#[test]
fn test_incremental_tracker_record_full_snapshot() {
let mut tracker = IncrementalTracker::new();
tracker.mark_dirty(1);
tracker.mark_dirty(2);
let now = std::time::SystemTime::now();
tracker.record_full_snapshot(now);
assert!(!tracker.is_dirty());
assert_eq!(tracker.version(), 1);
assert!(tracker.last_full_snapshot().is_some());
}
#[test]
fn test_apply_incremental_empty_delta() {
let mut base = make_base_snapshot();
let original_len = base.entries.len();
let delta = IncrementalSnapshot {
base_version: 0,
delta_version: 1,
changed_entries: vec![],
deleted_ids: vec![],
created_at: 0,
};
IndexPersistence::apply_incremental(&mut base, &delta).expect("apply empty delta");
assert_eq!(base.entries.len(), original_len);
}
#[test]
fn test_incremental_snapshot_only_dirty() {
use crate::hnsw::{DistanceMetric, VectorIndex};
const DIM: usize = 8;
const TOTAL: usize = 100;
const DIRTY_COUNT: usize = 5;
let mut index =
VectorIndex::new(DIM, DistanceMetric::L2, 8, 50).expect("create VectorIndex");
let mut cids = Vec::with_capacity(TOTAL);
for i in 0..TOTAL {
let data = bytes::Bytes::from(format!("embed-data-{}", i));
let block = ipfrs_core::Block::new(data).expect("create block");
let cid = *block.cid();
cids.push(cid);
let mut v = vec![0.0f32; DIM];
v[i % DIM] = 1.0 + (i as f32) * 0.001; index.add_embedding(&cid, &v).expect("add_embedding");
}
index.record_full_snapshot();
let dirty_start = TOTAL;
for i in dirty_start..(dirty_start + DIRTY_COUNT) {
let data = bytes::Bytes::from(format!("new-embed-{}", i));
let block = ipfrs_core::Block::new(data).expect("create block");
let cid = *block.cid();
let mut v = vec![0.0f32; DIM];
v[i % DIM] = 2.0 + (i as f32) * 0.001;
index.add_embedding(&cid, &v).expect("add_embedding dirty");
}
assert_eq!(
index.dirty_count(),
DIRTY_COUNT,
"expected exactly {} dirty entries after insertions",
DIRTY_COUNT
);
let base_version = index.tracker_version();
let delta = index
.snapshot_incremental(base_version)
.expect("snapshot_incremental");
assert_eq!(
delta.changed_entries.len(),
DIRTY_COUNT,
"incremental delta must contain exactly {} changed entries",
DIRTY_COUNT
);
assert!(
delta.deleted_ids.is_empty(),
"no deletions should be recorded"
);
}
}