weavatrix-search-vector 0.3.1

use super::format::{
    checksum, decode_f32s, decode_kind, decode_metric, decode_u64s, read_u32, read_u64, to_usize,
};
use super::storage::decode_storage;
use super::{HEADER_LEN, MAGIC, VERSION};
use crate::config::{DistanceMetric, IndexConfig};
use crate::error::SearchError;
use crate::quantized::{QuantizationKind, QuantizedIndex};

pub(crate) fn decode(bytes: &[u8]) -> Result<QuantizedIndex, SearchError> {
    validate_header(bytes)?;
    let dimensions = to_usize(read_u64(bytes, 32)?)?;
    let count = to_usize(read_u64(bytes, 40)?)?;
    let metric = decode_metric(read_u32(bytes, 48)?)?;
    let kind = decode_kind(read_u32(bytes, 52)?)?;
    let config = decode_config(bytes, dimensions, metric)?;
    validate_kind_metric(kind, metric)?;
    let words_per_vector = to_usize(read_u64(bytes, 112)?)?;
    let storage_len = to_usize(read_u64(bytes, 120)?)?;
    validate_storage_shape(kind, count, dimensions, words_per_vector, storage_len)?;
    let mut cursor = HEADER_LEN;
    let keys = decode_u64s(bytes, &mut cursor, count)?;
    if keys.windows(2).any(|pair| pair[0] >= pair[1]) {
        return Err(SearchError::CorruptSnapshot(
            "quantized snapshot keys are not strictly ordered",
        ));
    }
    let squared_norms = decode_f32s(bytes, &mut cursor, count)?;
    if squared_norms.iter().any(|norm| {
        !norm.is_finite() || *norm < 0.0 || (metric == DistanceMetric::Cosine && *norm == 0.0)
    }) {
        return Err(SearchError::CorruptSnapshot(
            "quantized snapshot norms are invalid",
        ));
    }
    let storage = decode_storage(kind, bytes, &mut cursor, storage_len)?;
    if cursor != bytes.len() {
        return Err(SearchError::CorruptSnapshot(
            "quantized snapshot contains trailing bytes",
        ));
    }
    Ok(QuantizedIndex {
        config,
        kind,
        keys,
        storage,
        squared_norms,
        words_per_vector,
    })
}

fn validate_header(bytes: &[u8]) -> Result<(), SearchError> {
    if bytes.len() < HEADER_LEN || bytes.get(..8) != Some(MAGIC.as_slice()) {
        return Err(SearchError::CorruptSnapshot(
            "quantized 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(
            "quantized snapshot header length does not match",
        ));
    }
    if to_usize(read_u64(bytes, 24)?)? != bytes.len() {
        return Err(SearchError::CorruptSnapshot(
            "quantized snapshot file length does not match",
        ));
    }
    if checksum(&bytes[HEADER_LEN..]) != read_u64(bytes, 16)? {
        return Err(SearchError::CorruptSnapshot(
            "quantized snapshot checksum does not match",
        ));
    }
    Ok(())
}

fn decode_config(
    bytes: &[u8],
    dimensions: usize,
    metric: DistanceMetric,
) -> Result<IndexConfig, SearchError> {
    let config = IndexConfig {
        dimensions,
        metric,
        connectivity: to_usize(read_u64(bytes, 56)?)?,
        expansion_build: to_usize(read_u64(bytes, 64)?)?,
        expansion_query: to_usize(read_u64(bytes, 72)?)?,
        replicas: to_usize(read_u64(bytes, 80)?)?,
        build_threads: to_usize(read_u64(bytes, 88)?)?,
        query_threads: to_usize(read_u64(bytes, 96)?)?,
        seed: read_u64(bytes, 104)?,
    };
    config.validate()?;
    Ok(config)
}

fn validate_kind_metric(kind: QuantizationKind, metric: DistanceMetric) -> Result<(), SearchError> {
    if matches!(kind, QuantizationKind::Int8 | QuantizationKind::Binary)
        && metric != DistanceMetric::Cosine
    {
        return Err(SearchError::CorruptSnapshot(
            "quantized snapshot kind and metric are incompatible",
        ));
    }
    Ok(())
}

fn validate_storage_shape(
    kind: QuantizationKind,
    count: usize,
    dimensions: usize,
    words_per_vector: usize,
    storage_len: usize,
) -> Result<(), SearchError> {
    if words_per_vector != dimensions.div_ceil(64) {
        return Err(SearchError::CorruptSnapshot(
            "quantized snapshot binary stride does not match dimensions",
        ));
    }
    let expected_storage = if kind == QuantizationKind::Binary {
        count
            .checked_mul(words_per_vector)
            .ok_or(SearchError::CapacityOverflow)?
    } else {
        count
            .checked_mul(dimensions)
            .ok_or(SearchError::CapacityOverflow)?
    };
    if storage_len != expected_storage {
        return Err(SearchError::CorruptSnapshot(
            "quantized snapshot storage length does not match",
        ));
    }
    Ok(())
}