weavatrix-search-vector 0.3.1

use super::{MAX_LEVEL, NodeLinks};
use crate::config::IndexConfig;
use crate::error::SearchError;
use crate::vector::{Candidate, VectorStore, splitmix64};

pub(super) fn select_neighbors(
    vectors: &VectorStore,
    target: usize,
    mut candidates: Vec<Candidate>,
    limit: usize,
) -> Vec<Candidate> {
    candidates.sort_unstable();
    candidates.dedup_by_key(|candidate| candidate.index());
    candidates.retain(|candidate| candidate.index() != target);
    let mut selected = Vec::with_capacity(limit);
    let mut deferred = Vec::new();
    for candidate in candidates {
        let diverse = selected.iter().all(|neighbor: &Candidate| {
            vectors.distance_indices(candidate.index(), neighbor.index()) >= candidate.distance
        });
        if diverse && selected.len() < limit {
            selected.push(candidate);
        } else {
            deferred.push(candidate);
        }
    }
    if selected.len() < limit {
        selected.extend(deferred.into_iter().take(limit - selected.len()));
    }
    selected
}

pub(super) fn max_degree(config: &IndexConfig, level: usize) -> Result<usize, SearchError> {
    if level == 0 {
        config
            .connectivity
            .checked_mul(2)
            .ok_or(SearchError::CapacityOverflow)
    } else {
        Ok(config.connectivity)
    }
}

pub(super) fn level_for(seed: u64, replica: usize, key: u64, connectivity: usize) -> usize {
    let replica = u64::try_from(replica).unwrap_or(u64::MAX);
    let divisor = u64::try_from(connectivity).unwrap_or(u64::MAX).max(2);
    let mut random = splitmix64(seed ^ key ^ replica.wrapping_mul(0xa076_1d64_78bd_642f));
    let mut level = 0;
    while level < MAX_LEVEL && random.is_multiple_of(divisor) {
        level += 1;
        random = splitmix64(random);
    }
    level
}

pub(super) fn empty_node(level: usize) -> Result<NodeLinks, SearchError> {
    let mut layers = Vec::new();
    layers
        .try_reserve_exact(level + 1)
        .map_err(|_| SearchError::AllocationFailed)?;
    layers.extend(std::iter::repeat_with(Vec::new).take(level + 1));
    Ok(NodeLinks { layers })
}

pub(super) fn to_node(index: usize) -> u32 {
    u32::try_from(index).expect("vector count was validated before graph construction")
}

pub(super) fn from_node(index: u32) -> usize {
    usize::try_from(index).expect("u32 node index fits usize on supported platforms")
}