weavatrix-search-vector 0.3.1

use crate::error::SearchError;
use crate::hnsw::VectorIndex;
use crate::multi::{MultiVectorIndex, MultiVectorKey};
use std::io::{Read, Write};
use std::path::Path;

const MAGIC: &[u8; 8] = b"WVMULT03";
const VERSION: u32 = 1;
const HEADER_LEN: usize = 56;
const HEADER_LEN_U32: u32 = 56;
const FNV_OFFSET: u64 = 0xcbf2_9ce4_8422_2325;
const FNV_PRIME: u64 = 0x0000_0100_0000_01b3;

impl MultiVectorIndex {
    /// Serializes the graph and caller-visible `(key, vector_id)` identities.
    ///
    /// # Errors
    ///
    /// Returns a capacity or allocation error.
    pub fn to_bytes(&self) -> Result<Vec<u8>, SearchError> {
        let index = self.index.to_bytes()?;
        let identity_bytes = self
            .identities
            .len()
            .checked_mul(16)
            .ok_or(SearchError::CapacityOverflow)?;
        let file_len = HEADER_LEN
            .checked_add(index.len())
            .and_then(|value| value.checked_add(identity_bytes))
            .ok_or(SearchError::CapacityOverflow)?;
        let mut bytes = Vec::new();
        bytes
            .try_reserve_exact(file_len)
            .map_err(|_| SearchError::AllocationFailed)?;
        bytes.extend_from_slice(MAGIC);
        bytes.resize(HEADER_LEN, 0);
        bytes.extend_from_slice(&index);
        for identity in &self.identities {
            bytes.extend_from_slice(&identity.key.to_le_bytes());
            bytes.extend_from_slice(&identity.vector_id.to_le_bytes());
        }
        let payload_checksum = checksum(&bytes[HEADER_LEN..]);
        put_u32(&mut bytes, 8, VERSION);
        put_u32(&mut bytes, 12, HEADER_LEN_U32);
        put_u64(&mut bytes, 16, payload_checksum);
        put_u64(&mut bytes, 24, to_u64(file_len)?);
        put_u64(&mut bytes, 32, to_u64(index.len())?);
        put_u64(&mut bytes, 40, to_u64(self.identities.len())?);
        Ok(bytes)
    }

    /// Restores and validates a graph and its multi-vector identities.
    ///
    /// # Errors
    ///
    /// Returns a format, integrity, graph, or allocation error.
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, SearchError> {
        if bytes.len() < HEADER_LEN || bytes.get(..8) != Some(MAGIC.as_slice()) {
            return Err(SearchError::CorruptSnapshot(
                "multi-vector snapshot header is invalid",
            ));
        }
        let version = read_u32(bytes, 8)?;
        if version != VERSION {
            return Err(SearchError::UnsupportedSnapshotVersion(version));
        }
        if read_u32(bytes, 12)? as usize != HEADER_LEN {
            return Err(SearchError::CorruptSnapshot(
                "multi-vector snapshot header length does not match",
            ));
        }
        if to_usize(read_u64(bytes, 24)?)? != bytes.len() {
            return Err(SearchError::CorruptSnapshot(
                "multi-vector snapshot file length does not match",
            ));
        }
        if checksum(&bytes[HEADER_LEN..]) != read_u64(bytes, 16)? {
            return Err(SearchError::CorruptSnapshot(
                "multi-vector snapshot checksum does not match",
            ));
        }
        let index_len = to_usize(read_u64(bytes, 32)?)?;
        let count = to_usize(read_u64(bytes, 40)?)?;
        let index_end = HEADER_LEN
            .checked_add(index_len)
            .ok_or(SearchError::CapacityOverflow)?;
        let index = VectorIndex::from_bytes(bytes.get(HEADER_LEN..index_end).ok_or(
            SearchError::CorruptSnapshot("multi-vector graph is truncated"),
        )?)?;
        if index.len() != count
            || !index
                .keys()
                .enumerate()
                .all(|(position, key)| key == position as u64)
        {
            return Err(SearchError::CorruptSnapshot(
                "multi-vector graph identities do not match",
            ));
        }
        let identity_len = count.checked_mul(16).ok_or(SearchError::CapacityOverflow)?;
        let identity_bytes = bytes
            .get(index_end..)
            .filter(|value| value.len() == identity_len)
            .ok_or(SearchError::CorruptSnapshot(
                "multi-vector identities are truncated",
            ))?;
        let mut identities = Vec::new();
        identities
            .try_reserve_exact(count)
            .map_err(|_| SearchError::AllocationFailed)?;
        for pair in identity_bytes.chunks_exact(16) {
            let mut key = [0_u8; 8];
            key.copy_from_slice(&pair[..8]);
            let mut vector_id = [0_u8; 8];
            vector_id.copy_from_slice(&pair[8..]);
            identities.push(MultiVectorKey {
                key: u64::from_le_bytes(key),
                vector_id: u64::from_le_bytes(vector_id),
            });
        }
        if identities
            .windows(2)
            .any(|pair| (pair[0].key, pair[0].vector_id) >= (pair[1].key, pair[1].vector_id))
        {
            return Err(SearchError::CorruptSnapshot(
                "multi-vector identities are not strictly ordered",
            ));
        }
        Ok(Self { index, identities })
    }

