weavatrix-search-vector 0.2.0

use crate::vector::{Candidate, VectorStore, splitmix64};
use crate::{IndexConfig, SearchError, SearchHit};
use std::cmp::Reverse;
use std::collections::BinaryHeap;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};

const MAX_LEVEL: usize = 32;

/// Immutable deterministic HNSW candidate index.
#[derive(Debug)]
pub struct VectorIndex {
    config: IndexConfig,
    vectors: Arc<VectorStore>,
    graphs: Vec<Graph>,
    routing: RoutingIndex,
}

impl VectorIndex {
    /// Builds independently seeded HNSW replicas over validated normalized
    /// vectors.
    ///
    /// Input order does not affect the resulting graph when keys, vectors,
    /// config, and seed are unchanged.
    ///
    /// # Errors
    ///
    /// Returns typed configuration, vector, capacity, allocation, or worker
    /// failures.
    pub fn build(config: IndexConfig, vectors: &[(u64, &[f32])]) -> Result<Self, SearchError> {
        config.validate()?;
        let vectors = Arc::new(VectorStore::build(config.dimensions, vectors)?);
        let graphs = build_graphs(&vectors, &config)?;
        let routing = RoutingIndex::build(&vectors)?;
        Ok(Self {
            config,
            vectors,
            graphs,
            routing,
        })
    }

    pub(crate) fn from_parts(
        config: IndexConfig,
        vectors: VectorStore,
        graphs: Vec<Graph>,
        routing: RoutingIndex,
    ) -> Result<Self, SearchError> {
        config.validate()?;
        if graphs.len() != config.replicas {
            return Err(SearchError::CorruptSnapshot(
                "graph replica count does not match index config",
            ));
        }
        if graphs
            .iter()
            .any(|graph| graph.nodes.len() != vectors.len())
        {
            return Err(SearchError::CorruptSnapshot(
                "graph node count does not match vector count",
            ));
        }
        Ok(Self {
            config,
            vectors: Arc::new(vectors),
            graphs,
            routing,
        })
    }

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

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

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

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

