weavatrix-search-vector 0.3.1

use super::support::current_len;
use super::{MutableSnapshot, MutableState, MutableVectorIndex, VectorRecord};
use crate::config::IndexConfig;
use crate::error::SearchError;
use crate::hnsw::VectorIndex;
use crate::metadata::MetadataIndex;
use crate::simd::DistanceKernel;
use std::collections::{BTreeMap, BTreeSet};
use std::sync::{Arc, RwLock};

impl MutableVectorIndex {
    /// Builds a mutable index and preserves supplied metadata.
    ///
    /// # Errors
    ///
    /// Returns typed config, vector, allocation, capacity, or duplicate-key
    /// errors.
    pub fn build(config: IndexConfig, records: &[VectorRecord]) -> Result<Self, SearchError> {
        let vectors = records
            .iter()
            .map(|record| (record.key, record.vector.as_slice()))
            .collect::<Vec<_>>();
        let base = Arc::new(VectorIndex::build(config.clone(), &vectors)?);
        let mut metadata = MetadataIndex::new();
        for record in records {
            if !record.metadata.is_empty() {
                metadata.insert(record.key, record.metadata.clone());
            }
        }
        Ok(Self {
            config,
            state: RwLock::new(MutableState {
                base,
                sealed: None,
                pending: BTreeMap::new(),
                deleted: BTreeSet::new(),
                metadata,
                generation: 0,
            }),
            distance_kernel: DistanceKernel::detect(),
        })
    }

    /// Wraps an existing immutable index with an empty mutable delta.
    #[must_use]
    pub fn from_index(index: VectorIndex) -> Self {
        let config = index.config().clone();
        Self {
            config,
            state: RwLock::new(MutableState {
                base: Arc::new(index),
                sealed: None,
                pending: BTreeMap::new(),
                deleted: BTreeSet::new(),
                metadata: MetadataIndex::new(),
                generation: 0,
            }),
            distance_kernel: DistanceKernel::detect(),
        }
    }

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

    #[must_use]
    pub fn len(&self) -> usize {
        let state = self
            .state
            .read()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        current_len(&state)
    }

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

    #[must_use]
    pub fn delta_len(&self) -> usize {
        let state = self
            .state
            .read()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        state
            .sealed
            .as_ref()
            .map_or(0, |index| index.len())
            .saturating_add(state.pending.len())
            .saturating_add(state.deleted.len())
    }

    /// Number of writes still using exact delta search.
    #[must_use]
    pub fn staged_len(&self) -> usize {
        let state = self
            .state
            .read()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        state.pending.len()
    }

    /// Number of live vectors in the sealed HNSW delta.
    #[must_use]
    pub fn sealed_len(&self) -> usize {
        let state = self
            .state
            .read()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        state.sealed.as_ref().map_or(0, |index| index.len())
    }

    #[must_use]
    pub fn should_compact(&self, maximum_delta: usize) -> bool {
        self.delta_len() >= maximum_delta
    }

    pub(crate) fn snapshot(&self) -> MutableSnapshot {
        let state = self
            .state
            .read()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        MutableSnapshot {
            config: self.config.clone(),
            base: Arc::clone(&state.base),
            sealed: state.sealed.clone(),
            pending: state.pending.clone(),
            deleted: state.deleted.clone(),
            metadata: state.metadata.clone(),
        }
    }

    pub(crate) fn from_snapshot(snapshot: MutableSnapshot) -> Self {
        Self {
            config: snapshot.config,
            state: RwLock::new(MutableState {
                base: snapshot.base,
                sealed: snapshot.sealed,
                pending: snapshot.pending,
                deleted: snapshot.deleted,
                metadata: snapshot.metadata,
                generation: 0,
            }),
            distance_kernel: DistanceKernel::detect(),
        }
    }
}