velesdb-core 7.0.0

//! HNSW insert operations.

use super::super::distance::DistanceEngine;
use super::super::layer::{Layer, NodeId};
use super::{NativeHnsw, NO_ENTRY_POINT};
use crate::perf_optimizations::ContiguousVectors;
use std::sync::atomic::Ordering;

/// Headroom the slow path of `expand_layers` adds past the node that needs
/// it, as a fraction of the capacity already covered: growing by an eighth
/// keeps the number of `layers.write` growths logarithmic in the index size,
/// and the slots it adds that no node uses yet to an eighth of them, or
/// `MIN_CAPACITY_GROWTH`.
const CAPACITY_GROWTH_DIVISOR: usize = 8;

/// Floor of that headroom, so a small index does not take `layers.write`
/// every few inserts.
const MIN_CAPACITY_GROWTH: usize = 256;

impl<D: DistanceEngine> NativeHnsw<D> {
    /// Allocates vector storage if needed and pushes the vector, returning its node ID.
    ///
    /// # Errors
    ///
    /// Returns an error if storage allocation or push fails.
    // The write guard spans lazy arena creation and the push because
    // `storage` borrows out of it, and because a concurrent search takes the
    // same lock for reading: holding it is what makes "arena exists" and
    // "vector is in it" indivisible to every reader.
    #[expect(clippy::significant_drop_tightening)]
    fn allocate_and_store_vector(&self, vector: &[f32]) -> crate::error::Result<NodeId> {
        let mut guard = self.vectors.write();
        if guard.is_none() {
            // Through `new_arena`, not `ContiguousVectors::new`: a graph whose
            // arena belongs on disk must get a mapped one here too, or the
            // lazy path would silently hand back a heap arena (#2112).
            *guard = Some(Self::new_arena(
                self.arena_home.as_ref(),
                vector.len(),
                ContiguousVectors::MIN_ARENA_CAPACITY,
            )?);
        }
        let storage = guard.as_mut().ok_or_else(|| {
            crate::error::Error::Internal("Vector storage missing after init".to_string())
        })?;
        let id = storage.len();
        storage.push(vector)?;
        Ok(id)
    }

