weavatrix-search-vector 0.3.1

use super::format::{checksum, kind_code, metric_code, put_u32, put_u64, to_u64};
use super::{HEADER_LEN, MAGIC, VERSION};
use crate::error::SearchError;
use crate::quantized::QuantizedIndex;

pub(crate) fn encode(index: &QuantizedIndex) -> Result<Vec<u8>, SearchError> {
    let storage_len = index.storage.len();
    let payload_len = index
        .keys
        .len()
        .checked_mul(8)
        .and_then(|value| value.checked_add(index.squared_norms.len().checked_mul(4)?))
        .and_then(|value| value.checked_add(index.storage.byte_len()?))
        .ok_or(SearchError::CapacityOverflow)?;
    let file_len = HEADER_LEN
        .checked_add(payload_len)
        .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);
    for key in &index.keys {
        bytes.extend_from_slice(&key.to_le_bytes());
    }
    for norm in &index.squared_norms {
        bytes.extend_from_slice(&norm.to_bits().to_le_bytes());
    }
    index.storage.encode_into(&mut bytes);
    let payload_checksum = checksum(&bytes[HEADER_LEN..]);
    put_u32(&mut bytes, 8, VERSION);
    put_u32(
        &mut bytes,
        12,
        u32::try_from(HEADER_LEN).expect("fixed header length fits 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.config.dimensions)?);
    put_u64(&mut bytes, 40, to_u64(index.keys.len())?);
    put_u32(&mut bytes, 48, metric_code(index.config.metric));
    put_u32(&mut bytes, 52, kind_code(index.kind));
    put_u64(&mut bytes, 56, to_u64(index.config.connectivity)?);
    put_u64(&mut bytes, 64, to_u64(index.config.expansion_build)?);
    put_u64(&mut bytes, 72, to_u64(index.config.expansion_query)?);
    put_u64(&mut bytes, 80, to_u64(index.config.replicas)?);
    put_u64(&mut bytes, 88, to_u64(index.config.build_threads)?);
    put_u64(&mut bytes, 96, to_u64(index.config.query_threads)?);
    put_u64(&mut bytes, 104, index.config.seed);
    put_u64(&mut bytes, 112, to_u64(index.words_per_vector)?);
    put_u64(&mut bytes, 120, to_u64(storage_len)?);
    Ok(bytes)
}