use crate::{AnnConfig, AnnIndex, AnnQuery};
fn build_index() -> AnnIndex {
let mut index = AnnIndex::new(AnnConfig::for_dimensions(4)).unwrap();
let records: Vec<(String, Vec<f32>)> = (0..500)
.map(|i| {
let angle = i as f32 * 0.041;
(
format!("record-{i:05}"),
vec![
angle.sin(),
angle.cos(),
(angle * 3.0).sin(),
(angle * 3.0).cos(),
],
)
})
.collect();
index.bulk_load(records, 4).unwrap();
index.delete("record-00042");
index
}
fn probe(index: &AnnIndex, seed: f32) -> Vec<crate::AnnCandidate> {
let angle = seed * 0.29;
index
.query(AnnQuery {
vector: &[
angle.sin(),
angle.cos(),
(angle * 3.0).sin(),
(angle * 3.0).cos(),
],
top_k: 10,
ef_search: Some(64),
filter: None,
})
.unwrap()
}
#[test]
fn snapshot_roundtrip_preserves_the_graph() {
let index = build_index();
let mut bytes = Vec::new();
index.write_to(&mut bytes).unwrap();
let restored = AnnIndex::read_from(bytes.as_slice()).unwrap();
assert_eq!(restored.len(), index.len());
for i in 0..25 {
assert_eq!(
probe(&index, i as f32),
probe(&restored, i as f32),
"snapshot roundtrip must preserve query results exactly"
);
}
}
#[test]
fn snapshot_roundtrip_preserves_tombstones() {
let index = build_index();
let mut bytes = Vec::new();
index.write_to(&mut bytes).unwrap();
let mut restored = AnnIndex::read_from(bytes.as_slice()).unwrap();
assert!(!restored.delete("record-00042"));
assert!(restored.delete("record-00043"));
}
#[test]
fn corrupted_snapshots_are_rejected() {
let index = build_index();
let mut bytes = Vec::new();
index.write_to(&mut bytes).unwrap();
for position in [0, 8, bytes.len() / 2, bytes.len() - 5] {
let mut corrupted = bytes.clone();
corrupted[position] ^= 0x40;
assert!(
AnnIndex::read_from(corrupted.as_slice()).is_err(),
"flipped byte at {position} must be detected"
);
}
let truncated = &bytes[..bytes.len() - 12];
assert!(AnnIndex::read_from(truncated).is_err());
}
#[test]
fn snapshot_saves_and_loads_from_disk() {
let index = build_index();
let path = std::env::temp_dir().join(format!("feox-ann-test-{}.fxa", std::process::id()));
index.save_to(&path).unwrap();
let restored = AnnIndex::load_from(&path).unwrap();
std::fs::remove_file(&path).ok();
assert_eq!(restored.len(), index.len());
assert_eq!(probe(&index, 1.0), probe(&restored, 1.0));
}