use uqa_core::DocId;
use super::metric::{deterministic_level, normalize_with_norm};
use super::prepare::{check, Control};
use super::search::{Candidate, Query};
use super::types::{active_key, HNSWIndex, HNSWNode};
use crate::{StorageBackendError, StorageBackendResult};
impl HNSWIndex {
pub(super) fn insert_vector(
&mut self,
doc_id: DocId,
vector_ordinal: u32,
raw_vector: Vec<f32>,
control: Control<'_>,
) -> StorageBackendResult<()> {
check(control)?;
let node_id = self.next_node_id;
self.next_node_id = self
.next_node_id
.checked_add(1)
.ok_or_else(|| StorageBackendError::Other("HNSW node id space exhausted".into()))?;
let level = deterministic_level(self.params.seed, node_id, self.params.m);
let (normalized_vector, norm) = normalize_with_norm(&raw_vector);
let previous_entry = self.entry_point;
let previous_max_level = self.max_level;
self.put_vectors(
node_id,
raw_vector,
normalized_vector.clone(),
norm,
control,
)?;
self.put_node(
HNSWNode {
id: node_id,
doc_id,
vector_ordinal,
level,
deleted: false,
neighbors: vec![Vec::new(); level + 1],
},
control,
)?;
self.active
.insert(active_key(doc_id, vector_ordinal), node_id, control)?;
let Some(mut entry) = previous_entry else {
self.entry_point = Some(node_id);
self.max_level = level;
return Ok(());
};
let mut query = Query::new(&normalized_vector, &self.memory);
if previous_max_level > level {
for layer in ((level + 1)..=previous_max_level).rev() {
entry = self.greedy_search_layer(&mut query, entry, layer, control)?;
}
}
for layer in (0..=level.min(previous_max_level)).rev() {
check(control)?;
let candidates = self.search_layer(
&mut query,
&[entry],
self.params.ef_construction,
layer,
control,
control,
)?;
let mut selected = self.select_candidates(
candidates.iter(),
self.max_connections(layer),
Some(node_id),
control,
)?;
if layer == 0 {
self.ensure_layer_zero_backbone(node_id, &mut selected)?;
}
self.modify_node(node_id, control, |node| {
node.neighbors[layer].clone_from(&selected);
})?;
for neighbor_id in selected {
check(control)?;
self.connect_and_prune_node(neighbor_id, node_id, layer, control)?;
}
self.prune_node(node_id, layer, control)?;
if let Some(Candidate { node_id, .. }) = candidates.iter().next().transpose()? {
entry = node_id;
}
}
if level > previous_max_level {
self.entry_point = Some(node_id);
self.max_level = level;
}
Ok(())
}
}