weavatrix-search-vector 0.3.1

use crate::error::SearchError;
use crate::hit::SearchHit;
use crate::hnsw::VectorIndex;
use crate::parallel;

/// Directed nearest-neighbor candidate edge.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct NeighborEdge {
    pub source: u64,
    pub target: u64,
    pub distance: f32,
}

/// Deterministic directed KNN candidate graph in compressed-row form.
///
/// This is deliberately not a semantic or domain graph: consumers decide how
/// to interpret, threshold, symmetrize, or add provenance to candidate edges.
#[derive(Debug, Clone, PartialEq)]
pub struct KnnGraph {
    nodes: Vec<u64>,
    offsets: Vec<usize>,
    edges: Vec<NeighborEdge>,
}

impl KnnGraph {
    /// Builds an approximate directed KNN graph with bounded query workers.
    ///
    /// # Errors
    ///
    /// Returns a typed query, worker, capacity, or allocation error.
    pub fn build(index: &VectorIndex, neighbors: usize) -> Result<Self, SearchError> {
        Self::build_with(index, neighbors, false)
    }

    /// Builds an exact directed KNN graph for oracle and small-corpus use.
    ///
    /// # Errors
    ///
    /// Returns a typed query, worker, capacity, or allocation error.
    pub fn build_exact(index: &VectorIndex, neighbors: usize) -> Result<Self, SearchError> {
        Self::build_with(index, neighbors, true)
    }

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

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

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

    #[must_use]
    pub fn nodes(&self) -> impl ExactSizeIterator<Item = u64> + '_ {
        self.nodes.iter().copied()
    }

    #[must_use]
    pub fn neighbors(&self, source: u64) -> Option<&[NeighborEdge]> {
        let index = self.nodes.binary_search(&source).ok()?;
        Some(&self.edges[self.offsets[index]..self.offsets[index + 1]])
    }

    #[must_use]
    pub fn edges(&self) -> impl ExactSizeIterator<Item = NeighborEdge> + '_ {
        self.edges.iter().copied()
    }

    fn build_with(index: &VectorIndex, neighbors: usize, exact: bool) -> Result<Self, SearchError> {
        let nodes = index.keys().collect::<Vec<_>>();
        let per_node = neighbors.min(index.len().saturating_sub(1));
        let queries = nodes
            .iter()
            .map(|key| {
                index
                    .vector(*key)
                    .expect("index key iterator resolves to a vector")
            })
            .collect::<Vec<_>>();
        let rows = if exact {
            parallel::search_batch(&queries, index.config().query_threads, |query| {
                index.search_exact(query, per_node.saturating_add(1))
            })?
        } else {
            index.search_batch(&queries, per_node.saturating_add(1))?
        };
        Self::from_rows(nodes, rows, per_node)
    }

    fn from_rows(
        nodes: Vec<u64>,
        rows: Vec<Vec<SearchHit>>,
        per_node: usize,
    ) -> Result<Self, SearchError> {
        let capacity = nodes
            .len()
            .checked_mul(per_node)
            .ok_or(SearchError::CapacityOverflow)?;
        let mut offsets = Vec::new();
        offsets
            .try_reserve_exact(nodes.len().saturating_add(1))
            .map_err(|_| SearchError::AllocationFailed)?;
        let mut edges = Vec::new();
        edges
            .try_reserve_exact(capacity)
            .map_err(|_| SearchError::AllocationFailed)?;
        offsets.push(0);
        for (source, row) in nodes.iter().copied().zip(rows) {
            edges.extend(
                row.into_iter()
                    .filter(|hit| hit.key != source)
                    .take(per_node)
                    .map(|hit| NeighborEdge {
                        source,
                        target: hit.key,
                        distance: hit.distance,
                    }),
            );
            offsets.push(edges.len());
        }
        Ok(Self {
            nodes,
            offsets,
            edges,
        })
    }
}