    /// Iterates stable caller-provided keys in ascending order.
    #[must_use]
    pub fn keys(&self) -> impl ExactSizeIterator<Item = u64> + '_ {
        self.vectors.keys().iter().copied()
    }

    /// Returns the normalized vector stored for `key`.
    #[must_use]
    pub fn vector(&self, key: u64) -> Option<&[f32]> {
        self.vectors
            .find_index(key)
            .map(|index| self.vectors.vector(index))
    }

    pub(crate) fn vectors(&self) -> &VectorStore {
        &self.vectors
    }

    pub(crate) fn graphs(&self) -> &[Graph] {
        &self.graphs
    }

    pub(crate) fn routing(&self) -> &RoutingIndex {
        &self.routing
    }

    /// Returns an allocation-based estimate of resident vector and graph
    /// storage. It excludes allocator metadata and temporary query scratch.
    #[must_use]
    pub fn estimated_memory_bytes(&self) -> usize {
        self.vectors
            .estimated_bytes()
            .saturating_add(self.graphs.iter().map(Graph::estimated_bytes).sum())
            .saturating_add(self.routing.estimated_bytes())
    }

    /// Returns approximate top-K hits ordered by exact cosine 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> {
        let mut scratch = SearchScratch::new(self.len())?;
        self.search_with_scratch(query, count, &mut scratch)
    }

    /// Returns exact top-K results over the same normalized vectors.
    ///
    /// This is intended as a correctness and recall oracle.
    ///
    /// # Errors
    ///
    /// Returns a typed error for an invalid query.
    pub fn search_exact(&self, query: &[f32], count: usize) -> Result<Vec<SearchHit>, SearchError> {
        self.vectors.exact(query, count)
    }

    /// Returns exact top-K results accepted by `filter`.
    ///
    /// # Errors
    ///
    /// Returns a typed error for an invalid query.
    pub fn search_exact_filtered<F>(
        &self,
        query: &[f32],
        count: usize,
        filter: F,
    ) -> Result<Vec<SearchHit>, SearchError>
    where
        F: Fn(u64) -> bool,
    {
        self.vectors.exact_filtered(query, count, filter)
    }

    /// Searches HNSW candidates accepted by `filter`, with an exact fallback
    /// when selective filters do not yield enough approximate candidates.
    ///
    /// # Errors
    ///
    /// Returns a typed error for an invalid query or scratch allocation.
    pub fn search_filtered<F>(
        &self,
        query: &[f32],
        count: usize,
        filter: F,
    ) -> Result<Vec<SearchHit>, SearchError>
    where
        F: Fn(u64) -> bool,
    {
        let requested = count.min(self.len());
        if requested == 0 {
            self.vectors.query_inverse_norm(query)?;
            return Ok(Vec::new());
        }
        let candidate_count = requested
            .saturating_mul(8)
            .max(self.config.expansion_query.saturating_mul(2))
            .min(self.len());
        let mut hits = self.search(query, candidate_count)?;
        hits.retain(|hit| filter(hit.key));
        hits.truncate(requested);
        if hits.len() == requested {
            return Ok(hits);
        }
        self.vectors.exact_filtered(query, requested, filter)
    }

    /// Searches independent queries using bounded scoped workers. Each worker
    /// reuses its visited-set and heaps across queries.
    ///
    /// # Errors
    ///
    /// Returns the first query error in input order or a worker-panic error.
    pub fn search_batch(
        &self,
        queries: &[&[f32]],
        count: usize,
    ) -> Result<Vec<Vec<SearchHit>>, SearchError> {
        if queries.is_empty() {
            return Ok(Vec::new());
        }
        let workers = self.config.query_threads.min(queries.len()).max(1);
        let chunk_size = queries.len().div_ceil(workers);
        let mut output = Vec::new();
        output
            .try_reserve_exact(queries.len())
            .map_err(|_| SearchError::AllocationFailed)?;
        output.extend(std::iter::repeat_with(|| None).take(queries.len()));
        let panicked = std::thread::scope(|scope| {
            let handles = queries
                .chunks(chunk_size)
                .enumerate()
                .map(|(chunk_index, chunk)| {
                    let start = chunk_index * chunk_size;
                    scope.spawn(move || {
                        let mut scratch = SearchScratch::new(self.len());
                        let mut local = Vec::with_capacity(chunk.len());
                        for (offset, query) in chunk.iter().enumerate() {
                            let result = match &mut scratch {
                                Ok(scratch) => self.search_with_scratch(query, count, scratch),
                                Err(error) => Err(error.clone()),
                            };
                            local.push((start + offset, result));
                        }
                        local
                    })
                })
                .collect::<Vec<_>>();
            let mut panicked = false;
            for handle in handles {
                match handle.join() {
                    Ok(local) => {
                        for (index, result) in local {
                            output[index] = Some(result);
                        }
                    }
                    Err(_) => panicked = true,
                }
            }
            panicked
        });
        if panicked {
            return Err(SearchError::WorkerPanic);
        }
        output
            .into_iter()
            .map(|slot| slot.ok_or(SearchError::WorkerPanic)?)
            .collect()
    }

    fn search_with_scratch(
        &self,
        query: &[f32],
        count: usize,
        scratch: &mut SearchScratch,
    ) -> Result<Vec<SearchHit>, SearchError> {
        let inverse_norm = self.vectors.query_inverse_norm(query)?;
        let limit = count.min(self.len());
        if limit == 0 {
            return Ok(Vec::new());
        }
        let expansion = self.config.expansion_query.max(limit);
        let mut merged = std::mem::take(&mut scratch.merged);
        merged.clear();
        merged
            .try_reserve(self.graphs.len().saturating_mul(expansion))
            .map_err(|_| SearchError::AllocationFailed)?;
        for graph in &self.graphs {
            graph.search_into(
                &self.vectors,
                query,
                inverse_norm,
                expansion,
                scratch,
                &mut merged,
            );
        }
        let mut routing_nodes = std::mem::take(&mut scratch.routing_nodes);
        routing_nodes.clear();
        self.routing
            .append_candidates(&self.vectors.routing_probes(query), &mut routing_nodes);
        merged
            .try_reserve(routing_nodes.len())
            .map_err(|_| SearchError::AllocationFailed)?;
        merged.extend(routing_nodes.iter().copied().map(from_node).map(|index| {
            Candidate::new(
                self.vectors.distance_query(index, query, inverse_norm),
                index,
            )
        }));
        routing_nodes.clear();
        scratch.routing_nodes = routing_nodes;
        merged.sort_unstable();
        let mut hits = Vec::new();
        hits.try_reserve_exact(limit)
            .map_err(|_| SearchError::AllocationFailed)?;
        for candidate in merged.iter().copied() {
            let key = self.vectors.key(candidate.index());
            if hits.iter().any(|hit: &SearchHit| hit.key == key) {
                continue;
            }
            hits.push(SearchHit {
                key,
                distance: candidate.distance,
            });
            if hits.len() == limit {
                break;
            }
        }
        merged.clear();
        scratch.merged = merged;
        Ok(hits)
    }
}

