weavatrix-search-vector 0.3.1

use super::format::checksum;
use super::mapped_validation::validate_sections;
use super::{HEADER_LEN, MappedVectorIndex, SnapshotValidation};
use crate::config::{DistanceMetric, IndexConfig};
use crate::error::SearchError;
use crate::mmap::Mapping;
use crate::simd::DistanceKernel;
use crate::vector::{distance, splitmix64, squared_norm};
use std::path::Path;

impl MappedVectorIndex {
    /// Opens and fully validates a read-only memory-mapped snapshot.
    ///
    /// # Errors
    ///
    /// Returns a typed storage, version, integrity, or config error.
    pub fn open(path: impl AsRef<Path>) -> Result<Self, SearchError> {
        Self::open_with_validation(path, SnapshotValidation::Full)
    }

    /// Opens a read-only snapshot with an explicit validation policy.
    ///
    /// # Errors
    ///
    /// Returns a typed storage, version, integrity, or config error.
    pub fn open_with_validation(
        path: impl AsRef<Path>,
        validation: SnapshotValidation,
    ) -> Result<Self, SearchError> {
        let mapping =
            Mapping::open(path.as_ref()).map_err(|error| SearchError::storage("map", &error))?;
        let header = super::Header::parse(mapping.bytes())?;
        if header.file_len != mapping.len() {
            return Err(SearchError::CorruptSnapshot(
                "header file length does not match mapped file",
            ));
        }
        if validation == SnapshotValidation::Full {
            let actual = checksum(&mapping.bytes()[HEADER_LEN..]);
            if actual != header.checksum {
                return Err(SearchError::CorruptSnapshot(
                    "snapshot payload checksum does not match",
                ));
            }
        }
        let graphs = validate_sections(&mapping, &header, validation)?;
        Ok(Self {
            mapping,
            header,
            graphs,
            distance_kernel: DistanceKernel::detect(),
        })
    }

    #[must_use]
    pub fn len(&self) -> usize {
        self.header.count
    }

    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.header.count == 0
    }

    #[must_use]
    pub fn dimensions(&self) -> usize {
        self.header.config.dimensions
    }

    #[must_use]
    pub fn config(&self) -> &IndexConfig {
        &self.header.config
    }

    /// Returns the mapped file length. Query scratch and the small graph-layout
    /// table are excluded.
    #[must_use]
    pub fn mapped_bytes(&self) -> usize {
        self.mapping.len()
    }

    /// Iterates stable keys in ascending order.
    #[must_use]
    pub fn keys(&self) -> impl ExactSizeIterator<Item = u64> + '_ {
        self.key_slice().iter().copied()
    }

    /// Returns the normalized mapped vector for `key`.
    #[must_use]
    pub fn vector(&self, key: u64) -> Option<&[f32]> {
        let index = self.key_slice().binary_search(&key).ok()?;
        Some(self.vector_at(index))
    }

    pub(super) fn key_slice(&self) -> &[u64] {
        self.mapping
            .u64_slice(self.header.keys_offset, self.len())
            .expect("validated key section")
    }

    pub(super) fn vector_slice(&self) -> &[f32] {
        self.mapping
            .f32_slice(self.header.vectors_offset, self.len() * self.dimensions())
            .expect("validated vector section")
    }

    pub(super) fn routing_codes(&self) -> &[u16] {
        self.mapping
            .u16_slice(self.header.routing_codes_offset, self.len())
            .expect("validated routing-code section")
    }

    pub(super) fn routing_nodes(&self) -> &[u32] {
        self.mapping
            .u32_slice(self.header.routing_nodes_offset, self.len())
            .expect("validated routing-node section")
    }

    pub(super) fn vector_at(&self, index: usize) -> &[f32] {
        let start = index * self.dimensions();
        &self.vector_slice()[start..start + self.dimensions()]
    }

    pub(super) fn query_squared_norm(&self, query: &[f32]) -> Result<f32, SearchError> {
        if query.len() != self.dimensions() {
            return Err(SearchError::DimensionMismatch {
                expected: self.dimensions(),
                actual: query.len(),
                vector: None,
            });
        }
        let norm = squared_norm(query, None)?;
        if self.header.config.metric == DistanceMetric::Cosine && norm == 0.0 {
            return Err(SearchError::ZeroVector { vector: None });
        }
        Ok(norm)
    }

    pub(super) fn distance_query(
        &self,
        index: usize,
        query: &[f32],
        query_squared_norm: f32,
    ) -> f32 {
        let vector = self.vector_at(index);
        let vector_squared_norm = if self.header.config.metric == DistanceMetric::Cosine {
            1.0
        } else {
            squared_norm(vector, Some(index)).expect("validated mapped vector")
        };
        distance(
            self.distance_kernel,
            self.header.config.metric,
            vector,
            vector_squared_norm,
            query,
            query_squared_norm,
        )
    }

    pub(super) fn routing_signs(&self) -> impl Iterator<Item = u16> + '_ {
        (0..self.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")
        })
    }
}