core_rules/hnsw.rs
1//! In-tree HNSW approximate nearest-neighbor index.
2//!
3//! Implements the Malkov & Yashunin (2018) Hierarchical Navigable Small World
4//! algorithm, including a first-rejection short-cut of the §3.5 diverse-neighbour heuristic (not Algorithm 4; see `select_neighbors_first_rejection`) that
5//! decides which links survive a prune. The index shape lives in
6//! [`HnswParams`]; [`hnsw_params`] reads it once from the environment.
7//!
8//! Parameters are set for high-dimensional text embeddings (768-d to 2048-d).
9//! At these dimensionalities the nearest-neighbour distribution is flat, so it
10//! is neighbour *diversity* rather than neighbour *count* that lets the beam
11//! route through the hierarchy and recover the true k-NN.
12//!
13//! All vectors are L2-normalized at insert time; cosine similarity reduces to
14//! dot product for unit vectors, which is faster and numerically stable.
15//!
16//! The index keeps its own copy of every vector as `f32` in one contiguous
17//! [`VecSlab`] addressed by slot — the store keeps the `f64`s — and the dot
18//! product is [`dot_f32`], summed in eight independent accumulators so the
19//! compiler can emit parallel lanes. The index's distances choose *candidates*;
20//! every score a caller sees is recomputed from the `f64` properties, so the
21//! narrower type costs a tie-break and nothing observable. One slab means one
22//! stride, which is why an embedding whose dimension differs from the first one
23//! indexed is skipped rather than truncated.
24//!
25//! **Determinism**, and its one sharp edge: the level assigned to each node is
26//! derived from a seeded PRNG (`splitmix64`) seeded with
27//! `FNV-1a(rule_name) XOR (node_id × PHI)`, so insertion order and seed fix the
28//! levels. They do **not** fix the graph on their own — the edges also depend on
29//! [`HnswParams`], which [`hnsw_params`] reads from `MUSHROOMDB_HNSW_PARAMS`.
30//!
31//! So: one WAL replayed by two processes with different `MUSHROOMDB_HNSW_PARAMS`
32//! produces two *different* graphs. That is sound, because the graph is derived
33//! state — every edge it yields is recomputable and the blob is never compared
34//! byte-for-byte across replicas — but it does mean the variable belongs with
35//! the deployment's configuration, not with a single node's environment. Set it
36//! identically across replicas, or accept that their approximate answers differ
37//! (each still above the recall floor its own shape was gated at).
38
39use serde::{Deserialize, Serialize};
40use std::cell::RefCell;
41use std::cmp::Reverse;
42use std::collections::{BTreeMap, BTreeSet, BinaryHeap};
43
44/// Maximum layer cap: prevents pathological depth on tiny graphs.
45const MAX_LEVEL: usize = 16;
46
47// ---------------------------------------------------------------------------
48// Index shape
49// ---------------------------------------------------------------------------
50
51/// Index shape. Store-wide, not per rule: nothing asked to tune one rule
52/// differently from another, and a per-rule parameter set would have to be
53/// persisted in `RuleDef` and versioned with it.
54///
55/// Not persisted, either — the parameters bound how a graph is *built*, never
56/// how a built graph is *read*, so a blob written under one shape is read
57/// correctly under any other. That is why the v2 blob's meaning is unchanged by
58/// this struct's existence.
59#[derive(Clone, Copy, Debug, PartialEq, Eq)]
60pub struct HnswParams {
61 /// Max connections per layer above layer 0.
62 pub m: usize,
63 /// Max connections at layer 0.
64 pub m0: usize,
65 /// Beam width during insertion.
66 pub ef_construction: usize,
67 /// Beam width during search; the floor under a query's own `k`.
68 pub ef_search: usize,
69 /// How far the §3.5 diverse-neighbour heuristic reaches.
70 pub prune: Prune,
71}
72
73/// Which side of a new link the §3.5 diverse-neighbour heuristic decides.
74///
75/// The heuristic always chooses the new node's **own** neighbours — that is
76/// where diversity is cheap, because the candidate list is built once per
77/// insert. Applying it a second time on the **neighbour** side, to re-decide
78/// each over-connected neighbour's list, is what costs: that call is O(m₀²)
79/// distance computations and it runs once per neighbour per insert.
80#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
81pub enum Prune {
82 /// Heuristic for the new node's own neighbours; keep-the-m-nearest on the
83 /// neighbour side.
84 ///
85 /// **Below the recall floor on clustered corpora at the default `m0`**: it
86 /// scores 0.8581 on `approximate_recall_5k_timing`, whose floor is 0.90. Opt
87 /// in only with a raised `m0` or a corpus measured to tolerate it.
88 ///
89 /// The speed-up depends entirely on the corpus, so quote both numbers: **5×**
90 /// on a raw 5,000 × 1,536-D index of uniformly random vectors (45.8 s against
91 /// 251 s), but only **1.5×** on the clustered corpus the rule engine actually
92 /// derives edges over (123.1 s against 180.5 s of backfill) — and that is the
93 /// corpus where it misses the floor.
94 Own,
95 /// Heuristic on both sides — the paper's Algorithm 4 applied in full. The
96 /// default, because it is the cheapest shape that passes every committed
97 /// recall gate: see `docs/site/rules.md` for the three-way comparison
98 /// against `own` and against raising `m0` instead.
99 #[default]
100 Both,
101}
102
103impl Prune {
104 /// Parse the `prune` field of `MUSHROOMDB_HNSW_PARAMS`.
105 fn parse(s: &str) -> Option<Self> {
106 match s.trim() {
107 "own" => Some(Self::Own),
108 "both" => Some(Self::Both),
109 _ => None,
110 }
111 }
112}
113
114impl Default for HnswParams {
115 /// The shipped shape, chosen by the recall gates (`hnsw_5k_1536_recall`,
116 /// `approximate_recall_above_floor_1536dim_1k`, `approximate_recall_5k_timing`)
117 /// and not by argument.
118 ///
119 /// `m` falls 32 → 16 and `m0` falls 128 → 64, because M₀ = 128 was standing
120 /// in for the §3.5 diverse-neighbour heuristic the prune now applies; the
121 /// prune's cost is quadratic in `m0`, so the per-insert constant falls with
122 /// it. `ef_construction` falls 400 → 200 because it buys almost nothing once
123 /// the neighbours are diverse.
124 ///
125 /// `ef_search` does **not** fall. At 1,536 dimensions the nearest-neighbour
126 /// distribution is flat enough that recall is a beam-width problem, and
127 /// `hnsw_5k_1536_recall` refused every narrower beam that was tried: at
128 /// `ef_search` = 128 the min recall is 0.30 at `m0` = 32 and 0.70 at
129 /// `m0` = 64, against a floor of 0.90.
130 ///
131 /// `prune` is [`Prune::Both`] because [`Prune::Own`], which builds 5× faster
132 /// on uniform vectors but only 1.5× faster on clustered ones, scores **0.8581** on
133 /// `approximate_recall_5k_timing` against a floor of 0.90 — that corpus's
134 /// clusters are 100 members wide against an `m0` of 64, and the neighbour
135 /// side is where those links get thrown away. The two ways to fix it are
136 /// this and `m0` = 128; this one is both faster end to end (180 s against
137 /// 269 s of backfill) and half the adjacency memory, so it is the default.
138 fn default() -> Self {
139 Self {
140 m: 16,
141 m0: 64,
142 ef_construction: 200,
143 ef_search: 400,
144 prune: Prune::Both,
145 }
146 }
147}
148
149impl HnswParams {
150 /// Parse `m,m0,ef_construction,ef_search[,prune]`. `None` if the shape is
151 /// wrong, or a numeric field is absent, unparseable or zero — a zero would
152 /// produce an index with no edges or a search with no beam. The `prune`
153 /// field is optional and defaults to [`Prune::Both`]; it is the one field
154 /// that is a word rather than a number, so it cannot be confused with the
155 /// four ahead of it.
156 fn parse(s: &str) -> Option<Self> {
157 let fields: Vec<&str> = s.split(',').collect();
158 if fields.len() < 4 || fields.len() > 5 {
159 return None;
160 }
161 let mut num = fields
162 .iter()
163 .take(4)
164 .map(|f| f.trim().parse::<usize>().ok().filter(|&v| v > 0));
165 let mut next = || num.next().flatten();
166 let (m, m0, ef_construction, ef_search) = (next()?, next()?, next()?, next()?);
167 let prune = match fields.get(4) {
168 Some(f) => Prune::parse(f)?,
169 None => Prune::default(),
170 };
171 Some(Self {
172 m,
173 m0,
174 ef_construction,
175 ef_search,
176 prune,
177 })
178 }
179}
180
181/// The index shape this process builds and searches with.
182///
183/// Read once from `MUSHROOMDB_HNSW_PARAMS`, formatted
184/// `m,m0,ef_construction,ef_search[,prune]`, where `prune` is `own` or `both`
185/// and defaults to `both`. Unset or unparseable →
186/// [`HnswParams::default`]. Read once rather than per call so a graph cannot be
187/// half-built under one shape and half under another, and so the value is a
188/// pointer chase on the insert path.
189///
190/// For benchmarks and for an operator who has measured their own corpus; not a
191/// per-rule knob.
192pub fn hnsw_params() -> HnswParams {
193 static PARAMS: std::sync::OnceLock<HnswParams> = std::sync::OnceLock::new();
194 *PARAMS.get_or_init(|| {
195 let Ok(raw) = std::env::var("MUSHROOMDB_HNSW_PARAMS") else {
196 return HnswParams::default();
197 };
198 match HnswParams::parse(&raw) {
199 Some(p) => p,
200 None => {
201 // Silence here means an operator who mistyped the variable gets
202 // the defaults and a graph built under a shape they did not
203 // ask for, with nothing to tell them apart from a shape they
204 // did. One line, once per process.
205 eprintln!(
206 "mushroomdb: MUSHROOMDB_HNSW_PARAMS={raw:?} is not \
207 `m,m0,ef_construction,ef_search[,own|both]` with non-zero numbers; \
208 using the defaults {:?}",
209 HnswParams::default()
210 );
211 HnswParams::default()
212 }
213 }
214 })
215}
216
217/// `M` — max connections per layer (except layer 0).
218#[deprecated(note = "read hnsw_params() instead")]
219pub const M: usize = 16;
220/// `M₀` — max connections at layer 0.
221#[deprecated(note = "read hnsw_params() instead")]
222pub const M0: usize = 64;
223/// Beam width for insertion.
224#[deprecated(note = "read hnsw_params() instead")]
225pub const EF_CONSTRUCTION: usize = 200;
226/// Beam width for search.
227#[deprecated(note = "read hnsw_params() instead")]
228pub const EF_SEARCH: usize = 400;
229
230// ---------------------------------------------------------------------------
231// Test hooks
232// ---------------------------------------------------------------------------
233//
234// A thread-local count of vectors pushed into any `HnswIndex` on this thread,
235// in the shape `with_ivf_drift_rebuild` (`index.rs:29-50`) already uses. It
236// exists so an integration test can assert that opening a store inserts
237// *nothing* — the persisted graph is adopted, not rebuilt. Gated on
238// `test-hooks` because `#[cfg(test)]` items in this crate are invisible to
239// `mushroomdb`'s integration tests.
240
241#[cfg(any(test, feature = "test-hooks"))]
242thread_local! {
243 static HNSW_INSERT_COUNT: std::cell::Cell<u64> = const { std::cell::Cell::new(0) };
244 static HNSW_REMOVE_SCANNED: std::cell::Cell<u64> = const { std::cell::Cell::new(0) };
245 static HNSW_SEARCH_COUNT: std::cell::Cell<u64> = const { std::cell::Cell::new(0) };
246 static HNSW_DIST_EVALS: std::cell::Cell<u64> = const { std::cell::Cell::new(0) };
247 static HNSW_DIST_EVALS_PAIRWISE: std::cell::Cell<u64> = const { std::cell::Cell::new(0) };
248 static HNSW_BEAM_SCRATCH_GROWS: std::cell::Cell<u64> = const { std::cell::Cell::new(0) };
249}
250
251/// Count one vector actually indexed. Called *after* the zero-vector early
252/// return, so the count is "vectors the graph took", not "calls made".
253/// Compiles away without `test-hooks`.
254#[inline]
255fn note_insert() {
256 #[cfg(any(test, feature = "test-hooks"))]
257 HNSW_INSERT_COUNT.with(|c| c.set(c.get().saturating_add(1)));
258}
259
260/// Count `n` nodes whose adjacency a `remove` had to touch. This is the
261/// operation count the O(in-degree) claim is made about — no wall clock.
262#[inline]
263fn note_remove_scanned(n: usize) {
264 #[cfg(any(test, feature = "test-hooks"))]
265 HNSW_REMOVE_SCANNED.with(|c| c.set(c.get().saturating_add(n as u64)));
266 #[cfg(not(any(test, feature = "test-hooks")))]
267 let _ = n;
268}
269
270/// Count one distance evaluation. Called from the two distance functions and
271/// nowhere else, so the count is "dot products the graph asked for".
272///
273/// One thread-local add per `dim` multiply-adds is under 0.1 % of the kernel at
274/// any dimension worth indexing, and a counter a benchmark cannot read is not a
275/// gate. Compiles away without `test-hooks`.
276#[inline]
277fn note_dist() {
278 #[cfg(any(test, feature = "test-hooks"))]
279 HNSW_DIST_EVALS.with(|c| c.set(c.get().saturating_add(1)));
280}
281
282/// Count one distance evaluation between two *indexed* vectors — the §3.5
283/// prune's diversity test and the neighbour-side scoring loop. Always
284/// accompanied by a [`note_dist`], so pairwise is a subset of the total and
285/// `total − pairwise` is the beam plus the descent.
286#[inline]
287fn note_dist_pairwise() {
288 #[cfg(any(test, feature = "test-hooks"))]
289 HNSW_DIST_EVALS_PAIRWISE.with(|c| c.set(c.get().saturating_add(1)));
290}
291
292/// Distance evaluations on this thread since the last reset.
293#[doc(hidden)]
294#[cfg(any(test, feature = "test-hooks"))]
295pub fn hnsw_dist_evals() -> u64 {
296 HNSW_DIST_EVALS.with(|c| c.get())
297}
298
299/// The subset of [`hnsw_dist_evals`] that compared two indexed vectors rather
300/// than a vector against a query.
301#[doc(hidden)]
302#[cfg(any(test, feature = "test-hooks"))]
303pub fn hnsw_dist_evals_pairwise() -> u64 {
304 HNSW_DIST_EVALS_PAIRWISE.with(|c| c.get())
305}
306
307/// Reset both distance-evaluation counters to zero.
308#[doc(hidden)]
309#[cfg(any(test, feature = "test-hooks"))]
310pub fn hnsw_dist_evals_reset() {
311 HNSW_DIST_EVALS.with(|c| c.set(0));
312 HNSW_DIST_EVALS_PAIRWISE.with(|c| c.set(0));
313}
314
315/// Count one growth of the reused beam-search visited buffer.
316#[inline]
317fn note_beam_scratch_grow() {
318 #[cfg(any(test, feature = "test-hooks"))]
319 HNSW_BEAM_SCRATCH_GROWS.with(|c| c.set(c.get().saturating_add(1)));
320}
321
322/// Times the beam-search visited buffer grew on this thread since the last reset.
323#[doc(hidden)]
324#[cfg(any(test, feature = "test-hooks"))]
325pub fn hnsw_beam_scratch_grows() -> u64 {
326 HNSW_BEAM_SCRATCH_GROWS.with(|c| c.get())
327}
328
329/// Reset the beam-search visited-buffer growth counter to zero.
330#[doc(hidden)]
331#[cfg(any(test, feature = "test-hooks"))]
332pub fn hnsw_beam_scratch_grows_reset() {
333 HNSW_BEAM_SCRATCH_GROWS.with(|c| c.set(0));
334}
335
336/// Count one query answered by the graph itself (past the empty-index guards).
337#[inline]
338fn note_search() {
339 #[cfg(any(test, feature = "test-hooks"))]
340 HNSW_SEARCH_COUNT.with(|c| c.set(c.get().saturating_add(1)));
341}
342
343/// Nodes whose adjacency lists `HnswIndex::remove` has touched on this thread
344/// since the last reset.
345#[doc(hidden)]
346#[cfg(any(test, feature = "test-hooks"))]
347pub fn hnsw_remove_scanned() -> u64 {
348 HNSW_REMOVE_SCANNED.with(|c| c.get())
349}
350
351/// Reset this thread's removal-scan counter to zero.
352#[doc(hidden)]
353#[cfg(any(test, feature = "test-hooks"))]
354pub fn hnsw_remove_scanned_reset() {
355 HNSW_REMOVE_SCANNED.with(|c| c.set(0));
356}
357
358/// Vectors indexed by any `HnswIndex` on this thread since the last reset.
359#[doc(hidden)]
360#[cfg(any(test, feature = "test-hooks"))]
361pub fn hnsw_insert_count() -> u64 {
362 HNSW_INSERT_COUNT.with(|c| c.get())
363}
364
365/// Reset this thread's indexed-vector counter to zero.
366#[doc(hidden)]
367#[cfg(any(test, feature = "test-hooks"))]
368pub fn hnsw_insert_count_reset() {
369 HNSW_INSERT_COUNT.with(|c| c.set(0));
370}
371
372/// Queries answered by an `HnswIndex` graph on this thread since the last
373/// reset. Zero means every approximate query fell back to a full scan.
374#[doc(hidden)]
375#[cfg(any(test, feature = "test-hooks"))]
376pub fn hnsw_search_count() -> u64 {
377 HNSW_SEARCH_COUNT.with(|c| c.get())
378}
379
380/// Reset this thread's graph-query counter to zero.
381#[doc(hidden)]
382#[cfg(any(test, feature = "test-hooks"))]
383pub fn hnsw_search_count_reset() {
384 HNSW_SEARCH_COUNT.with(|c| c.set(0));
385}
386
387// ---------------------------------------------------------------------------
388// PRNG helpers
389// ---------------------------------------------------------------------------
390
391/// splitmix64 step — one round of the splitmix64 PRNG.
392#[inline]
393fn splitmix64(x: u64) -> u64 {
394 let x = x.wrapping_add(0x9E3779B97F4A7C15);
395 let x = (x ^ (x >> 30)).wrapping_mul(0xBF58476D1CE4E5B9);
396 let x = (x ^ (x >> 27)).wrapping_mul(0x94D049BB133111EB);
397 x ^ (x >> 31)
398}
399
400/// Generate the insertion level for `node_id` using the rule's `base_seed`.
401///
402/// Level follows the geometric distribution used by HNSW:
403/// `l = floor(-ln(uniform) / ln(M))`
404/// where `uniform` is deterministically derived from the seed.
405fn gen_level(base_seed: u64, node_id: u32) -> usize {
406 // Fibonacci-hash the node id to spread seeds uniformly.
407 let mixed = base_seed ^ (node_id as u64).wrapping_mul(0x9E3779B97F4A7C15);
408 let rng = splitmix64(mixed);
409 // Map upper 53 bits to (0, 1] — avoids ln(0).
410 let bits = (rng >> 11) | 1; // ensure non-zero
411 let uniform = bits as f64 / (1u64 << 53) as f64;
412 let ml = 1.0 / (hnsw_params().m as f64).ln();
413 let level = (-uniform.ln() * ml).floor() as usize;
414 level.min(MAX_LEVEL)
415}
416
417// ---------------------------------------------------------------------------
418// f64 ordering wrapper (for BinaryHeap)
419// ---------------------------------------------------------------------------
420
421/// f64 wrapper implementing total order (NaN sorts last).
422#[derive(Debug, Clone, Copy, PartialEq)]
423struct OrdF64(f64);
424
425impl Eq for OrdF64 {}
426impl PartialOrd for OrdF64 {
427 fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
428 Some(self.cmp(other))
429 }
430}
431impl Ord for OrdF64 {
432 fn cmp(&self, other: &Self) -> std::cmp::Ordering {
433 self.0
434 .partial_cmp(&other.0)
435 .unwrap_or(std::cmp::Ordering::Greater)
436 }
437}
438
439// ---------------------------------------------------------------------------
440// The distance kernel
441// ---------------------------------------------------------------------------
442
443/// The index's own copy of every vector: one contiguous `f32` allocation, `dim`
444/// floats per slot, addressed by slot.
445///
446/// The store keeps `f64`. This copy exists only to **choose candidates** —
447/// `index.rs::hnsw_candidates` throws away the similarity a search returns and
448/// keeps the ids, and every score a rule or a query reports is recomputed from
449/// the `f64` properties. So `f32` here costs a tie-break at the 1e-7 level and
450/// nothing a caller can observe, and it buys half the bytes and twice the lanes.
451///
452/// One slab rather than one `Vec<f64>` per node also means a distance is a
453/// pointer offset into a known stride instead of a chase through an
454/// independently allocated 12 KB block, and that the pair-distance path needs no
455/// copy at all.
456///
457/// `dim` is 0 until the first vector arrives, and fixed from then on — with one
458/// correction: an election made on a *single* sample can be undone, because the
459/// first vector indexed may be the odd one out (see [`HnswIndex::insert`]).
460///
461/// `data` grows through `Vec::resize`, so its capacity grows geometrically: a
462/// slab can hold up to about twice the bytes its rows need, transiently, and
463/// [`HnswIndex::memory_stats`] does not see that — it counts rows, which is why
464/// it documents itself as a floor on resident size. One amortised doubling is
465/// the price of not copying the whole slab on every insert.
466#[derive(Serialize, Deserialize, Clone, Debug, Default)]
467struct VecSlab {
468 dim: usize,
469 data: Vec<f32>,
470}
471
472impl VecSlab {
473 /// The row for `slot`, or an empty slice when the slab has no row there.
474 /// Empty is unreachable for a live slot — every `alloc_slot` writes one —
475 /// and returning it rather than panicking keeps a corrupt blob from taking
476 /// the process down.
477 #[inline]
478 fn get(&self, slot: u32) -> &[f32] {
479 if self.dim == 0 {
480 return &[];
481 }
482 let start = slot as usize * self.dim;
483 self.data.get(start..start + self.dim).unwrap_or(&[])
484 }
485
486 /// Write `v` into `slot`'s row, converting from the caller's unit `f64`
487 /// vector. Fixes `dim` on the first call. `false` — nothing written — when
488 /// `v.len()` disagrees with an established `dim`.
489 fn put(&mut self, slot: u32, v: &[f64]) -> bool {
490 if self.dim == 0 {
491 self.dim = v.len();
492 }
493 if self.dim == 0 || v.len() != self.dim {
494 return false;
495 }
496 let start = slot as usize * self.dim;
497 let end = start + self.dim;
498 if self.data.len() < end {
499 self.data.resize(end, 0.0);
500 }
501 for (d, s) in self.data[start..end].iter_mut().zip(v) {
502 *d = *s as f32;
503 }
504 true
505 }
506
507 /// Floats the slab holds for `live` nodes — an element count, so the
508 /// caller decides what a float costs.
509 #[inline]
510 fn floats_for(&self, live: usize) -> usize {
511 live * self.dim
512 }
513}
514
515/// Dot product of two equal-length `f32` slices, summed in **eight independent
516/// accumulators**.
517///
518/// IEEE addition is not associative, so LLVM may not reassociate a single
519/// accumulator: `a.iter().zip(b).map(|(x, y)| x * y).sum()` is a serial chain of
520/// `dim` dependent multiply-adds, `dim` × the 3–4 cycle latency of one add. At
521/// 1,536 dimensions that is ~1.5 µs of a machine that could have done the work
522/// in a fraction of it. Choosing the summation order here — eight partial sums
523/// over `chunks_exact(8)`, which hands LLVM a known-length slice — is what lets
524/// it emit parallel FMAs (NEON is baseline on `aarch64`; SSE2 on x86-64, which
525/// is 4-wide multiply and add without FMA and so a smaller win).
526///
527/// `std::simd` would say this declaratively and is nightly-only; the toolchain
528/// is pinned stable, so this is the portable way to say it. No `unsafe`, no
529/// `cfg`, one code path.
530#[inline]
531fn dot_f32(a: &[f32], b: &[f32]) -> f32 {
532 debug_assert_eq!(a.len(), b.len(), "a dot product needs equal lengths");
533 let mut acc = [0.0f32; 8];
534 let mut ca = a.chunks_exact(8);
535 let mut cb = b.chunks_exact(8);
536 for (x, y) in ca.by_ref().zip(cb.by_ref()) {
537 for i in 0..8 {
538 acc[i] += x[i] * y[i];
539 }
540 }
541 let tail: f32 = ca
542 .remainder()
543 .iter()
544 .zip(cb.remainder())
545 .map(|(x, y)| x * y)
546 .sum();
547 acc.iter().sum::<f32>() + tail
548}
549
550/// Cosine distance between the two vectors in `a` and `b` — the prune's
551/// diversity test and the neighbour-side scoring loop. Both are already unit, so
552/// this is `1 − dot`, clamped to [0, 2].
553///
554/// Replaces the 12 KB `nb_vec.clone()` the neighbour-side prune used to make
555/// once per over-connected neighbour per insert.
556#[inline]
557fn dist_slots(slab: &VecSlab, a: u32, b: u32) -> f64 {
558 note_dist();
559 note_dist_pairwise();
560 let dot = dot_f32(slab.get(a), slab.get(b)) as f64;
561 (1.0 - dot.clamp(-1.0, 1.0)).max(0.0)
562}
563
564/// Cosine distance from the vector in `slot` to the already-unit `f32` query
565/// `q`. Returns `1 - dot` clamped to [0, 2] (0 = identical, 2 = opposite).
566#[inline]
567fn dist_to(slab: &VecSlab, slot: u32, q: &[f32]) -> f64 {
568 note_dist();
569 let dot = dot_f32(slab.get(slot), q) as f64;
570 (1.0 - dot.clamp(-1.0, 1.0)).max(0.0)
571}
572
573/// Build a slab from vectors decoded out of an older blob, `vectors[slot]` being
574/// the vector for that slot.
575///
576/// The stride is the first non-empty vector's length — a freed slot decodes as an
577/// empty one, and so does a node a 0.6.5 writer left without a vector. A vector
578/// of some *other* non-zero length is a mixed-dimension index, which older
579/// builds accepted and whose distances they computed over the shorter of the two
580/// vectors; it is copied as far as it goes and zero-filled beyond, because the
581/// node is already in the graph and dropping it would leave adjacency naming a
582/// slot that is not there. One log line names how many.
583fn slab_of_decoded(vectors: &[Vec<f64>]) -> (VecSlab, u64) {
584 let dim = vectors
585 .iter()
586 .map(|v| v.len())
587 .find(|&l| l != 0)
588 .unwrap_or(0);
589 let mut slab = VecSlab {
590 dim,
591 data: vec![0.0; vectors.len() * dim],
592 };
593 if dim == 0 {
594 return (slab, 0);
595 }
596 let mut odd = 0usize;
597 for (slot, v) in vectors.iter().enumerate() {
598 if v.len() == dim {
599 slab.put(slot as u32, v);
600 continue;
601 }
602 if v.is_empty() {
603 continue; // a freed slot: its row stays zero and nothing reads it
604 }
605 odd += 1;
606 let start = slot * dim;
607 for (d, s) in slab.data[start..start + dim].iter_mut().zip(v) {
608 *d = *s as f32;
609 }
610 }
611 if odd > 0 {
612 eprintln!(
613 "mushroomdb: HNSW loaded {odd} node(s) whose embedding is not {dim} \
614 dimensions; they were padded to the index's stride. Re-embed the \
615 collection with one model — their distances were already meaningless. \
616 This index will not claim the fast path while they are in it."
617 );
618 }
619 (slab, odd as u64)
620}
621
622/// The `f32` form of a unit `f64` vector — the query shape every search path
623/// uses, and bit-for-bit what [`VecSlab::put`] stored. Search/insert fill a
624/// thread-local buffer via [`as_f32_into`]; this owned helper remains for tests.
625#[cfg(test)]
626#[inline]
627fn as_f32(v: &[f64]) -> Vec<f32> {
628 let mut out = Vec::with_capacity(v.len());
629 as_f32_into(&mut out, v);
630 out
631}
632
633#[inline]
634fn as_f32_into(dst: &mut Vec<f32>, v: &[f64]) {
635 dst.clear();
636 dst.extend(v.iter().map(|&x| x as f32));
637}
638
639/// Reused beam-search buffers. Thread-local so `&self` search can share them
640/// without interior mutability on the index (cloned and serialized).
641struct BeamScratch {
642 visited: Vec<bool>,
643 c_heap: BinaryHeap<Reverse<(OrdF64, u32)>>,
644 w_heap: BinaryHeap<(OrdF64, u32)>,
645}
646
647impl BeamScratch {
648 const fn empty() -> Self {
649 Self {
650 visited: Vec::new(),
651 c_heap: BinaryHeap::new(),
652 w_heap: BinaryHeap::new(),
653 }
654 }
655
656 fn reset_beam(&mut self, n_slots: usize) {
657 if self.visited.len() == n_slots {
658 self.visited.fill(false);
659 } else {
660 if self.visited.capacity() < n_slots {
661 note_beam_scratch_grow();
662 }
663 self.visited.clear();
664 self.visited.resize(n_slots, false);
665 }
666 self.c_heap.clear();
667 self.w_heap.clear();
668 }
669}
670
671thread_local! {
672 static BEAM_SCRATCH: RefCell<BeamScratch> = const { RefCell::new(BeamScratch::empty()) };
673 static QUERY_F32: RefCell<Vec<f32>> = const { RefCell::new(Vec::new()) };
674}
675
676fn with_beam_scratch<R>(f: impl FnOnce(&mut BeamScratch) -> R) -> R {
677 BEAM_SCRATCH.with(|cell| f(&mut cell.borrow_mut()))
678}
679
680fn take_query_f32(v: &[f64]) -> Vec<f32> {
681 QUERY_F32.with(|cell| {
682 let mut q = std::mem::take(&mut *cell.borrow_mut());
683 as_f32_into(&mut q, v);
684 q
685 })
686}
687
688fn stash_query_f32(q: Vec<f32>) {
689 QUERY_F32.with(|cell| {
690 *cell.borrow_mut() = q;
691 });
692}
693
694// ---------------------------------------------------------------------------
695// Node storage
696// ---------------------------------------------------------------------------
697
698/// A node's place in the hierarchy. The vector lives in the index's [`VecSlab`],
699/// addressed by the same slot, which is why this struct has no vector field and
700/// why the blob is version 3.
701#[derive(Serialize, Deserialize, Clone, Debug, Default)]
702struct HnswNode {
703 /// Assigned layer level (inclusive; node has layers 0..=level).
704 level: usize,
705 /// `layers[l]` = neighbor **slots** at layer `l`.
706 ///
707 /// Slots, not node ids: a distance is then a `Vec` index rather than a
708 /// `BTreeMap<u32, HnswNode>` descent, and a beam step is a pointer offset.
709 /// Adjacency therefore only ever names live slots — a stale slot would name
710 /// whichever node reused it, so `remove` must strip every reference to a
711 /// slot before freeing it. `back_refs` is what makes that affordable.
712 layers: Vec<Vec<u32>>,
713}
714
715// ---------------------------------------------------------------------------
716// HnswIndex
717// ---------------------------------------------------------------------------
718
719/// In-tree HNSW approximate nearest-neighbor index for cosine similarity.
720///
721/// Stores L2-normalized vectors and answers approximate k-NN queries using the
722/// Malkov & Yashunin hierarchical graph.
723#[derive(Serialize, Deserialize, Clone, Debug, Default)]
724pub struct HnswIndex {
725 /// Base seed derived from `fnv1a(rule_name)`. Mixed with each node's id
726 /// to generate deterministic per-node levels.
727 base_seed: u64,
728 /// Dense node storage, indexed by slot. Freed slots are kept (blanked) and
729 /// reused, so a slot index is stable for as long as the node lives.
730 slots: Vec<HnswNode>,
731 /// node id → slot. The authority on which node ids are in the index.
732 slot_of: BTreeMap<u32, u32>,
733 /// slot → node id. `u32::MAX` marks a freed slot; that doubles as the
734 /// liveness check the search paths need.
735 id_of: Vec<u32>,
736 /// Freed slots, reused by the next insert (LIFO).
737 free: Vec<u32>,
738 /// Every indexed vector, `f32`, addressed by the same slot as `slots`.
739 ///
740 /// A freed slot keeps its row until another node takes the slot over: the
741 /// bytes stay resident where a per-node `Vec` would have freed them. In
742 /// exchange the index loses one allocator header and one fragmentation risk
743 /// per node, and `memory_stats` counts live nodes only, so the figure it
744 /// reports stays a floor on resident size rather than a measurement of it.
745 slab: VecSlab,
746 /// Vectors **refused** because their dimension disagreed with the slab's
747 /// settled stride. Any refusal means the index is missing a vector it was
748 /// offered, so [`HnswIndex::can_answer`] stops claiming the fast path — the
749 /// caller's exhaustive scan is the correct answer and this one is not.
750 ///
751 /// A re-elected stride (see [`HnswIndex::insert`]) **evicts** rather than
752 /// refuses, and does not count here: after it the index holds every vector
753 /// it was offered at the stride it now has.
754 ///
755 /// Not persisted: a loaded index has refused nothing, and the vectors it
756 /// holds are whatever the writer put in it. The first refusal on an index
757 /// logs one line, and later ones are silent, so an ingest pointed at the
758 /// wrong model cannot print a line per node.
759 #[serde(skip)]
760 dim_mismatches: u64,
761 /// The node ids behind [`Self::dim_mismatches`], so a reopen can skip
762 /// re-offering what it already knows it refuses.
763 ///
764 /// Not the same thing as the counter, and deliberately not derived from it:
765 /// the counter is monotone (a refusal never becomes un-counted, which keeps
766 /// [`Self::can_answer`] conservative), while this set is *maintained* — an
767 /// id leaves it the moment the index accepts or removes that node. A blob
768 /// may therefore carry a count larger than this set, and that is valid: an
769 /// upgraded v1/v2 index counts padded rows whose nodes are in the graph and
770 /// were never refused at all.
771 ///
772 /// Persisted from blob v4. Before v4 the open-time node scan re-derived it
773 /// by re-offering every vector and collecting the same refusals.
774 #[serde(skip)]
775 refused: BTreeSet<u32>,
776 /// Vectors **evicted** by a stride re-election, kept `(id, unit vector)` so
777 /// that a later re-election to their dimension can put them back.
778 ///
779 /// A re-election drops the one vector standing behind the old stride, and
780 /// that vector may be the *real* corpus: in the order
781 /// `[real, stray, real, …]` the first real vector is evicted by the stray and
782 /// the stray is evicted by the second real one. Parking is what makes the
783 /// first case recoverable — the second real vector re-elects that dimension
784 /// and the parked vector is re-inserted — and [`Self::can_answer`] is what
785 /// makes the gap safe while it lasts.
786 ///
787 /// Bounded by the number of distinct dimensions ever elected, not by the
788 /// corpus: a re-election needs the index to hold at most one vector, and a
789 /// stream that changes dimension on every insert parks one entry per
790 /// change. Every insert and remove pays an `O(parked)` scan. Never persisted, and holding
791 /// an `f64` copy of a vector the index is not indexing, which
792 /// [`Self::memory_stats`] does not count.
793 #[serde(skip)]
794 parked: Vec<(u32, Vec<f64>)>,
795 /// Reverse adjacency: slot → the slots that list it as a neighbour on *any*
796 /// layer. Maintained in lockstep with `HnswNode::layers` by insert, remove
797 /// and the prune. Turns removal from O(live nodes × M₀) into O(in-degree).
798 ///
799 /// Derivable from `slots`, so it is never serialized: a load rebuilds it
800 /// from the decoded adjacency lists (see `rebuild_back_refs`).
801 #[serde(skip)]
802 back_refs: BTreeMap<u32, BTreeSet<u32>>,
803 /// True when this graph was decoded from a blob a sliced build had not
804 /// finished writing, so it holds a prefix of its rule's corpus.
805 ///
806 /// A partial graph answers a search perfectly well — it just answers about
807 /// the wrong set, and nothing in the result says so. That is the one shape
808 /// of wrongness this index is not allowed to have, so
809 /// [`Self::can_answer`] refuses and the caller takes its exhaustive path.
810 ///
811 /// Only the **read** path ever sees this set: a live handle's authority on
812 /// an unfinished build is `RuleEngine::pending_builds`, and
813 /// [`SideIndex::adopt_hnsw`] clears the flag. An incomplete blob's remainder
814 /// is sliced by `pump_index_build`; a complete blob's missing nodes (writes
815 /// after the snapshot) are still supplied by the open-time scan. It is the
816 /// lazily-decoded copy on a clean open — which no scan follows — that needs
817 /// evidence carried in the bytes.
818 ///
819 /// Not serialized as part of the index: it is a property of the *blob*, and
820 /// [`HnswBlob::complete`] is where it lives on disk.
821 #[serde(skip)]
822 incomplete: bool,
823 /// Entry-point **slot**.
824 entry_point: Option<u32>,
825 max_level: usize,
826 /// Overrides [`hnsw_params`]'s `prune` for this index only.
827 ///
828 /// Exists because `hnsw_params()` reads the environment through a
829 /// `OnceLock`: one process gets one shape, so a test that wants to compare
830 /// both prune strategies cannot get there through the env. Compiled out
831 /// entirely without `test-hooks`, and `#[serde(skip)]` regardless, so it can
832 /// reach neither a production build nor a blob.
833 #[doc(hidden)]
834 #[cfg(any(test, feature = "test-hooks"))]
835 #[serde(skip)]
836 pub prune_override_for_test: Option<Prune>,
837}
838
839impl HnswIndex {
840 /// The prune this index builds with: the per-index test override when one is
841 /// set, otherwise the process-wide shape.
842 #[inline]
843 fn resolved_prune(&self, params: &HnswParams) -> Prune {
844 #[cfg(any(test, feature = "test-hooks"))]
845 if let Some(p) = self.prune_override_for_test {
846 return p;
847 }
848 params.prune
849 }
850}
851
852/// A freed slot's id sentinel.
853const DEAD: u32 = u32::MAX;
854
855/// What an [`HnswIndex`] is holding, in countable units rather than bytes.
856///
857/// Produced by [`HnswIndex::memory_stats`]. Payload only: see that method for
858/// what is deliberately not counted.
859#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
860pub struct HnswMemoryStats {
861 /// Vectors currently indexed.
862 pub live_nodes: usize,
863 /// Total neighbour entries across every live node's every layer.
864 pub neighbour_slots: usize,
865 /// Total entries in the reverse adjacency index.
866 pub back_ref_entries: usize,
867 /// Total floats stored across every live node's vector. `f32`s since 0.6.6:
868 /// the index's own copy is the slab's, and the `f64`s stay in the store.
869 pub vector_floats: usize,
870}
871
872impl HnswMemoryStats {
873 /// Payload bytes per indexed vector: adjacency (4 B per entry, forward and
874 /// reverse) plus the index's own copy of the vector (4 B per dimension — the
875 /// store's `f64` copy is not the index's business). Zero for an empty index.
876 pub fn bytes_per_node(&self) -> f64 {
877 if self.live_nodes == 0 {
878 return 0.0;
879 }
880 let bytes = (self.neighbour_slots + self.back_ref_entries) * 4
881 + self.vector_floats * std::mem::size_of::<f32>();
882 bytes as f64 / self.live_nodes as f64
883 }
884
885 /// The adjacency half of [`Self::bytes_per_node`] — the half the index
886 /// shape controls. The vector half is fixed by the embedding's dimension.
887 pub fn adjacency_bytes_per_node(&self) -> f64 {
888 if self.live_nodes == 0 {
889 return 0.0;
890 }
891 ((self.neighbour_slots + self.back_ref_entries) * 4) as f64 / self.live_nodes as f64
892 }
893}
894
895impl HnswIndex {
896 /// Create a new empty HNSW index seeded by `base_seed`.
897 ///
898 /// Typically `base_seed = fnv1a(rule_name.as_bytes())` so WAL replay
899 /// with the same rule name always produces the same graph structure.
900 pub fn new(base_seed: u64) -> Self {
901 Self {
902 base_seed,
903 ..Self::default()
904 }
905 }
906
907 /// Number of indexed vectors.
908 pub fn len(&self) -> usize {
909 self.slot_of.len()
910 }
911
912 /// True when `id` currently has a vector in this graph.
913 pub fn contains(&self, id: u32) -> bool {
914 self.slot_of.contains_key(&id)
915 }
916
917 /// True when no vectors are indexed.
918 pub fn is_empty(&self) -> bool {
919 self.slot_of.is_empty()
920 }
921
922 /// **Ask this before searching.** True when the index can answer a query of
923 /// `q_len` dimensions *completely* — meaning a caller may use its answer
924 /// instead of an exhaustive scan.
925 ///
926 /// Four things have to hold, and each of them is a way the index can be
927 /// useless rather than wrong:
928 ///
929 /// * It holds something. An empty index answers nothing.
930 /// * Its stride **is** the query's dimension. A slab has one stride, so an
931 /// index of 1,536-D vectors cannot compare a 3-D query to anything, and —
932 /// the case that matters — an index that elected a 3-D stride from a stray
933 /// first vector cannot answer the 1,536-D queries the rule actually makes.
934 /// * It has refused nothing (`dim_mismatches == 0`). A refused vector is one
935 /// the caller asked to index and the index does not hold, so its candidate
936 /// set is incomplete and the caller's own scan is the correct answer.
937 /// * Nothing of `q_len` dimensions is **parked** — evicted by a stride
938 /// re-election and not yet put back. A re-election revives every parked
939 /// vector of the dimension it elects, so this clause should always hold for
940 /// the current stride; it is the assertion that makes that a guarantee
941 /// rather than a claim. A parked vector of some *other* dimension is not a
942 /// gap in this answer: it is not comparable with anything at this stride,
943 /// by the same rule (`def.rs`'s `VectorSimilar` refuses a pair of unequal
944 /// length) that makes it unable to be an edge.
945 ///
946 /// When this is false the caller must take its exhaustive path —
947 /// `index.rs::hnsw_candidates` returns `hnsw_tracked`, and
948 /// `engine.rs::hnsw_search_dst`/`_any_dst` return `None` so
949 /// `db.rs::find_similar_vector` brute-forces. Correct and slower beats fast
950 /// and short.
951 pub fn can_answer(&self, q_len: usize) -> bool {
952 !self.is_empty()
953 && !self.incomplete
954 && self.dim_mismatches == 0
955 && self.slab.dim == q_len
956 && !self.parked.iter().any(|(_, v)| v.len() == q_len)
957 }
958
959 /// Declare this graph whole, clearing the [`Self::incomplete`] flag a
960 /// partial blob set.
961 ///
962 /// Called by [`SideIndex::adopt_hnsw`]. A live handle's authority on an
963 /// unfinished build is `RuleEngine::pending_builds` (see `hnsw_search_dst`);
964 /// this flag is what the lazily-decoded read-path copy uses, which has no
965 /// scan and no pending-build map behind it.
966 pub fn mark_complete(&mut self) {
967 self.incomplete = false;
968 }
969
970 /// True when this graph came from a blob written mid-build.
971 #[doc(hidden)]
972 pub fn is_incomplete(&self) -> bool {
973 self.incomplete
974 }
975
976 /// Returns all node ids currently in the index.
977 pub fn node_ids(&self) -> BTreeSet<u32> {
978 self.slot_of.keys().copied().collect()
979 }
980
981 /// Every id this index can account for: indexed, parked, or refused.
982 ///
983 /// This is what an open-time scan must skip re-offering. [`Self::node_ids`]
984 /// is the narrower "actually in the graph" answer and is still what the
985 /// exhaustive-fallback candidate set wants — re-offering a parked vector
986 /// destroys it (`insert` supersedes the parked copy, then refuses the
987 /// vector against the elected stride), and re-offering a refused one counts
988 /// the same mismatch twice.
989 ///
990 /// On a graph decoded from a blob older than v4 the parked and refused sets
991 /// are empty, so this equals `node_ids()` and the pre-v4 behaviour — scan,
992 /// re-offer, re-derive — is unchanged.
993 pub fn accounted_ids(&self) -> BTreeSet<u32> {
994 let mut out = self.node_ids();
995 out.extend(self.parked.iter().map(|(id, _)| *id));
996 out.extend(self.refused.iter().copied());
997 out
998 }
999
1000 /// True when [`Self::accounted_ids`] holds `id`, without building the set.
1001 pub fn accounts_for(&self, id: u32) -> bool {
1002 self.slot_of.contains_key(&id)
1003 || self.refused.contains(&id)
1004 || self.parked.iter().any(|(pid, _)| *pid == id)
1005 }
1006
1007 /// The refusal count [`Self::can_answer`] consults.
1008 #[doc(hidden)]
1009 pub fn dim_mismatches(&self) -> u64 {
1010 self.dim_mismatches
1011 }
1012
1013 /// The ids this index refused for a dimension disagreement.
1014 #[doc(hidden)]
1015 pub fn refused_ids(&self) -> BTreeSet<u32> {
1016 self.refused.clone()
1017 }
1018
1019 /// Parked entries in the slab's unit, which is what the blob persists.
1020 ///
1021 /// Narrowing to `f32` is lossless here and not by luck of the values: the
1022 /// only site that parks a vector reads it back out of the `f32` slab and
1023 /// widens it, so every `f64` in `parked` already has an exact `f32`
1024 /// preimage. A future park that stored a caller's raw `f64` would break
1025 /// that, and the round trip would start perturbing the low mantissa bits.
1026 fn parked_rows_f32(&self) -> Vec<(u32, Vec<f32>)> {
1027 self.parked
1028 .iter()
1029 .map(|(id, v)| (*id, v.iter().map(|&x| x as f32).collect()))
1030 .collect()
1031 }
1032
1033 /// Restore the state blob v4 carries. Pre-v4 blobs leave it all default.
1034 fn restore_side_state(
1035 &mut self,
1036 dim_mismatches: u64,
1037 refused: BTreeSet<u32>,
1038 parked: Vec<(u32, Vec<f32>)>,
1039 ) {
1040 self.dim_mismatches = dim_mismatches;
1041 // A blob claiming an id is both indexed and refused is inconsistent;
1042 // trusting it would make the scan skip a node the graph does not hold.
1043 self.refused = refused
1044 .into_iter()
1045 .filter(|id| !self.slot_of.contains_key(id))
1046 .collect();
1047 // A parked vector whose length is the elected stride would have been
1048 // revived rather than parked, so the blob disagrees with itself. Drop
1049 // it and let the scan re-offer that id.
1050 self.parked = parked
1051 .into_iter()
1052 .filter(|(id, v)| {
1053 !self.slot_of.contains_key(id) && (self.slab.dim == 0 || v.len() != self.slab.dim)
1054 })
1055 .map(|(id, v)| (id, v.into_iter().map(f64::from).collect()))
1056 .collect();
1057 }
1058
1059 // -----------------------------------------------------------------------
1060 // Slot bookkeeping
1061 // -----------------------------------------------------------------------
1062
1063 /// True when `slot` names a live node.
1064 #[inline]
1065 fn is_live(id_of: &[u32], slot: u32) -> bool {
1066 id_of.get(slot as usize).is_some_and(|&i| i != DEAD)
1067 }
1068
1069 /// Bind `id` to a slot (reusing a freed one when available), store `node`
1070 /// there and write `unit` into the slab's row for it. The caller owns linking
1071 /// it into the graph, and must already have checked that `unit` matches the
1072 /// slab's dimension — `debug_assert`ed here, because a refused write would
1073 /// leave a live slot pointing at another node's stale row.
1074 fn alloc_slot(&mut self, id: u32, node: HnswNode, unit: &[f64]) -> u32 {
1075 let slot = match self.free.pop() {
1076 Some(s) => {
1077 self.slots[s as usize] = node;
1078 self.id_of[s as usize] = id;
1079 s
1080 }
1081 None => {
1082 self.slots.push(node);
1083 self.id_of.push(id);
1084 (self.slots.len() - 1) as u32
1085 }
1086 };
1087 let written = self.slab.put(slot, unit);
1088 debug_assert!(written, "alloc_slot was handed a vector the slab refused");
1089 self.slot_of.insert(id, slot);
1090 slot
1091 }
1092
1093 // -----------------------------------------------------------------------
1094 // Reverse adjacency
1095 // -----------------------------------------------------------------------
1096
1097 /// Replace `slot`'s layer-`lc` adjacency with `next`, keeping `back_refs`
1098 /// in lockstep. The single write path for an adjacency list, so the
1099 /// invariant `back_refs[t] == {s : t ∈ slots[s].layers[*]}` holds by
1100 /// construction.
1101 fn set_layer(&mut self, slot: u32, lc: usize, next: Vec<u32>) {
1102 let next_set: BTreeSet<u32> = next.iter().copied().collect();
1103 let prev = std::mem::replace(&mut self.slots[slot as usize].layers[lc], next);
1104 let prev_set: BTreeSet<u32> = prev.into_iter().collect();
1105
1106 for &old in prev_set.difference(&next_set) {
1107 // Still listed on another layer? Then the back-ref stands.
1108 if self.slots[slot as usize]
1109 .layers
1110 .iter()
1111 .any(|l| l.contains(&old))
1112 {
1113 continue;
1114 }
1115 if let Some(refs) = self.back_refs.get_mut(&old) {
1116 refs.remove(&slot);
1117 if refs.is_empty() {
1118 self.back_refs.remove(&old);
1119 }
1120 }
1121 }
1122 for &added in next_set.difference(&prev_set) {
1123 self.back_refs.entry(added).or_default().insert(slot);
1124 }
1125 }
1126
1127 /// Record that `from` now lists `target` as a neighbour.
1128 fn link_back_ref(&mut self, target: u32, from: u32) {
1129 self.back_refs.entry(target).or_default().insert(from);
1130 }
1131
1132 /// The slots that list `slot` as a neighbour on any layer.
1133 fn referrers_of(&self, slot: u32) -> Vec<u32> {
1134 self.back_refs
1135 .get(&slot)
1136 .map(|s| s.iter().copied().collect())
1137 .unwrap_or_default()
1138 }
1139
1140 /// Forget who lists `slot`; the caller has just stripped them all.
1141 fn clear_back_refs(&mut self, slot: u32) {
1142 self.back_refs.remove(&slot);
1143 }
1144
1145 /// Up-convert a decoded 0.6.5 index. Slots are assigned in ascending node-id
1146 /// order — the order a fresh build over ascending ids would have produced —
1147 /// so the result is deterministic, and adjacency ids are remapped onto
1148 /// them. Ids naming a node the map does not hold are dropped, which is what
1149 /// the old `nodes.contains_key` guard in the search paths did anyway.
1150 fn from_v1(v1: HnswIndexV1) -> Self {
1151 let id_of: Vec<u32> = v1.nodes.keys().copied().collect();
1152 let slot_of: BTreeMap<u32, u32> = id_of
1153 .iter()
1154 .enumerate()
1155 .map(|(s, &id)| (id, s as u32))
1156 .collect();
1157 let mut vectors: Vec<Vec<f64>> = Vec::with_capacity(id_of.len());
1158 let slots: Vec<HnswNode> = v1
1159 .nodes
1160 .into_values()
1161 .map(|n| {
1162 vectors.push(n.vector);
1163 HnswNode {
1164 level: n.level,
1165 layers: n
1166 .layers
1167 .into_iter()
1168 .map(|l| l.iter().filter_map(|id| slot_of.get(id).copied()).collect())
1169 .collect(),
1170 }
1171 })
1172 .collect();
1173
1174 let (slab, odd) = slab_of_decoded(&vectors);
1175 let mut out = Self {
1176 base_seed: v1.base_seed,
1177 entry_point: v1.entry_point.and_then(|e| slot_of.get(&e).copied()),
1178 max_level: v1.max_level,
1179 slab,
1180 // A padded position is a vector the index does not really hold, so
1181 // it counts exactly as a refusal does: `can_answer` goes false and
1182 // every caller takes its exhaustive path. Without this an upgraded
1183 // mixed-dimension index claims the fast path over fabricated
1184 // coordinates, which is the one thing the BREAKING note promises it
1185 // will not do.
1186 dim_mismatches: odd,
1187 slots,
1188 slot_of,
1189 id_of,
1190 free: Vec::new(),
1191 ..Self::default()
1192 };
1193 // A dangling entry point would have panicked the old `dist_to`; elect a
1194 // live one instead.
1195 if out.entry_point.is_none() && !out.slot_of.is_empty() {
1196 let (_, &ep) = out
1197 .slot_of
1198 .iter()
1199 .max_by_key(|(_, &s)| out.slots[s as usize].level)
1200 .expect("slot_of is non-empty");
1201 out.entry_point = Some(ep);
1202 out.max_level = out.slots[ep as usize].level;
1203 }
1204 out.rebuild_back_refs();
1205 out
1206 }
1207
1208 /// Adopt a decoded 0.6.6-pre-3b (blob v2) index: the slot layout is already
1209 /// this one's, so only the vectors move — out of the per-node `Vec<f64>`s and
1210 /// into the slab, converted and nothing else. No distance is computed and no
1211 /// vector is re-inserted, so the graph is adopted exactly as it was built.
1212 fn from_v2(v2: HnswIndexV2) -> Self {
1213 let mut vectors: Vec<Vec<f64>> = Vec::with_capacity(v2.slots.len());
1214 let slots: Vec<HnswNode> = v2
1215 .slots
1216 .into_iter()
1217 .map(|n| {
1218 vectors.push(n.vector);
1219 HnswNode {
1220 level: n.level,
1221 layers: n.layers,
1222 }
1223 })
1224 .collect();
1225 let (slab, odd) = slab_of_decoded(&vectors);
1226 let mut out = Self {
1227 base_seed: v2.base_seed,
1228 slab,
1229 // See `from_v1`: a padded position counts as a refusal.
1230 dim_mismatches: odd,
1231 slots,
1232 slot_of: v2.slot_of,
1233 id_of: v2.id_of,
1234 free: v2.free,
1235 entry_point: v2.entry_point,
1236 max_level: v2.max_level,
1237 ..Self::default()
1238 };
1239 out.rebuild_back_refs();
1240 out
1241 }
1242
1243 /// Derive `back_refs` from the adjacency lists. Used after a load, where
1244 /// the field is not serialized. No distance is computed.
1245 fn rebuild_back_refs(&mut self) {
1246 let mut refs: BTreeMap<u32, BTreeSet<u32>> = BTreeMap::new();
1247 for (s, node) in self.slots.iter().enumerate() {
1248 if !Self::is_live(&self.id_of, s as u32) {
1249 continue;
1250 }
1251 for layer in &node.layers {
1252 for &t in layer {
1253 refs.entry(t).or_default().insert(s as u32);
1254 }
1255 }
1256 }
1257 self.back_refs = refs;
1258 }
1259
1260 /// Release `slot`. The caller must already have stripped every adjacency
1261 /// reference to it — a reused slot names a different node.
1262 ///
1263 /// The slab's row is left as it is: nothing reads a dead slot's vector
1264 /// (`is_live` gates every search path), and the next occupant overwrites it.
1265 /// The cost is that the row stays resident until then.
1266 fn free_slot(&mut self, slot: u32) {
1267 let id = std::mem::replace(&mut self.id_of[slot as usize], DEAD);
1268 self.slot_of.remove(&id);
1269 self.slots[slot as usize] = HnswNode::default();
1270 self.free.push(slot);
1271 }
1272
1273 // -----------------------------------------------------------------------
1274 // Internal helpers
1275 // -----------------------------------------------------------------------
1276
1277 /// Beam search on a single layer.
1278 ///
1279 /// Returns a list of `(slot, cosine_distance)` — the `ef` nearest
1280 /// candidates found starting from `ep`. Ascending distance order is not
1281 /// guaranteed (callers sort as needed).
1282 fn beam_search(
1283 slots: &[HnswNode],
1284 id_of: &[u32],
1285 slab: &VecSlab,
1286 q: &[f32],
1287 ep: u32,
1288 layer: usize,
1289 ef: usize,
1290 ) -> Vec<(u32, f64)> {
1291 with_beam_scratch(|scratch| {
1292 // visited: avoid re-expanding a node. A `Vec<bool>` indexed by slot,
1293 // not a `BTreeSet`: this is probed once per candidate edge — order
1294 // 10⁴ times per insert — and each probe was an O(log V) chase
1295 // through separately allocated tree nodes. One memset per call buys
1296 // O(1) probes. It changes neither which nodes are expanded nor the
1297 // order they are pushed in, which is what keeps the graph a
1298 // function of the WAL. The buffer is thread-local and cleared per
1299 // call so a search allocates once per thread, not once per beam.
1300 scratch.reset_beam(slots.len());
1301 if let Some(v) = scratch.visited.get_mut(ep as usize) {
1302 *v = true;
1303 }
1304
1305 let ep_dist = dist_to(slab, ep, q);
1306
1307 // c_heap: min-heap of (dist, slot) — candidates to expand
1308 scratch.c_heap.push(Reverse((OrdF64(ep_dist), ep)));
1309
1310 // w_heap: max-heap of (dist, slot) — ef-best results (worst on top for eviction)
1311 scratch.w_heap.push((OrdF64(ep_dist), ep));
1312
1313 while let Some(&Reverse((OrdF64(c_dist), c))) = scratch.c_heap.peek() {
1314 // furthest in result set
1315 let f_dist = scratch
1316 .w_heap
1317 .peek()
1318 .map(|(OrdF64(d), _)| *d)
1319 .unwrap_or(f64::MAX);
1320 if c_dist > f_dist {
1321 break; // all remaining candidates are farther than our worst result
1322 }
1323 scratch.c_heap.pop();
1324
1325 let neighbors: &[u32] = slots
1326 .get(c as usize)
1327 .and_then(|n| n.layers.get(layer))
1328 .map(|l| l.as_slice())
1329 .unwrap_or_default();
1330
1331 for &e in neighbors {
1332 if scratch.visited.get(e as usize).copied().unwrap_or(true) {
1333 continue;
1334 }
1335 if !Self::is_live(id_of, e) {
1336 continue; // defensive: a freed slot is never a candidate
1337 }
1338 scratch.visited[e as usize] = true;
1339
1340 let e_dist = dist_to(slab, e, q);
1341 let f_dist = scratch
1342 .w_heap
1343 .peek()
1344 .map(|(OrdF64(d), _)| *d)
1345 .unwrap_or(f64::MAX);
1346 if e_dist < f_dist || scratch.w_heap.len() < ef {
1347 scratch.c_heap.push(Reverse((OrdF64(e_dist), e)));
1348 scratch.w_heap.push((OrdF64(e_dist), e));
1349 if scratch.w_heap.len() > ef {
1350 scratch.w_heap.pop(); // evict furthest
1351 }
1352 }
1353 }
1354 }
1355
1356 scratch
1357 .w_heap
1358 .drain()
1359 .map(|(OrdF64(d), s)| (s, d))
1360 .collect()
1361 })
1362 }
1363
1364 /// The **first-rejection prune** — a short-cut of HNSW Algorithm 4, not
1365 /// Algorithm 4 itself. Read this before changing it: the difference is
1366 /// deliberate, measured, and visible in the graph.
1367 ///
1368 /// `candidates` must be `(slot, distance-to-base)` sorted nearest-first; the
1369 /// base vector itself is never needed again, only those distances. Walking
1370 /// nearest-first, a candidate is **kept** when it is closer to the base than
1371 /// to every neighbour already kept — a candidate sitting behind an
1372 /// already-kept neighbour is reachable *through* it, so the link would spend
1373 /// a slot without adding a route. That much is Algorithm 4's diversity test,
1374 /// and it is what `m0` = 128 was paying for before 0.6.6.
1375 ///
1376 /// **Where it stops being Algorithm 4.** Once `rejections ≥ candidates.len()
1377 /// − m`, every candidate still unseen is taken without testing it, and the
1378 /// walk ends. Consequences, both real:
1379 ///
1380 /// * The tail it takes is the **farthest** candidates, untested. Algorithm 4
1381 /// would have gone on testing, and then filled any shortfall with the
1382 /// **nearest** rejects (`keepPrunedConnections`). Those are different sets,
1383 /// so this builds a different graph — it is a short-cut, not an
1384 /// optimisation.
1385 /// * On the neighbour side the candidate list is always exactly `m + 1`
1386 /// long, so `candidates.len() − m` is 1 and the walk stops at the **first**
1387 /// rejection. That prune therefore performs exactly one diversity
1388 /// rejection and keeps the rest as it found them, rather than re-deciding
1389 /// the whole adjacency list.
1390 /// * Because the walk only ever ends with `kept.len() ≥ m` (or with the tail
1391 /// taken), `keepPrunedConnections` would never fire. There is no backfill
1392 /// here; if you restore the full algorithm you must restore it too, or
1393 /// nodes will under-connect into inbound-only leaves that beam search
1394 /// cannot route through.
1395 ///
1396 /// **Why the short-cut ships, and it is not only speed.** Full Algorithm 4
1397 /// was implemented and measured against this, both at `prune = both`,
1398 /// `m0` = 64:
1399 ///
1400 /// | gate | full Algorithm 4 | this |
1401 /// |---|---|---|
1402 /// | `hnsw_5k_1536_recall` (5 000 × 1 536-D uniform) | 1.0000 / 1.0000 in **369 s** | 1.0000 / 1.0000 in **254 s** |
1403 /// | `clustered_…_wider_than_m0` (40 × 120, 128-D) | min **0.5000** / mean 0.9725 | min **0.8000** / mean 0.9950 |
1404 /// | `approximate_recall_5k_timing` (5 000 clustered) | 1.0000, backfill **226.6 s** | 1.0000, backfill **181.5 s** |
1405 ///
1406 /// So the short-cut is 1.25–1.45× faster *and* strictly better on the corpus
1407 /// whose clusters are wider than `m0`. The reason is `keepPrunedConnections`
1408 /// itself: it backfills with the **nearest** rejects, and on a wide cluster
1409 /// those are all crowded in the one direction the diversity test just
1410 /// rejected. Taking the untested far tail instead keeps longer-range links,
1411 /// which is what makes the cluster reachable from outside. The paper's
1412 /// algorithm is the more principled one; on this corpus and at this `m0` it
1413 /// is measurably the worse one, which is why the deviation is deliberate
1414 /// rather than a bug to fix later. `docs/site/rules.md` carries the same
1415 /// table for operators.
1416 fn select_neighbors_first_rejection(
1417 slab: &VecSlab,
1418 id_of: &[u32],
1419 candidates: &[(u32, f64)],
1420 m: usize,
1421 ) -> Vec<u32> {
1422 let mut kept: Vec<u32> = Vec::with_capacity(m);
1423
1424 for (i, &(cand, d_base)) in candidates.iter().enumerate() {
1425 if kept.len() >= m {
1426 break;
1427 }
1428 // `rejections >= candidates.len() - m`, rearranged to avoid an
1429 // underflow when the list is shorter than `m`. Everything from here
1430 // on is taken untested: see the doc comment for what that costs.
1431 if kept.len() + (candidates.len() - i) <= m {
1432 kept.extend(
1433 candidates[i..]
1434 .iter()
1435 .map(|&(c, _)| c)
1436 .filter(|&c| Self::is_live(id_of, c)),
1437 );
1438 break;
1439 }
1440 if !Self::is_live(id_of, cand) {
1441 continue;
1442 }
1443 // Closer to the base than to anything already kept?
1444 let diverse = kept.iter().all(|&k| d_base < dist_slots(slab, k, cand));
1445 if diverse {
1446 kept.push(cand);
1447 }
1448 // A rejection is not recorded: nothing downstream reads it, because
1449 // there is no backfill. The `kept.len() + remaining <= m` test above
1450 // is the same condition as `rejections >= candidates.len() - m`.
1451 }
1452 debug_assert!(kept.len() <= m, "the prune must respect its allowance");
1453 kept
1454 }
1455
1456 /// Greedy 1-NN descent from `ep` at `layer`. Returns the nearest slot
1457 /// found (used for upper-layer descent during insert/search).
1458 fn greedy_step(
1459 slots: &[HnswNode],
1460 id_of: &[u32],
1461 slab: &VecSlab,
1462 q: &[f32],
1463 ep: u32,
1464 layer: usize,
1465 ) -> u32 {
1466 let mut curr = ep;
1467 let mut curr_dist = dist_to(slab, ep, q);
1468 loop {
1469 let mut improved = false;
1470 let neighbors: &[u32] = slots
1471 .get(curr as usize)
1472 .and_then(|n| n.layers.get(layer))
1473 .map(|l| l.as_slice())
1474 .unwrap_or_default();
1475 for &nb in neighbors {
1476 if !Self::is_live(id_of, nb) {
1477 continue;
1478 }
1479 let d = dist_to(slab, nb, q);
1480 if d < curr_dist {
1481 curr_dist = d;
1482 curr = nb;
1483 improved = true;
1484 }
1485 }
1486 if !improved {
1487 break;
1488 }
1489 }
1490 curr
1491 }
1492
1493 // -----------------------------------------------------------------------
1494 // Public API
1495 // -----------------------------------------------------------------------
1496
1497 /// Insert vector `v` for node `id`.
1498 ///
1499 /// Zero vectors are silently skipped (cosine is undefined for them).
1500 /// If `id` already exists it is replaced (remove + re-insert semantics).
1501 ///
1502 /// A vector whose dimension differs from the one this index settled on is
1503 /// **skipped too**, and the first such skip is logged. The slab has one
1504 /// stride, so there is nowhere to put it; before 0.6.6 the distance
1505 /// `zip`-truncated to the shorter of the two vectors and produced a number
1506 /// that meant nothing, which is a worse answer than no answer. A skip makes
1507 /// [`Self::can_answer`] false for good, so the rule falls back to its
1508 /// exhaustive scan rather than answering from an index that is missing a
1509 /// vector.
1510 ///
1511 /// **The stride is re-elected when it was elected from a single sample, and
1512 /// the vector it displaces is parked rather than lost.**
1513 ///
1514 /// The first vector an index takes sets the stride, and that vector may be
1515 /// the odd one out — one 3-element stray ingested ahead of a corpus of
1516 /// 1,536-D embeddings would otherwise refuse every real vector and leave a
1517 /// non-empty index that can answer nothing. So when the index holds at most
1518 /// one vector and the incoming one disagrees with it, the stride is
1519 /// re-elected to the incoming dimension, and the one vector standing behind
1520 /// the old stride is **evicted and parked**: removed from the graph, kept as
1521 /// `(id, unit vector)` in [`Self::parked`], and re-inserted the moment a
1522 /// re-election elects its dimension again.
1523 ///
1524 /// Parking is not a detail. In the order `[real, stray, real, …]` the first
1525 /// real vector is evicted by the stray, and the stray is evicted by the second
1526 /// real one — so without parking the first real vector would be gone for good
1527 /// from an index that believes itself complete, and the rule would silently
1528 /// lose its edges. With it, the second real vector's re-election puts the
1529 /// first one back, and [`Self::can_answer`] refuses the fast path for any
1530 /// dimension still sitting in the parked list.
1531 ///
1532 /// An eviction is not a refusal. A refusal discards a vector the index has no
1533 /// record of; an eviction keeps it, so the index can say precisely what it is
1534 /// missing and stop claiming only that.
1535 pub fn insert(&mut self, id: u32, v: &[f64]) {
1536 // An explicit insert supersedes any parked copy of the same node: the
1537 // caller is telling us this node's vector, and a stale parked one must
1538 // never be revived over it — including when the new vector is the zero
1539 // vector the index will not hold.
1540 self.parked.retain(|(pid, _)| *pid != id);
1541 // The caller is naming this node's vector, so whatever this index
1542 // refused for it before is no longer what it holds. Clearing here (and
1543 // in `remove`, and on revival) is what keeps the skip set from
1544 // suppressing a legitimate re-offer forever. `dim_mismatches` stays
1545 // where it is: the counter is monotone on purpose.
1546 self.refused.remove(&id);
1547 let Some(unit) = l2_normalize(v) else {
1548 return; // zero vector — skip, and do not count it as indexed
1549 };
1550 if self.slab.dim != 0 && unit.len() != self.slab.dim {
1551 // At most one vector in, so the stride was elected on a sample of
1552 // one — or on a node that has since been removed, leaving a stride
1553 // with nothing behind it. Either way the election is not evidence
1554 // against the incoming vector: re-elect, park the single earlier
1555 // vector if there is one, and bring back anything parked at the
1556 // dimension now being elected.
1557 if self.len() <= 1 {
1558 let was = self.slab.dim;
1559 if let Some((&evicted, &slot)) = self.slot_of.iter().next() {
1560 // From the slab rather than from the caller's original `f64`:
1561 // an `f32` widened to `f64` is exact, so this is the vector
1562 // the index was holding.
1563 let kept: Vec<f64> = self.slab.get(slot).iter().map(|&x| x as f64).collect();
1564 eprintln!(
1565 "mushroomdb: HNSW re-elected its embedding dimension from {was} to {} \
1566 at node {id}, and parked node {evicted}: the first vector indexed \
1567 set the dimension and was the odd one out. Node {evicted} returns \
1568 to the index if {was} dimensions are elected again; until then \
1569 this index answers {was}-dimension queries through the full scan.",
1570 unit.len()
1571 );
1572 self.remove(evicted);
1573 self.parked.push((evicted, kept));
1574 }
1575 self.slab = VecSlab::default();
1576 self.slab.dim = unit.len();
1577 // The stride moved, so every earlier refusal was judged against
1578 // a dimension this index no longer holds. Their vectors are
1579 // gone — a refusal discards them — so the most this can do is
1580 // stop skipping them, and let the next scan offer them again.
1581 // `dim_mismatches` stays put: `can_answer` must not recover on
1582 // its own.
1583 self.refused.clear();
1584 // Whatever was parked at this dimension belongs in the index
1585 // again. Their own inserts cannot re-enter this branch — their
1586 // length is the stride — so the recursion is one level deep.
1587 let mut revive: Vec<(u32, Vec<f64>)> = Vec::new();
1588 self.parked.retain(|entry| {
1589 if entry.1.len() == unit.len() {
1590 revive.push(entry.clone());
1591 false
1592 } else {
1593 true
1594 }
1595 });
1596 for (pid, pv) in revive {
1597 self.insert(pid, &pv);
1598 }
1599 } else {
1600 // A refusal still has to honour this method's contract that an
1601 // existing id is *replaced*: the `remove` below is past the
1602 // early return, so without this the index would keep the old
1603 // vector under an id whose new embedding it just rejected.
1604 if self.slot_of.contains_key(&id) {
1605 self.remove(id);
1606 }
1607 self.dim_mismatches += 1;
1608 self.refused.insert(id);
1609 if self.dim_mismatches == 1 {
1610 eprintln!(
1611 "mushroomdb: HNSW skipped node {id}: its embedding has {} dimensions \
1612 and this index holds {}. A mixed-dimension index cannot be \
1613 searched, so this index will now answer through the full scan \
1614 instead; re-embed the collection with one model. Further skips \
1615 on this index are silent.",
1616 unit.len(),
1617 self.slab.dim
1618 );
1619 }
1620 return;
1621 }
1622 }
1623 note_insert();
1624
1625 // Remove existing entry if any (handles update = remove + re-insert).
1626 if self.slot_of.contains_key(&id) {
1627 self.remove(id);
1628 }
1629
1630 let level = gen_level(self.base_seed, id);
1631
1632 // Allocate the slot up front. Nothing lists it yet, so beam search
1633 // cannot reach it; scoring it during the prune below then needs no
1634 // special case for "the node not in the graph yet".
1635 let slot = self.alloc_slot(
1636 id,
1637 HnswNode {
1638 level,
1639 layers: vec![vec![]; level + 1],
1640 },
1641 &unit,
1642 );
1643
1644 let Some(ep) = self.entry_point else {
1645 // First node ever inserted.
1646 self.entry_point = Some(slot);
1647 self.max_level = level;
1648 return;
1649 };
1650
1651 // The query every distance on this insert path is taken against: the
1652 // same `f32` values the slab now holds for `slot`, so a distance to the
1653 // new node is exactly a distance between two slab rows. Filled into a
1654 // thread-local buffer rather than borrowed from the slab because the
1655 // graph is mutated below.
1656 let q = take_query_f32(&unit);
1657
1658 let params = hnsw_params();
1659 let prune = self.resolved_prune(¶ms);
1660 let max_level = self.max_level;
1661 let mut curr_ep = ep;
1662
1663 // Phase 1: greedy descent from max_level to level+1 (ef=1).
1664 for lc in ((level + 1)..=max_level).rev() {
1665 curr_ep = Self::greedy_step(&self.slots, &self.id_of, &self.slab, &q, curr_ep, lc);
1666 }
1667
1668 // Phase 2: beam-search + connect at each layer from min(level, max_level)
1669 // down to 0.
1670 for lc in (0..=level.min(max_level)).rev() {
1671 let m_lc = if lc == 0 { params.m0 } else { params.m };
1672
1673 // Beam search to collect ef_construction nearest candidates.
1674 let mut candidates = Self::beam_search(
1675 &self.slots,
1676 &self.id_of,
1677 &self.slab,
1678 &q,
1679 curr_ep,
1680 lc,
1681 params.ef_construction,
1682 );
1683
1684 sort_by_distance(&mut candidates);
1685
1686 // Advance curr_ep to the nearest candidate before the heuristic
1687 // thins the list: the descent wants the nearest node, not the most
1688 // diverse one.
1689 if let Some(&(nearest, _)) = candidates.first() {
1690 curr_ep = nearest;
1691 }
1692
1693 // The new node's own neighbours go through the heuristic too. This
1694 // is where the diversity buys recall: taking the m nearest leaves
1695 // every node in a dense cluster pointing back into the same
1696 // cluster, and the beam never crosses out of it.
1697 let neighbors =
1698 Self::select_neighbors_first_rejection(&self.slab, &self.id_of, &candidates, m_lc);
1699 self.set_layer(slot, lc, neighbors.clone());
1700
1701 // Add bidirectional links and prune over-connected neighbors.
1702 for &nb in &neighbors {
1703 let nb_node = &mut self.slots[nb as usize];
1704 while nb_node.layers.len() <= lc {
1705 nb_node.layers.push(vec![]);
1706 }
1707 if !nb_node.layers[lc].contains(&slot) {
1708 nb_node.layers[lc].push(slot);
1709 self.link_back_ref(slot, nb);
1710 }
1711
1712 if self.slots[nb as usize].layers[lc].len() <= m_lc {
1713 continue;
1714 }
1715 // Over-connected. `set_layer` keeps `back_refs` in step for
1716 // every link the prune drops, whichever prune that is.
1717 //
1718 // Scoring is two slab rows per candidate — where 0.6.5 cloned
1719 // the neighbour's whole 12 KB vector first, once per
1720 // over-connected neighbour and so up to `m0` times per insert.
1721 let current: Vec<u32> = self.slots[nb as usize].layers[lc].clone();
1722 let mut scored: Vec<(u32, f64)> = current
1723 .iter()
1724 .filter(|&&s| Self::is_live(&self.id_of, s))
1725 .map(|&s| (s, dist_slots(&self.slab, s, nb)))
1726 .collect();
1727 sort_by_distance(&mut scored);
1728 let kept: Vec<u32> = match prune {
1729 // Keep the m nearest. One distance per candidate, already
1730 // computed above.
1731 Prune::Own => scored.iter().take(m_lc).map(|(s, _)| *s).collect(),
1732 // Re-decide the whole list by diversity: O(m₀²) distances,
1733 // once per over-connected neighbour per insert.
1734 Prune::Both => Self::select_neighbors_first_rejection(
1735 &self.slab,
1736 &self.id_of,
1737 &scored,
1738 m_lc,
1739 ),
1740 };
1741 self.set_layer(nb, lc, kept);
1742 }
1743 }
1744
1745 // Update entry point if new node has a higher level.
1746 if level > max_level {
1747 self.entry_point = Some(slot);
1748 self.max_level = level;
1749 }
1750 stash_query_f32(q);
1751 }
1752
1753 /// Remove node `id` from the index.
1754 ///
1755 /// Every reference to its slot is stripped — the slot is about to be reused
1756 /// by another node, so a leftover link would silently name the wrong
1757 /// vector. `back_refs` is what makes finding those references O(in-degree)
1758 /// rather than a scan of the whole index. If `id` was the entry point, a
1759 /// new entry point is elected (highest remaining level).
1760 pub fn remove(&mut self, id: u32) {
1761 // A removed node must not come back through a later stride re-election.
1762 // Done before the early return, because the node may be parked rather
1763 // than indexed — which is exactly the state a removal has to clear.
1764 self.parked.retain(|(pid, _)| *pid != id);
1765 // Same reason, for the same early return: a refused node is not in
1766 // `slot_of`, so a removal that stopped at the return below would leave
1767 // its id in the skip set and a later re-offer would never be tried.
1768 self.refused.remove(&id);
1769 let Some(slot) = self.slot_of.get(&id).copied() else {
1770 return;
1771 };
1772
1773 // Drop the node's own out-edges, and with them its back-ref claims.
1774 // Done directly rather than through `set_layer` because every layer is
1775 // cleared at once: no target can still be listed on another layer.
1776 let mut targets: BTreeSet<u32> = BTreeSet::new();
1777 for layer in &self.slots[slot as usize].layers {
1778 targets.extend(layer.iter().copied());
1779 }
1780 for layer in self.slots[slot as usize].layers.iter_mut() {
1781 layer.clear();
1782 }
1783 for t in targets {
1784 if let Some(refs) = self.back_refs.get_mut(&t) {
1785 refs.remove(&slot);
1786 if refs.is_empty() {
1787 self.back_refs.remove(&t);
1788 }
1789 }
1790 }
1791
1792 // Walk every node that lists this slot and strip it. Asymmetric pruning
1793 // means these are NOT just the nodes it listed back, which is why the
1794 // reverse index exists.
1795 let referrers = self.referrers_of(slot);
1796 note_remove_scanned(referrers.len());
1797 for referrer in referrers {
1798 for layer in self.slots[referrer as usize].layers.iter_mut() {
1799 layer.retain(|&x| x != slot);
1800 }
1801 }
1802 self.clear_back_refs(slot);
1803
1804 self.free_slot(slot);
1805
1806 // Update entry point if needed.
1807 if self.entry_point == Some(slot) {
1808 if self.slot_of.is_empty() {
1809 self.entry_point = None;
1810 self.max_level = 0;
1811 } else {
1812 // Iterate in node-id order and keep the last maximum, exactly
1813 // as the BTreeMap scan this replaced did.
1814 let (_, &new_ep) = self
1815 .slot_of
1816 .iter()
1817 .max_by_key(|(_, &s)| self.slots[s as usize].level)
1818 .expect("slot_of is non-empty");
1819 self.entry_point = Some(new_ep);
1820 self.max_level = self.slots[new_ep as usize].level;
1821 }
1822 }
1823 }
1824
1825 /// Approximate k-nearest-neighbor search by cosine similarity.
1826 ///
1827 /// Returns up to `k` results as `(node_id, cosine_similarity)` pairs,
1828 /// sorted descending by similarity.
1829 ///
1830 /// **This method never falls back.** Zero-norm queries return empty, and so
1831 /// does a query whose dimension is not this index's stride — there is
1832 /// nothing to compare it against here. The fall-back to an exhaustive scan
1833 /// belongs to the caller, and [`Self::can_answer`] is how a caller knows it
1834 /// is needed: ask it first, and take your own path when it says no. Every
1835 /// caller in this workspace does (`index.rs::hnsw_candidates`,
1836 /// `engine.rs::hnsw_search_dst`/`_any_dst`).
1837 ///
1838 /// **The similarity is for ordering, not for reporting.** It is computed
1839 /// from the index's `f32` copies, so it is accurate to ~1e-6 and an exact
1840 /// duplicate scores 0.9999999 rather than 1.0. `hnsw_candidates` discards it
1841 /// and keeps the ids; `db.rs::find_similar_vector` re-scores every candidate
1842 /// against the `f64` property vectors before it applies `min`, orders, or
1843 /// reports anything. A new caller must do the same.
1844 pub fn search(&self, q: &[f64], k: usize) -> Vec<(u32, f64)> {
1845 self.search_with_ef(q, k, self.ef_for(k))
1846 }
1847
1848 /// The beam width [`HnswIndex::search`] uses for `k` results.
1849 ///
1850 /// Exposed so a caller that widens the beam itself — an exact
1851 /// `VectorSimilar` rule looking for *every* hit above a floor — can start
1852 /// from the same place `search` would have. Reads [`hnsw_params`], so the
1853 /// widening loop in `index.rs` inherits whatever shape this process was
1854 /// configured with and never names a constant of its own.
1855 pub fn ef_for(&self, k: usize) -> usize {
1856 k.max(hnsw_params().ef_search)
1857 }
1858
1859 /// [`HnswIndex::search`], with the layer-0 beam width set independently of
1860 /// the result count.
1861 ///
1862 /// `ef` below `k` is raised to `k`: a beam narrower than the answer cannot
1863 /// produce the answer.
1864 pub fn search_with_ef(&self, q: &[f64], k: usize, ef: usize) -> Vec<(u32, f64)> {
1865 let Some(unit_q) = l2_normalize(q) else {
1866 return vec![];
1867 };
1868 if self.slab.dim == 0 || unit_q.len() != self.slab.dim {
1869 return vec![];
1870 }
1871 let Some(ep) = self.entry_point.filter(|&s| Self::is_live(&self.id_of, s)) else {
1872 // No entry point, or one left dangling by a bug: answer nothing
1873 // rather than walk from a freed slot and hand back a `u32::MAX` id.
1874 return vec![];
1875 };
1876 if k == 0 {
1877 return vec![];
1878 }
1879
1880 note_search();
1881 // The caller's width, floored at `k`: a beam narrower than the answer
1882 // cannot produce the answer. `search` passes `ef_for(k)`, which is the
1883 // `hnsw_params()` width this function used before the width became a
1884 // parameter, so its behaviour is unchanged.
1885 let ef = ef.max(k);
1886 let unit_q = take_query_f32(&unit_q);
1887 let mut curr_ep = ep;
1888
1889 // Greedy descent from max_level to layer 1.
1890 for lc in (1..=self.max_level).rev() {
1891 curr_ep = Self::greedy_step(&self.slots, &self.id_of, &self.slab, &unit_q, curr_ep, lc);
1892 }
1893
1894 // Beam search at layer 0 with ef candidates.
1895 let candidates = Self::beam_search(
1896 &self.slots,
1897 &self.id_of,
1898 &self.slab,
1899 &unit_q,
1900 curr_ep,
1901 0,
1902 ef,
1903 );
1904
1905 // Convert slots → node ids and distances → cosine similarities; sort
1906 // descending; take k.
1907 let mut results: Vec<(u32, f64)> = candidates
1908 .into_iter()
1909 .map(|(s, dist)| (self.id_of[s as usize], (1.0 - dist).clamp(-1.0, 1.0)))
1910 .collect();
1911 results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
1912 results.truncate(k);
1913 stash_query_f32(unit_q);
1914 results
1915 }
1916
1917 /// Count the graph's own storage, so a benchmark can report bytes per
1918 /// indexed vector instead of guessing at them.
1919 ///
1920 /// Counts payload only — the `u32`s in every adjacency list, the `u32`s in
1921 /// the reverse index, and the `f32`s of every live node's slab row.
1922 /// Allocator headers and the per-`Vec`/per-`BTreeSet` bookkeeping are not
1923 /// counted, and neither are the slab rows of freed slots, so this is a floor
1924 /// on resident size, not a measurement of it.
1925 pub fn memory_stats(&self) -> HnswMemoryStats {
1926 let mut neighbour_slots = 0usize;
1927 for (s, node) in self.slots.iter().enumerate() {
1928 if !Self::is_live(&self.id_of, s as u32) {
1929 continue;
1930 }
1931 neighbour_slots += node.layers.iter().map(|l| l.len()).sum::<usize>();
1932 }
1933 let live_nodes = self.slot_of.len();
1934 HnswMemoryStats {
1935 live_nodes,
1936 neighbour_slots,
1937 back_ref_entries: self.back_refs.values().map(|s| s.len()).sum(),
1938 vector_floats: self.slab.floats_for(live_nodes),
1939 }
1940 }
1941
1942 // -----------------------------------------------------------------------
1943 // Test-only introspection
1944 // -----------------------------------------------------------------------
1945
1946 /// The slots `back_refs` claims list `id`'s slot as a neighbour.
1947 #[cfg(any(test, feature = "test-hooks"))]
1948 pub fn back_refs_for_test(&self, id: u32) -> BTreeSet<u32> {
1949 let Some(&slot) = self.slot_of.get(&id) else {
1950 return BTreeSet::new();
1951 };
1952 self.back_refs.get(&slot).cloned().unwrap_or_default()
1953 }
1954
1955 /// The same set, derived by walking every live node's adjacency lists.
1956 /// The oracle `back_refs_for_test` must agree with.
1957 #[cfg(any(test, feature = "test-hooks"))]
1958 pub fn scan_back_refs_for_test(&self, id: u32) -> BTreeSet<u32> {
1959 let Some(&slot) = self.slot_of.get(&id) else {
1960 return BTreeSet::new();
1961 };
1962 let mut out = BTreeSet::new();
1963 for (s, node) in self.slots.iter().enumerate() {
1964 if !Self::is_live(&self.id_of, s as u32) {
1965 continue;
1966 }
1967 if node.layers.iter().any(|l| l.contains(&slot)) {
1968 out.insert(s as u32);
1969 }
1970 }
1971 out
1972 }
1973
1974 /// Number of adjacency slots ever allocated, live or freed. A removal that
1975 /// reuses a slot does not grow it.
1976 #[cfg(any(test, feature = "test-hooks"))]
1977 pub fn slot_capacity_for_test(&self) -> usize {
1978 self.slots.len()
1979 }
1980}
1981
1982// ---------------------------------------------------------------------------
1983// Persisted blob: magic + version
1984// ---------------------------------------------------------------------------
1985
1986/// Magic bytes at the head of every versioned (0.6.6 and later) HNSW blob.
1987pub const HNSW_BLOB_MAGIC: [u8; 4] = *b"MHNS";
1988/// Highest blob version this build can read, and the one it writes.
1989///
1990/// **3** since the distance kernel: the index no longer holds `f64` vectors per
1991/// node, it holds one `f32` slab, so the serialized shape changed and the blob
1992/// halved. Version 2 (0.6.6 before the kernel) and the bare 0.6.5 shape both
1993/// still decode, with no vector re-inserted and no distance computed. A reader
1994/// older than this one meeting a v3 blob fails its version check and leaves the
1995/// side on its `hnsw_tracked` full scan — slower, never wrong.
1996/// Version 4 (0.6.9) appends the state `HnswIndex` does not serialize —
1997/// the refusal count, the refused ids, and the parked vectors — so a reopen
1998/// restores them instead of re-deriving them by re-offering every vector. The
1999/// index body inside is byte-identical to v3's, and `complete` stays the last
2000/// field so [`hnsw_blob_complete`] can still peek it without decoding.
2001pub const HNSW_BLOB_VERSION: u16 = 4;
2002
2003/// Bytes of the wrapper's header: `magic` (4 raw bytes) then `version` (a
2004/// little-endian `u16`) under bincode's fixed-int encoding. Read directly rather
2005/// than through a deserialize, because the rest of the wrapper's shape depends on
2006/// the version it carries.
2007const HNSW_BLOB_HEADER_LEN: usize = 6;
2008
2009/// On-disk wrapper for a persisted HNSW graph.
2010///
2011/// `magic` + `version` make a 0.6.5 blob and a 0.6.6 blob distinguishable
2012/// without bumping the snapshot format: section 6 carries the index as two
2013/// opaque `Vec<u8>` per rule, so only the bytes inside change.
2014///
2015/// The reverse direction is safe by construction: an older binary meeting one of
2016/// these either fails its version check or fails its
2017/// `bincode::deserialize::<HnswIndex>` against the shape it knows, and falls
2018/// back to the `hnsw_tracked` full scan — slower, never wrong.
2019#[derive(Serialize, Deserialize, Clone, Debug)]
2020pub struct HnswBlob {
2021 pub magic: [u8; 4],
2022 pub version: u16,
2023 pub index: HnswIndex,
2024 /// `HnswIndex::dim_mismatches`, which the index itself does not serialize.
2025 pub dim_mismatches: u64,
2026 /// The ids behind that count. May be shorter than the count — see the field
2027 /// doc on `HnswIndex::refused`.
2028 pub refused: BTreeSet<u32>,
2029 /// `HnswIndex::parked`, in the slab's `f32` unit rather than a second `f64`
2030 /// copy.
2031 pub parked: Vec<(u32, Vec<f32>)>,
2032 /// False when the rule's sliced build had not finished when this blob was
2033 /// written, so the graph inside holds a prefix of the corpus.
2034 ///
2035 /// `RuleEngine::pending_builds` is not persisted, so a snapshot taken
2036 /// mid-build has to carry its own evidence: without this, a reader opening
2037 /// that snapshot decodes the partial graph, finds it perfectly answerable
2038 /// and serves `find_similar` from a fraction of the corpus with no signal.
2039 /// A `false` here makes [`HnswIndex::can_answer`] refuse, which sends every
2040 /// such query to the exhaustive scan until a write, `mushroomdb build-index`
2041 /// or `serve`'s pump finishes the build. Open also reads it to register the
2042 /// rule in `pending_builds`, so a restarted `serve` has something to pump
2043 /// without waiting for a write.
2044 pub complete: bool,
2045}
2046
2047/// Serialize-only twin of [`HnswBlob`] so encoding never clones the index.
2048///
2049/// bincode is positional: this must stay field-for-field in lockstep with
2050/// [`HnswBlob`], and `complete` must stay last.
2051#[derive(Serialize)]
2052struct HnswBlobRef<'a> {
2053 magic: [u8; 4],
2054 version: u16,
2055 index: &'a HnswIndex,
2056 dim_mismatches: u64,
2057 refused: BTreeSet<u32>,
2058 parked: Vec<(u32, Vec<f32>)>,
2059 complete: bool,
2060}
2061
2062/// The v3 wrapper: the same index body with no side state after it.
2063///
2064/// Read-only — nothing writes it any more. Keeping it as its own struct is what
2065/// lets the v4 body append fields without giving v3 a second index shape to
2066/// decode.
2067#[derive(Deserialize)]
2068struct HnswBlobV3 {
2069 #[allow(dead_code)]
2070 magic: [u8; 4],
2071 #[allow(dead_code)]
2072 version: u16,
2073 index: HnswIndex,
2074 complete: bool,
2075}
2076
2077/// The 0.6.5 on-disk shape: node ids key the map *and* name the adjacency
2078/// lists. Read-only — nothing writes it any more.
2079#[derive(Deserialize)]
2080struct HnswIndexV1 {
2081 base_seed: u64,
2082 nodes: BTreeMap<u32, HnswNodeV1>,
2083 entry_point: Option<u32>,
2084 max_level: usize,
2085}
2086
2087#[derive(Deserialize)]
2088struct HnswNodeV1 {
2089 level: usize,
2090 vector: Vec<f64>,
2091 /// Neighbour node **ids**.
2092 layers: Vec<Vec<u32>>,
2093}
2094
2095/// The blob-v2 shape — 0.6.6 before the distance kernel. Slot-keyed already, but
2096/// with an `f64` vector inside every node instead of a slab beside them.
2097/// Read-only: nothing writes it any more.
2098///
2099/// The field order is the v2 `HnswIndex`'s, and it must stay that way: bincode is
2100/// positional, so this struct *is* the old format's definition.
2101#[derive(Deserialize)]
2102struct HnswIndexV2 {
2103 base_seed: u64,
2104 slots: Vec<HnswNodeV2>,
2105 slot_of: BTreeMap<u32, u32>,
2106 id_of: Vec<u32>,
2107 free: Vec<u32>,
2108 entry_point: Option<u32>,
2109 max_level: usize,
2110}
2111
2112#[derive(Deserialize)]
2113struct HnswNodeV2 {
2114 level: usize,
2115 vector: Vec<f64>,
2116 /// Neighbour **slots**, as in v3.
2117 layers: Vec<Vec<u32>>,
2118}
2119
2120/// The v2 wrapper: the same magic and version, a different index shape.
2121#[derive(Deserialize)]
2122struct HnswBlobV2 {
2123 #[allow(dead_code)]
2124 magic: [u8; 4],
2125 #[allow(dead_code)]
2126 version: u16,
2127 index: HnswIndexV2,
2128}
2129
2130/// Serialize `index` as a versioned blob. `None` only if bincode fails.
2131///
2132/// `complete` is false when the rule's sliced build is still owed vectors, and
2133/// is what stops a reader over this snapshot answering from a prefix of the
2134/// corpus. Every caller that cannot be mid-build passes `true`.
2135pub fn encode_hnsw_blob(index: &HnswIndex, complete: bool) -> Option<Vec<u8>> {
2136 bincode::serialize(&HnswBlobRef {
2137 magic: HNSW_BLOB_MAGIC,
2138 version: HNSW_BLOB_VERSION,
2139 index,
2140 dim_mismatches: index.dim_mismatches(),
2141 refused: index.refused_ids(),
2142 parked: index.parked_rows_f32(),
2143 complete,
2144 })
2145 .ok()
2146}
2147
2148/// Whether a persisted blob was written as a finished graph.
2149///
2150/// v3 carries `complete` as its last field, so this peeks without decoding the
2151/// index. v1 and v2 have no flag and were always whole. `None` if the bytes are
2152/// empty or not a blob this build can read.
2153pub fn hnsw_blob_complete(blob: &[u8]) -> Option<bool> {
2154 if blob.is_empty() {
2155 return None;
2156 }
2157 if blob.len() >= HNSW_BLOB_HEADER_LEN && blob[..4] == HNSW_BLOB_MAGIC {
2158 let version = u16::from_le_bytes([blob[4], blob[5]]);
2159 return match version {
2160 // v4 appends its side state *before* `complete`, so `complete` is
2161 // still the last byte and this peek is unchanged. Leaving v4 out of
2162 // this arm would return `None`, which the open path reads as
2163 // "complete" — a mid-build blob would silently stop registering its
2164 // pending build and nothing would fail loudly.
2165 3 | 4 => {
2166 if blob.len() < HNSW_BLOB_HEADER_LEN + 1 {
2167 return None;
2168 }
2169 match blob[blob.len() - 1] {
2170 0 => Some(false),
2171 1 => Some(true),
2172 _ => None,
2173 }
2174 }
2175 1 | 2 => Some(true),
2176 _ => None,
2177 };
2178 }
2179 // No wrapper: a 0.6.5 v1 blob. Sliced builds did not exist.
2180 Some(true)
2181}
2182
2183/// Decode a persisted HNSW blob.
2184///
2185/// Accepts, in this order:
2186///
2187/// * **v3** — this build's shape, slot-keyed with an `f32` slab.
2188/// * **v2** — 0.6.6 before the distance kernel, slot-keyed with an `f64` vector
2189/// per node. The vectors move into a slab and nothing else changes.
2190/// * **v1** — a bare bincoded `HnswIndex` from 0.6.5, id-keyed, up-converted by
2191/// remapping the decoded adjacency lists onto slots.
2192///
2193/// No distance is computed and no vector is re-inserted on any of those paths:
2194/// the graph is adopted as it was built.
2195///
2196/// The version is read from the header rather than inferred from a successful
2197/// deserialize, because v2 and v3 differ *inside* the wrapper: a v2 blob fed to
2198/// v3's shape could in principle decode into nonsense rather than fail. A blob
2199/// whose magic matches but whose version this build does not know is rejected
2200/// exactly as a corrupt one is — the caller leaves the side on its
2201/// `hnsw_tracked` full-scan fallback rather than risk misreading it.
2202/// The checks and repairs every wrapper version owes its decoded index.
2203///
2204/// Shared by the v3 and v4 arms so the two cannot drift: a guard added for one
2205/// version and forgotten in the other is exactly how a truncated blob would
2206/// reach `dot_f32`.
2207fn finish_decoded_index(
2208 mut index: HnswIndex,
2209 complete: bool,
2210 version: u16,
2211) -> Result<HnswIndex, String> {
2212 // bincode will happily decode a `Vec<f32>` shorter than the slots claim —
2213 // it reads the length prefix it is given. A short slab then hands `dot_f32`
2214 // two slices of unequal length, which is a `debug_assert` in a debug build
2215 // and a garbage distance in a release one. Refuse instead: the caller keeps
2216 // its `hnsw_tracked` scan.
2217 if index.slab.dim != 0 && index.slab.data.len() < index.slots.len() * index.slab.dim {
2218 return Err(format!(
2219 "HNSW v{version} blob is truncated: the slab holds {} floats, {} slots \
2220 of {} dimensions need {}",
2221 index.slab.data.len(),
2222 index.slots.len(),
2223 index.slab.dim,
2224 index.slots.len() * index.slab.dim
2225 ));
2226 }
2227 if index.id_of.len() != index.slots.len() {
2228 return Err(format!(
2229 "HNSW v{version} blob is inconsistent: {} slots against {} id entries",
2230 index.slots.len(),
2231 index.id_of.len()
2232 ));
2233 }
2234 index.rebuild_back_refs();
2235 index.incomplete = !complete;
2236 Ok(index)
2237}
2238
2239pub fn decode_hnsw_blob(blob: &[u8]) -> Result<HnswIndex, String> {
2240 if blob.is_empty() {
2241 return Err("empty blob".to_string());
2242 }
2243 if blob.len() >= HNSW_BLOB_HEADER_LEN && blob[..4] == HNSW_BLOB_MAGIC {
2244 let version = u16::from_le_bytes([blob[4], blob[5]]);
2245 return match version {
2246 4 => bincode::deserialize::<HnswBlob>(blob)
2247 .map_err(|e| format!("HNSW v4 blob did not decode ({e})"))
2248 .and_then(|b| {
2249 let mut index = finish_decoded_index(b.index, b.complete, 4)?;
2250 index.restore_side_state(b.dim_mismatches, b.refused, b.parked);
2251 Ok(index)
2252 }),
2253 3 => bincode::deserialize::<HnswBlobV3>(blob)
2254 .map_err(|e| format!("HNSW v3 blob did not decode ({e})"))
2255 // No side state on the wire: parked and refused stay empty and
2256 // the counter stays zero, so the open-time scan re-offers every
2257 // vector and re-derives them exactly as it did before v4.
2258 .and_then(|b| finish_decoded_index(b.index, b.complete, 3)),
2259 2 => bincode::deserialize::<HnswBlobV2>(blob)
2260 .map_err(|e| format!("HNSW v2 blob did not decode ({e})"))
2261 .map(|b| HnswIndex::from_v2(b.index)),
2262 v => Err(format!(
2263 "HNSW blob version {v} is not readable by this build (reads up to \
2264 {HNSW_BLOB_VERSION})"
2265 )),
2266 };
2267 }
2268 // No wrapper, or foreign magic: try the 0.6.5 shape.
2269 match bincode::deserialize::<HnswIndexV1>(blob) {
2270 Ok(v1) => Ok(HnswIndex::from_v1(v1)),
2271 Err(v1_err) => Err(format!(
2272 "not a versioned HNSW blob (magic {:?}) and not a v1 one ({v1_err})",
2273 &blob[..4.min(blob.len())]
2274 )),
2275 }
2276}
2277
2278// ---------------------------------------------------------------------------
2279// Archived search helper
2280// ---------------------------------------------------------------------------
2281
2282/// Deserialize an `HnswIndex` from a bincoded blob and search it.
2283///
2284/// Returns an empty vec if the blob is corrupt or `q` is the zero vector.
2285/// Blobs are produced by `engine.rs` when it calls `bincode::serialize` on the
2286/// index before handing it to the V8 encoder.
2287///
2288/// # Caller note
2289///
2290/// This function has no current caller in the codebase. It is a Task-3 /
2291/// future-use primitive: external callers with direct access to a V8 HNSW blob
2292/// (e.g. snapshot introspection tools or the upcoming MCP search path) can use
2293/// this to run an ANN query without a live `RuleEngine`.
2294pub fn search_hnsw_blob(blob: &[u8], q: &[f64], k: usize) -> Vec<(u32, f64)> {
2295 let Ok(idx) = decode_hnsw_blob(blob) else {
2296 return vec![];
2297 };
2298 idx.search(q, k)
2299}
2300
2301// ---------------------------------------------------------------------------
2302// Vector helpers
2303// ---------------------------------------------------------------------------
2304
2305/// Sort `(slot, distance)` pairs nearest-first, breaking ties on the slot so
2306/// the order is a function of the graph and not of heap iteration order. The
2307/// §3.5 heuristic reads this order, so a tie decided differently on two
2308/// machines would be two different graphs from the same WAL.
2309fn sort_by_distance(v: &mut [(u32, f64)]) {
2310 v.sort_by(|a, b| {
2311 a.1.partial_cmp(&b.1)
2312 .unwrap_or(std::cmp::Ordering::Equal)
2313 .then_with(|| a.0.cmp(&b.0))
2314 });
2315}
2316
2317/// L2-normalize `v`. Returns `None` for the zero vector.
2318pub(crate) fn l2_normalize(v: &[f64]) -> Option<Vec<f64>> {
2319 let norm = v.iter().map(|x| x * x).sum::<f64>().sqrt();
2320 if norm == 0.0 {
2321 return None;
2322 }
2323 Some(v.iter().map(|x| x / norm).collect())
2324}
2325
2326/// Build `n` deterministic unit vectors in `dim` dimensions from a splitmix64
2327/// stream seeded with `seed`. Each vector is L2-normalized.
2328///
2329/// Public — and not `#[cfg(test)]` — so the in-crate recall gate, the
2330/// `tests/hnsw_scale.rs` benchmark and any external measurement all draw from
2331/// one generator: a recall number and a build time are only comparable when the
2332/// vectors behind them are the same vectors.
2333#[doc(hidden)]
2334pub fn make_unit_vecs(n: usize, dim: usize, seed: u64) -> Vec<Vec<f64>> {
2335 let mut state = seed;
2336 (0..n)
2337 .map(|_| {
2338 let raw: Vec<f64> = (0..dim)
2339 .map(|_| {
2340 state = splitmix64(state);
2341 // Map to [-1, 1]
2342 (state as i64 as f64) / (i64::MAX as f64)
2343 })
2344 .collect();
2345 l2_normalize(&raw).unwrap_or_else(|| vec![1.0; dim])
2346 })
2347 .collect()
2348}
2349
2350/// Build `clusters × per_cluster` deterministic unit vectors, each a cluster
2351/// centre plus a small perturbation, in cluster order.
2352///
2353/// Uniformly random high-dimensional vectors are all nearly orthogonal, so they
2354/// exercise beam width and nothing else. Real embeddings cluster, and a cluster
2355/// wider than the layer-0 allowance is the case the §3.5 heuristic exists for:
2356/// without it every member of the cluster spends all its links inside the
2357/// cluster and the graph stops being navigable between clusters.
2358#[doc(hidden)]
2359pub fn make_clustered_unit_vecs(
2360 clusters: usize,
2361 per_cluster: usize,
2362 dim: usize,
2363 seed: u64,
2364) -> Vec<Vec<f64>> {
2365 let centres = make_unit_vecs(clusters, dim, seed);
2366 let mut state = seed ^ 0xA5A5_5A5A_1234_9876;
2367 let mut out = Vec::with_capacity(clusters * per_cluster);
2368 for centre in ¢res {
2369 for _ in 0..per_cluster {
2370 let raw: Vec<f64> = centre
2371 .iter()
2372 .map(|c| {
2373 state = splitmix64(state);
2374 // A perturbation small enough that cluster membership is
2375 // unambiguous, large enough that members are distinct.
2376 c + 0.12 * (state as i64 as f64) / (i64::MAX as f64) / (dim as f64).sqrt()
2377 })
2378 .collect();
2379 out.push(l2_normalize(&raw).unwrap_or_else(|| centre.clone()));
2380 }
2381 }
2382 out
2383}
2384
2385// ---------------------------------------------------------------------------
2386// Tests
2387// ---------------------------------------------------------------------------
2388
2389#[cfg(test)]
2390mod tests {
2391 use super::*;
2392
2393 /// Exact brute-force k-NN by cosine (dot product for unit vecs).
2394 fn exact_knn(vecs: &[Vec<f64>], q: &[f64], k: usize) -> Vec<usize> {
2395 let unit_q = l2_normalize(q).unwrap();
2396 let mut scores: Vec<(usize, f64)> = vecs
2397 .iter()
2398 .enumerate()
2399 .map(|(i, v)| {
2400 let dot: f64 = v.iter().zip(unit_q.iter()).map(|(a, b)| a * b).sum();
2401 (i, dot)
2402 })
2403 .collect();
2404 scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
2405 scores.into_iter().take(k).map(|(i, _)| i).collect()
2406 }
2407
2408 #[test]
2409 fn hnsw_recalls_near_duplicate() {
2410 // 200 deterministic unit vectors in dim 32.
2411 let vecs = make_unit_vecs(200, 32, 0xDEAD_BEEF_1234_5678);
2412 let seed = crate::index::fnv1a_u64(b"test-rule");
2413 let mut idx = HnswIndex::new(seed);
2414 for (i, v) in vecs.iter().enumerate() {
2415 idx.insert(i as u32, v);
2416 }
2417
2418 // Query: vec[7] + tiny noise (so nearest is definitely 7).
2419 let mut noisy = vecs[7].clone();
2420 noisy[0] += 1e-4;
2421 noisy[1] -= 1e-4;
2422
2423 let results = idx.search(&noisy, 1);
2424 assert!(!results.is_empty(), "HNSW must return at least one result");
2425 assert_eq!(
2426 results[0].0, 7,
2427 "nearest to vec[7]+noise must be 7, got {} (cos={:.6})",
2428 results[0].0, results[0].1
2429 );
2430 }
2431
2432 /// The distance kernel is the insert path's whole cost, so the count of
2433 /// evaluations is the measurement this task is judged on. A counter a test
2434 /// cannot read is not a gate.
2435 #[test]
2436 fn a_single_insert_reports_its_distance_evaluations() {
2437 let vecs = make_unit_vecs(201, 32, 0x0D15_7A17);
2438 let mut idx = HnswIndex::new(crate::index::fnv1a_u64(b"dist-evals"));
2439 for (i, v) in vecs.iter().take(200).enumerate() {
2440 idx.insert(i as u32, v);
2441 }
2442
2443 hnsw_dist_evals_reset();
2444 idx.insert(200, &vecs[200]);
2445 let total = hnsw_dist_evals();
2446 let pairwise = hnsw_dist_evals_pairwise();
2447 assert!(
2448 total > 0,
2449 "an insert into a 200-node index must evaluate distances"
2450 );
2451 assert!(
2452 pairwise > 0,
2453 "the prune compares candidates against each other, so some \
2454 evaluations must be pairwise"
2455 );
2456 assert!(
2457 pairwise <= total,
2458 "pairwise ({pairwise}) is a subset of total ({total})"
2459 );
2460
2461 // A one-node index: the only node the beam can reach is the entry point,
2462 // so a search scores it once per layer it descends through plus once in
2463 // the beam, and none of those evaluations is pairwise.
2464 let mut one = HnswIndex::new(7);
2465 one.insert(0, &vecs[0]);
2466 hnsw_dist_evals_reset();
2467 assert_eq!(one.search(&vecs[1], 5).len(), 1);
2468 assert_eq!(
2469 hnsw_dist_evals(),
2470 1 + one.max_level as u64,
2471 "a search of a one-node graph scores the entry point once per \
2472 descended layer and once in the beam"
2473 );
2474 assert_eq!(
2475 hnsw_dist_evals_pairwise(),
2476 0,
2477 "a search compares the query against nodes, never two nodes"
2478 );
2479 }
2480
2481 /// Scratch reuse (visited bitset, heaps, query f32) must not change the
2482 /// neighbour set, order, or distance-eval count on a fixed seed.
2483 #[test]
2484 fn beam_search_scratch_does_not_change_hits() {
2485 let vecs = make_unit_vecs(80, 16, 0x5C12_A7C4);
2486 let mut idx = HnswIndex::new(crate::index::fnv1a_u64(b"scratch-hits"));
2487 for (i, v) in vecs.iter().enumerate() {
2488 idx.insert(i as u32, v);
2489 }
2490
2491 let q = &vecs[13];
2492 let k = 10;
2493
2494 hnsw_dist_evals_reset();
2495 let baseline = idx.search(q, k);
2496 let baseline_evals = hnsw_dist_evals();
2497 assert_eq!(baseline.len(), k, "fixture must return k hits");
2498
2499 hnsw_dist_evals_reset();
2500 hnsw_beam_scratch_grows_reset();
2501 let again = idx.search(q, k);
2502 let again_evals = hnsw_dist_evals();
2503
2504 assert_eq!(
2505 again, baseline,
2506 "scratch reuse must not change hits or order"
2507 );
2508 assert_eq!(
2509 again_evals, baseline_evals,
2510 "scratch reuse must not change dist-eval count"
2511 );
2512 assert_eq!(
2513 hnsw_beam_scratch_grows(),
2514 0,
2515 "a second search on a warm index must reuse the beam visited buffer"
2516 );
2517
2518 // Comparing the implementation with itself proves it is deterministic,
2519 // not that it is right: a scratch buffer that leaked state between
2520 // queries would corrupt both runs identically and satisfy every
2521 // assertion above. Anchor it to an answer computed without the index at
2522 // all. At this size the beam visits the whole graph, so the approximate
2523 // path owes the exact one its ids in order.
2524 let mut exact: Vec<(u32, f64)> = vecs
2525 .iter()
2526 .enumerate()
2527 .map(|(i, v)| {
2528 let dot: f64 = v.iter().zip(q.iter()).map(|(a, b)| a * b).sum();
2529 (i as u32, dot)
2530 })
2531 .collect();
2532 exact.sort_by(|a, b| {
2533 b.1.partial_cmp(&a.1)
2534 .unwrap_or(std::cmp::Ordering::Equal)
2535 .then_with(|| a.0.cmp(&b.0))
2536 });
2537 let want: Vec<u32> = exact.iter().take(k).map(|(id, _)| *id).collect();
2538 let got: Vec<u32> = baseline.iter().map(|(id, _)| *id).collect();
2539 assert_eq!(
2540 got, want,
2541 "the beam must return the exact top-k for a graph this small; a \
2542 scratch buffer carrying state between queries would show up here \
2543 and nowhere above"
2544 );
2545 }
2546
2547 #[test]
2548 fn hnsw_empty_returns_empty() {
2549 let idx = HnswIndex::new(42);
2550 assert!(idx.search(&[1.0, 0.0], 5).is_empty());
2551 }
2552
2553 #[test]
2554 fn hnsw_zero_vector_skipped() {
2555 let seed = 1;
2556 let mut idx = HnswIndex::new(seed);
2557 idx.insert(0, &[0.0, 0.0]); // zero vector — skipped
2558 idx.insert(1, &[1.0, 0.0]);
2559 // Only node 1 was actually inserted.
2560 assert_eq!(idx.len(), 1);
2561 let r = idx.search(&[1.0, 0.0], 5);
2562 assert_eq!(r.len(), 1);
2563 assert_eq!(r[0].0, 1);
2564 }
2565
2566 #[test]
2567 fn hnsw_remove_works() {
2568 let seed = crate::index::fnv1a_u64(b"rm-test");
2569 let mut idx = HnswIndex::new(seed);
2570 idx.insert(0, &[1.0, 0.0]);
2571 idx.insert(1, &[0.0, 1.0]);
2572 idx.insert(2, &[1.0, 0.0]); // same direction as 0
2573 idx.remove(0);
2574 // Search for [1,0] — 0 is gone, 2 is the nearest remaining.
2575 let r = idx.search(&[1.0, 0.0], 1);
2576 assert!(!r.is_empty());
2577 assert_eq!(r[0].0, 2, "after removing 0, nearest must be 2");
2578 }
2579
2580 #[test]
2581 fn hnsw_cosine_order_preserved() {
2582 let seed = 99;
2583 let mut idx = HnswIndex::new(seed);
2584 // node 0: [1,0] (cos=1.0 with query)
2585 // node 1: [0.6, 0.8] (cos=0.6 with [1,0] query)
2586 // node 2: [0,1] (cos=0.0 with [1,0] query)
2587 idx.insert(0, &[1.0, 0.0]);
2588 idx.insert(1, &[0.6, 0.8]);
2589 idx.insert(2, &[0.0, 1.0]);
2590 let r = idx.search(&[1.0, 0.0], 3);
2591 assert_eq!(r.len(), 3);
2592 // Results must be descending by cosine.
2593 assert!(r[0].1 >= r[1].1);
2594 assert!(r[1].1 >= r[2].1);
2595 assert_eq!(r[0].0, 0, "node 0 must be nearest");
2596 }
2597
2598 // -----------------------------------------------------------------------
2599 // The slab and the kernel
2600 // -----------------------------------------------------------------------
2601
2602 /// The index's own copy of a vector is one `f32` per dimension, in a slab
2603 /// indexed by slot — half the bytes of the `f64` store and one contiguous
2604 /// allocation rather than one per node. And because the slab has a single
2605 /// stride, a vector of some other dimension cannot be stored in it at all:
2606 /// it is skipped, where the `zip`-truncating distance of 0.6.5 would have
2607 /// taken it and produced meaningless distances.
2608 #[test]
2609 fn the_index_holds_one_f32_per_dimension() {
2610 const DIM: usize = 64;
2611 let vecs = make_unit_vecs(300, DIM, 0x51AB_1234);
2612 let mut idx = HnswIndex::new(crate::index::fnv1a_u64(b"slab"));
2613 for (i, v) in vecs.iter().enumerate() {
2614 idx.insert(i as u32, v);
2615 }
2616
2617 let mem = idx.memory_stats();
2618 assert_eq!(mem.live_nodes, 300);
2619 assert_eq!(
2620 mem.vector_floats,
2621 300 * DIM,
2622 "the slab holds one float per dimension per live node"
2623 );
2624 let expected = mem.adjacency_bytes_per_node() + (DIM * 4) as f64;
2625 assert!(
2626 (mem.bytes_per_node() - expected).abs() < 1.0,
2627 "bytes per node is {:.1}, expected {expected:.1} = adjacency + {DIM} f32s",
2628 mem.bytes_per_node()
2629 );
2630
2631 // A vector of another dimension cannot go in: the slab's stride is the
2632 // dimension of the first vector indexed.
2633 let odd = make_unit_vecs(1, DIM - 1, 0x0DD);
2634 hnsw_insert_count_reset();
2635 idx.insert(1_000, &odd[0]);
2636 assert_eq!(idx.len(), 300, "a 63-D vector must not enter a 64-D index");
2637 assert_eq!(
2638 hnsw_insert_count(),
2639 0,
2640 "a skipped vector must not be counted as indexed"
2641 );
2642 assert!(
2643 idx.search(&odd[0], 5).is_empty(),
2644 "a query of the wrong dimension cannot be answered"
2645 );
2646
2647 // ...and one of the right dimension still can.
2648 let more = make_unit_vecs(1, DIM, 0xF00D);
2649 idx.insert(1_001, &more[0]);
2650 assert_eq!(idx.len(), 301, "a 64-D vector is still accepted");
2651 assert_eq!(hnsw_insert_count(), 1);
2652 }
2653
2654 /// The first vector indexed elects the stride, and it may be the odd one
2655 /// out. One 3-element stray ahead of a real corpus must not void the index:
2656 /// the stride is re-elected and the stray evicted, and the index stays the
2657 /// fast path because it then holds everything it was offered at that stride.
2658 #[test]
2659 fn a_stray_first_vector_re_elects_the_stride() {
2660 const DIM: usize = 64;
2661 let vecs = make_unit_vecs(50, DIM, 0x5712_A140);
2662 let mut idx = HnswIndex::new(crate::index::fnv1a_u64(b"re-elect"));
2663
2664 // The stray arrives first and elects 3.
2665 idx.insert(900, &[1.0, 2.0, 3.0]);
2666 assert_eq!(idx.len(), 1);
2667 assert!(idx.can_answer(3), "a 3-D index can answer a 3-D query");
2668 assert!(!idx.can_answer(DIM), "...and not a 64-D one");
2669
2670 for (i, v) in vecs.iter().enumerate() {
2671 idx.insert(i as u32, v);
2672 }
2673
2674 assert_eq!(idx.len(), 50, "every real vector must be indexed");
2675 assert!(
2676 !idx.node_ids().contains(&900),
2677 "the stray must have been evicted, not kept"
2678 );
2679 assert!(
2680 idx.can_answer(DIM),
2681 "an index that re-elected its stride has refused nothing and must \
2682 stay the fast path"
2683 );
2684 assert!(!idx.can_answer(3), "the old stride is gone");
2685 assert_eq!(
2686 idx.search(&vecs[7], 1).first().map(|&(id, _)| id),
2687 Some(7),
2688 "and it must still answer correctly"
2689 );
2690 for &id in idx.node_ids().iter() {
2691 assert_eq!(
2692 idx.back_refs_for_test(id),
2693 idx.scan_back_refs_for_test(id),
2694 "back_refs[{id}] disagrees with a full scan after the eviction"
2695 );
2696 }
2697
2698 // A stride with nothing behind it is not evidence either: drain the
2699 // index and it will take whatever dimension arrives next. Re-ingesting a
2700 // collection under a new embedding model must not need a new rule.
2701 for id in idx.node_ids() {
2702 idx.remove(id);
2703 }
2704 let sevens = make_unit_vecs(3, 7, 0x5E7E_0007);
2705 for (i, v) in sevens.iter().enumerate() {
2706 idx.insert(i as u32, v);
2707 }
2708 assert_eq!(idx.len(), 3, "an emptied index re-elects its stride");
2709 assert!(idx.can_answer(7) && !idx.can_answer(DIM));
2710 }
2711
2712 /// The order that broke the first attempt at re-election: a real vector, a
2713 /// stray, then the rest of the corpus. The stray's re-election displaces the
2714 /// first real vector, and the second real vector's re-election displaces the
2715 /// stray — so the first one has to come back, and until it does the index
2716 /// must not claim it can answer for its dimension.
2717 #[test]
2718 fn a_real_vector_evicted_by_a_stray_comes_back() {
2719 const DIM: usize = 8;
2720 let vecs = make_unit_vecs(6, DIM, 0x0E01_C7ED);
2721 let mut idx = HnswIndex::new(crate::index::fnv1a_u64(b"parked"));
2722
2723 idx.insert(0, &vecs[0]);
2724 assert!(idx.can_answer(DIM));
2725
2726 // The stray re-elects to 3 and parks node 0.
2727 idx.insert(900, &[1.0, 2.0, 3.0]);
2728 assert_eq!(idx.node_ids(), [900].into_iter().collect());
2729 assert!(
2730 !idx.can_answer(DIM),
2731 "while an 8-D vector is parked the index must not answer 8-D queries \
2732 — that is the window in which node 0 is missing"
2733 );
2734
2735 // The next real vector re-elects to 8, parks the stray, and brings node 0
2736 // back.
2737 idx.insert(1, &vecs[1]);
2738 assert_eq!(
2739 idx.node_ids(),
2740 [0, 1].into_iter().collect(),
2741 "the vector the stray displaced must be re-inserted"
2742 );
2743 assert!(
2744 idx.can_answer(DIM),
2745 "with nothing of this dimension parked the index is complete again"
2746 );
2747 assert!(
2748 !idx.can_answer(3),
2749 "the stray's dimension is not the stride"
2750 );
2751
2752 for (i, v) in vecs.iter().enumerate().skip(2) {
2753 idx.insert(i as u32, v);
2754 }
2755 assert_eq!(idx.len(), 6, "every real vector is indexed");
2756 assert_eq!(
2757 idx.search(&vecs[0], 1).first().map(|&(id, _)| id),
2758 Some(0),
2759 "and the re-inserted one is reachable"
2760 );
2761 for &id in idx.node_ids().iter() {
2762 assert_eq!(
2763 idx.back_refs_for_test(id),
2764 idx.scan_back_refs_for_test(id),
2765 "back_refs[{id}] disagrees with a full scan after a re-insertion"
2766 );
2767 }
2768 }
2769
2770 /// A node removed while parked must stay removed: a later re-election of its
2771 /// dimension must not resurrect a vector the caller deleted.
2772 #[test]
2773 fn a_removed_node_does_not_return_from_the_parked_list() {
2774 const DIM: usize = 8;
2775 let vecs = make_unit_vecs(3, DIM, 0xDE1E_7ED0);
2776 let mut idx = HnswIndex::new(crate::index::fnv1a_u64(b"parked-rm"));
2777
2778 idx.insert(0, &vecs[0]);
2779 idx.insert(900, &[1.0, 2.0, 3.0]); // parks node 0
2780 idx.remove(0); // the caller deletes it while it is parked
2781 idx.insert(1, &vecs[1]); // re-elects 8 — node 0 must not come back
2782 assert_eq!(
2783 idx.node_ids(),
2784 [1].into_iter().collect(),
2785 "a removed node must not be revived by a re-election"
2786 );
2787 assert!(idx.can_answer(DIM));
2788
2789 // An explicit insert supersedes a parked copy, so a re-election can never
2790 // revive a vector the node no longer has.
2791 let mut idx = HnswIndex::new(crate::index::fnv1a_u64(b"parked-stale"));
2792 idx.insert(0, &vecs[0]); // stride 8
2793 idx.insert(900, &[1.0, 2.0, 3.0]); // parks node 0's 8-D vector
2794 idx.insert(0, &[7.0, 8.0, 9.0]); // node 0 again, now 3-D
2795 assert_eq!(idx.len(), 2, "nodes 900 and 0, both at the 3-D stride");
2796 idx.remove(900); // back to one vector, so a re-election is possible
2797 idx.insert(2, &vecs[2]); // re-elects 8 and parks node 0's *3-D* vector
2798 assert_eq!(
2799 idx.node_ids(),
2800 [2].into_iter().collect(),
2801 "node 0's stale 8-D copy must not be revived — its vector is 3-D now"
2802 );
2803 assert!(
2804 idx.can_answer(8),
2805 "nothing of 8 dimensions is parked, so the index is complete at its \
2806 stride"
2807 );
2808 }
2809
2810 /// A vector refused *after* the stride has settled leaves the index missing
2811 /// something it was offered, so it must stop claiming it can answer and let
2812 /// the caller scan. This is the property `hnsw_candidates`,
2813 /// `hnsw_search_dst` and `find_similar_vector` all hang their fallback on.
2814 #[test]
2815 fn an_index_that_refused_a_vector_will_not_claim_to_answer() {
2816 const DIM: usize = 64;
2817 let vecs = make_unit_vecs(30, DIM, 0xBADD_14E0);
2818 let mut idx = HnswIndex::new(crate::index::fnv1a_u64(b"refused"));
2819 for (i, v) in vecs.iter().enumerate() {
2820 idx.insert(i as u32, v);
2821 }
2822 assert!(idx.can_answer(DIM), "a clean index answers");
2823
2824 hnsw_insert_count_reset();
2825 idx.insert(900, &[1.0, 2.0, 3.0]);
2826 assert_eq!(idx.len(), 30, "the stray must not be indexed");
2827 assert_eq!(hnsw_insert_count(), 0, "nor counted");
2828 assert!(
2829 !idx.can_answer(DIM),
2830 "an index that refused a vector must send the caller to its scan"
2831 );
2832 assert!(
2833 !idx.search(&vecs[3], 5).is_empty(),
2834 "`search` itself still answers — the fallback is the caller's \
2835 decision, taken on `can_answer`"
2836 );
2837 }
2838
2839 /// The kernel sums in eight accumulators rather than one, which is a
2840 /// different summation order and therefore a different answer. This bounds
2841 /// how different: well inside the granularity at which candidate order can
2842 /// change, at every dimension including the awkward ones either side of the
2843 /// chunk width.
2844 #[test]
2845 fn the_dot_kernel_agrees_with_an_f64_reference() {
2846 for dim in [1usize, 7, 8, 15, 64, 1_536] {
2847 let vecs = make_unit_vecs(6, dim, 0x4047_0000 ^ dim as u64);
2848 for i in 0..vecs.len() {
2849 let a32: Vec<f32> = vecs[i].iter().map(|&x| x as f32).collect();
2850 let self_dot = dot_f32(&a32, &a32);
2851 assert!(
2852 (self_dot as f64 - 1.0).abs() < 1e-5,
2853 "dim {dim}: a unit vector dotted with itself is {self_dot}, not 1.0"
2854 );
2855 for j in 0..vecs.len() {
2856 let b32: Vec<f32> = vecs[j].iter().map(|&x| x as f32).collect();
2857 let reference: f64 =
2858 vecs[i].iter().zip(vecs[j].iter()).map(|(a, b)| a * b).sum();
2859 let got = dot_f32(&a32, &b32) as f64;
2860 assert!(
2861 (got - reference).abs() < 2e-5,
2862 "dim {dim}, pair ({i},{j}): kernel {got} against reference \
2863 {reference}"
2864 );
2865 }
2866 }
2867 }
2868 }
2869
2870 // -----------------------------------------------------------------------
2871 // Parameters
2872 // -----------------------------------------------------------------------
2873
2874 /// The deprecated constants are the default shape's fields. They are kept
2875 /// so an external reference still compiles, and this is what stops them
2876 /// drifting away from the values the index actually uses.
2877 #[test]
2878 #[allow(deprecated)]
2879 fn the_deprecated_constants_still_name_the_default_shape() {
2880 let d = HnswParams::default();
2881 assert_eq!(
2882 (M, M0, EF_CONSTRUCTION, EF_SEARCH),
2883 (d.m, d.m0, d.ef_construction, d.ef_search)
2884 );
2885 }
2886
2887 /// The documented format, and every way of getting it wrong. A bad value
2888 /// falls back to the default rather than building a degenerate index: an
2889 /// `m0` of 0 is a graph with no edges, an `ef_search` of 0 is a beam with
2890 /// no width.
2891 #[test]
2892 fn params_parse_accepts_the_documented_format_and_nothing_else() {
2893 assert_eq!(
2894 HnswParams::parse("16,64,200,400"),
2895 Some(HnswParams::default())
2896 );
2897 assert_eq!(
2898 HnswParams::parse(" 8 , 64 , 300 , 96 "),
2899 Some(HnswParams {
2900 m: 8,
2901 m0: 64,
2902 ef_construction: 300,
2903 ef_search: 96,
2904 prune: Prune::Both
2905 }),
2906 "whitespace around a field must not defeat the override"
2907 );
2908 assert_eq!(
2909 HnswParams::parse("16,64,200,400,both"),
2910 Some(HnswParams::default()),
2911 "the prune field is optional, and `both` is what omitting it means"
2912 );
2913 assert_eq!(
2914 HnswParams::parse("16,64,200,400, own "),
2915 Some(HnswParams {
2916 prune: Prune::Own,
2917 ..HnswParams::default()
2918 }),
2919 "`own` opts out of the neighbour-side heuristic"
2920 );
2921 for bad in [
2922 "",
2923 "16",
2924 "16,64,200",
2925 "16,64,200,400,64",
2926 "16,64,200,400,neither",
2927 "16,64,200,400,own,own",
2928 "16,64,200,x",
2929 "0,64,200,400",
2930 "16,0,200,400",
2931 "16,64,0,400",
2932 "16,64,200,0",
2933 "-16,64,200,400",
2934 ] {
2935 assert_eq!(HnswParams::parse(bad), None, "{bad:?} must not parse");
2936 }
2937 }
2938
2939 /// The prune strategy changes the graph, so it must change what the
2940 /// neighbour-side prune does — and nothing else. Both shapes must respect
2941 /// the degree bound and both must answer.
2942 #[test]
2943 fn both_prune_strategies_build_a_searchable_bounded_graph() {
2944 let params = hnsw_params();
2945 let vecs = make_unit_vecs(600, 24, 0x9121_5EED);
2946 for prune in [Prune::Own, Prune::Both] {
2947 let mut idx = HnswIndex::new(crate::index::fnv1a_u64(b"prune-shapes"));
2948 idx.prune_override_for_test = Some(prune);
2949 for (i, v) in vecs.iter().enumerate() {
2950 idx.insert(i as u32, v);
2951 }
2952 assert_eq!(idx.len(), 600, "{prune:?} lost nodes");
2953 for (s, node) in idx.slots.iter().enumerate() {
2954 if !HnswIndex::is_live(&idx.id_of, s as u32) {
2955 continue;
2956 }
2957 for (lc, layer) in node.layers.iter().enumerate() {
2958 let allowed = if lc == 0 { params.m0 } else { params.m };
2959 assert!(
2960 layer.len() <= allowed,
2961 "{prune:?} slot {s} layer {lc} holds {} links for an allowance of \
2962 {allowed}",
2963 layer.len()
2964 );
2965 }
2966 }
2967 let hits = idx.search(&vecs[42], 5);
2968 assert_eq!(
2969 hits.first().map(|&(id, _)| id),
2970 Some(42),
2971 "{prune:?} search"
2972 );
2973 for &id in idx.node_ids().iter() {
2974 assert_eq!(
2975 idx.back_refs_for_test(id),
2976 idx.scan_back_refs_for_test(id),
2977 "{prune:?} back_refs[{id}] disagrees with a full scan"
2978 );
2979 }
2980 }
2981 }
2982
2983 // -----------------------------------------------------------------------
2984 // The §3.5 diverse-neighbour heuristic
2985 // -----------------------------------------------------------------------
2986
2987 /// Build a throwaway slab holding exactly `vecs`, normalized, one per slot,
2988 /// so a heuristic call can be made against known geometry. The prune reads
2989 /// the slab and the liveness map and nothing else.
2990 fn slab_of(vecs: &[Vec<f64>]) -> (VecSlab, Vec<u32>) {
2991 let mut slab = VecSlab::default();
2992 for (s, v) in vecs.iter().enumerate() {
2993 assert!(slab.put(s as u32, &l2_normalize(v).unwrap()));
2994 }
2995 let id_of = (0..vecs.len() as u32).collect();
2996 (slab, id_of)
2997 }
2998
2999 /// The prune's whole point: given candidates at nearly the same distance,
3000 /// some of which sit behind a neighbour already kept, it keeps the ones that
3001 /// open a new direction. Sized so the diversity test genuinely runs rather
3002 /// than being short-circuited by the tail rule.
3003 #[test]
3004 fn the_prune_prefers_a_new_direction_over_a_redundant_neighbour() {
3005 // Six candidates, `m` = 2, so the tail short-cut cannot fire until four
3006 // rejections have been made: the diversity test runs at least three
3007 // times and its verdicts decide the answer. Candidate 0 opens one
3008 // direction; 1, 2 and 3 sit behind it at increasing distance; 4 opens a
3009 // second direction; 5 is far and behind 0. Keeping the m *nearest* would
3010 // answer [0, 1]; deleting the diversity test would also answer [0, 1].
3011 let base = l2_normalize(&[1.0, 0.0, 0.0]).unwrap();
3012 let vecs = vec![
3013 vec![0.80, 0.0, 0.60], // 0: nearest, direction A
3014 vec![0.78, 0.0, 0.63], // 1: behind 0
3015 vec![0.76, 0.0, 0.65], // 2: behind 0
3016 vec![0.74, 0.0, 0.67], // 3: behind 0
3017 vec![0.70, 0.71, 0.0], // 4: direction B — diverse
3018 vec![0.60, 0.0, 0.80], // 5: far, behind 0
3019 ];
3020 let (slab, id_of) = slab_of(&vecs);
3021 let base = as_f32(&base);
3022 let mut cands: Vec<(u32, f64)> = (0..6u32).map(|s| (s, dist_to(&slab, s, &base))).collect();
3023 sort_by_distance(&mut cands);
3024 assert_eq!(
3025 cands.iter().map(|&(s, _)| s).collect::<Vec<_>>(),
3026 vec![0, 1, 2, 3, 4, 5],
3027 "fixture: candidates must arrive in this nearest-first order"
3028 );
3029
3030 let kept = HnswIndex::select_neighbors_first_rejection(&slab, &id_of, &cands, 2);
3031 assert_eq!(
3032 kept,
3033 vec![0, 4],
3034 "the prune must reject 1, 2 and 3 as reachable through 0 and keep 4, \
3035 which opens a direction 0 does not cover"
3036 );
3037 }
3038
3039 /// The documented deviation from Algorithm 4, pinned so it cannot change by
3040 /// accident. Once `rejections >= candidates.len() - m` the prune takes the
3041 /// remaining candidates **untested** — the farthest ones — where Algorithm 4
3042 /// would keep testing and then backfill with the **nearest** reject. This
3043 /// fixture is a case where those two answers differ, and it asserts ours.
3044 #[test]
3045 fn the_prune_keeps_the_untested_tail_where_algorithm_4_would_backfill() {
3046 // `m` = 2 over three candidates, so one rejection triggers the tail.
3047 // 0 opens a direction; 1 is near but behind 0; 2 is far and behind 0.
3048 let base = l2_normalize(&[1.0, 0.0, 0.0]).unwrap();
3049 let vecs = vec![
3050 vec![0.80, 0.0, 0.60], // 0: nearest
3051 vec![0.78, 0.0, 0.63], // 1: behind 0, near
3052 vec![0.55, 0.0, 0.84], // 2: behind 0, far
3053 ];
3054 let (slab, id_of) = slab_of(&vecs);
3055 let base = as_f32(&base);
3056 let mut cands: Vec<(u32, f64)> = (0..3u32).map(|s| (s, dist_to(&slab, s, &base))).collect();
3057 sort_by_distance(&mut cands);
3058 assert_eq!(
3059 cands.iter().map(|&(s, _)| s).collect::<Vec<_>>(),
3060 vec![0, 1, 2],
3061 "fixture: nearest-first order"
3062 );
3063
3064 let kept = HnswIndex::select_neighbors_first_rejection(&slab, &id_of, &cands, 2);
3065 assert_eq!(
3066 kept,
3067 vec![0, 2],
3068 "the short-cut keeps the untested far candidate; full Algorithm 4 \
3069 would have rejected it too and backfilled with the nearer reject 1, \
3070 answering [0, 1]. If this ever reads [0, 1] the prune has become \
3071 Algorithm 4 and `select_neighbors_first_rejection`'s name, its doc \
3072 comment and docs/site/rules.md are all now wrong."
3073 );
3074 }
3075
3076 /// A dead slot is never linked to. `insert` filters them before scoring,
3077 /// but the heuristic is the last gate before an adjacency list is written.
3078 #[test]
3079 fn the_prune_never_keeps_a_freed_slot() {
3080 let vecs = make_unit_vecs(4, 8, 0x7EA0_1234);
3081 let (slab, mut id_of) = slab_of(&vecs);
3082 id_of[1] = DEAD;
3083 let base = as_f32(&l2_normalize(&vecs[0]).unwrap());
3084 let mut cands: Vec<(u32, f64)> = (0..4u32).map(|s| (s, dist_to(&slab, s, &base))).collect();
3085 sort_by_distance(&mut cands);
3086 let kept = HnswIndex::select_neighbors_first_rejection(&slab, &id_of, &cands, 4);
3087 assert!(
3088 !kept.contains(&1),
3089 "a freed slot must not become a neighbour"
3090 );
3091 }
3092
3093 /// The memory win, guarded on every `cargo test`: degree bounds per layer
3094 /// and payload bytes per node against a stated ceiling.
3095 ///
3096 /// The whole argument for `m0` = 64 is that adjacency halves, and the 5,000 ×
3097 /// 1,536-D measurement that backs it is an `#[ignore]`d release run. This is
3098 /// the cheap version that actually runs: a regression that lets degree drift
3099 /// — a prune that stops pruning, a `set_layer` that leaks — shows up here
3100 /// first.
3101 #[test]
3102 fn degree_and_adjacency_bytes_stay_within_the_shape() {
3103 let params = hnsw_params();
3104 let vecs = make_unit_vecs(1_200, 32, 0xDE6E_E5EE);
3105 let mut idx = HnswIndex::new(crate::index::fnv1a_u64(b"degree-bound"));
3106 for (i, v) in vecs.iter().enumerate() {
3107 idx.insert(i as u32, v);
3108 }
3109
3110 let mut worst_layer0 = 0usize;
3111 let mut worst_upper = 0usize;
3112 for (s, node) in idx.slots.iter().enumerate() {
3113 if !HnswIndex::is_live(&idx.id_of, s as u32) {
3114 continue;
3115 }
3116 for (lc, layer) in node.layers.iter().enumerate() {
3117 let allowed = if lc == 0 { params.m0 } else { params.m };
3118 assert!(
3119 layer.len() <= allowed,
3120 "slot {s} layer {lc} holds {} links for an allowance of {allowed}",
3121 layer.len()
3122 );
3123 if lc == 0 {
3124 worst_layer0 = worst_layer0.max(layer.len());
3125 } else {
3126 worst_upper = worst_upper.max(layer.len());
3127 }
3128 }
3129 }
3130
3131 // Payload per node. Forward entries are bounded by `m0` on layer 0 plus
3132 // `m` on each layer above, and the reverse index holds at most one entry
3133 // per (source, target) pair, so it can never exceed the forward count.
3134 // `2 * (m0 + m) * 4` bytes is therefore a true ceiling with room for the
3135 // upper layers, and it is tight enough to catch `m0` doubling.
3136 let mem = idx.memory_stats();
3137 let ceiling = 2.0 * (params.m0 + params.m) as f64 * 4.0;
3138 eprintln!(
3139 "degree: layer0 <= {worst_layer0} (allowance {}), upper <= {worst_upper} \
3140 (allowance {}); adjacency {:.1} B/node against a ceiling of {ceiling:.1}; \
3141 vector {:.1} B/node",
3142 params.m0,
3143 params.m,
3144 mem.adjacency_bytes_per_node(),
3145 (mem.vector_floats * 4) as f64 / mem.live_nodes as f64,
3146 );
3147 assert_eq!(mem.live_nodes, 1_200);
3148 assert!(
3149 mem.back_ref_entries <= mem.neighbour_slots,
3150 "the reverse index ({}) cannot hold more pairs than the forward one ({})",
3151 mem.back_ref_entries,
3152 mem.neighbour_slots
3153 );
3154 assert!(
3155 mem.adjacency_bytes_per_node() <= ceiling,
3156 "adjacency is {:.1} B/node against a ceiling of {ceiling:.1} — the \
3157 index shape grew",
3158 mem.adjacency_bytes_per_node()
3159 );
3160 // 32 f32s per vector, exactly, or the fixture is not what it says.
3161 assert_eq!(mem.vector_floats, 1_200 * 32);
3162 }
3163
3164 /// Recall on a corpus whose clusters are wider than the layer-0 allowance.
3165 ///
3166 /// This is the case `M₀ = 128` was raised for in v0.4.2 and the case the
3167 /// §3.5 heuristic replaces it for: 40 clusters of 120 members each, against
3168 /// an `m0` of 64. Without diversity in the selection, every member of a
3169 /// cluster spends all 64 of its layer-0 links on other members of the same
3170 /// cluster, the cluster becomes a closed component, and a query that enters
3171 /// the graph elsewhere never reaches it.
3172 ///
3173 /// Queries are drawn from the same clustered distribution, so the top-10
3174 /// are inside one cluster and finding them means the beam must have got
3175 /// into that cluster.
3176 ///
3177 /// Run with:
3178 /// `cargo test --release -p mushroomdb-rules -- clustered_recall_survives_clusters_wider_than_m0 --ignored --nocapture`
3179 #[test]
3180 #[ignore = "slow: builds a 4,800-vector clustered index"]
3181 fn clustered_recall_survives_clusters_wider_than_m0() {
3182 const CLUSTERS: usize = 40;
3183 const PER_CLUSTER: usize = 120;
3184 const DIM: usize = 128;
3185 const K: usize = 10;
3186
3187 assert!(
3188 PER_CLUSTER > hnsw_params().m0,
3189 "the fixture only proves anything when a cluster is wider than m0 \
3190 ({PER_CLUSTER} vs {})",
3191 hnsw_params().m0
3192 );
3193
3194 let vecs = make_clustered_unit_vecs(CLUSTERS, PER_CLUSTER, DIM, 0xC1_05_7E_12_34_56_78_9A);
3195
3196 // Queries: one member of each cluster, nudged. The true top-10 is then
3197 // that member and its nine nearest cluster-mates — a well-defined set,
3198 // unlike a query at the cluster centre, where 120 near-equidistant
3199 // members make "the top 10" a coin toss and recall measures noise.
3200 let queries: Vec<Vec<f64>> = (0..CLUSTERS)
3201 .map(|c| {
3202 let j = c * PER_CLUSTER + 17;
3203 let mut q = vecs[j].clone();
3204 q[0] += 1e-6;
3205 q
3206 })
3207 .collect();
3208 let exact: Vec<BTreeSet<usize>> = queries
3209 .iter()
3210 .map(|q| exact_knn(&vecs, q, K).into_iter().collect())
3211 .collect();
3212
3213 // Both prune strategies, because this fixture is the only one that can
3214 // tell them apart, and the difference is the whole reason `Prune::Both`
3215 // exists as an option. Floors are per strategy: they are what was
3216 // measured, and the gap between them is the trade-off `rules.md`
3217 // documents.
3218 let mut scored: Vec<(Prune, f64, f64)> = Vec::new();
3219 for prune in [Prune::Own, Prune::Both] {
3220 let mut idx = HnswIndex::new(crate::index::fnv1a_u64(b"clustered-recall"));
3221 idx.prune_override_for_test = Some(prune);
3222 for (i, v) in vecs.iter().enumerate() {
3223 idx.insert(i as u32, v);
3224 }
3225 let recalls: Vec<f64> = queries
3226 .iter()
3227 .zip(exact.iter())
3228 .map(|(q, truth)| {
3229 let found = idx
3230 .search(q, K)
3231 .into_iter()
3232 .filter(|(id, _)| truth.contains(&(*id as usize)))
3233 .count();
3234 found as f64 / K as f64
3235 })
3236 .collect();
3237 let min = recalls.iter().cloned().fold(f64::MAX, f64::min);
3238 let mean = recalls.iter().sum::<f64>() / recalls.len() as f64;
3239 eprintln!(
3240 "clustered recall@{K} ({CLUSTERS}x{PER_CLUSTER}, dim {DIM}, m0={}, \
3241 prune={prune:?}): min={min:.4} mean={mean:.4}",
3242 hnsw_params().m0
3243 );
3244 scored.push((prune, min, mean));
3245 }
3246
3247 // Measured: Own min 0.5000 / mean 0.9350, Both min 0.8000 / mean 0.9950.
3248 // A query here sits inside a 120-member cluster whose members are
3249 // genuinely close together, so ranks 8 to 12 are separated by very
3250 // little and the tail of the top-10 is the hardest thing this index is
3251 // ever asked for. The mean says the beam reached the right cluster; the
3252 // min says no single query got stuck in one member's neighbourhood.
3253 for &(prune, min, mean) in &scored {
3254 let (min_floor, mean_floor) = match prune {
3255 Prune::Own => (0.40, 0.90),
3256 Prune::Both => (0.70, 0.95),
3257 };
3258 assert!(
3259 min >= min_floor,
3260 "{prune:?} min clustered recall@{K} = {min:.4} < {min_floor}"
3261 );
3262 assert!(
3263 mean >= mean_floor,
3264 "{prune:?} mean clustered recall@{K} = {mean:.4} < {mean_floor}"
3265 );
3266 }
3267
3268 // The reason the option exists: on a corpus whose clusters are wider
3269 // than `m0`, the full heuristic must actually be better. If this ever
3270 // stops holding, `Prune::Both` is paying 5x the build time for nothing.
3271 let own = scored[0];
3272 let both = scored[1];
3273 assert!(
3274 both.1 >= own.1 && both.2 >= own.2,
3275 "Prune::Both (min {:.4} mean {:.4}) is not better than Prune::Own \
3276 (min {:.4} mean {:.4}) on clusters wider than m0 — the option has no \
3277 justification left",
3278 both.1,
3279 both.2,
3280 own.1,
3281 own.2
3282 );
3283 }
3284
3285 /// Recall after churn. Insert 5,000 1,536-D vectors, remove a fifth,
3286 /// re-insert a tenth of them, remove a disjoint seventh, then measure
3287 /// recall@10 against brute force over exactly the surviving set. A graph
3288 /// that healed badly answers from the wrong neighbourhood; this is the test
3289 /// that would have caught a prune that orphaned nodes on removal.
3290 ///
3291 /// The corpus is the shape `hnsw_5k_1536_recall` uses, and for the same
3292 /// reason: `ef_search` is 400, so a graph of a few hundred reachable nodes
3293 /// is searched exhaustively and its recall is 1.0 whatever the churn did.
3294 /// Only at thousands of survivors does the number mean anything. That makes
3295 /// it a release gate rather than a `cargo test` one.
3296 ///
3297 /// Run with:
3298 /// `cargo test --release -p mushroomdb-rules -- recall_survives_insert_remove_churn --ignored --nocapture`
3299 #[test]
3300 #[ignore = "slow: builds a 5k x 1536-D index and churns it"]
3301 fn recall_survives_insert_remove_churn() {
3302 const N: usize = 5_000;
3303 const DIM: usize = 1_536;
3304 const K: usize = 10;
3305
3306 let vecs = make_unit_vecs(N, DIM, 0xC0FF_EE00_5EED_1234);
3307 let mut idx = HnswIndex::new(crate::index::fnv1a_u64(b"churn-recall"));
3308 for (i, v) in vecs.iter().enumerate() {
3309 idx.insert(i as u32, v);
3310 }
3311
3312 // Churn: drop every 5th, then bring back every 10th of those, then drop
3313 // a disjoint slice. Slots are reused throughout, so a stale adjacency
3314 // reference would now name the wrong vector.
3315 let mut live: BTreeSet<usize> = (0..N).collect();
3316 for i in (0..N).step_by(5) {
3317 idx.remove(i as u32);
3318 live.remove(&i);
3319 }
3320 for i in (0..N).step_by(10) {
3321 idx.insert(i as u32, &vecs[i]);
3322 live.insert(i);
3323 }
3324 for i in (3..N).step_by(7) {
3325 idx.remove(i as u32);
3326 live.remove(&i);
3327 }
3328 assert_eq!(
3329 idx.len(),
3330 live.len(),
3331 "the index and the oracle disagree on size"
3332 );
3333
3334 // back_refs must still be exact — recall means nothing on a corrupt
3335 // graph. Spot-checked rather than verified for every node: the full
3336 // check is O(n² · m₀), which at n = 5,000 costs more than the whole
3337 // rest of this test, and `back_refs_match_a_full_scan` already runs it
3338 // exhaustively over a churned 500-node index on every `cargo test`.
3339 for &id in idx.node_ids().iter().step_by(53) {
3340 assert_eq!(
3341 idx.back_refs_for_test(id),
3342 idx.scan_back_refs_for_test(id),
3343 "back_refs[{id}] disagrees with a full scan after churn"
3344 );
3345 }
3346
3347 // Ground truth over the survivors only.
3348 let survivors: Vec<usize> = live.iter().copied().collect();
3349 let queries = make_unit_vecs(40, DIM, 0x9111_0BED);
3350 let mut recalls = Vec::with_capacity(queries.len());
3351 for q in &queries {
3352 let mut scored: Vec<(usize, f64)> = survivors
3353 .iter()
3354 .map(|&i| {
3355 let dot: f64 = vecs[i].iter().zip(q.iter()).map(|(a, b)| a * b).sum();
3356 (i, dot)
3357 })
3358 .collect();
3359 scored.sort_by(|a, b| {
3360 b.1.partial_cmp(&a.1)
3361 .unwrap_or(std::cmp::Ordering::Equal)
3362 .then_with(|| a.0.cmp(&b.0))
3363 });
3364 let exact: BTreeSet<usize> = scored.into_iter().take(K).map(|(i, _)| i).collect();
3365
3366 let hits = idx.search(q, K);
3367 for (id, _) in &hits {
3368 assert!(
3369 live.contains(&(*id as usize)),
3370 "search returned {id}, which was removed"
3371 );
3372 }
3373 let found = hits
3374 .iter()
3375 .filter(|(id, _)| exact.contains(&(*id as usize)))
3376 .count();
3377 recalls.push(found as f64 / K as f64);
3378 }
3379 let min = recalls.iter().cloned().fold(f64::MAX, f64::min);
3380 let mean = recalls.iter().sum::<f64>() / recalls.len() as f64;
3381 eprintln!(
3382 "churn recall@{K}: min={min:.4} mean={mean:.4} over {} survivors",
3383 live.len()
3384 );
3385 // The same floor `hnsw_5k_1536_recall` holds the un-churned graph to.
3386 // Everything here is seeded, so these are fixed numbers, not a sample.
3387 assert!(min >= 0.90, "min recall@{K} after churn = {min:.4} < 0.90");
3388 assert!(
3389 mean >= 0.95,
3390 "mean recall@{K} after churn = {mean:.4} < 0.95"
3391 );
3392 }
3393
3394 // -----------------------------------------------------------------------
3395 // Reverse adjacency
3396 // -----------------------------------------------------------------------
3397
3398 /// `back_refs` must agree with a full scan of every adjacency list, through
3399 /// inserts, removes and re-inserts. This is the invariant that makes it
3400 /// safe to reuse a slot: if a stale reference survived a removal it would
3401 /// silently name whichever node took the slot over.
3402 #[test]
3403 fn back_refs_match_a_full_scan() {
3404 let vecs = make_unit_vecs(500, 32, 0x5EED_1234);
3405 let mut idx = HnswIndex::new(crate::index::fnv1a_u64(b"backrefs"));
3406 for (i, v) in vecs.iter().enumerate() {
3407 idx.insert(i as u32, v);
3408 }
3409 for i in (0..500).step_by(5) {
3410 idx.remove(i as u32);
3411 }
3412 for i in (0..500).step_by(5) {
3413 idx.insert(i as u32, &vecs[i]);
3414 }
3415 assert_eq!(idx.len(), 500, "every node is back in the index");
3416 for &id in idx.node_ids().iter() {
3417 assert_eq!(
3418 idx.back_refs_for_test(id),
3419 idx.scan_back_refs_for_test(id),
3420 "back_refs[{id}] disagrees with a full scan"
3421 );
3422 }
3423 // Freed slots were reused rather than appended to.
3424 assert_eq!(
3425 idx.slot_capacity_for_test(),
3426 500,
3427 "a remove + re-insert must recycle the slot"
3428 );
3429 }
3430
3431 /// An insert that replaces an existing id is a remove plus an insert, and
3432 /// must leave the graph exactly as if the id had never been there.
3433 #[test]
3434 fn insert_then_remove_equals_never_inserted() {
3435 let vecs = make_unit_vecs(200, 16, 0xABCD_0001);
3436 let seed = crate::index::fnv1a_u64(b"rm-equiv");
3437
3438 // Reference: ids 0..199 except 42, 77 and 150.
3439 let skipped = [42usize, 77, 150];
3440 let mut reference = HnswIndex::new(seed);
3441 for (i, v) in vecs.iter().enumerate() {
3442 if skipped.contains(&i) {
3443 continue;
3444 }
3445 reference.insert(i as u32, v);
3446 }
3447
3448 // Subject: every id, then the three removed.
3449 let mut subject = HnswIndex::new(seed);
3450 for (i, v) in vecs.iter().enumerate() {
3451 subject.insert(i as u32, v);
3452 }
3453 for &i in &skipped {
3454 subject.remove(i as u32);
3455 }
3456
3457 assert_eq!(
3458 subject.node_ids(),
3459 reference.node_ids(),
3460 "the removed ids must be gone from the index"
3461 );
3462 for &id in subject.node_ids().iter() {
3463 assert_eq!(
3464 subject.back_refs_for_test(id),
3465 subject.scan_back_refs_for_test(id),
3466 "back_refs[{id}] disagrees with a full scan after the removals"
3467 );
3468 }
3469 // A removed id is unreachable: query at its own vector and it must not
3470 // come back, however many neighbours are asked for.
3471 for &i in &skipped {
3472 let hits = subject.search(&vecs[i], 20);
3473 assert!(
3474 !hits.iter().any(|&(id, _)| id == i as u32),
3475 "removed id {i} is still reachable through the graph"
3476 );
3477 }
3478 }
3479
3480 /// Removing the entry point — the node that owns the top layer — must
3481 /// re-elect a live one and leave the graph searchable. A dangling entry
3482 /// point would either panic in `dist_to` or hand back a freed slot's id.
3483 #[test]
3484 fn removing_the_entry_point_re_elects_and_still_searches() {
3485 let vecs = make_unit_vecs(300, 16, 0xE47E_9001);
3486 let mut idx = HnswIndex::new(crate::index::fnv1a_u64(b"ep-churn"));
3487 for (i, v) in vecs.iter().enumerate() {
3488 idx.insert(i as u32, v);
3489 }
3490
3491 // Peel off the entry point repeatedly: each removal must re-elect the
3492 // highest-level node still present, and the top layer eventually
3493 // collapses onto a lower one.
3494 let mut seen_levels = Vec::new();
3495 for _ in 0..12 {
3496 let ep_slot = idx
3497 .entry_point
3498 .expect("a non-empty index has an entry point");
3499 let ep_id = idx.id_of[ep_slot as usize];
3500 assert_ne!(ep_id, DEAD, "the entry point must name a live node");
3501 seen_levels.push(idx.max_level);
3502
3503 idx.remove(ep_id);
3504
3505 let new_ep = idx.entry_point.expect("re-election must find a live node");
3506 assert!(
3507 HnswIndex::is_live(&idx.id_of, new_ep),
3508 "the re-elected entry point is a freed slot"
3509 );
3510 assert_ne!(new_ep, ep_slot, "the removed slot is still the entry point");
3511 assert_eq!(
3512 idx.max_level, idx.slots[new_ep as usize].level,
3513 "max_level must follow the re-elected entry point"
3514 );
3515 let highest = idx
3516 .slot_of
3517 .values()
3518 .map(|&s| idx.slots[s as usize].level)
3519 .max()
3520 .unwrap();
3521 assert_eq!(
3522 idx.max_level, highest,
3523 "the entry point must be a highest-level node"
3524 );
3525 assert!(
3526 !idx.node_ids().contains(&ep_id),
3527 "the old entry point lingers"
3528 );
3529
3530 // The graph still answers, and never with a removed id.
3531 let hits = idx.search(&vecs[200], 5);
3532 assert!(
3533 !hits.is_empty(),
3534 "the graph stopped answering after re-election"
3535 );
3536 for (id, _) in &hits {
3537 assert!(
3538 idx.node_ids().contains(id),
3539 "search returned {id}, which is not in the index"
3540 );
3541 }
3542 for &id in idx.node_ids().iter() {
3543 assert_eq!(
3544 idx.back_refs_for_test(id),
3545 idx.scan_back_refs_for_test(id),
3546 "back_refs[{id}] disagrees with a full scan after an entry-point removal"
3547 );
3548 }
3549 }
3550 assert!(
3551 seen_levels.iter().any(|&l| l > 0),
3552 "the fixture never had a multi-layer entry point, so this proved nothing"
3553 );
3554
3555 // Drain it entirely: the last removal must clear the entry point.
3556 for id in idx.node_ids() {
3557 idx.remove(id);
3558 }
3559 assert!(idx.is_empty());
3560 assert_eq!(idx.entry_point, None);
3561 assert_eq!(idx.max_level, 0);
3562 assert!(idx.search(&vecs[0], 5).is_empty());
3563 }
3564
3565 /// A removal must not touch every node in the index. Asserted on the
3566 /// operation count, not the wall clock: `remove` visits exactly the nodes
3567 /// that list the victim, and nothing else.
3568 #[test]
3569 fn remove_touches_only_the_nodes_that_list_it() {
3570 let vecs = make_unit_vecs(1_500, 32, 0xD00D_0007);
3571 let mut idx = HnswIndex::new(crate::index::fnv1a_u64(b"rm-cost"));
3572 for (i, v) in vecs.iter().enumerate() {
3573 idx.insert(i as u32, v);
3574 }
3575
3576 let mut worst = 0u64;
3577 for i in (0..1_500).step_by(100) {
3578 let in_degree = idx.back_refs_for_test(i as u32).len() as u64;
3579 hnsw_remove_scanned_reset();
3580 idx.remove(i as u32);
3581 let scanned = hnsw_remove_scanned();
3582 assert_eq!(
3583 scanned, in_degree,
3584 "removing {i} visited {scanned} nodes for an in-degree of {in_degree}"
3585 );
3586 worst = worst.max(scanned);
3587 }
3588 let live = idx.len() as u64;
3589 assert!(
3590 worst * 2 < live,
3591 "worst removal touched {worst} of {live} live nodes — the O(N·M₀) scan \
3592 is still there"
3593 );
3594 }
3595
3596 /// The wall-clock form of the same claim, for the record. Ignored by
3597 /// default because it builds a 5k index.
3598 ///
3599 /// Run with:
3600 /// `cargo test --release -p mushroomdb-rules -- remove_is_not_a_full_scan --ignored --nocapture`
3601 #[test]
3602 #[ignore = "slow: builds a 5k index"]
3603 fn remove_is_not_a_full_scan() {
3604 let vecs = make_unit_vecs(5_000, 64, 0xD00D);
3605 let mut idx = HnswIndex::new(crate::index::fnv1a_u64(b"rm-cost"));
3606 for (i, v) in vecs.iter().enumerate() {
3607 idx.insert(i as u32, v);
3608 }
3609 let t = std::time::Instant::now();
3610 for i in (0..5_000).step_by(100) {
3611 idx.remove(i as u32);
3612 }
3613 let mean = t.elapsed() / 50;
3614 eprintln!("mean removal at n=5000: {mean:?}");
3615 assert!(
3616 mean < std::time::Duration::from_millis(5),
3617 "mean removal {mean:?} at n=5000 — the O(N·M₀) scan is still there"
3618 );
3619 }
3620
3621 // -----------------------------------------------------------------------
3622 // Blob magic and version
3623 // -----------------------------------------------------------------------
3624
3625 /// The 0.6.5 on-disk shape, serialize side, so a test can write one.
3626 #[derive(Serialize)]
3627 struct V1Node {
3628 level: usize,
3629 vector: Vec<f64>,
3630 layers: Vec<Vec<u32>>,
3631 }
3632
3633 #[derive(Serialize)]
3634 struct V1Index {
3635 base_seed: u64,
3636 nodes: BTreeMap<u32, V1Node>,
3637 entry_point: Option<u32>,
3638 max_level: usize,
3639 }
3640
3641 /// The blob-v2 on-disk shape, serialize side, so a test can write one.
3642 #[derive(Serialize)]
3643 struct V2Node {
3644 level: usize,
3645 vector: Vec<f64>,
3646 layers: Vec<Vec<u32>>,
3647 }
3648
3649 #[derive(Serialize)]
3650 struct V2Index {
3651 base_seed: u64,
3652 slots: Vec<V2Node>,
3653 slot_of: BTreeMap<u32, u32>,
3654 id_of: Vec<u32>,
3655 free: Vec<u32>,
3656 entry_point: Option<u32>,
3657 max_level: usize,
3658 }
3659
3660 #[derive(Serialize)]
3661 struct V2Blob {
3662 magic: [u8; 4],
3663 version: u16,
3664 index: V2Index,
3665 }
3666
3667 /// A slot's vector back in `f64`, as the older shapes stored it. An `f32`
3668 /// widened to `f64` and narrowed again is bit-exact, so a round trip through
3669 /// either older blob loses nothing — which is what lets the upgrade tests
3670 /// compare `search` results exactly.
3671 fn vector_of(idx: &HnswIndex, slot: u32) -> Vec<f64> {
3672 idx.slab.get(slot).iter().map(|&x| x as f64).collect()
3673 }
3674
3675 /// Re-express a live index in the id-keyed 0.6.5 shape.
3676 fn as_v1_blob(idx: &HnswIndex) -> Vec<u8> {
3677 as_v1_blob_with(idx, &[])
3678 }
3679
3680 /// [`as_v1_blob`], but `by_id[id]` overrides the vector written for `id`
3681 /// when it is present — the only way to build a mixed-dimension blob, since
3682 /// this build refuses one on the way in and an older build did not.
3683 fn as_v1_blob_with(idx: &HnswIndex, by_id: &[Vec<f64>]) -> Vec<u8> {
3684 let nodes: BTreeMap<u32, V1Node> = idx
3685 .slot_of
3686 .iter()
3687 .map(|(&id, &s)| {
3688 let n = &idx.slots[s as usize];
3689 (
3690 id,
3691 V1Node {
3692 level: n.level,
3693 vector: by_id
3694 .get(id as usize)
3695 .cloned()
3696 .unwrap_or_else(|| vector_of(idx, s)),
3697 layers: n
3698 .layers
3699 .iter()
3700 .map(|l| l.iter().map(|&t| idx.id_of[t as usize]).collect())
3701 .collect(),
3702 },
3703 )
3704 })
3705 .collect();
3706 bincode::serialize(&V1Index {
3707 base_seed: idx.base_seed,
3708 nodes,
3709 entry_point: idx.entry_point.map(|s| idx.id_of[s as usize]),
3710 max_level: idx.max_level,
3711 })
3712 .unwrap()
3713 }
3714
3715 /// Re-express a live index in the slot-keyed blob-v2 shape — 0.6.6 before the
3716 /// distance kernel, with an `f64` vector inside every node.
3717 fn as_v2_blob(idx: &HnswIndex) -> Vec<u8> {
3718 as_v2_blob_with(idx, &[])
3719 }
3720
3721 /// [`as_v2_blob`], with the same `by_id` override as [`as_v1_blob_with`].
3722 fn as_v2_blob_with(idx: &HnswIndex, by_id: &[Vec<f64>]) -> Vec<u8> {
3723 let slots: Vec<V2Node> = idx
3724 .slots
3725 .iter()
3726 .enumerate()
3727 .map(|(s, n)| V2Node {
3728 level: n.level,
3729 vector: by_id
3730 .get(*idx.id_of.get(s).unwrap_or(&DEAD) as usize)
3731 .cloned()
3732 .unwrap_or_else(|| vector_of(idx, s as u32)),
3733 layers: n.layers.clone(),
3734 })
3735 .collect();
3736 bincode::serialize(&V2Blob {
3737 magic: HNSW_BLOB_MAGIC,
3738 version: 2,
3739 index: V2Index {
3740 base_seed: idx.base_seed,
3741 slots,
3742 slot_of: idx.slot_of.clone(),
3743 id_of: idx.id_of.clone(),
3744 free: idx.free.clone(),
3745 entry_point: idx.entry_point,
3746 max_level: idx.max_level,
3747 },
3748 })
3749 .unwrap()
3750 }
3751
3752 /// The v3 wrapper as 0.6.6b-0.6.8 wrote it: this build's index body with
3753 /// `complete` straight after it and no side state.
3754 fn as_v3_blob(idx: &HnswIndex, complete: bool) -> Vec<u8> {
3755 #[derive(Serialize)]
3756 struct V3Ref<'a> {
3757 magic: [u8; 4],
3758 version: u16,
3759 index: &'a HnswIndex,
3760 complete: bool,
3761 }
3762 bincode::serialize(&V3Ref {
3763 magic: HNSW_BLOB_MAGIC,
3764 version: 3,
3765 index: idx,
3766 complete,
3767 })
3768 .expect("v3 encode")
3769 }
3770
3771 fn blob_fixture() -> (Vec<Vec<f64>>, HnswIndex) {
3772 let vecs = make_unit_vecs(120, 24, 0x0B10_B0B0);
3773 let mut idx = HnswIndex::new(crate::index::fnv1a_u64(b"blob-rt"));
3774 for (i, v) in vecs.iter().enumerate() {
3775 idx.insert(i as u32, v);
3776 }
3777 (vecs, idx)
3778 }
3779
3780 fn assert_matches(loaded: &HnswIndex, original: &HnswIndex, q: &[f64]) {
3781 assert_eq!(loaded.node_ids(), original.node_ids(), "node ids differ");
3782 for &id in loaded.node_ids().iter() {
3783 assert_eq!(
3784 loaded.back_refs_for_test(id),
3785 loaded.scan_back_refs_for_test(id),
3786 "back_refs[{id}] disagrees with a full scan after the load"
3787 );
3788 }
3789 assert_eq!(
3790 loaded.search(q, 10),
3791 original.search(q, 10),
3792 "the loaded graph answers differently"
3793 );
3794 }
3795
3796 /// A blob written by 0.6.5 — a bare bincoded `HnswIndex` with id-keyed
3797 /// adjacency — is up-converted in place, with no vector re-inserted.
3798 #[test]
3799 fn a_v1_blob_upgrades_in_place() {
3800 let (vecs, idx) = blob_fixture();
3801 let blob = as_v1_blob(&idx);
3802
3803 hnsw_insert_count_reset();
3804 let loaded = decode_hnsw_blob(&blob).expect("a 0.6.5 blob must still load");
3805 assert_eq!(
3806 hnsw_insert_count(),
3807 0,
3808 "up-converting a v1 blob must not re-insert a single vector"
3809 );
3810 assert_matches(&loaded, &idx, &vecs[3]);
3811 }
3812
3813 /// A mixed-dimension v1 or v2 blob upgrades into an index that **declines**
3814 /// the fast path.
3815 ///
3816 /// 0.6.5 accepted vectors of unequal length, so this is exactly where a
3817 /// mixed-dimension graph can come from. `slab_of_decoded` pads the odd ones
3818 /// to the elected stride, which puts a fabricated position in the graph;
3819 /// counting each as a refusal is what makes the release's promise — "an
3820 /// index that skipped a vector stops claiming the fast path" — true on the
3821 /// upgrade path too. The odd node stays in the graph, because dropping it
3822 /// would leave adjacency naming a slot that is not there.
3823 #[test]
3824 fn a_mixed_dimension_upgrade_declines_the_fast_path() {
3825 for shape in ["v1", "v2"] {
3826 let mut vecs = make_unit_vecs(40, 24, 0x0D1D_0DDD);
3827 let mut idx = HnswIndex::new(crate::index::fnv1a_u64(b"mixed"));
3828 for (i, v) in vecs.iter().enumerate() {
3829 idx.insert(i as u32, v);
3830 }
3831 // Shorten one vector *in the encoded bytes*, which is the only way
3832 // a mixed-dimension graph can exist: this build refuses one on the
3833 // way in, and an older build did not.
3834 vecs[7].truncate(23);
3835 let blob = match shape {
3836 "v1" => as_v1_blob_with(&idx, &vecs),
3837 _ => as_v2_blob_with(&idx, &vecs),
3838 };
3839
3840 let loaded = decode_hnsw_blob(&blob).expect("a mixed blob still loads");
3841 assert_eq!(
3842 loaded.node_ids(),
3843 idx.node_ids(),
3844 "{shape}: the odd node must stay in the graph"
3845 );
3846 assert!(
3847 !loaded.can_answer(24),
3848 "{shape}: an index holding a padded position must not claim the fast path"
3849 );
3850 // And it is still a usable graph — the caller's scan is what answers,
3851 // but nothing here may panic.
3852 let _ = loaded.search(&vecs[3], 5);
3853 }
3854 }
3855
3856 /// A truncated v3 slab is refused, not decoded into a graph whose distance
3857 /// kernel reads past the rows it has.
3858 #[test]
3859 fn a_truncated_v3_slab_is_refused() {
3860 let (_vecs, mut idx) = blob_fixture();
3861
3862 // A *well-formed* blob whose slab is short of what its slots claim.
3863 // bincode reads the length prefix it is given, so this decodes happily
3864 // and `VecSlab::get` then hands `dot_f32` an empty slice against a
3865 // 24-element query — a `debug_assert` in debug, a garbage distance in
3866 // release. The header check has to catch it before that.
3867 let full = idx.slab.data.len();
3868 idx.slab.data.truncate(full - idx.slab.dim);
3869 let blob = encode_hnsw_blob(&idx, true).expect("encode");
3870
3871 let err = decode_hnsw_blob(&blob).expect_err("a truncated slab must be refused");
3872 assert!(
3873 err.contains("truncated"),
3874 "the error must name the problem; got {err:?}"
3875 );
3876
3877 // An inconsistent slot/id pairing is refused by the same gate.
3878 let (_v, mut bad) = blob_fixture();
3879 bad.id_of.pop();
3880 let err = decode_hnsw_blob(&encode_hnsw_blob(&bad, true).expect("encode"))
3881 .expect_err("a slot/id mismatch must be refused");
3882 assert!(
3883 err.contains("inconsistent"),
3884 "the error must name the problem; got {err:?}"
3885 );
3886 }
3887
3888 /// A blob written mid-build says so, and the graph it decodes to refuses to
3889 /// answer until something finishes the build.
3890 #[test]
3891 fn an_incomplete_blob_refuses_to_answer() {
3892 let (vecs, idx) = blob_fixture();
3893
3894 let whole = decode_hnsw_blob(&encode_hnsw_blob(&idx, true).expect("encode"))
3895 .expect("a complete blob loads");
3896 assert!(whole.can_answer(24), "a complete blob must answer");
3897 assert!(!whole.is_incomplete());
3898
3899 let partial = decode_hnsw_blob(&encode_hnsw_blob(&idx, false).expect("encode"))
3900 .expect("an incomplete blob still loads");
3901 assert!(
3902 partial.is_incomplete(),
3903 "the flag must survive the round trip"
3904 );
3905 assert!(
3906 !partial.can_answer(24),
3907 "a graph holding a prefix of its corpus must not claim the fast path"
3908 );
3909 // It is a real graph, not a broken one: the flag is about completeness,
3910 // not about validity, and `mark_complete` is what the adopting side
3911 // calls once the open-time scan has filled it in.
3912 assert_eq!(partial.node_ids(), idx.node_ids());
3913 let mut adopted = partial;
3914 adopted.mark_complete();
3915 assert!(adopted.can_answer(24));
3916 assert_eq!(adopted.search(&vecs[3], 5), whole.search(&vecs[3], 5));
3917 }
3918
3919 /// v3 writes `complete` as the last byte, so open can peek it without
3920 /// decoding the graph.
3921 #[test]
3922 fn hnsw_blob_complete_peeks_the_last_byte() {
3923 let (_vecs, idx) = blob_fixture();
3924 let whole = encode_hnsw_blob(&idx, true).expect("encode");
3925 let partial = encode_hnsw_blob(&idx, false).expect("encode");
3926 assert_eq!(hnsw_blob_complete(&whole), Some(true));
3927 assert_eq!(hnsw_blob_complete(&partial), Some(false));
3928 assert_eq!(&whole[..whole.len() - 1], &partial[..partial.len() - 1]);
3929 assert_eq!(whole[whole.len() - 1], 1);
3930 assert_eq!(partial[partial.len() - 1], 0);
3931 assert_eq!(hnsw_blob_complete(&[]), None);
3932 }
3933
3934 /// A refused insert still replaces the id it was offered for.
3935 ///
3936 /// `insert`'s contract is that an existing id is replaced; the refusal arm
3937 /// returns before the `remove` that implements it, so without the paired
3938 /// remove the index keeps the *old* vector under an id whose new embedding
3939 /// it just rejected — a stale answer under a live key. The engine removes
3940 /// first today, so this is a guard on the method, not a live defect.
3941 #[test]
3942 fn a_refused_insert_does_not_leave_a_stale_vector() {
3943 let mut idx = HnswIndex::new(crate::index::fnv1a_u64(b"stale"));
3944 let vecs = make_unit_vecs(4, 8, 0x57A1_E000);
3945 for (i, v) in vecs.iter().enumerate() {
3946 idx.insert(i as u32, v);
3947 }
3948 assert!(idx.node_ids().contains(&0));
3949
3950 // Re-embed node 0 at the wrong dimension: the index cannot hold it.
3951 idx.insert(0, &make_unit_vecs(1, 4, 0x57A1_E001)[0]);
3952 assert!(
3953 !idx.node_ids().contains(&0),
3954 "the refused id must not keep its old vector"
3955 );
3956 assert!(
3957 !idx.can_answer(8),
3958 "a refusal means the index is missing a vector it was offered"
3959 );
3960 }
3961
3962 /// A set-but-malformed `MUSHROOMDB_HNSW_PARAMS` falls back to the defaults.
3963 ///
3964 /// The warning itself is an `eprintln!` and not asserted here; what is
3965 /// asserted is that no malformed spelling silently becomes a *different*
3966 /// shape, which is the property the graph depends on.
3967 #[test]
3968 fn a_malformed_params_string_yields_the_defaults() {
3969 for bad in [
3970 "",
3971 "16,64,200",
3972 "16,64,200,400,neither",
3973 "a,b,c,d",
3974 "16,0,200,400",
3975 ] {
3976 assert_eq!(
3977 HnswParams::parse(bad),
3978 None,
3979 "{bad:?} must not parse to a shape"
3980 );
3981 }
3982 }
3983
3984 /// This build's wrapper round-trips, and it is version 4.
3985 #[test]
3986 fn a_v4_blob_round_trips() {
3987 let (vecs, idx) = blob_fixture();
3988 let blob = encode_hnsw_blob(&idx, true).expect("encode");
3989 assert_eq!(
3990 &blob[..4],
3991 &HNSW_BLOB_MAGIC,
3992 "the blob must carry its magic"
3993 );
3994 assert_eq!(
3995 u16::from_le_bytes([blob[4], blob[5]]),
3996 4,
3997 "this build writes blob version 4"
3998 );
3999
4000 hnsw_insert_count_reset();
4001 let loaded = decode_hnsw_blob(&blob).expect("a v4 blob must load");
4002 assert_eq!(hnsw_insert_count(), 0, "a load must not re-insert vectors");
4003 assert_matches(&loaded, &idx, &vecs[3]);
4004 assert_eq!(
4005 loaded.slab.dim, idx.slab.dim,
4006 "the slab's stride must survive the round trip"
4007 );
4008 }
4009
4010 /// A v3 blob — this build's own index body with no side state after it —
4011 /// still loads, and comes back with nothing restored, so the open-time scan
4012 /// re-offers every vector and re-derives the refusals exactly as before v4.
4013 #[test]
4014 fn v3_blob_still_loads_and_reoffers() {
4015 let (vecs, idx) = blob_fixture();
4016 let blob = as_v3_blob(&idx, true);
4017 assert_eq!(
4018 u16::from_le_bytes([blob[4], blob[5]]),
4019 3,
4020 "the fixture must actually be a v3 blob"
4021 );
4022
4023 hnsw_insert_count_reset();
4024 let loaded = decode_hnsw_blob(&blob).expect("a v3 blob must still load");
4025 assert_eq!(hnsw_insert_count(), 0, "a load must not re-insert vectors");
4026 assert_matches(&loaded, &idx, &vecs[3]);
4027 assert_eq!(
4028 loaded.accounted_ids(),
4029 loaded.node_ids(),
4030 "a v3 blob carries no parked or refused state, so the scan's skip \
4031 set is exactly the graph — the pre-v4 behaviour"
4032 );
4033 assert_eq!(loaded.dim_mismatches(), 0);
4034 assert!(loaded.refused_ids().is_empty());
4035 }
4036
4037 /// v4 carries the refusal count, the refused ids and the parked vectors, so
4038 /// a reopen restores them instead of re-deriving them — and the ids it
4039 /// restored are the ids the open-time scan must skip.
4040 #[test]
4041 fn a_v4_blob_restores_parked_and_refused() {
4042 // Stride 2 elected by the corpus, then one 3-D stray refused, and a
4043 // re-election that parks a vector at the old stride.
4044 let mut idx = HnswIndex::new(crate::index::fnv1a_u64(b"v4-state"));
4045 idx.insert(1, &[1.0, 0.0]);
4046 idx.insert(2, &[0.0, 1.0]);
4047 idx.insert(3, &[1.0, 1.0, 1.0]); // refused: stride is 2, len > 1
4048 assert_eq!(idx.dim_mismatches(), 1, "the stray must be refused");
4049 assert_eq!(idx.refused_ids(), BTreeSet::from([3]));
4050
4051 let blob = encode_hnsw_blob(&idx, true).expect("encode");
4052 hnsw_insert_count_reset();
4053 let loaded = decode_hnsw_blob(&blob).expect("decode");
4054 assert_eq!(hnsw_insert_count(), 0, "a load must not re-insert vectors");
4055
4056 assert_eq!(
4057 loaded.dim_mismatches(),
4058 idx.dim_mismatches(),
4059 "the refusal count survives the round trip rather than starting at zero"
4060 );
4061 assert_eq!(
4062 loaded.refused_ids(),
4063 idx.refused_ids(),
4064 "and so do the ids behind it"
4065 );
4066 assert!(
4067 loaded.accounts_for(3),
4068 "the refused id is accounted for, so the open-time scan skips it \
4069 instead of refusing it a second time"
4070 );
4071 assert!(
4072 !loaded.node_ids().contains(&3),
4073 "accounted for is not the same as in the graph"
4074 );
4075 }
4076
4077 /// A parked vector comes back parked, and is not re-offered. Re-offering it
4078 /// is what would destroy it: `insert` supersedes the parked copy and then
4079 /// refuses the vector against the elected stride, so the one copy that made
4080 /// it recoverable is gone.
4081 #[test]
4082 fn a_v4_blob_restores_a_parked_vector() {
4083 // `[stray, real, real]`: the first vector elects its own dimension and
4084 // is parked by the re-election the second one triggers.
4085 let mut idx = HnswIndex::new(crate::index::fnv1a_u64(b"v4-parked"));
4086 idx.insert(7, &[1.0, 1.0, 1.0]);
4087 idx.insert(1, &[1.0, 0.0]);
4088 idx.insert(2, &[0.0, 1.0]);
4089 assert!(
4090 idx.accounts_for(7) && !idx.node_ids().contains(&7),
4091 "node 7 must be parked, not indexed"
4092 );
4093
4094 let blob = encode_hnsw_blob(&idx, true).expect("encode");
4095 let loaded = decode_hnsw_blob(&blob).expect("decode");
4096 assert!(
4097 loaded.accounts_for(7),
4098 "the parked entry survives the round trip"
4099 );
4100 assert!(
4101 !loaded.node_ids().contains(&7),
4102 "and is still parked rather than indexed"
4103 );
4104 assert!(
4105 loaded.accounted_ids().contains(&7),
4106 "so the open-time scan skips it"
4107 );
4108 }
4109
4110 /// A blob written by 0.6.6 before the distance kernel — slot-keyed, with an
4111 /// `f64` vector per node — loads into the slab shape with no vector
4112 /// re-inserted and answers identically.
4113 #[test]
4114 fn a_v2_blob_upgrades_in_place() {
4115 let (vecs, idx) = blob_fixture();
4116 let blob = as_v2_blob(&idx);
4117
4118 hnsw_insert_count_reset();
4119 let loaded = decode_hnsw_blob(&blob).expect("a v2 blob must still load");
4120 assert_eq!(
4121 hnsw_insert_count(),
4122 0,
4123 "up-converting a v2 blob must not re-insert a single vector"
4124 );
4125 assert_eq!(
4126 loaded.slab.dim, idx.slab.dim,
4127 "the slab's stride comes from the decoded vectors"
4128 );
4129 assert_matches(&loaded, &idx, &vecs[3]);
4130 }
4131
4132 /// A version this build does not know is treated as corrupt, not guessed
4133 /// at. The caller's fallback is the full scan — slower, never wrong.
4134 #[test]
4135 fn an_unknown_version_is_rejected() {
4136 let (_, idx) = blob_fixture();
4137 let mut blob = encode_hnsw_blob(&idx, true).expect("encode");
4138 blob[4] = HNSW_BLOB_VERSION as u8 + 1; // bump the version's low byte
4139 let err = decode_hnsw_blob(&blob).expect_err("a future version must not be read");
4140 assert!(
4141 err.contains("version"),
4142 "the refusal must name the version: {err}"
4143 );
4144 }
4145
4146 /// The downgrade this release actually ships into: a reader whose ceiling is
4147 /// version 2 — 0.6.6 up to Task 1 — meets a v3 blob and refuses it, so its
4148 /// caller keeps the `hnsw_tracked` full scan.
4149 ///
4150 /// The old reader is reconstructed here rather than asserted about: it read
4151 /// the wrapper against the v2 index shape and then refused any version above
4152 /// its own ceiling, so [`v2_era_decode`] is those two steps with
4153 /// [`HnswIndexV2`] — which this build still carries to *read* v2 blobs — in
4154 /// place of the live shape. A v2 blob proves the reconstruction works; a v3
4155 /// blob is then refused by it, on whichever of the two steps fires first.
4156 #[test]
4157 fn a_v2_reader_refuses_a_v3_blob_and_still_reads_a_v2_one() {
4158 /// `HNSW_BLOB_VERSION` as 0.6.6 shipped it before the distance kernel.
4159 const V2_CEILING: u16 = 2;
4160
4161 /// The pre-3b decoder, in the two decisions it made: deserialize the
4162 /// wrapper against the shape it knew, then refuse an unknown version.
4163 fn v2_era_decode(blob: &[u8]) -> Result<usize, String> {
4164 match bincode::deserialize::<HnswBlobV2>(blob) {
4165 Ok(b) if b.magic == HNSW_BLOB_MAGIC => {
4166 if b.version == 0 || b.version > V2_CEILING {
4167 Err(format!("version {} is not readable", b.version))
4168 } else {
4169 Ok(b.index.slots.len())
4170 }
4171 }
4172 Ok(b) => Err(format!("unrecognised magic {:?}", b.magic)),
4173 Err(e) => Err(format!("not a v2 blob ({e})")),
4174 }
4175 }
4176
4177 let (_, idx) = blob_fixture();
4178
4179 // Positive control: the reconstruction really does read a v2 blob, so its
4180 // refusal below is about the version and not about being broken.
4181 let v2 = as_v2_blob(&idx);
4182 assert_eq!(
4183 v2_era_decode(&v2),
4184 Ok(idx.slots.len()),
4185 "the reconstructed v2 reader must read a v2 blob"
4186 );
4187
4188 // And it refuses this build's blob.
4189 let v3 = encode_hnsw_blob(&idx, true).expect("encode");
4190 assert_eq!(&v3[..4], &HNSW_BLOB_MAGIC, "same magic, new version");
4191 assert_eq!(u16::from_le_bytes([v3[4], v3[5]]), HNSW_BLOB_VERSION);
4192 let err = v2_era_decode(&v3).expect_err(
4193 "a build that reads up to v2 must refuse a v3 blob rather than \
4194 misread it — its caller then keeps the full scan",
4195 );
4196 eprintln!("a v2-era reader on a v3 blob: {err}");
4197
4198 // The same gate in this build, for the version after this one: the
4199 // refusal names the version, which is what the caller logs before it
4200 // falls back.
4201 let mut future = v3.clone();
4202 future[4] = HNSW_BLOB_VERSION as u8 + 1;
4203 let err = decode_hnsw_blob(&future).expect_err("a future version must not be read");
4204 assert!(err.contains("version"), "{err}");
4205 }
4206
4207 /// Foreign magic is rejected too, and does not fall through to a garbage
4208 /// v1 read.
4209 #[test]
4210 fn a_foreign_magic_is_rejected() {
4211 let (_, idx) = blob_fixture();
4212 let mut blob = encode_hnsw_blob(&idx, true).expect("encode");
4213 blob[0] = b'X';
4214 assert!(
4215 decode_hnsw_blob(&blob).is_err(),
4216 "a blob with foreign magic must not be read"
4217 );
4218 }
4219
4220 /// The downgrade direction: a 0.6.5 binary meeting a v3 blob fails its
4221 /// `bincode::deserialize::<HnswIndex>` against the old shape rather than
4222 /// misreading it. `V1ReadIndex` is that old shape's read side.
4223 #[test]
4224 fn a_v3_blob_is_not_readable_as_a_v1_index() {
4225 #[derive(Deserialize)]
4226 #[allow(dead_code)]
4227 struct V1ReadNode {
4228 level: usize,
4229 vector: Vec<f64>,
4230 layers: Vec<Vec<u32>>,
4231 }
4232 #[derive(Deserialize)]
4233 #[allow(dead_code)]
4234 struct V1ReadIndex {
4235 base_seed: u64,
4236 nodes: BTreeMap<u32, V1ReadNode>,
4237 entry_point: Option<u32>,
4238 max_level: usize,
4239 }
4240
4241 let (_, idx) = blob_fixture();
4242 let blob = encode_hnsw_blob(&idx, true).expect("encode");
4243 assert!(
4244 bincode::deserialize::<V1ReadIndex>(&blob).is_err(),
4245 "a 0.6.5 reader must reject a 0.6.6 blob, not misread it"
4246 );
4247 }
4248
4249 /// Recall probe: 5 000 vectors × dim 1536, 50 queries, min recall@10 ≥ 0.90.
4250 ///
4251 /// Run with: `cargo test --release -p mushroomdb-rules -- hnsw_5k_1536_recall --ignored`
4252 ///
4253 /// Fixed seed — never random per run. Asserts are gates for the CI report.
4254 #[test]
4255 #[ignore]
4256 fn hnsw_5k_1536_recall() {
4257 const N: usize = 5_000;
4258 const DIM: usize = 1_536;
4259 const N_QUERIES: usize = 50;
4260 const K: usize = 10;
4261
4262 let seed = crate::index::fnv1a_u64(b"recall-probe-5k-1536");
4263 let vecs = make_unit_vecs(N, DIM, seed);
4264
4265 let mut idx = HnswIndex::new(seed);
4266 for (i, v) in vecs.iter().enumerate() {
4267 idx.insert(i as u32, v);
4268 }
4269
4270 // Use a different seed for query vectors so they differ from index vecs.
4271 let q_seed = crate::index::fnv1a_u64(b"recall-queries");
4272 let queries = make_unit_vecs(N_QUERIES, DIM, q_seed);
4273
4274 let mut recalls = Vec::with_capacity(N_QUERIES);
4275 for q in &queries {
4276 let exact_set: std::collections::BTreeSet<usize> =
4277 exact_knn(&vecs, q, K).into_iter().collect();
4278 let approx_ids: Vec<usize> = idx
4279 .search(q, K)
4280 .into_iter()
4281 .map(|(id, _)| id as usize)
4282 .collect();
4283 let hits = approx_ids
4284 .iter()
4285 .filter(|id| exact_set.contains(id))
4286 .count();
4287 recalls.push(hits as f64 / K as f64);
4288 }
4289
4290 let min_recall = recalls.iter().cloned().fold(f64::MAX, f64::min);
4291 let mean_recall = recalls.iter().sum::<f64>() / recalls.len() as f64;
4292
4293 eprintln!("HNSW 5k/1536 recall@{K}: min={min_recall:.4} mean={mean_recall:.4}");
4294
4295 assert!(
4296 min_recall >= 0.90,
4297 "min recall@{K} = {min_recall:.4} < 0.90"
4298 );
4299 assert!(
4300 mean_recall >= 0.95,
4301 "mean recall@{K} = {mean_recall:.4} < 0.95"
4302 );
4303 }
4304}