weavatrix-search-vector 0.3.1

use super::graph_helpers::{empty_node, level_for};
use super::{Graph, SearchScratch};
use crate::config::IndexConfig;
use crate::error::SearchError;
use crate::vector::{VectorStore, splitmix64};
use std::sync::Mutex;
use std::sync::atomic::{AtomicUsize, Ordering};

impl Graph {
    fn build(
        vectors: &VectorStore,
        config: &IndexConfig,
        replica: usize,
        workers: usize,
    ) -> Result<Self, SearchError> {
        let mut nodes = Vec::new();
        nodes
            .try_reserve_exact(vectors.len())
            .map_err(|_| SearchError::AllocationFailed)?;
        for index in 0..vectors.len() {
            let level = level_for(
                config.seed,
                replica,
                vectors.key(index),
                config.connectivity,
            );
            nodes.push(empty_node(level)?);
        }
        let mut graph = Self {
            entry: None,
            max_level: 0,
            nodes,
        };
        let mut order = (0..vectors.len()).collect::<Vec<_>>();
        let replica_seed = config.seed
            ^ u64::try_from(replica)
                .unwrap_or(u64::MAX)
                .wrapping_mul(0xd6e8_feb8_6659_fd93);
        order.sort_unstable_by_key(|index| {
            (
                splitmix64(replica_seed ^ vectors.key(*index)),
                vectors.key(*index),
            )
        });
        let mut scratch = SearchScratch::new(vectors.len())?;
        let workers = workers.min(order.len()).max(1);
        if workers == 1 {
            for index in order {
                graph.insert(vectors, config, index, &mut scratch)?;
            }
            return Ok(graph);
        }
        let seed_count = order.len().min((workers * 8).clamp(64, 256));
        for &index in &order[..seed_count] {
            graph.insert(vectors, config, index, &mut scratch)?;
        }
        let wave_size = (workers * 32).clamp(64, 512);
        for wave in order[seed_count..].chunks(wave_size) {
            graph.insert_wave(vectors, config, wave, workers)?;
        }
        Ok(graph)
    }

    pub(super) fn estimated_bytes(&self) -> usize {
        self.nodes
            .capacity()
            .saturating_mul(std::mem::size_of::<super::NodeLinks>())
            .saturating_add(
                self.nodes
                    .iter()
                    .map(|node| {
                        node.layers
                            .capacity()
                            .saturating_mul(std::mem::size_of::<Vec<u32>>())
                            .saturating_add(
                                node.layers
                                    .iter()
                                    .map(|layer| {
                                        layer.capacity().saturating_mul(std::mem::size_of::<u32>())
                                    })
                                    .sum(),
                            )
                    })
                    .sum(),
            )
    }
}

pub(super) fn build_graphs(
    vectors: &VectorStore,
    config: &IndexConfig,
) -> Result<Vec<Graph>, SearchError> {
    if vectors.is_empty() {
        let mut graphs = Vec::new();
        graphs
            .try_reserve_exact(config.replicas)
            .map_err(|_| SearchError::AllocationFailed)?;
        graphs.extend((0..config.replicas).map(|_| Graph {
            entry: None,
            max_level: 0,
            nodes: Vec::new(),
        }));
        return Ok(graphs);
    }
    let workers = config.build_threads.min(config.replicas).max(1);
    let graph_workers = (config.build_threads / workers).max(1);
    let next = AtomicUsize::new(0);
    let mut slots = Vec::new();
    slots
        .try_reserve_exact(config.replicas)
        .map_err(|_| SearchError::AllocationFailed)?;
    slots.extend(std::iter::repeat_with(|| None).take(config.replicas));
    let slots = Mutex::new(slots);
    let panicked = std::thread::scope(|scope| {
        let handles = (0..workers)
            .map(|_| {
                let slots = &slots;
                let next = &next;
                scope.spawn(move || {
                    loop {
                        let replica = next.fetch_add(1, Ordering::Relaxed);
                        if replica >= config.replicas {
                            break;
                        }
                        let graph = Graph::build(vectors, config, replica, graph_workers);
                        slots
                            .lock()
                            .unwrap_or_else(std::sync::PoisonError::into_inner)[replica] =
                            Some(graph);
                    }
                })
            })
            .collect::<Vec<_>>();
        let mut panicked = false;
        for handle in handles {
            panicked |= handle.join().is_err();
        }
        panicked
    });
    if panicked {
        return Err(SearchError::WorkerPanic);
    }
    slots
        .into_inner()
        .unwrap_or_else(std::sync::PoisonError::into_inner)
        .into_iter()
        .map(|slot| slot.ok_or(SearchError::WorkerPanic)?)
        .collect()
}