use std::sync::atomic::Ordering;
use super::super::NativeHnsw;
use crate::distance::DistanceMetric;
use crate::index::hnsw::native::distance::CpuDistance;
fn graph(max_elements: usize) -> NativeHnsw<CpuDistance> {
NativeHnsw::new(
CpuDistance::new(DistanceMetric::Euclidean),
16,
100,
max_elements,
)
}
fn shortest_layer(hnsw: &NativeHnsw<CpuDistance>) -> usize {
hnsw.layers
.read()
.iter()
.map(|layer| layer.neighbors.len())
.min()
.unwrap_or(0)
}
#[test]
fn test_growing_past_the_capacity_restores_the_fast_path() {
let hnsw = graph(8);
for i in 0..64_u8 {
hnsw.insert(&[f32::from(i), 0.0, 0.0, 0.0]).expect("insert");
}
let covered = hnsw.pre_allocated_capacity.load(Ordering::Relaxed);
assert!(
covered > hnsw.len(),
"capacity {covered} does not cover the next insert past {} nodes",
hnsw.len()
);
}
#[test]
fn test_a_layer_added_after_pre_expansion_covers_the_capacity() {
let hnsw = graph(0);
hnsw.pre_expand_layers(1000);
let above_top = hnsw.layers.read().len();
hnsw.expand_layers(10, above_top);
hnsw.expand_layers(500, above_top);
assert!(
shortest_layer(&hnsw) > 500,
"a layer holds {} slots, below node 500 the fast path let through",
shortest_layer(&hnsw)
);
}
#[test]
fn test_a_smaller_pre_expansion_keeps_the_capacity() {
let hnsw = graph(0);
hnsw.pre_expand_layers(1000);
hnsw.pre_expand_layers(10);
let covered = hnsw.pre_allocated_capacity.load(Ordering::Relaxed);
assert!(covered >= 1000, "capacity fell to {covered}");
assert!(shortest_layer(&hnsw) >= covered);
}
#[test]
fn test_a_layer_added_by_a_smaller_pre_expansion_covers_the_capacity() {
let hnsw = graph(0);
hnsw.expand_layers(0, 0);
let covered = hnsw.pre_allocated_capacity.load(Ordering::Relaxed);
let layers_before = hnsw.layers.read().len();
hnsw.pre_expand_layers(100);
assert!(
hnsw.layers.read().len() > layers_before,
"no layer was added"
);
assert!(
shortest_layer(&hnsw) >= covered,
"a layer holds {} slots, below the capacity {covered}",
shortest_layer(&hnsw)
);
}
#[test]
fn test_a_small_index_keeps_its_upper_layers_small() {
let hnsw = graph(100_000);
for i in 0..200_u8 {
hnsw.insert(&[f32::from(i), 0.0, 0.0, 0.0]).expect("insert");
}
let layers = hnsw.layers.read();
for (level, layer) in layers.iter().enumerate().skip(1) {
assert!(
layer.neighbors.len() < 1_000,
"layer {level} holds {} slots for 200 nodes",
layer.neighbors.len()
);
}
}
#[test]
fn test_growth_past_the_floor_stays_within_an_eighth() {
let hnsw = graph(0);
for i in 0..3_000_u16 {
hnsw.insert(&[f32::from(i), 0.0, 0.0, 0.0]).expect("insert");
}
let held = hnsw.len();
let bound = held + held / 8 + 1;
for (level, layer) in hnsw.layers.read().iter().enumerate() {
assert!(
layer.neighbors.len() <= bound,
"layer {level} holds {} slots for {held} nodes, past {bound}",
layer.neighbors.len()
);
}
}