#[derive(Debug)]
pub(crate) struct RoutingIndex {
    pub(crate) entries: Vec<(u16, u32)>,
}

impl RoutingIndex {
    pub(crate) fn build(vectors: &VectorStore) -> Result<Self, SearchError> {
        let mut entries = Vec::new();
        entries
            .try_reserve_exact(vectors.len())
            .map_err(|_| SearchError::AllocationFailed)?;
        entries.extend((0..vectors.len()).map(|index| {
            (
                vectors.stored_routing_code(index),
                u32::try_from(index).expect("vector count is bounded by u32::MAX"),
            )
        }));
        entries.sort_unstable();
        Ok(Self { entries })
    }

    fn estimated_bytes(&self) -> usize {
        self.entries
            .capacity()
            .saturating_mul(std::mem::size_of::<(u16, u32)>())
    }

    fn append_candidates(&self, codes: &[u16], output: &mut Vec<u32>) {
        for code in codes {
            self.append_bucket(*code, output);
        }
    }

    fn append_bucket(&self, code: u16, output: &mut Vec<u32>) {
        let start = self.entries.partition_point(|entry| entry.0 < code);
        let end = self.entries.partition_point(|entry| entry.0 <= code);
        output.extend(self.entries[start..end].iter().map(|entry| entry.1));
    }
}

#[derive(Debug)]
pub(crate) struct NodeLinks {
    pub(crate) layers: Vec<Vec<u32>>,
}

