use super::graph_helpers::{max_degree, select_neighbors, to_node};
use super::{Graph, InsertionPlan, SearchScratch};
use crate::config::IndexConfig;
use crate::error::SearchError;
use crate::vector::{Candidate, VectorStore};
use std::sync::atomic::{AtomicUsize, Ordering};
impl Graph {
pub(super) 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(())
}
pub(super) 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(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, 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);
}
}
}