weavatrix-search-vector 0.3.1

use super::{FNV_OFFSET, FNV_PRIME};
use crate::config::DistanceMetric;
use crate::error::SearchError;
use crate::quantized::QuantizationKind;

pub(super) fn decode_u16s(
    bytes: &[u8],
    cursor: &mut usize,
    length: usize,
) -> Result<Vec<u16>, SearchError> {
    let raw = take(
        bytes,
        cursor,
        length.checked_mul(2).ok_or(SearchError::CapacityOverflow)?,
    )?;
    Ok(raw
        .chunks_exact(2)
        .map(|chunk| u16::from_le_bytes(chunk.try_into().expect("two-byte chunk")))
        .collect())
}

pub(super) fn decode_u64s(
    bytes: &[u8],
    cursor: &mut usize,
    length: usize,
) -> Result<Vec<u64>, SearchError> {
    let raw = take(
        bytes,
        cursor,
        length.checked_mul(8).ok_or(SearchError::CapacityOverflow)?,
    )?;
    Ok(raw
        .chunks_exact(8)
        .map(|chunk| u64::from_le_bytes(chunk.try_into().expect("eight-byte chunk")))
        .collect())
}

pub(super) fn decode_f32s(
    bytes: &[u8],
    cursor: &mut usize,
    length: usize,
) -> Result<Vec<f32>, SearchError> {
    let raw = take(
        bytes,
        cursor,
        length.checked_mul(4).ok_or(SearchError::CapacityOverflow)?,
    )?;
    Ok(raw
        .chunks_exact(4)
        .map(|chunk| {
            f32::from_bits(u32::from_le_bytes(
                chunk.try_into().expect("four-byte chunk"),
            ))
        })
        .collect())
}

pub(super) fn take<'a>(
    bytes: &'a [u8],
    cursor: &mut usize,
    length: usize,
) -> Result<&'a [u8], SearchError> {
    let end = cursor
        .checked_add(length)
        .ok_or(SearchError::CapacityOverflow)?;
    let value = bytes.get(*cursor..end).ok_or(SearchError::CorruptSnapshot(
        "quantized snapshot payload is truncated",
    ))?;
    *cursor = end;
    Ok(value)
}

pub(super) const fn metric_code(metric: DistanceMetric) -> u32 {
    match metric {
        DistanceMetric::Cosine => 0,
        DistanceMetric::Dot => 1,
        DistanceMetric::SquaredEuclidean => 2,
    }
}

pub(super) fn decode_metric(code: u32) -> Result<DistanceMetric, SearchError> {
    match code {
        0 => Ok(DistanceMetric::Cosine),
        1 => Ok(DistanceMetric::Dot),
        2 => Ok(DistanceMetric::SquaredEuclidean),
        _ => Err(SearchError::CorruptSnapshot(
            "quantized snapshot metric is unknown",
        )),
    }
}

pub(super) const fn kind_code(kind: QuantizationKind) -> u32 {
    match kind {
        QuantizationKind::BFloat16 => 0,
        QuantizationKind::Float16 => 1,
        QuantizationKind::Float8E4M3 => 2,
        QuantizationKind::Int8 => 3,
        QuantizationKind::Binary => 4,
    }
}

pub(super) fn decode_kind(code: u32) -> Result<QuantizationKind, SearchError> {
    match code {
        0 => Ok(QuantizationKind::BFloat16),
        1 => Ok(QuantizationKind::Float16),
        2 => Ok(QuantizationKind::Float8E4M3),
        3 => Ok(QuantizationKind::Int8),
        4 => Ok(QuantizationKind::Binary),
        _ => Err(SearchError::CorruptSnapshot(
            "quantized snapshot kind is unknown",
        )),
    }
}

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

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

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

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

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

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

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