weavatrix-search-vector 0.3.1

use super::{VectorStore, splitmix64};
use crate::config::DistanceMetric;
use crate::error::SearchError;
use crate::simd::DistanceKernel;
use crate::vector::squared_norm;

impl VectorStore {
    pub(crate) fn build(
        dimensions: usize,
        metric: DistanceMetric,
        vectors: &[(u64, &[f32])],
    ) -> Result<Self, SearchError> {
        if vectors.len() > u32::MAX as usize {
            return Err(SearchError::CapacityOverflow);
        }
        let elements = dimensions
            .checked_mul(vectors.len())
            .ok_or(SearchError::CapacityOverflow)?;
        let mut order = (0..vectors.len()).collect::<Vec<_>>();
        order.sort_unstable_by_key(|index| vectors[*index].0);
        for pair in order.windows(2) {
            if vectors[pair[0]].0 == vectors[pair[1]].0 {
                return Err(SearchError::DuplicateKey(vectors[pair[0]].0));
            }
        }

        let mut keys = Vec::new();
        keys.try_reserve_exact(vectors.len())
            .map_err(|_| SearchError::AllocationFailed)?;
        let mut values = Vec::new();
        values
            .try_reserve_exact(elements)
            .map_err(|_| SearchError::AllocationFailed)?;
        let mut squared_norms = Vec::new();
        squared_norms
            .try_reserve_exact(vectors.len())
            .map_err(|_| SearchError::AllocationFailed)?;
        let routing_signs = routing_signs(dimensions)?;
        for (sorted_index, source_index) in order.into_iter().enumerate() {
            let (key, vector) = vectors[source_index];
            if vector.len() != dimensions {
                return Err(SearchError::DimensionMismatch {
                    expected: dimensions,
                    actual: vector.len(),
                    vector: Some(source_index),
                });
            }
            let squared_norm = squared_norm(vector, Some(source_index))?;
            keys.push(key);
            if metric == DistanceMetric::Cosine {
                if squared_norm == 0.0 {
                    return Err(SearchError::ZeroVector {
                        vector: Some(source_index),
                    });
                }
                let inverse_norm = squared_norm.sqrt().recip();
                values.extend(vector.iter().map(|value| value * inverse_norm));
                squared_norms.push(1.0);
            } else {
                values.extend_from_slice(vector);
                squared_norms.push(squared_norm);
            }
            debug_assert_eq!(keys.len(), sorted_index + 1);
        }
        Ok(Self {
            dimensions,
            metric,
            keys,
            values,
            squared_norms,
            routing_signs,
            distance_kernel: DistanceKernel::detect(),
        })
    }

    pub(crate) fn from_values(
        dimensions: usize,
        metric: DistanceMetric,
        keys: Vec<u64>,
        values: Vec<f32>,
    ) -> Result<Self, SearchError> {
        let expected = dimensions
            .checked_mul(keys.len())
            .ok_or(SearchError::CapacityOverflow)?;
        if values.len() != expected {
            return Err(SearchError::CorruptSnapshot(
                "vector payload length does not match dimensions and key count",
            ));
        }
        if keys.windows(2).any(|pair| pair[0] >= pair[1]) {
            return Err(SearchError::CorruptSnapshot(
                "vector keys are not strictly increasing",
            ));
        }
        if values.iter().any(|value| !value.is_finite()) {
            return Err(SearchError::CorruptSnapshot(
                "vector payload contains a non-finite value",
            ));
        }
        let routing_signs = routing_signs(dimensions)?;
        let mut squared_norms = Vec::new();
        squared_norms
            .try_reserve_exact(keys.len())
            .map_err(|_| SearchError::AllocationFailed)?;
        for (index, vector) in values.chunks_exact(dimensions).enumerate() {
            let norm = squared_norm(vector, Some(index))?;
            if metric == DistanceMetric::Cosine && norm == 0.0 {
                return Err(SearchError::CorruptSnapshot(
                    "cosine snapshot contains a zero vector",
                ));
            }
            squared_norms.push(norm);
        }
        Ok(Self {
            dimensions,
            metric,
            keys,
            values,
            squared_norms,
            routing_signs,
            distance_kernel: DistanceKernel::detect(),
        })
    }
}

fn routing_signs(dimensions: usize) -> Result<Vec<u16>, SearchError> {
    let mut signs = Vec::new();
    signs
        .try_reserve_exact(dimensions)
        .map_err(|_| SearchError::AllocationFailed)?;
    signs.extend((0..dimensions).map(|dimension| {
        let mixed =
            splitmix64(u64::try_from(dimension).unwrap_or(u64::MAX) ^ 0xa076_1d64_78bd_642f);
        u16::try_from(mixed & u64::from(u16::MAX)).expect("masked routing signs fit u16")
    }));
    Ok(signs)
}