use std::collections::HashMap;
use std::fs::File;
use std::io::{BufWriter, Read, Write};
use std::path::{Path, PathBuf};
use std::thread;
use crate::graph::Node;
use crate::index::MAX_ENTRY_POINTS;
use crate::{AnnConfig, AnnError, AnnIndex, Result};
const MAGIC: &[u8; 4] = b"FXA1";
const VERSION: u16 = 1;
const NO_ENTRY: u64 = u64::MAX;
impl AnnIndex {
pub fn save_to(&self, path: impl AsRef<Path>) -> Result<()> {
let file = File::create(path)?;
let mut writer = BufWriter::new(file);
self.write_to(&mut writer)?;
writer.flush()?;
writer.get_ref().sync_all()?;
Ok(())
}
pub fn load_from(path: impl AsRef<Path>) -> Result<Self> {
Self::read_from(File::open(path)?)
}
pub fn save_in_background(&self, path: PathBuf) -> thread::JoinHandle<Result<()>> {
let snapshot = self.clone();
thread::spawn(move || snapshot.save_to(path))
}
pub fn write_to<W: Write>(&self, writer: W) -> Result<()> {
let mut writer = CrcWriter::new(writer);
writer.write_all(MAGIC)?;
write_u16(&mut writer, VERSION)?;
write_u16(&mut writer, 0)?;
write_u32(&mut writer, self.config.dimensions as u32)?;
write_u32(&mut writer, self.config.max_neighbors as u32)?;
write_u32(&mut writer, self.config.max_base_neighbors as u32)?;
write_u32(&mut writer, self.config.ef_construction as u32)?;
write_u32(&mut writer, self.config.ef_search as u32)?;
write_u32(&mut writer, self.config.max_level as u32)?;
write_u64(
&mut writer,
self.entry.map(|e| e as u64).unwrap_or(NO_ENTRY),
)?;
write_u32(&mut writer, self.max_level as u32)?;
write_u64(&mut writer, self.active as u64)?;
write_u64(&mut writer, self.nodes.len() as u64)?;
write_u32(&mut writer, self.entry_points.len() as u32)?;
for &entry_point in &self.entry_points {
write_u64(&mut writer, entry_point as u64)?;
}
for node in &self.nodes {
let id_bytes = node.id.as_bytes();
if id_bytes.len() > u16::MAX as usize {
return Err(AnnError::InvalidVector(
"record id exceeds 65535 bytes and cannot be serialized".to_string(),
));
}
write_u16(&mut writer, id_bytes.len() as u16)?;
writer.write_all(id_bytes)?;
writer.write_all(&[u8::from(node.deleted)])?;
write_u16(&mut writer, node.neighbors.len() as u16)?;
for value in &node.vector {
writer.write_all(&value.to_le_bytes())?;
}
for layer in &node.neighbors {
write_u32(&mut writer, layer.len() as u32)?;
for &neighbor in layer {
write_u64(&mut writer, neighbor as u64)?;
}
}
}
let crc = writer.crc();
writer.into_inner().write_all(&crc.to_le_bytes())?;
Ok(())
}
pub fn read_from<R: Read>(mut reader: R) -> Result<Self> {
let mut bytes = Vec::new();
reader.read_to_end(&mut bytes)?;
if bytes.len() < MAGIC.len() + 4 {
return Err(corrupted("snapshot is too short"));
}
let (payload, crc_bytes) = bytes.split_at(bytes.len() - 4);
let stored_crc = u32::from_le_bytes(crc_bytes.try_into().expect("crc is four bytes"));
if crc32(payload) != stored_crc {
return Err(corrupted("checksum mismatch"));
}
let mut cursor = Cursor::new(payload);
if cursor.take_bytes(4)? != MAGIC {
return Err(corrupted("bad magic"));
}
let version = cursor.take_u16()?;
if version != VERSION {
return Err(corrupted(&format!("unsupported version {version}")));
}
cursor.take_u16()?;
let config = AnnConfig {
dimensions: cursor.take_u32()? as usize,
max_neighbors: cursor.take_u32()? as usize,
max_base_neighbors: cursor.take_u32()? as usize,
ef_construction: cursor.take_u32()? as usize,
ef_search: cursor.take_u32()? as usize,
max_level: cursor.take_u32()? as usize,
};
let entry_raw = cursor.take_u64()?;
let max_level = cursor.take_u32()? as usize;
let active = cursor.take_u64()? as usize;
let node_count = cursor.take_u64()? as usize;
let entry_points_len = cursor.take_u32()? as usize;
if entry_points_len > MAX_ENTRY_POINTS {
return Err(corrupted("entry point count exceeds maximum"));
}
let mut entry_points = Vec::with_capacity(entry_points_len);
for _ in 0..entry_points_len {
let entry_point = cursor.take_u64()? as usize;
if entry_point >= node_count {
return Err(corrupted("entry point out of bounds"));
}
entry_points.push(entry_point);
}
let mut index = AnnIndex::new(config)?;
index.reserve(node_count);
let mut ids: HashMap<String, usize> = HashMap::with_capacity(node_count);
let mut visible = 0_usize;
for position in 0..node_count {
let id_len = cursor.take_u16()? as usize;
let id = String::from_utf8(cursor.take_bytes(id_len)?.to_vec())
.map_err(|_| corrupted("record id is not utf-8"))?;
let deleted = match cursor.take_bytes(1)?[0] {
0 => false,
1 => true,
_ => return Err(corrupted("invalid deleted flag")),
};
let levels = cursor.take_u16()? as usize;
if levels == 0 || levels > index.config.max_level + 1 {
return Err(corrupted("invalid level count"));
}
let mut vector = Vec::with_capacity(index.config.dimensions);
for _ in 0..index.config.dimensions {
vector.push(f32::from_le_bytes(
cursor.take_bytes(4)?.try_into().expect("f32 is four bytes"),
));
}
let mut neighbors = Vec::with_capacity(levels);
for _ in 0..levels {
let count = cursor.take_u32()? as usize;
let mut layer = Vec::with_capacity(count);
for _ in 0..count {
let neighbor = cursor.take_u64()? as usize;
if neighbor >= node_count {
return Err(corrupted("neighbor index out of bounds"));
}
layer.push(neighbor);
}
neighbors.push(layer);
}
if !deleted {
if ids.insert(id.clone(), position).is_some() {
return Err(corrupted("duplicate visible record id"));
}
visible += 1;
}
index.nodes.push(Node {
id,
vector,
deleted,
neighbors,
});
}
if !cursor.is_empty() {
return Err(corrupted("trailing bytes after last node"));
}
if visible != active {
return Err(corrupted("active count does not match visible records"));
}
index.ids = ids;
index.active = active;
index.max_level = max_level;
index.entry_points = entry_points;
index.entry = if entry_raw == NO_ENTRY {
None
} else {
let entry = entry_raw as usize;
if entry >= node_count {
return Err(corrupted("entry index out of bounds"));
}
Some(entry)
};
Ok(index)
}
}
fn corrupted(message: &str) -> AnnError {
AnnError::Corrupted(message.to_string())
}
struct Cursor<'a> {
bytes: &'a [u8],
}
impl<'a> Cursor<'a> {
fn new(bytes: &'a [u8]) -> Self {
Self { bytes }
}
fn take_bytes(&mut self, len: usize) -> Result<&'a [u8]> {
if self.bytes.len() < len {
return Err(corrupted("unexpected end of snapshot"));
}
let (taken, rest) = self.bytes.split_at(len);
self.bytes = rest;
Ok(taken)
}
fn take_u16(&mut self) -> Result<u16> {
Ok(u16::from_le_bytes(
self.take_bytes(2)?.try_into().expect("u16 is two bytes"),
))
}
fn take_u32(&mut self) -> Result<u32> {
Ok(u32::from_le_bytes(
self.take_bytes(4)?.try_into().expect("u32 is four bytes"),
))
}
fn take_u64(&mut self) -> Result<u64> {
Ok(u64::from_le_bytes(
self.take_bytes(8)?.try_into().expect("u64 is eight bytes"),
))
}
fn is_empty(&self) -> bool {
self.bytes.is_empty()
}
}
struct CrcWriter<W: Write> {
inner: W,
crc: u32,
}
impl<W: Write> CrcWriter<W> {
fn new(inner: W) -> Self {
Self {
inner,
crc: 0xffff_ffff,
}
}
fn crc(&self) -> u32 {
self.crc ^ 0xffff_ffff
}
fn into_inner(self) -> W {
self.inner
}
}
impl<W: Write> Write for CrcWriter<W> {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
let written = self.inner.write(buf)?;
self.crc = crc32_update(self.crc, &buf[..written]);
Ok(written)
}
fn flush(&mut self) -> std::io::Result<()> {
self.inner.flush()
}
}
fn write_u16<W: Write>(writer: &mut W, value: u16) -> Result<()> {
writer.write_all(&value.to_le_bytes())?;
Ok(())
}
fn write_u32<W: Write>(writer: &mut W, value: u32) -> Result<()> {
writer.write_all(&value.to_le_bytes())?;
Ok(())
}
fn write_u64<W: Write>(writer: &mut W, value: u64) -> Result<()> {
writer.write_all(&value.to_le_bytes())?;
Ok(())
}
fn crc32(bytes: &[u8]) -> u32 {
crc32_update(0xffff_ffff, bytes) ^ 0xffff_ffff
}
fn crc32_update(mut crc: u32, bytes: &[u8]) -> u32 {
static TABLE: std::sync::OnceLock<[u32; 256]> = std::sync::OnceLock::new();
let table = TABLE.get_or_init(|| {
let mut table = [0_u32; 256];
let mut i = 0;
while i < 256 {
let mut value = i as u32;
let mut bit = 0;
while bit < 8 {
value = if value & 1 == 1 {
0xedb8_8320 ^ (value >> 1)
} else {
value >> 1
};
bit += 1;
}
table[i] = value;
i += 1;
}
table
});
for &byte in bytes {
crc = table[((crc ^ u32::from(byte)) & 0xff) as usize] ^ (crc >> 8);
}
crc
}