feox-ann 0.1.0

use crate::graph::{compare_scored_desc, Node, ScoredNode, SearchScratch};
use crate::math::{dot, normalize_in_place};
use crate::{AnnIndex, Result};

pub struct AnnInsertCursor<'a> {
    index: &'a mut AnnIndex,
    scratch: SearchScratch,
}

impl AnnIndex {
    pub fn insert_cursor(&mut self) -> AnnInsertCursor<'_> {
        AnnInsertCursor {
            index: self,
            scratch: SearchScratch::default(),
        }
    }

    pub fn upsert(&mut self, id: String, vector: &[f32]) -> Result<()> {
        let mut scratch = SearchScratch::default();
        self.insert_vector(id, vector.to_vec(), &mut scratch)
    }

    pub fn upsert_owned(&mut self, id: String, vector: Vec<f32>) -> Result<()> {
        let mut scratch = SearchScratch::default();
        self.insert_vector(id, vector, &mut scratch)
    }

    pub(crate) fn insert_vector(
        &mut self,
        id: String,
        mut vector: Vec<f32>,
        scratch: &mut SearchScratch,
    ) -> Result<()> {
        self.validate_vector(&vector)?;
        if let Some(index) = self.ids.remove(&id) {
            if !self.nodes[index].deleted {
                self.nodes[index].deleted = true;
                self.active = self.active.saturating_sub(1);
            }
        }

        normalize_in_place(&mut vector);
        let level = self.level_for_id(&id);
        let index = self.nodes.len();
        self.nodes.push(Node {
            id: id.clone(),
            vector,
            deleted: false,
            neighbors: vec![Vec::new(); level + 1],
        });
        self.ids.insert(id, index);
        self.add_entry_point(index);
        self.active += 1;

        let Some(mut entry) = self.entry else {
            self.entry = Some(index);
            self.max_level = level;
            return Ok(());
        };

        for layer in ((level + 1)..=self.max_level).rev() {
            entry = self.greedy_search(index, entry, layer);
        }

        let top_layer = level.min(self.max_level);
        for layer in (0..=top_layer).rev() {
            self.search_layer(index, entry, layer, self.config.ef_construction, scratch);
            let selected = self.select_neighbors(index, layer, &mut scratch.found);
            for &neighbor in &selected {
                self.nodes[neighbor].neighbors[layer].push(index);
                self.prune_neighbors(neighbor, layer);
            }
            self.nodes[index].neighbors[layer] = selected;
            if let Some(best) = self.best_candidate(index, layer) {
                entry = best;
            }
        }

        if level > self.max_level {
            self.entry = Some(index);
            self.max_level = level;
        }

        Ok(())
    }

    pub(crate) fn select_neighbors(
        &self,
        query_index: usize,
        layer: usize,
        candidates: &mut Vec<ScoredNode>,
    ) -> Vec<usize> {
        let limit = self.neighbor_limit(layer);
        candidates.retain(|candidate| candidate.index != query_index);
        candidates.sort_by(compare_scored_desc);
        candidates.dedup_by_key(|candidate| candidate.index);

        let mut selected: Vec<usize> = Vec::with_capacity(limit);
        let mut pruned: Vec<usize> = Vec::new();
        for candidate in candidates.iter() {
            if selected.len() >= limit {
                break;
            }
            let vector = &self.nodes[candidate.index].vector;
            let diverse = selected
                .iter()
                .all(|&kept| candidate.score >= dot(vector, &self.nodes[kept].vector));
            if diverse {
                selected.push(candidate.index);
            } else {
                pruned.push(candidate.index);
            }
        }
        for index in pruned {
            if selected.len() >= limit {
                break;
            }
            selected.push(index);
        }
        selected
    }

    pub(crate) fn prune_neighbors(&mut self, index: usize, layer: usize) {
        let limit = self.neighbor_limit(layer);
        if self.nodes[index].neighbors[layer].len() <= limit {
            return;
        }
        let mut neighbors = self.nodes[index].neighbors[layer].clone();
        neighbors.sort_unstable();
        neighbors.dedup();
        if neighbors.len() <= limit {
            self.nodes[index].neighbors[layer] = neighbors;
            return;
        }

        let source = &self.nodes[index].vector;
        let mut candidates = Vec::with_capacity(neighbors.len());
        for neighbor in neighbors {
            candidates.push(ScoredNode {
                index: neighbor,
                score: dot(source, &self.nodes[neighbor].vector),
            });
        }
        self.nodes[index].neighbors[layer] = self.select_neighbors(index, layer, &mut candidates);
    }
}

impl AnnInsertCursor<'_> {
    pub fn reserve(&mut self, additional: usize) {
        self.index.reserve(additional);
    }

    pub fn upsert_owned(&mut self, id: String, vector: Vec<f32>) -> Result<()> {
        self.index.insert_vector(id, vector, &mut self.scratch)
    }
}