    /// Writes a complete multi-vector snapshot to a stream.
    ///
    /// # Errors
    ///
    /// Returns a serialization or stream error.
    pub fn write_to(&self, mut writer: impl Write) -> Result<(), SearchError> {
        writer
            .write_all(&self.to_bytes()?)
            .map_err(|error| SearchError::storage("write multi-vector snapshot", &error))
    }

    /// Reads a complete multi-vector snapshot from a stream.
    ///
    /// # Errors
    ///
    /// Returns a stream, format, or integrity error.
    pub fn read_from(mut reader: impl Read) -> Result<Self, SearchError> {
        let mut bytes = Vec::new();
        reader
            .read_to_end(&mut bytes)
            .map_err(|error| SearchError::storage("read multi-vector snapshot", &error))?;
        Self::from_bytes(&bytes)
    }

    /// Atomically saves a complete multi-vector snapshot.
    ///
    /// # Errors
    ///
    /// Returns a serialization or filesystem error.
    pub fn save(&self, path: impl AsRef<Path>) -> Result<(), SearchError> {
        crate::atomic_file::atomic_write(path.as_ref(), &self.to_bytes()?)
    }

    /// Loads a complete multi-vector snapshot.
    ///
    /// # Errors
    ///
    /// Returns a filesystem, format, or integrity error.
    pub fn load(path: impl AsRef<Path>) -> Result<Self, SearchError> {
        let bytes = std::fs::read(path.as_ref())
            .map_err(|error| SearchError::storage("read multi-vector snapshot", &error))?;
        Self::from_bytes(&bytes)
    }
}

fn checksum(bytes: &[u8]) -> u64 {
    bytes.iter().fold(FNV_OFFSET, |hash, byte| {
        (hash ^ u64::from(*byte)).wrapping_mul(FNV_PRIME)
    })
}

fn read_u32(bytes: &[u8], offset: usize) -> Result<u32, SearchError> {
    Ok(u32::from_le_bytes(
        bytes
            .get(offset..offset + 4)
            .ok_or(SearchError::CorruptSnapshot(
                "multi-vector header is truncated",
            ))?
            .try_into()
            .expect("checked four-byte header value"),
    ))
}

fn read_u64(bytes: &[u8], offset: usize) -> Result<u64, SearchError> {
    Ok(u64::from_le_bytes(
        bytes
            .get(offset..offset + 8)
            .ok_or(SearchError::CorruptSnapshot(
                "multi-vector header is truncated",
            ))?
            .try_into()
            .expect("checked eight-byte header value"),
    ))
}

fn put_u32(bytes: &mut [u8], offset: usize, value: u32) {
    bytes[offset..offset + 4].copy_from_slice(&value.to_le_bytes());
}

fn put_u64(bytes: &mut [u8], offset: usize, value: u64) {
    bytes[offset..offset + 8].copy_from_slice(&value.to_le_bytes());
}

fn to_u64(value: usize) -> Result<u64, SearchError> {
    u64::try_from(value).map_err(|_| SearchError::CapacityOverflow)
}

fn to_usize(value: u64) -> Result<usize, SearchError> {
    usize::try_from(value).map_err(|_| SearchError::CapacityOverflow)
}