#[derive(Debug)]
pub(crate) struct Graph {
    pub(crate) entry: Option<usize>,
    pub(crate) max_level: usize,
    pub(crate) nodes: Vec<NodeLinks>,
}

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,
            );
            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));
            nodes.push(NodeLinks { layers });
        }
        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)
    }

    fn estimated_bytes(&self) -> usize {
        self.nodes
            .capacity()
            .saturating_mul(std::mem::size_of::<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(),
            )
    }

    fn insert(
        &mut self,
        vectors: &VectorStore,
        config: &IndexConfig,
        index: usize,
        scratch: &mut SearchScratch,
    ) -> Result<(), SearchError> {
        let level = self.nodes[index].layers.len() - 1;
        let Some(mut entry) = self.entry else {
            self.entry = Some(index);
            self.max_level = level;
            return Ok(());
        };

        for current_level in ((level + 1)..=self.max_level).rev() {
            entry = self.greedy_index(vectors, index, entry, current_level);
        }
        for current_level in (0..=level.min(self.max_level)).rev() {
            let candidates = self.search_layer(
                vectors,
                entry,
                config.expansion_build,
                current_level,
                scratch,
                |candidate| vectors.distance_indices(index, candidate),
            );
            if let Some(best) = candidates.first() {
                entry = best.index();
            }
            let degree = max_degree(config, current_level)?;
            let selected = select_neighbors(vectors, index, candidates, degree);
            self.nodes[index].layers[current_level]
                .extend(selected.iter().map(|candidate| to_node(candidate.index())));
            for candidate in selected {
                self.add_backlink(candidate.index(), index, current_level);
            }
        }
        if level > self.max_level {
            self.entry = Some(index);
            self.max_level = level;
        }
        Ok(())
    }

    fn insert_wave(
        &mut self,
        vectors: &VectorStore,
        config: &IndexConfig,
        wave: &[usize],
        workers: usize,
    ) -> Result<(), SearchError> {
        if wave.is_empty() {
            return Ok(());
        }
        let workers = workers.min(wave.len()).max(1);
        let next = AtomicUsize::new(0);
        let mut slots = Vec::new();
        slots
            .try_reserve_exact(wave.len())
            .map_err(|_| SearchError::AllocationFailed)?;
        slots.extend(std::iter::repeat_with(|| None).take(wave.len()));
        let mut panicked = false;
        std::thread::scope(|scope| {
            let handles = (0..workers)
                .map(|_| {
                    let next = &next;
                    let graph = &*self;
                    scope.spawn(move || {
                        let mut scratch = SearchScratch::new(vectors.len());
                        let mut local = Vec::new();
                        loop {
                            let position = next.fetch_add(1, Ordering::Relaxed);
                            let Some(&index) = wave.get(position) else {
                                break;
                            };
                            let plan = match &mut scratch {
                                Ok(scratch) => graph.plan_insert(vectors, config, index, scratch),
                                Err(error) => Err(error.clone()),
                            };
                            local.push((position, plan));
                        }
                        local
                    })
                })
                .collect::<Vec<_>>();
            for handle in handles {
                match handle.join() {
                    Ok(local) => {
                        for (position, plan) in local {
                            slots[position] = Some(plan);
                        }
                    }
                    Err(_) => panicked = true,
                }
            }
        });
        if panicked {
            return Err(SearchError::WorkerPanic);
        }
        for slot in slots {
            let plan = slot.ok_or(SearchError::WorkerPanic)??;
            self.apply_plan(vectors, plan);
        }
        Ok(())
    }

    fn plan_insert(
        &self,
        vectors: &VectorStore,
        config: &IndexConfig,
        index: usize,
        scratch: &mut SearchScratch,
    ) -> Result<InsertionPlan, SearchError> {
        let level = self.nodes[index].layers.len() - 1;
        let Some(mut entry) = self.entry else {
            return Ok(InsertionPlan {
                index,
                level,
                links: Vec::new(),
            });
        };
        for current_level in ((level + 1)..=self.max_level).rev() {
            entry = self.greedy_index(vectors, index, entry, current_level);
        }
        let link_count = level.min(self.max_level) + 1;
        let mut links = Vec::new();
        links
            .try_reserve_exact(link_count)
            .map_err(|_| SearchError::AllocationFailed)?;
        for current_level in (0..=level.min(self.max_level)).rev() {
            let candidates = self.search_layer(
                vectors,
                entry,
                config.expansion_build,
                current_level,
                scratch,
                |candidate| vectors.distance_indices(index, candidate),
            );
            if let Some(best) = candidates.first() {
                entry = best.index();
            }
            let degree = max_degree(config, current_level)?;
            links.push((
                current_level,
                select_neighbors(vectors, index, candidates, degree)
                    .into_iter()
                    .map(Candidate::index)
                    .collect(),
            ));
        }
        Ok(InsertionPlan {
            index,
            level,
            links,
        })
    }

    fn apply_plan(&mut self, _vectors: &VectorStore, plan: InsertionPlan) {
        for (level, selected) in plan.links {
            self.nodes[plan.index].layers[level].extend(selected.iter().copied().map(to_node));
            for candidate in selected {
                self.add_backlink(candidate, plan.index, level);
            }
        }
        if plan.level > self.max_level {
            self.entry = Some(plan.index);
            self.max_level = plan.level;
        }
    }

    fn add_backlink(&mut self, node: usize, backlink: usize, level: usize) {
        if node == backlink || self.nodes[node].layers.len() <= level {
            return;
        }
        let links = &mut self.nodes[node].layers[level];
        let backlink = to_node(backlink);
        if !links.contains(&backlink) {
            links.push(backlink);
        }
    }

    fn greedy_index(
        &self,
        vectors: &VectorStore,
        target: usize,
        entry: usize,
        level: usize,
    ) -> usize {
        self.greedy(entry, level, |candidate| {
            vectors.distance_indices(target, candidate)
        })
    }

    fn greedy_query(
        &self,
        vectors: &VectorStore,
        query: &[f32],
        inverse_norm: f32,
        entry: usize,
        level: usize,
    ) -> usize {
        self.greedy(entry, level, |candidate| {
            vectors.distance_query(candidate, query, inverse_norm)
        })
    }

    fn greedy<F>(&self, mut best: usize, level: usize, mut distance: F) -> usize
    where
        F: FnMut(usize) -> f32,
    {
        let mut best_candidate = Candidate::new(distance(best), best);
        loop {
            let mut improved = false;
            for neighbor in &self.nodes[best].layers[level] {
                let neighbor = from_node(*neighbor);
                let candidate = Candidate::new(distance(neighbor), neighbor);
                if candidate < best_candidate {
                    best = neighbor;
                    best_candidate = candidate;
                    improved = true;
                }
            }
            if !improved {
                return best;
            }
        }
    }

    fn search_into(
        &self,
        vectors: &VectorStore,
        query: &[f32],
        inverse_norm: f32,
        expansion: usize,
        scratch: &mut SearchScratch,
        output: &mut Vec<Candidate>,
    ) {
        let Some(mut entry) = self.entry else {
            return;
        };
        for level in (1..=self.max_level).rev() {
            entry = self.greedy_query(vectors, query, inverse_norm, entry, level);
        }
        self.explore_layer(entry, expansion, 0, scratch, |candidate| {
            vectors.distance_query(candidate, query, inverse_norm)
        });
        output.extend(scratch.results.iter().copied());
    }

    fn search_layer<F>(
        &self,
        _vectors: &VectorStore,
        entry: usize,
        expansion: usize,
        level: usize,
        scratch: &mut SearchScratch,
        mut distance: F,
    ) -> Vec<Candidate>
    where
        F: FnMut(usize) -> f32,
    {
        self.explore_layer(entry, expansion, level, scratch, &mut distance);
        let mut results = scratch.results.clone().into_vec();
        results.sort_unstable();
        results
    }

    fn explore_layer<F>(
        &self,
        entry: usize,
        expansion: usize,
        level: usize,
        scratch: &mut SearchScratch,
        mut distance: F,
    ) where
        F: FnMut(usize) -> f32,
    {
        scratch.begin();
        let first = Candidate::new(distance(entry), entry);
        scratch.mark(entry);
        scratch.candidates.push(Reverse(first));
        scratch.results.push(first);

        while let Some(Reverse(current)) = scratch.candidates.pop() {
            if scratch.results.len() >= expansion
                && scratch.results.peek().is_some_and(|worst| current > *worst)
            {
                break;
            }
            for neighbor in &self.nodes[current.index()].layers[level] {
                let neighbor = from_node(*neighbor);
                if !scratch.mark(neighbor) {
                    continue;
                }
                let candidate = Candidate::new(distance(neighbor), neighbor);
                if scratch.results.len() < expansion
                    || scratch
                        .results
                        .peek()
                        .is_some_and(|worst| candidate < *worst)
                {
                    scratch.candidates.push(Reverse(candidate));
                    scratch.results.push(candidate);
                    if scratch.results.len() > expansion {
                        scratch.results.pop();
                    }
                }
            }
        }
    }
}

struct InsertionPlan {
    index: usize,
    level: usize,
    links: Vec<(usize, Vec<usize>)>,
}

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
}

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)
    }
}

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
}

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()
}

struct SearchScratch {
    marks: Vec<u32>,
    generation: u32,
    candidates: BinaryHeap<Reverse<Candidate>>,
    results: BinaryHeap<Candidate>,
    routing_nodes: Vec<u32>,
    merged: Vec<Candidate>,
}

impl SearchScratch {
    fn new(nodes: usize) -> Result<Self, SearchError> {
        let mut marks = Vec::new();
        marks
            .try_reserve_exact(nodes)
            .map_err(|_| SearchError::AllocationFailed)?;
        marks.resize(nodes, 0);
        Ok(Self {
            marks,
            generation: 0,
            candidates: BinaryHeap::new(),
            results: BinaryHeap::new(),
            routing_nodes: Vec::new(),
            merged: Vec::new(),
        })
    }

    fn begin(&mut self) {
        self.candidates.clear();
        self.results.clear();
        if self.generation == u32::MAX {
            self.marks.fill(0);
            self.generation = 1;
        } else {
            self.generation += 1;
        }
    }

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

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

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