    /// Pre-creates layers and allocates node capacity for an upcoming batch.
    ///
    /// Uses a statistical upper bound for the max expected layer:
    /// `ceil(log_M(total_nodes)) + 2`, capped at 15.
    // Reason: cast_precision_loss acceptable for statistical bound calculation
    // Reason: cast_possible_truncation result is capped at 15
    // Reason: cast_sign_loss log of positive numbers is positive
    #[allow(
        clippy::cast_precision_loss,
        clippy::cast_possible_truncation,
        clippy::cast_sign_loss
    )]
    pub(in crate::index::hnsw::native) fn pre_expand_layers(&self, total_nodes: usize) {
        let max_layer = if self.max_connections > 1 && total_nodes > 1 {
            let log_m = (total_nodes as f64).ln() / (self.max_connections as f64).ln();
            (log_m.ceil() as usize + 2).min(15)
        } else {
            15
        };

        self.grow_layers(max_layer, total_nodes);
    }

    /// Ensures all layers up to `node_layer` exist and have capacity for `node_id`.
    fn expand_layers(&self, node_id: NodeId, node_layer: usize) {
        // Fast path: skip write lock if pre-allocation covers this insert
        if node_id < self.pre_allocated_capacity.load(Ordering::Relaxed)
            && node_layer < self.layers.read().len()
        {
            return;
        }
        // Slow path: past the capacity, grow it with headroom, so the inserts
        // that follow return to the fast path instead of each taking
        // `layers.write` (#2306); a missing layer alone keeps the capacity.
        let covered = self.pre_allocated_capacity.load(Ordering::Relaxed);
        let capacity = if node_id < covered {
            covered
        } else {
            node_id + 1 + (covered / CAPACITY_GROWTH_DIVISOR).max(MIN_CAPACITY_GROWTH)
        };
        self.grow_layers(node_layer, capacity);
    }

    /// Adds the layers up to `top_layer` and grows every layer to at least
    /// `capacity` slots, then publishes the capacity to the fast path of
    /// [`Self::expand_layers`].
    ///
    /// Keeps the invariant that fast path relies on: every layer holds at
    /// least `pre_allocated_capacity` slots. A layer added here therefore
    /// covers what the fast path already trusts, not only the node that
    /// created it: `Layer`'s accessors skip an id past a layer's end in
    /// silence, so a shorter layer would drop the neighbour writes of every
    /// later node sent to it. The capacity only grows, whatever `capacity`
    /// asks for.
    fn grow_layers(&self, top_layer: usize, capacity: usize) {
        let mut layers = self.layers.write();
        let capacity = capacity.max(self.pre_allocated_capacity.load(Ordering::Relaxed));
        while layers.len() <= top_layer {
            let level = layers.len();
            layers.push(Layer::at_level(level, capacity));
        }
        if let Some(last) = capacity.checked_sub(1) {
            for layer in layers.iter_mut() {
                layer.ensure_capacity(last);
            }
        }
        // Stored under the write guard: a fast-path reader that sees the new
        // capacity takes `layers.read()` afterwards, which cannot be granted
        // before this guard is released.
        self.pre_allocated_capacity
            .fetch_max(capacity, Ordering::Relaxed);
    }

    /// Inserts a vector into the index.
    ///
    /// Accepts a borrowed slice to avoid forcing callers to clone. For cosine
    /// metric with pre-normalization, a temporary copy is made internally;
    /// for all other metrics the slice is used directly (zero-copy).
    ///
    /// # Errors
    ///
    /// Returns an error if vector storage allocation or insertion fails.
    pub fn insert(&self, vector: &[f32]) -> crate::error::Result<NodeId> {
        let query = self.prepare_query(vector);

        let node_id = self.allocate_and_store_vector(&query)?;
        self.link_node(node_id, &query);

        // Invalidate GPU caches — topology and vectors both changed.
        // `vectors.write()` is already released at this point
        // (`allocate_and_store_vector` drops its own write guard before
        // returning), so the `gpu_vectors_snapshot` acquisition
        // inside the helper does not nest inside the vectors lock and
        // respects the declared order
        // (`GpuVectorsSnapshot` rank 5 → `Vectors` rank 10).
        #[cfg(feature = "gpu")]
        self.invalidate_gpu_caches();

        Ok(node_id)
    }

    /// Links `node_id`, whose vector the arena already holds, with `query`,
    /// the prepared form of that vector: draws the node's layer, makes room
    /// for it in the layers, and connects it through the entry point.
    pub(in crate::index::hnsw::native) fn link_node(&self, node_id: NodeId, query: &[f32]) {
        let node_layer = self.random_layer();
        self.expand_layers(node_id, node_layer);
        let entry_point = self.entry_point.load(Ordering::Acquire);
        self.link_new_node(node_id, query, node_layer, entry_point);
        self.count.fetch_add(1, Ordering::Relaxed);
    }

    /// Links a node just placed into the graph, given the entry point it saw.
    ///
    /// A node that saw none claims the empty graph's entry point. Losing that
    /// claim means another node won it in between: this one then connects
    /// through the winner like any other insert, instead of staying with no
    /// edge at all (#2259).
    fn link_new_node(&self, node_id: NodeId, query: &[f32], node_layer: usize, seen: NodeId) {
        let entry_point = if seen == NO_ENTRY_POINT {
            match self.claim_first_entry_point(node_id, node_layer) {
                None => return,
                Some(winner) => winner,
            }
        } else {
            seen
        };
        self.insert_with_entry_point(node_id, query, node_layer, entry_point);
        self.anchor_one(node_id);
        self.promote_entry_point(node_id, node_layer);
    }

    /// Claims an empty graph's entry point for `node_id`.
    ///
    /// Returns `None` when `node_id` is now the entry point, or the node
    /// that claimed it first. The claim and `max_layer` move together under
    /// the promotion lock, so no promotion lands between them; a graph that
    /// already has an entry point returns before taking the lock.
    pub(in crate::index::hnsw::native) fn claim_first_entry_point(
        &self,
        node_id: NodeId,
        node_layer: usize,
    ) -> Option<NodeId> {
        let seen = self.entry_point.load(Ordering::Acquire);
        if seen != NO_ENTRY_POINT {
            return Some(seen);
        }
        self.with_promotion_lock(|| {
            match self.entry_point.compare_exchange(
                NO_ENTRY_POINT,
                node_id,
                Ordering::AcqRel,
                Ordering::Acquire,
            ) {
                Ok(_) => {
                    // An empty graph's `max_layer` is 0, and only this lock
                    // writes it.
                    self.max_layer.store(node_layer, Ordering::Release);
                    None
                }
                Err(winner) => Some(winner),
            }
        })
    }

    /// `insert` for a node that saw an empty graph although another node has
    /// claimed its entry point since: the losing side of the first-insert
    /// race, which no schedule reproduces on demand.
    #[cfg(test)]
    pub(in crate::index::hnsw::native) fn insert_seeing_no_entry_point(
        &self,
        vector: &[f32],
    ) -> crate::error::Result<NodeId> {
        let query = self.prepare_query(vector);
        let node_id = self.allocate_and_store_vector(&query)?;
        let node_layer = self.random_layer();
        self.expand_layers(node_id, node_layer);
        self.link_new_node(node_id, &query, node_layer, NO_ENTRY_POINT);
        self.count.fetch_add(1, Ordering::Relaxed);
        Ok(node_id)
    }

    /// Makes `node_id` the entry point when the index is empty, or when the
    /// node reaches a higher layer than the current maximum.
    ///
    /// Readers load `entry_point` and `max_layer` without a lock. Writers take
    /// the promotion lock, and the check, the two stores and
    /// [`Self::reparent_entry_point`] run in it as one step: two promotions
    /// that overlapped could reparent in the wrong order, the later one's
    /// clearing the anchor the earlier one had just given its node, and strand
    /// the tree below it (#2259). Promotion is rare (O(log_M(N)) times per
    /// index lifetime), and a node at or below the maximum layer, the common
    /// case, returns before taking the lock.
    ///
    /// Between the `max_layer` store and the swap, a concurrent reader may see
    /// the new `max_layer` with the old `entry_point`. This is safe:
    /// `search_layer_single` finds no neighbour of the old EP at a layer where
    /// it has no edges, and returns it unchanged: a no-op greedy descent.
    pub(in crate::index::hnsw::native) fn promote_entry_point(
        &self,
        node_id: NodeId,
        node_layer: usize,
    ) {
        if self.entry_point.load(Ordering::Acquire) == NO_ENTRY_POINT
            && self.claim_first_entry_point(node_id, node_layer).is_none()
        {
            return;
        }
        if node_layer <= self.max_layer.load(Ordering::Acquire) {
            return;
        }
        self.with_promotion_lock(|| {
            if node_layer <= self.max_layer.load(Ordering::Acquire) {
                return;
            }
            self.max_layer.store(node_layer, Ordering::Release);
            let previous = self.entry_point.swap(node_id, Ordering::AcqRel);
            self.reparent_entry_point(previous, node_id);
        });
    }

    /// Batch-allocates vectors and assigns random layers.
    ///
    /// Returns the `(NodeId, layer)` assignments for each vector, in input
    /// order (`assignments[i]` corresponds to `vectors[i]`).
    ///
    /// This is Phase A of the two-phase batch insertion: all vectors are stored
    /// and layers expanded in single lock scopes. The caller is responsible for
    /// connecting nodes (Phase B) and updating `entry_point`/`count` (Phase C).
    ///
    /// # Allocation strategy (PERF2)
    ///
    /// Raw input slices are pushed directly into [`ContiguousVectors`]; for
    /// the pre-normalized cosine configuration the pushed range is then
    /// normalized **in place** inside the storage, under the same write lock
    /// (see [`Self::bulk_push_vectors`]). This avoids materializing one owned
    /// normalized buffer per vector (`Vec<Cow>` of N owned buffers held alive
    /// for the whole batch). Phase B re-derives the normalized query per node
    /// from a thread-local scratch (`with_prepared_query`), which is
    /// bit-identical to the stored form (same `normalize_inplace_native`
    /// kernel on the same input bits).
    ///
    /// # Lock Strategy (I2 optimization)
    ///
    /// The vector write lock is acquired in two separate scopes:
    /// 1. **Capacity reservation** — initializes storage and pre-grows the buffer
    ///    if needed. May trigger an expensive realloc+copy on the cold path.
    /// 2. **Bulk push** — copies vectors into pre-reserved space (guaranteed
    ///    fast memcpy with no reallocation), then normalizes the pushed range
    ///    in place for pre-normalized cosine.
    ///
    /// # Errors
    ///
    /// Returns an error if vector storage allocation fails.
    pub(in crate::index::hnsw::native) fn allocate_batch(
        &self,
        vectors: &[&[f32]],
    ) -> crate::error::Result<Vec<(NodeId, usize)>> {
        if vectors.is_empty() {
            return Ok(Vec::new());
        }

        // Pre-expand layers before vectors (lock ordering safe: layers only)
        let current_len = self
            .vectors
            .read()
            .as_ref()
            .map_or(0, ContiguousVectors::len);
        self.pre_expand_layers(current_len + vectors.len());

        // Scope 1: Initialize storage and reserve capacity (may resize — cold path)
        self.reserve_vector_capacity(vectors[0].len(), vectors.len())?;

        // Scope 2: Bulk push into pre-reserved space (fast memcpy — no resize)
        let first_id = self.bulk_push_vectors(vectors)?;

        let assignments: Vec<(NodeId, usize)> = (0..vectors.len())
            .map(|i| (first_id + i, self.random_layer()))
            .collect();

        Ok(assignments)
    }

    /// Phase A of a batch insert for nodes the arena already holds: makes room
    /// for every slot in the layers, as [`Self::allocate_batch`] does, and
    /// draws each node's layer, in input order.
    pub(in crate::index::hnsw::native) fn assign_layers(
        &self,
        nodes: &[NodeId],
    ) -> Vec<(NodeId, usize)> {
        self.pre_expand_layers(self.with_vectors_read(ContiguousVectors::len));
        nodes
            .iter()
            .map(|&node| (node, self.random_layer()))
            .collect()
    }

    /// Initializes vector storage if needed and pre-reserves capacity.
    ///
    /// Cold path: the write lock may be held during a buffer resize.
    // Same lazy-init-then-use shape as `allocate_and_store_vector`: the guard
    // spans arena creation and the reservation that depends on it.
    #[expect(clippy::significant_drop_tightening)]
    fn reserve_vector_capacity(
        &self,
        dimension: usize,
        batch_size: usize,
    ) -> crate::error::Result<()> {
        let mut guard = self.vectors.write();
        if guard.is_none() {
            *guard = Some(Self::new_arena(
                self.arena_home.as_ref(),
                dimension,
                batch_size,
            )?);
        }
        let storage = guard.as_mut().ok_or_else(|| {
            crate::error::Error::Internal("Vector storage missing after init".to_string())
        })?;
        storage.reserve_additional(batch_size)?;
        Ok(())
    }

    /// Pushes all raw vectors into pre-reserved storage, then normalizes the
    /// pushed range in place for the pre-normalized cosine configuration.
    ///
    /// Fast path: write lock held only for bulk memcpy (plus, for cosine
    /// pre-normalized, one in-place normalization pass), no reallocation.
    ///
    /// # Ordering invariant (PERF2)
    ///
    /// The in-place normalization runs under the **same** write lock as the
    /// push, so no reader can ever observe a raw (un-normalized) cosine
    /// vector: graph construction (Phase B) only starts after this method
    /// returns, and any concurrent search takes the vectors read lock, which
    /// is excluded while this write lock is held.
    // Held for the whole batch push deliberately -- the doc comment above
    // states the contract: a concurrent search takes the vectors read lock and
    // is excluded for exactly as long as this write lock is held, so the batch
    // becomes visible all at once.
    #[expect(clippy::significant_drop_tightening)]
    fn bulk_push_vectors(&self, vectors: &[&[f32]]) -> crate::error::Result<NodeId> {
        let mut guard = self.vectors.write();
        let storage = guard.as_mut().ok_or_else(|| {
            crate::error::Error::Internal("Vector storage missing after reserve".to_string())
        })?;
        self.push_stored_form(storage, vectors)
    }

    /// Pushes `vectors` into the arena without linking them into the graph,
    /// and returns the first slot: the direct writer's placement, which
    /// [`Self::link_placed`] links later, where it is. Each vector is stored
    /// as [`Self::allocate_batch`] stores it, under one write lock.
    ///
    /// # Errors
    ///
    /// Returns an error if the arena is not initialized or cannot grow.
    pub(in crate::index::hnsw) fn push_unlinked(
        &self,
        vectors: &[&[f32]],
    ) -> crate::error::Result<NodeId> {
        self.with_vectors_write(|storage| self.push_stored_form(storage, vectors))
    }

    /// Appends `vectors` to `storage` in the form this arena stores, and
    /// returns the first slot: unit-norm when [`Self::stores_unit_norm`], as
    /// given otherwise. The graph's batch insert and the direct writer both
    /// push through here, so neither leaves a raw slot in an arena whose
    /// `.vectors` declares unit-norm (#2246). The caller holds the vectors
    /// write lock across the call.
    fn push_stored_form(
        &self,
        storage: &mut ContiguousVectors,
        vectors: &[&[f32]],
    ) -> crate::error::Result<NodeId> {
        let first = storage.len();
        storage.push_batch(vectors)?;
        if self.stores_unit_norm() {
            for i in first..storage.len() {
                // `i` is in `first..storage.len()`, i.e. a slot we just pushed,
                // so `get_mut` is always `Some`. Assert it: a `None` here would
                // leave a raw (un-normalized) cosine vector in storage and
                // silently corrupt every distance computed against it.
                let stored = storage.get_mut(i);
                debug_assert!(
                    stored.is_some(),
                    "just-pushed vector slot {i} missing during in-place normalization"
                );
                if let Some(stored) = stored {
                    crate::simd_native::normalize_inplace_native(stored);
                }
            }
        }
        Ok(first)
    }

    /// Greedy descent through upper HNSW layers above `node_layer` to find
    /// the best entry point for the target layers.
    pub(in crate::index::hnsw::native) fn greedy_descent_upper_layers(
        &self,
        query: &[f32],
        node_layer: usize,
        mut entry_point: NodeId,
    ) -> NodeId {
        let max_layer = self.max_layer.load(Ordering::Relaxed);
        for layer_idx in (node_layer + 1..=max_layer).rev() {
            entry_point = self.search_layer_single(query, entry_point, layer_idx);
        }
        entry_point
    }

    /// Connects a node into the HNSW graph at layers 0..=`node_layer`.
    ///
    /// Searches for neighbors at each layer, selects the best candidates,
    /// and creates bidirectional connections.
    pub(in crate::index::hnsw::native) fn connect_node(
        &self,
        node_id: NodeId,
        query: &[f32],
        node_layer: usize,
        entry_point: NodeId,
    ) {
        self.connect_node_with_ef(
            node_id,
            query,
            node_layer,
            entry_point,
            self.ef_construction,
            0,
        );
    }

    /// Connects a node into the HNSW graph using a caller-specified ef budget.
    ///
    /// Used by `connect_batch_chunked` to apply adaptive `ef_construction`
    /// reduction during bulk insert (lower search budget for large batches)
    /// without affecting single-vector insert or the stored `ef_construction`.
    pub(in crate::index::hnsw::native) fn connect_node_with_ef(
        &self,
        node_id: NodeId,
        query: &[f32],
        node_layer: usize,
        mut entry_point: NodeId,
        effective_ef: usize,
        stagnation: usize,
    ) {
        for layer_idx in (0..=node_layer).rev() {
            let max_conn = if layer_idx == 0 {
                self.max_connections_0
            } else {
                self.max_connections
            };
            let neighbors = self.search_layer(
                query,
                &[entry_point],
                effective_ef,
                layer_idx,
                stagnation,
                None,
            );
            let selected = self.select_neighbors(&neighbors, max_conn);
            self.connect_neighbors_batch(node_id, &selected, layer_idx, max_conn);
            if !neighbors.is_empty() {
                entry_point = neighbors[0].0;
            }
        }
    }

    /// Performs the two-phase HNSW insertion when an entry point exists:
    /// 1. Greedy descent through upper layers above `node_layer`
    /// 2. Neighbor selection and bidirectional connection at layers 0..=`node_layer`
    #[inline]
    fn insert_with_entry_point(
        &self,
        node_id: NodeId,
        query: &[f32],
        node_layer: usize,
        ep: NodeId,
    ) {
        let current_ep = self.greedy_descent_upper_layers(query, node_layer, ep);
        self.connect_node(node_id, query, node_layer, current_ep);
    }
}

#[cfg(test)]
#[path = "insert_tests.rs"]
mod tests;