weavatrix-search-vector 0.3.1

use super::MappedVectorIndex;
use crate::error::SearchError;
use crate::hit::SearchHit;
use crate::hnsw::SearchPolicy;
use crate::vector::{Candidate, RoutingCandidate, routing_probes_from_signs};
use std::cmp::Reverse;
use std::collections::BinaryHeap;

pub(super) struct MappedScratch {
    pub(super) marks: Vec<u32>,
    pub(super) generation: u32,
    pub(super) candidates: BinaryHeap<Reverse<Candidate>>,
    pub(super) results: BinaryHeap<Candidate>,
    pub(super) routing_probes: Vec<u16>,
    pub(super) routing_probe_heap: BinaryHeap<Reverse<RoutingCandidate>>,
}

impl MappedScratch {
    pub(super) fn new(len: usize) -> Result<Self, SearchError> {
        let mut marks = Vec::new();
        marks
            .try_reserve_exact(len)
            .map_err(|_| SearchError::AllocationFailed)?;
        marks.resize(len, 0);
        Ok(Self {
            marks,
            generation: 0,
            candidates: BinaryHeap::new(),
            results: BinaryHeap::new(),
            routing_probes: Vec::new(),
            routing_probe_heap: BinaryHeap::new(),
        })
    }

    pub(super) fn begin(&mut self) {
        self.candidates.clear();
        self.results.clear();
        self.generation = self.generation.wrapping_add(1);
        if self.generation == 0 {
            self.marks.fill(0);
            self.generation = 1;
        }
    }

    pub(super) fn mark(&mut self, index: usize) -> bool {
        if self.marks[index] == self.generation {
            return false;
        }
        self.marks[index] = self.generation;
        true
    }
}

impl MappedVectorIndex {
    /// Returns approximate top-K hits ordered by exact distance and key.
    ///
    /// # Errors
    ///
    /// Returns a typed error for an invalid query or scratch allocation.
    pub fn search(&self, query: &[f32], count: usize) -> Result<Vec<SearchHit>, SearchError> {
        self.search_with_policy(
            query,
            count,
            SearchPolicy::new(self.header.config.expansion_query),
        )
    }

    /// Returns approximate top-K hits with a per-query recall policy.
    ///
    /// # Errors
    ///
    /// Returns a typed error for an invalid policy, query, or scratch
    /// allocation.
    pub fn search_with_policy(
        &self,
        query: &[f32],
        count: usize,
        policy: SearchPolicy,
    ) -> Result<Vec<SearchHit>, SearchError> {
        self.search_where(query, count, policy.validate()?, &|_| true)
    }

    pub(super) fn search_where<F>(
        &self,
        query: &[f32],
        count: usize,
        policy: SearchPolicy,
        accepts: &F,
    ) -> Result<Vec<SearchHit>, SearchError>
    where
        F: Fn(u64) -> bool,
    {
        let query_norm = self.query_squared_norm(query)?;
        let limit = count.min(self.len());
        if limit == 0 {
            return Ok(Vec::new());
        }
        let expansion = policy.expansion.max(limit);
        let mut scratch = MappedScratch::new(self.len())?;
        let mut merged = Vec::new();
        merged
            .try_reserve(self.graphs.len().saturating_mul(expansion))
            .map_err(|_| SearchError::AllocationFailed)?;
        for graph in &self.graphs {
            graph.search_into(
                self,
                query,
                query_norm,
                expansion,
                &mut scratch,
                &mut merged,
                accepts,
            );
        }
        self.append_routing(
            query,
            query_norm,
            policy,
            accepts,
            &mut scratch,
            &mut merged,
        )?;
        merged.sort_unstable();
        collect_hits(self, merged, limit)
    }

    fn append_routing<F>(
        &self,
        query: &[f32],
        query_norm: f32,
        policy: SearchPolicy,
        accepts: &F,
        scratch: &mut MappedScratch,
        merged: &mut Vec<Candidate>,
    ) -> Result<(), SearchError>
    where
        F: Fn(u64) -> bool,
    {
        routing_probes_from_signs(
            query,
            self.routing_signs(),
            policy.routing_probes,
            &mut scratch.routing_probes,
            &mut scratch.routing_probe_heap,
        )?;
        let codes = self.routing_codes();
        let nodes = self.routing_nodes();
        for code in scratch.routing_probes.iter().copied() {
            let start = codes.partition_point(|entry| *entry < code);
            let end = codes.partition_point(|entry| *entry <= code);
            merged
                .try_reserve(end.saturating_sub(start))
                .map_err(|_| SearchError::AllocationFailed)?;
            merged.extend(
                nodes[start..end]
                    .iter()
                    .copied()
                    .map(|node| node as usize)
                    .filter(|index| accepts(self.key_slice()[*index]))
                    .map(|index| {
                        Candidate::new(self.distance_query(index, query, query_norm), index)
                    }),
            );
        }
        Ok(())
    }
}

fn collect_hits(
    index: &MappedVectorIndex,
    candidates: Vec<Candidate>,
    limit: usize,
) -> Result<Vec<SearchHit>, SearchError> {
    let mut hits = Vec::new();
    hits.try_reserve_exact(limit)
        .map_err(|_| SearchError::AllocationFailed)?;
    for candidate in candidates {
        let key = index.key_slice()[candidate.index()];
        if hits.iter().any(|hit: &SearchHit| hit.key == key) {
            continue;
        }
        hits.push(SearchHit {
            key,
            distance: candidate.distance,
        });
        if hits.len() == limit {
            break;
        }
    }
    Ok(hits)
}