Skip to main content

writerslogic_warm/core/qgraph/
vertex.rs

1// Copyright 2024-2026 WritersLogic Contributors
2// SPDX-License-Identifier: Apache-2.0
3//! Graph index with one compact code per vertex instead of one code per edge.
4//!
5//! Every vertex stores its mean-centred vector as 8-bit integers with one scale (locally
6//! adaptive scalar quantization, LVQ; Aguerrebere et al., VLDB 2023). Traversal estimates
7//! `<q, x> - <q, mean> ~= s_x * s_q * <c_q, c_x>` with an int8 query code and one integer dot
8//! product per neighbour, so a neighbour costs one short code fetch (`dim` bytes, the scales
9//! live in their own small array) rather than a full-precision vector. An optional second code
10//! of the quantization residual (8-bit, LVQ-8x8, or 4-bit) re-ranks the best `rerank`
11//! candidates of the final pool with a float query.
12//!
13//! Memory per vertex is `4 * degree + dim + 4` bytes, plus `dim + 4` with the 8-bit residual
14//! (`dim / 2 + 4` with the 4-bit one), against `4 * dim` for the vectors alone in an HNSW
15//! index.
16
17use std::collections::BinaryHeap;
18use std::sync::atomic::{AtomicU64, Ordering};
19use std::sync::OnceLock;
20use std::time::Instant;
21
22use rayon::prelude::*;
23
24use super::adjacency::Adjacency;
25use super::fastscan::{pack32, scan32, BLOCK};
26use super::kernels::{
27    dot_i4_query, dot_i4_query_x4, dot_i8_query, dot_i8_query_x4, interleave_for_i4, prefetch,
28    prefetch_bytes, raw_one, QueryCode, Rotation, SplitMix,
29};
30use super::rerank::{dot_f32_i4, dot_f32_i8};
31use super::vnni::query_sum;
32use super::{invalid, prepare_rows, upper_bytes, BuildParams, Graph, GraphView, Metric, NONE};
33use crate::core::error::WarmError;
34
35mod file;
36mod gated;
37pub use file::Residency;
38use file::Slab;
39
40/// Scoring contract of the serving codes, independent of graph construction geometry.
41#[derive(Clone, Copy, Debug, PartialEq, Eq)]
42pub enum VMetric {
43    Angular,
44    Euclidean,
45    InnerProduct,
46}
47
48impl From<Metric> for VMetric {
49    fn from(metric: Metric) -> Self {
50        match metric {
51            Metric::Angular => Self::Angular,
52            Metric::Euclidean => Self::Euclidean,
53        }
54    }
55}
56
57impl VMetric {
58    fn code(self) -> u64 {
59        match self {
60            Self::Angular => 0,
61            Self::Euclidean => 1,
62            Self::InnerProduct => 2,
63        }
64    }
65}
66
67#[derive(Default, serde::Serialize)]
68pub struct QueryProfile {
69    pub query_encoding_ns: u64,
70    pub entry_ns: u64,
71    pub adjacency_visited_prefetch_ns: u64,
72    pub scoring_pool_ns: u64,
73    pub rerank_sort_ns: u64,
74    pub total_ns: u64,
75    pub expansions: usize,
76    pub neighbors_read: usize,
77    pub entry_evaluations: usize,
78    pub traversal_evaluations: usize,
79    pub rerank_evaluations: usize,
80    pub beam_retained: usize,
81    pub retained: usize,
82}
83
84/// Search parameters of a [`VGraph`].
85#[derive(Clone, Copy, Debug)]
86pub struct VSearchParams {
87    /// Candidate pool size (beam width).
88    pub ef: usize,
89    /// Independent result capacity, offered every row score before beam admission.
90    /// 0 uses the legacy final beam; otherwise bounded to `max(k, reservoir).min(n)`.
91    /// Traversal, including the `ef >= k` rule, is unchanged.
92    pub reservoir: usize,
93    /// Number of best pool candidates re-scored with the float query (and the residual code,
94    /// when built); 0 returns the pool order.
95    pub rerank: usize,
96    /// Early stop: end the traversal once this many consecutive expansions have left the
97    /// `k` best estimates unchanged (0 = the pool rule alone). The pool rule still applies.
98    pub patience: usize,
99    /// With a screen (see [`VGraph::from_graph_with`]): a fresh neighbour gets the 8-bit
100    /// estimate only if its 1-bit estimate minus this many standard deviations of the
101    /// screen's error could still enter the full pool. Larger values skip less.
102    pub screen_sigmas: f32,
103    /// Consult the screen when the index has one (false runs the plain search on it).
104    pub screen: bool,
105    /// With neighbour blocks: fuse the block estimates a vertex receives from every expanded
106    /// vertex that lists it (precision-weighted by the edge residual's scale) and re-admit
107    /// it to the pool when the fused estimate qualifies, instead of keeping its first one.
108    pub fuse: bool,
109    /// Per-row search: estimate only the first this many neighbours of each expanded vertex
110    /// (lists keep the build's order, nearest first); 0 = all.
111    pub neighbour_cap: usize,
112}
113
114impl Default for VSearchParams {
115    fn default() -> Self {
116        Self {
117            ef: 64,
118            reservoir: 0,
119            rerank: 0,
120            patience: 0,
121            screen_sigmas: 1.0,
122            screen: true,
123            fuse: false,
124            neighbour_cap: 0,
125        }
126    }
127}
128
129/// Runtime knobs of the search loop, each off by default so that the results are those of the
130/// plain search. [`VGraph::searcher`] reads them once per process from the environment
131/// ([`SearchKnobs::from_env`]); [`VGraph::searcher_with`] takes them explicitly.
132///
133/// Mechanisms, and what they may change:
134/// * `prefetch` (`WARM_QGRAPH_PREFETCH`, 0..=3, results never change): 0 prefetches the next
135///   candidate's neighbour list, the code row and estimate terms (scale, Euclidean norm) of
136///   every admitted neighbour, and the rows the re-rank reads; 2 adds the neighbour list of
137///   the second-nearest unexpanded candidate; 3 adds a lookahead that, after an expansion,
138///   prefetches the code rows (or screen rows) of the next candidate's unvisited neighbours.
139/// * `batch_est` (`WARM_QGRAPH_BATCH_EST`, results never change): the admitted neighbours'
140///   4-bit rows are estimated four at a time with shared query loads (`dot_i4_inline_x4` in
141///   the kernels module); the integer sums are the same.
142/// * `frontier` (`WARM_QGRAPH_FRONTIER`, `0.0..=1.0`, 1.0 = off): with a full pool, stop when
143///   the nearest unexpanded estimate is past the point `frontier` of the way from the `k`-th
144///   best to the `ef`-th best. Fewer expansions, so lower recall at a given `ef`; the
145///   recall/QPS curve over `ef` is what to compare.
146/// * `block_screen` (`WARM_QGRAPH_BLOCK_SCREEN`, needs [`VGraph::with_blocks`]): the neighbour
147///   block screens instead of ordering; survivors fetch their code row and the pool holds row
148///   estimates, as in the plain search. A fresh neighbour is skipped when its block estimate,
149///   improved by `screen_sigmas` deviations of the calibrated edge error, could not enter the
150///   full pool. Off, the block search is the existing one.
151#[derive(Clone, Copy, Debug, PartialEq)]
152pub struct SearchKnobs {
153    /// Prefetch depth, 0 to 3 (values above 3 act as 3).
154    pub prefetch: u8,
155    /// Estimate admitted neighbours four rows at a time.
156    pub batch_est: bool,
157    /// Frontier stop fraction in `0.0..=1.0`; 1.0 (or a non-finite value) is off.
158    pub frontier: f32,
159    /// Screen with the neighbour blocks and estimate survivors from their rows.
160    pub block_screen: bool,
161    /// Bound-gated re-rank: with `z > 0`, candidates are re-scored in estimate order and the
162    /// re-rank stops at the first whose estimate, corrected by the mean error and lowered by
163    /// `z` standard deviations (both measured on this query's re-scored rows), cannot beat
164    /// the current `k`-th exact distance. `rerank` stays the cap. 0 (or non-finite) is off.
165    pub rerank_z: f32,
166}
167
168impl SearchKnobs {
169    /// Every knob off: the plain search.
170    pub const OFF: Self = Self {
171        prefetch: 0,
172        batch_est: false,
173        frontier: 1.0,
174        block_screen: false,
175        rerank_z: 0.0,
176    };
177
178    /// The knobs set in the environment, read once per process; unset or invalid values
179    /// (reported once) take the [`SearchKnobs::OFF`] value.
180    pub fn from_env() -> Self {
181        static KNOBS: OnceLock<SearchKnobs> = OnceLock::new();
182        *KNOBS.get_or_init(|| Self {
183            prefetch: env_knob::<u8>("WARM_QGRAPH_PREFETCH").unwrap_or(0).min(3),
184            batch_est: env_knob::<u8>("WARM_QGRAPH_BATCH_EST").unwrap_or(0) != 0,
185            frontier: env_knob::<f32>("WARM_QGRAPH_FRONTIER")
186                .filter(|f| f.is_finite())
187                .map_or(1.0, |f| f.clamp(0.0, 1.0)),
188            block_screen: env_knob::<u8>("WARM_QGRAPH_BLOCK_SCREEN").unwrap_or(0) != 0,
189            rerank_z: env_knob::<f32>("WARM_QGRAPH_RERANK_Z")
190                .filter(|z| z.is_finite() && *z > 0.0)
191                .unwrap_or(0.0),
192        })
193    }
194}
195
196impl Default for SearchKnobs {
197    fn default() -> Self {
198        Self::OFF
199    }
200}
201
202/// `name` parsed as `T`: `None` when unset, and when the value does not parse, which is
203/// reported.
204fn env_knob<T: std::str::FromStr>(name: &str) -> Option<T> {
205    let value = crate::compatibility::var(name).ok()?;
206    match value.trim().parse::<T>() {
207        Ok(v) => Some(v),
208        Err(_) => {
209            eprintln!("{name}={value:?} is not a valid value; using the default");
210            None
211        }
212    }
213}
214
215/// On unless `WARM_QGRAPH_HUGEPAGES` is set to 0: the large arrays of an index built in this
216/// process or opened onto the heap (code rows, neighbour lists, scales, norm terms, residual
217/// rows, neighbour blocks) are advised into transparent huge pages before they are first
218/// touched, so the random row fetches of a search share one TLB entry per 2 MiB instead of
219/// needing one per 4 KiB (on x86 a page walk per fetch once the index outgrows the TLB's
220/// reach). Linux only; a no-op elsewhere, and where transparent huge pages are disabled.
221/// Read once.
222fn huge_pages() -> bool {
223    static HUGE: OnceLock<bool> = OnceLock::new();
224    *HUGE.get_or_init(|| env_knob::<u8>("WARM_QGRAPH_HUGEPAGES").unwrap_or(1) != 0)
225}
226
227/// Huge-page advice given by this process (see [`huge_page_status`]).
228static HUGE_ADVISED_BYTES: AtomicU64 = AtomicU64::new(0);
229static HUGE_ADVICE_FAILED: AtomicU64 = AtomicU64::new(0);
230/// Test hook: treat every `madvise` as failed, the result under THP `never`, a seccomp filter
231/// or a sandbox that rejects the call.
232#[cfg(test)]
233static FAIL_HUGE_ADVICE: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
234
235/// What the huge-page advice of this process achieved.
236#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, serde::Serialize)]
237pub struct HugePageStatus {
238    /// The advice is on (`WARM_QGRAPH_HUGEPAGES` unset or non-zero).
239    pub enabled: bool,
240    /// Bytes for which `madvise(MADV_HUGEPAGE)` succeeded.
241    pub advised_bytes: u64,
242    /// `madvise` calls that failed (each falls back to base pages; nothing else changes).
243    pub failed_calls: u64,
244    /// `AnonHugePages` of the whole process from `/proc/self/smaps_rollup`, whoever asked for
245    /// them (the allocator may hold huge pages of its own). `None` where the file is
246    /// unreadable (outside Linux, or a sandbox without it). [`VGraph::code_rows_huge_bytes`]
247    /// gives the figure for one index's code rows.
248    pub anon_huge_bytes: Option<u64>,
249}
250
251/// The huge-page advice given so far and the huge pages the process holds.
252pub fn huge_page_status() -> HugePageStatus {
253    let anon_huge_bytes = std::fs::read_to_string("/proc/self/smaps_rollup")
254        .ok()
255        .and_then(|text| {
256            let line = text.lines().find(|l| l.starts_with("AnonHugePages:"))?;
257            let kib: u64 = line.split_whitespace().nth(1)?.parse().ok()?;
258            Some(kib * 1024)
259        });
260    HugePageStatus {
261        enabled: huge_pages(),
262        advised_bytes: HUGE_ADVISED_BYTES.load(Ordering::Relaxed),
263        failed_calls: HUGE_ADVICE_FAILED.load(Ordering::Relaxed),
264        anon_huge_bytes,
265    }
266}
267
268/// `AnonHugePages` of the mappings of `/proc/self/smaps` that overlap `[ptr, ptr + len)`;
269/// `None` where the file is unreadable. A mapping may hold other allocations too, so this
270/// bounds the huge pages behind the range from above.
271fn huge_bytes_behind(ptr: *const u8, len: usize) -> Option<u64> {
272    let text = std::fs::read_to_string("/proc/self/smaps").ok()?;
273    let (lo, hi) = (ptr as u64, ptr as u64 + len as u64);
274    let mut total = 0u64;
275    let mut inside = false;
276    for line in text.lines() {
277        let head = line.split_whitespace().next().unwrap_or("");
278        if let Some((a, b)) = head.split_once('-') {
279            if let (Ok(a), Ok(b)) = (u64::from_str_radix(a, 16), u64::from_str_radix(b, 16)) {
280                inside = a < hi && lo < b;
281                continue;
282            }
283        }
284        if inside {
285            if let Some(rest) = line.strip_prefix("AnonHugePages:") {
286                let kib: u64 = rest.split_whitespace().next()?.parse().ok()?;
287                total += kib * 1024;
288            }
289        }
290    }
291    Some(total)
292}
293
294/// Asks the kernel to back the 2 MiB-aligned interior of `[ptr, ptr + len)` with transparent
295/// huge pages (Linux `MADV_HUGEPAGE`). Must run before the pages are first touched, so the
296/// page-fault path allocates them huge where THP is `always` or `madvise`; afterwards only
297/// `khugepaged` could collapse them, slowly. A no-op outside Linux. A failure (THP `never`,
298/// a kernel without THP, a filtered syscall) is counted and otherwise ignored: the advice
299/// changes paging policy only, never the mapping or its contents.
300fn advise_huge_pages(ptr: *const u8, len: usize) {
301    #[cfg(target_os = "linux")]
302    {
303        const HUGE: usize = 2 << 20;
304        const MADV_HUGEPAGE: i32 = 14;
305        extern "C" {
306            fn madvise(addr: *mut std::ffi::c_void, len: usize, advice: i32) -> i32;
307        }
308        let start = (ptr as usize).next_multiple_of(HUGE);
309        let end = (ptr as usize + len) & !(HUGE - 1);
310        if start < end {
311            #[cfg(test)]
312            if FAIL_HUGE_ADVICE.load(Ordering::Relaxed) {
313                HUGE_ADVICE_FAILED.fetch_add(1, Ordering::Relaxed);
314                return;
315            }
316            // SAFETY: `[start, end)` lies inside the live allocation `[ptr, ptr + len)` and
317            // `madvise` with `MADV_HUGEPAGE` reads or writes nothing in it.
318            let rc = unsafe { madvise(start as *mut std::ffi::c_void, end - start, MADV_HUGEPAGE) };
319            if rc == 0 {
320                HUGE_ADVISED_BYTES.fetch_add((end - start) as u64, Ordering::Relaxed);
321            } else {
322                HUGE_ADVICE_FAILED.fetch_add(1, Ordering::Relaxed);
323            }
324        }
325    }
326    #[cfg(not(target_os = "linux"))]
327    let _ = (ptr, len);
328}
329
330/// [`advise_huge_pages`] over the capacity of a freshly allocated, still untouched vector when
331/// [`huge_pages`] is on.
332fn huge_capacity<T>(ptr: *const T, capacity: usize) {
333    if huge_pages() {
334        advise_huge_pages(ptr.cast::<u8>(), capacity * std::mem::size_of::<T>());
335    }
336}
337
338/// One 128-byte cache line (the line size of Apple cores).
339#[repr(C, align(128))]
340#[derive(Clone, Copy)]
341struct Line([i8; 128]);
342
343/// Bytes starting on a cache-line boundary (zeroed when built here).
344struct Lines {
345    /// `len.div_ceil(128)` lines.
346    lines: Slab<Line>,
347    len: usize,
348}
349
350impl Lines {
351    fn zeroed(len: usize) -> Self {
352        let count = len.div_ceil(128);
353        let mut lines = Vec::with_capacity(count);
354        huge_capacity(lines.as_ptr(), count);
355        lines.resize(count, Line([0; 128]));
356        Self {
357            lines: lines.into(),
358            len,
359        }
360    }
361
362    #[inline(always)]
363    fn bytes(&self) -> &[i8] {
364        // SAFETY: `Line` is `repr(C)` around `[i8; 128]` (size 128, no padding), so `lines` is
365        // `128 * lines.len() >= len` initialized, contiguous bytes.
366        unsafe { std::slice::from_raw_parts(self.lines.as_ptr().cast::<i8>(), self.len) }
367    }
368
369    fn bytes_mut(&mut self) -> &mut [i8] {
370        let len = self.len;
371        let lines = self.lines.as_mut_slice();
372        // SAFETY: as in `bytes`, with the exclusive borrow of `lines`.
373        unsafe { std::slice::from_raw_parts_mut(lines.as_mut_ptr().cast::<i8>(), len) }
374    }
375}
376
377/// Logical payload bytes counted by [`VGraph::index_bytes`], including row padding.
378/// Excludes allocation capacity, slab rounding, object/allocator metadata, rotations and
379/// query scratch. These are not resident, unique, proportional, or peak memory measures.
380#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, serde::Serialize)]
381pub struct VGraphMemoryBreakdown {
382    pub adjacency_ids_bytes: usize,
383    pub adjacency_offsets_bytes: usize,
384    pub traversal_codes_bytes: usize,
385    pub traversal_scales_bytes: usize,
386    pub residual_codes_bytes: usize,
387    pub residual_scales_bytes: usize,
388    pub mean_bytes: usize,
389    pub norms_bytes: usize,
390    pub original_ids_bytes: usize,
391    pub screen_bytes: usize,
392    pub neighbor_blocks_bytes: usize,
393    pub upper_layers_bytes: usize,
394}
395
396impl VGraphMemoryBreakdown {
397    pub fn total_bytes(&self) -> usize {
398        self.adjacency_ids_bytes
399            + self.adjacency_offsets_bytes
400            + self.traversal_codes_bytes
401            + self.traversal_scales_bytes
402            + self.residual_codes_bytes
403            + self.residual_scales_bytes
404            + self.mean_bytes
405            + self.norms_bytes
406            + self.original_ids_bytes
407            + self.screen_bytes
408            + self.neighbor_blocks_bytes
409            + self.upper_layers_bytes
410    }
411}
412
413pub struct VGraph {
414    n: usize,
415    dim: usize,
416    metric: VMetric,
417    /// `dim` rounded up to a multiple of 16 (32 with 4-bit codes), the integer kernels' width.
418    padded: usize,
419    /// Bits per traversal code coordinate: 8, or 4 (two per byte).
420    bits: u8,
421    /// Bytes of one traversal code: `padded` or `padded / 2`.
422    cbytes: usize,
423    /// Bytes from one code row to the next: `cbytes`, or with aligned rows the smallest
424    /// power of two (up to 64) or multiple of 128 that holds it, so no row straddles a line.
425    stride: usize,
426    degree: usize,
427    mean: Vec<f32>,
428    /// Euclidean: `|x|^2 / 2` per stored vertex, the norm term of the distance order; empty
429    /// for angular.
430    half_norms: Slab<f32>,
431    /// Neighbour ids, `degree` per vertex padded with `NONE`; empty with 3-byte ids.
432    adj: Slab<u32>,
433    /// The same as 3-byte little-endian ids (padded with `NONE3`); empty with `u32` ids.
434    adj3: Slab<u8>,
435    /// In-memory exact-degree spans; empty for the persisted fixed-slot layout.
436    adj_off: Vec<u32>,
437    /// Delta-coded lists (`id_bytes` 0): per vertex, the sorted neighbour ids as LEB128
438    /// varints of the zigzag first difference to the vertex and the gaps less one after,
439    /// with `adjv_off[v]..adjv_off[v + 1]` the byte span of vertex `v`; empty otherwise.
440    adjv: Slab<u8>,
441    adjv_off: Slab<u32>,
442    codes: Lines,
443    scales: Slab<f32>,
444    /// 0 (none), 4 or 8.
445    res_bits: u8,
446    /// Bytes per residual row: `padded` (8-bit) or `padded / 2` (4-bit, two per byte).
447    res_stride: usize,
448    residual: Slab<i8>,
449    res_scales: Slab<f32>,
450    entry: u32,
451    upper_ids: Vec<u32>,
452    layers: Vec<Vec<Vec<u32>>>,
453    /// Original id of each stored vertex when the vertices were renumbered; empty otherwise.
454    ids: Slab<u32>,
455    /// 1-bit screen rows: `screen_words` u64 of sign bits of the rotated centred vector, then
456    /// one u64 holding the f32 factor (low half) and the popcount (high half); empty without
457    /// a screen.
458    screen: Slab<u64>,
459    screen_words: usize,
460    screen_rot: Option<Rotation>,
461    /// Standard deviation of (screen estimate - 8-bit estimate), in estimate units, measured
462    /// at build on index vectors used as queries.
463    screen_sigma: f32,
464    /// Bits per coordinate of the neighbour-block codes (0 = no blocks; 1, 2 or 4). See
465    /// [`VGraph::with_blocks`].
466    block_bits: u8,
467    /// Coordinates of a block code: the block rotation's width.
468    block_dim: usize,
469    /// Nibble groups per block code: `block_dim * block_bits / 4`.
470    block_m: usize,
471    /// Blocks per vertex: `degree.div_ceil(32)`.
472    blocks_per_vertex: usize,
473    /// Per vertex, `blocks_per_vertex` records of `16 * block_m` code bytes in the fast-scan
474    /// layout followed by the 32 neighbours' scales as little-endian f32.
475    blocks: Lines,
476    block_rot: Option<Rotation>,
477    /// Error model of the block estimates: `var = block_kappa * s^2 * |q - x_v|^2` for an
478    /// edge of scale `s` seen from the expanded vertex `v`, calibrated at build.
479    block_kappa: f32,
480}
481
482/// Seed of the screen's rotation (fixed: the screen is a property of the encoding).
483const SCREEN_SEED: u64 = 0x5C12_EE00;
484/// Seed of the neighbour-block rotation.
485const BLOCK_SEED: u64 = 0xB10C_5CA0;
486
487/// Bytes of one block record: codes, then 32 scales, then 32 offsets (little-endian f32).
488fn block_record_bytes(m: usize) -> usize {
489    16 * m + 8 * BLOCK
490}
491
492/// Level of coordinate `i` of a block code (see [`encode_block_code`]).
493#[inline]
494fn block_level(nibbles: &[u8], bits: u8, i: usize) -> f32 {
495    match bits {
496        1 => {
497            if nibbles[i / 4] >> (i % 4) & 1 == 1 {
498                1.0
499            } else {
500                -1.0
501            }
502        }
503        2 => 2.0 * f32::from(nibbles[i / 2] >> (2 * (i % 2)) & 3) - 3.0,
504        _ => f32::from(nibbles[i]) - 7.0,
505    }
506}
507
508/// Little-endian f32 number `i` of `bytes`.
509#[inline]
510fn le_f32(bytes: &[u8], i: usize) -> f32 {
511    f32::from_le_bytes([
512        bytes[4 * i],
513        bytes[4 * i + 1],
514        bytes[4 * i + 2],
515        bytes[4 * i + 3],
516    ])
517}
518
519/// Encodes a rotated, centred vector as `m` nibbles for a `bits`-per-coordinate block code
520/// and returns its scale, so that coordinate `i ~= scale * level(i)` with the levels
521/// `{-1, 1}` (1 bit), `{-3, -1, 1, 3}` (2 bits) or `-7..=7` (4 bits; nibble `level + 7`).
522/// The levels of the 2- and 4-bit codes are chosen with the scale among eight candidates
523/// that minimizes the squared reconstruction error. The returned scale, for every width, is
524/// `|x|^2 / <levels, x>`: it makes the code error orthogonal to `x` itself (for one bit this
525/// is the RaBitQ factor `|x|^2 / sum |x_i|`), so inner products with queries that are
526/// correlated with `x` are unbiased, where the least-squares scale underestimates them by
527/// the relative error energy.
528fn encode_block_code(x: &[f32], bits: u8, nibbles: &mut [u8]) -> f32 {
529    nibbles.fill(0);
530    let m = x.iter().fold(0f32, |a, &v| a.max(v.abs()));
531    if m == 0.0 {
532        return 0.0;
533    }
534    match bits {
535        1 => {
536            let (mut norm2, mut abs) = (0f32, 0f32);
537            for (i, &v) in x.iter().enumerate() {
538                norm2 += v * v;
539                abs += v.abs();
540                if v >= 0.0 {
541                    nibbles[i / 4] |= 1 << (i % 4);
542                }
543            }
544            norm2 / abs
545        }
546        2 => {
547            // Levels {-3, -1, 1, 3} * s; index t = (level + 3) / 2.
548            let err = |s: f32| -> f32 {
549                x.iter()
550                    .map(|&v| {
551                        let t = ((v / s + 3.0) / 2.0).round().clamp(0.0, 3.0);
552                        let d = v - s * (2.0 * t - 3.0);
553                        d * d
554                    })
555                    .sum()
556            };
557            let (mut best, mut best_err) = (m / 3.0, f32::INFINITY);
558            for k in 0..8 {
559                let s = m / 3.0 * (1.0 - 0.1 * k as f32);
560                let e = err(s);
561                if e < best_err {
562                    (best, best_err) = (s, e);
563                }
564            }
565            let mut dot = 0f32;
566            for (i, &v) in x.iter().enumerate() {
567                let t = ((v / best + 3.0) / 2.0).round().clamp(0.0, 3.0) as u8;
568                nibbles[i / 2] |= t << (2 * (i % 2));
569                dot += v * (2.0 * f32::from(t) - 3.0);
570            }
571            unbiased_scale(x, dot, best)
572        }
573        _ => {
574            let err = |s: f32| -> f32 {
575                x.iter()
576                    .map(|&v| {
577                        let d = v - s * (v / s).round().clamp(-7.0, 7.0);
578                        d * d
579                    })
580                    .sum()
581            };
582            let (mut best, mut best_err) = (m / 7.0, f32::INFINITY);
583            for k in 0..8 {
584                let s = m / 7.0 * (1.0 - 0.06 * k as f32);
585                let e = err(s);
586                if e < best_err {
587                    (best, best_err) = (s, e);
588                }
589            }
590            let mut dot = 0f32;
591            for (i, &v) in x.iter().enumerate() {
592                let level = (v / best).round().clamp(-7.0, 7.0);
593                nibbles[i] = (level as i8 + 7) as u8;
594                dot += v * level;
595            }
596            unbiased_scale(x, dot, best)
597        }
598    }
599}
600
601/// `|x|^2 / <levels, x>` (see [`encode_block_code`]), or the least-squares scale when the
602/// levels do not correlate with `x`.
603fn unbiased_scale(x: &[f32], dot: f32, fallback: f32) -> f32 {
604    let norm2: f32 = x.iter().map(|v| v * v).sum();
605    if dot > 0.0 {
606        norm2 / dot
607    } else {
608        fallback
609    }
610}
611
612/// Fills the float lookup table `t` (`16 * m` entries) of a rotated query for `bits`-per-
613/// coordinate block codes: `t[16 * j + v] = sum over the coordinates of nibble j of
614/// q_i * level_i(v)`.
615fn block_table(q: &[f32], bits: u8, m: usize, t: &mut [f32]) {
616    match bits {
617        1 => {
618            for j in 0..m {
619                let qs = &q[4 * j..4 * j + 4];
620                let total: f32 = qs.iter().sum();
621                for v in 0..16usize {
622                    let on: f32 = (0..4).filter(|b| v >> b & 1 == 1).map(|b| qs[b]).sum();
623                    t[16 * j + v] = 2.0 * on - total;
624                }
625            }
626        }
627        2 => {
628            for j in 0..m {
629                let (qa, qb) = (q[2 * j], q[2 * j + 1]);
630                for v in 0..16usize {
631                    let (ta, tb) = ((v & 3) as f32, (v >> 2) as f32);
632                    t[16 * j + v] = qa * (2.0 * ta - 3.0) + qb * (2.0 * tb - 3.0);
633                }
634            }
635        }
636        _ => {
637            for j in 0..m {
638                for v in 0..16usize {
639                    t[16 * j + v] = q[j] * (v as f32 - 7.0);
640                }
641            }
642        }
643    }
644}
645
646/// Per-query block lookup table: `lut` as the scan reads it, with the affine map
647/// `entry ~= delta * lut + min_j` so that a code's inner product is
648/// `scale * (delta * sum + min_sum)`.
649pub struct BlockQuery {
650    table: Vec<f32>,
651    lut: Vec<u8>,
652    /// Per group, the minimum of its table entries.
653    mins: Vec<f32>,
654    delta: f32,
655    min_sum: f32,
656    /// `1 / s_q`, which puts the block estimate in [`VGraph::est`] units.
657    inv_sq: f32,
658}
659
660impl BlockQuery {
661    fn new(m: usize) -> Self {
662        Self {
663            table: vec![0.0; 16 * m],
664            lut: vec![0; 16 * m],
665            mins: vec![0.0; m],
666            delta: 1.0,
667            min_sum: 0.0,
668            inv_sq: 1.0,
669        }
670    }
671
672    /// Quantizes the float table to bytes: per group, entries are offset by the group's
673    /// minimum; one step `delta` serves every group.
674    fn quantize(&mut self, m: usize) {
675        debug_assert_eq!(self.lut.len(), 16 * m);
676        debug_assert_eq!(self.mins.len(), m);
677        let mut range = 0f32;
678        let mut min_sum = 0f32;
679        let groups = self.table.as_chunks::<16>().0;
680        for (g, lo) in groups.iter().zip(self.mins.iter_mut()) {
681            *lo = g.iter().fold(f32::INFINITY, |a, &v| a.min(v));
682            let hi = g.iter().fold(f32::NEG_INFINITY, |a, &v| a.max(v));
683            min_sum += *lo;
684            range = range.max(hi - *lo);
685        }
686        let delta = if range > 0.0 { range / 255.0 } else { 1.0 };
687        let luts = self.lut.as_chunks_mut::<16>().0;
688        for ((g, &lo), lut) in groups.iter().zip(&self.mins).zip(luts) {
689            for (o, &e) in lut.iter_mut().zip(g) {
690                *o = ((e - lo) / delta).round().clamp(0.0, 255.0) as u8;
691            }
692        }
693        self.delta = delta;
694        self.min_sum = min_sum;
695    }
696}
697
698/// Inverse of [`key`] for the distance part.
699#[inline]
700fn key_dist(k: u64) -> f32 {
701    let o = (k >> 32) as u32;
702    f32::from_bits(if o >> 31 == 1 { o & 0x7FFF_FFFF } else { !o })
703}
704
705/// Writes the symmetric 8-bit code of `x` into `code` (zero padded) and returns its scale, so
706/// that `x_i ~= scale * code_i`.
707pub(super) fn quantize(x: &[f32], code: &mut [i8]) -> f32 {
708    let m = x.iter().fold(0f32, |a, &v| a.max(v.abs()));
709    code.fill(0);
710    if m == 0.0 {
711        return 0.0;
712    }
713    let scale = m / 127.0;
714    for (c, &v) in code.iter_mut().zip(x) {
715        *c = (v / scale).round().clamp(-127.0, 127.0) as i8;
716    }
717    scale
718}
719
720/// [`quantize`] with 4-bit codes in `-7..=7`, packed two per byte as [`dot_f32_i4`] reads them
721/// (`code` holds `code.len() * 2 >= x.len()` values).
722pub(super) fn quantize4(x: &[f32], code: &mut [i8]) -> f32 {
723    let m = x.iter().fold(0f32, |a, &v| a.max(v.abs()));
724    code.fill(0);
725    if m == 0.0 {
726        return 0.0;
727    }
728    let scale = m / 7.0;
729    let q = |i: usize| {
730        x.get(i)
731            .map_or(0, |&v| (v / scale).round().clamp(-7.0, 7.0) as i8)
732    };
733    for (j, c) in code.iter_mut().enumerate() {
734        *c = ((q(2 * j) as u8 & 0x0F) | ((q(2 * j + 1) as u8) << 4)) as i8;
735    }
736    scale
737}
738
739/// Breadth-first order of the bottom layer from the entry (unreached vertices follow in id
740/// order): neighbours get nearby positions, so a search touches fewer distinct pages and lines.
741pub(super) fn bfs_order(g: GraphView<'_>) -> Vec<u32> {
742    let mut seen = vec![false; g.n];
743    let mut order = Vec::with_capacity(g.n);
744    for s in std::iter::once(g.entry as usize).chain(0..g.n) {
745        if seen[s] {
746            continue;
747        }
748        seen[s] = true;
749        let mut head = order.len();
750        order.push(s as u32);
751        while head < order.len() {
752            let v = order[head] as usize;
753            head += 1;
754            for &u in g.neighbors(v) {
755                if !std::mem::replace(&mut seen[u as usize], true) {
756                    order.push(u);
757                }
758            }
759        }
760    }
761    order
762}
763
764/// Coordinate `i` of a code row of `bits` bits (0 past its end).
765pub(super) fn code_at(c: &[i8], bits: u8, i: usize) -> i8 {
766    if bits == 8 {
767        c.get(i).copied().unwrap_or(0)
768    } else {
769        c.get(i / 2).map_or(0, |&b| {
770            if i.is_multiple_of(2) {
771                (((b as u8) << 4) as i8) >> 4
772            } else {
773                b >> 4
774            }
775        })
776    }
777}
778
779/// Padding of a 3-byte neighbour list.
780const NONE3: u32 = 0x00FF_FFFF;
781
782#[inline]
783/// LEB128 varint (7 bits per byte, low bits first).
784fn put_varint(out: &mut Vec<u8>, mut v: u64) {
785    while v >= 0x80 {
786        out.push((v as u8) | 0x80);
787        v >>= 7;
788    }
789    out.push(v as u8);
790}
791
792/// Decodes one LEB128 varint; returns the value and the bytes used.
793#[inline]
794fn get_varint(bytes: &[u8]) -> (u64, usize) {
795    let mut v = 0u64;
796    for (i, &b) in bytes.iter().enumerate() {
797        v |= u64::from(b & 0x7f) << (7 * i);
798        if b & 0x80 == 0 {
799            return (v, i + 1);
800        }
801    }
802    (v, bytes.len())
803}
804
805fn zigzag(v: i64) -> u64 {
806    ((v << 1) ^ (v >> 63)) as u64
807}
808
809fn unzigzag(v: u64) -> i64 {
810    ((v >> 1) as i64) ^ -((v & 1) as i64)
811}
812
813fn prefetch_u8(bytes: &[u8]) {
814    // SAFETY: any initialized `u8` slice is also initialized `i8`s.
815    prefetch_bytes(unsafe { std::slice::from_raw_parts(bytes.as_ptr().cast::<i8>(), bytes.len()) });
816}
817
818#[inline]
819fn prefetch_f32(x: &f32) {
820    // SAFETY: an initialized `f32` is four initialized bytes.
821    prefetch_bytes(unsafe { std::slice::from_raw_parts((x as *const f32).cast::<i8>(), 4) });
822}
823
824#[inline]
825fn prefetch_ids(ids: &[u32]) {
826    // SAFETY: any initialized `u32` slice is also `4 * len` initialized bytes.
827    prefetch_bytes(unsafe { std::slice::from_raw_parts(ids.as_ptr().cast::<i8>(), ids.len() * 4) });
828}
829
830/// Orders by distance, then id; `f32` order is kept for every non-NaN distance.
831#[inline]
832fn key(dist: f32, id: u32) -> u64 {
833    let b = dist.to_bits();
834    let o = if b >> 31 == 1 { !b } else { b | 0x8000_0000 };
835    (u64::from(o) << 32) | u64::from(id)
836}
837
838#[inline]
839fn key_id(k: u64) -> u32 {
840    k as u32
841}
842
843struct VertexLayout {
844    padded: usize,
845    cbytes: usize,
846    stride: usize,
847    res_stride: usize,
848}
849
850impl VertexLayout {
851    fn new(
852        n: usize,
853        dim: usize,
854        bits: u8,
855        residual_bits: u8,
856        align: bool,
857        id_bytes: u8,
858    ) -> Result<Self, WarmError> {
859        if !matches!(bits, 4 | 8) {
860            return Err(invalid(format!(
861                "vertex code bits must be 4 or 8, not {bits}"
862            )));
863        }
864        if !(id_bytes == 4 || id_bytes == 0 || (id_bytes == 3 && (n as u64) < u64::from(NONE3))) {
865            return Err(invalid(format!("id_bytes must be 4, 0 (delta varints), or 3 below 2^24 - 1 vertices (got {id_bytes} for {n} vertices)")));
866        }
867        if !matches!(residual_bits, 0 | 4 | 8) {
868            return Err(invalid(format!(
869                "residual code bits must be 0, 4 or 8, not {residual_bits}"
870            )));
871        }
872        let padded = dim.next_multiple_of(if bits == 4 { 32 } else { 16 });
873        let cbytes = if bits == 4 { padded / 2 } else { padded };
874        let stride = match (align, cbytes) {
875            (false, _) => cbytes,
876            (true, c) if c <= 64 => c.next_power_of_two(),
877            (true, c) => c.next_multiple_of(128),
878        };
879        let res_stride = match residual_bits {
880            8 => padded,
881            4 => padded / 2,
882            _ => 0,
883        };
884        Ok(Self {
885            padded,
886            cbytes,
887            stride,
888            res_stride,
889        })
890    }
891}
892
893impl VGraph {
894    /// Encodes a built graph with the 8-bit residual code used by `rerank` (or none), unaligned
895    /// rows and the build's vertex order.
896    pub fn from_graph(g: &Graph, residual: bool) -> Result<Self, WarmError> {
897        Self::from_graph_with(g, 8, if residual { 8 } else { 0 }, false, false, 4, false)
898    }
899
900    /// Encodes a built graph. `bits` (8 or 4) sizes the traversal code and `residual_bits`
901    /// (0, 4 or 8) the re-rank code of its residual; `align` keeps every code row inside as
902    /// few 128-byte lines as its size allows; `reorder` renumbers the vertices in breadth-first
903    /// order (search still returns the original ids); `id_bytes` 3 stores neighbour ids in
904    /// three bytes (graphs below 2^24 - 1 vertices) instead of 4, and 0 stores each list
905    /// sorted and delta-coded as varints (sizes vary; `reorder` makes the deltas small); `screen` adds a 1-bit
906    /// sign code of the rotated centred vector per vertex (RaBitQ-style, with a per-vertex
907    /// factor) that the search consults before paying for the 8-bit estimate.
908    pub fn from_graph_with(
909        g: &Graph,
910        bits: u8,
911        residual_bits: u8,
912        align: bool,
913        reorder: bool,
914        id_bytes: u8,
915        screen: bool,
916    ) -> Result<Self, WarmError> {
917        Self::encode_graph(
918            g,
919            (g.dim, g.metric.into()),
920            (bits, residual_bits),
921            align,
922            reorder,
923            id_bytes,
924            screen,
925        )
926    }
927
928    /// Keeps a lifted Euclidean graph's topology while encoding only its original coordinates
929    /// for negative inner-product scoring. The final construction coordinate is omitted.
930    pub fn from_lifted_ip_graph(
931        g: &Graph,
932        bits: u8,
933        residual_bits: u8,
934        align: bool,
935        reorder: bool,
936        id_bytes: u8,
937    ) -> Result<Self, WarmError> {
938        if g.metric != Metric::Euclidean || g.dim < 2 {
939            return Err(invalid(
940                "direct IP serving requires a lifted Euclidean graph".into(),
941            ));
942        }
943        Self::encode_graph(
944            g,
945            (g.dim - 1, VMetric::InnerProduct),
946            (bits, residual_bits),
947            align,
948            reorder,
949            id_bytes,
950            false,
951        )
952    }
953
954    fn encode_graph(
955        g: &Graph,
956        geometry: (usize, VMetric),
957        precision: (u8, u8),
958        align: bool,
959        reorder: bool,
960        id_bytes: u8,
961        screen: bool,
962    ) -> Result<Self, WarmError> {
963        let (dim, metric) = geometry;
964        let (bits, residual_bits) = precision;
965        let layout = VertexLayout::new(g.n, dim, bits, residual_bits, align, id_bytes)?;
966        let euclidean = metric == VMetric::Euclidean;
967        if metric != VMetric::Angular && screen {
968            return Err(invalid("the 1-bit screen is angular only".into()));
969        }
970        let n = g.n;
971        let VertexLayout {
972            padded,
973            cbytes,
974            stride,
975            res_stride,
976        } = layout;
977        let mut sum = vec![0f64; dim];
978        for x in g.vectors.chunks_exact(g.dim) {
979            sum.iter_mut().zip(x).for_each(|(s, &v)| *s += f64::from(v));
980        }
981        let mut mean: Vec<f32> = sum.iter().map(|s| (s / n as f64) as f32).collect();
982        mean.resize(padded, 0.0);
983
984        let ids = if reorder {
985            bfs_order(g.view())
986        } else {
987            Vec::new()
988        };
989        let old = |v: usize| if reorder { ids[v] as usize } else { v };
990        let mut new_id: Vec<u32> = Vec::new();
991        if reorder {
992            new_id = vec![0; n];
993            for (v, &o) in ids.iter().enumerate() {
994                new_id[o as usize] = v as u32;
995            }
996        }
997        let new = |o: u32| if reorder { new_id[o as usize] } else { o };
998        let half_norms: Vec<f32> = if euclidean {
999            let mut half_norms = Vec::with_capacity(n);
1000            huge_capacity(half_norms.as_ptr(), n);
1001            half_norms.extend((0..n).map(|v| 0.5 * g.norms[old(v)]));
1002            half_norms
1003        } else {
1004            Vec::new()
1005        };
1006
1007        let screen_rot = screen.then(|| Rotation::new(dim, 0, SCREEN_SEED));
1008        let screen_words = screen_rot.as_ref().map_or(0, |r| r.padded() / 64);
1009        // The centred row of stored vertex `v`, which every code below encodes, written to a
1010        // per-worker buffer.
1011        let centred = |v: usize, out: &mut Vec<f32>| {
1012            let o = old(v);
1013            out.clear();
1014            out.extend(
1015                g.vectors[o * g.dim..o * g.dim + dim]
1016                    .iter()
1017                    .zip(&mean)
1018                    .map(|(a, m)| a - m),
1019            );
1020        };
1021        // Codes are written in place, so no per-vertex buffers sit next to the final arrays.
1022        let mut codes = Lines::zeroed(n * stride);
1023        let mut scales = Vec::with_capacity(n);
1024        huge_capacity(scales.as_ptr(), n);
1025        scales.resize(n, 0f32);
1026        codes
1027            .bytes_mut()
1028            .par_chunks_exact_mut(stride)
1029            .zip(scales.par_iter_mut())
1030            .enumerate()
1031            .for_each_init(Vec::<f32>::new, |x, (v, (dst, s))| {
1032                centred(v, x);
1033                *s = if bits == 8 {
1034                    quantize(x, &mut dst[..cbytes])
1035                } else {
1036                    quantize4(x, &mut dst[..cbytes])
1037                };
1038            });
1039        let mut residual = Vec::with_capacity(n * res_stride);
1040        huge_capacity(residual.as_ptr(), n * res_stride);
1041        residual.resize(n * res_stride, 0i8);
1042        let mut res_scales = vec![0f32; if residual_bits > 0 { n } else { 0 }];
1043        if residual_bits > 0 {
1044            let code_bytes = codes.bytes();
1045            residual
1046                .par_chunks_exact_mut(res_stride)
1047                .zip(res_scales.par_iter_mut())
1048                .enumerate()
1049                .for_each_init(Vec::<f32>::new, |x, (v, (r, sr))| {
1050                    centred(v, x);
1051                    let c = &code_bytes[v * stride..v * stride + cbytes];
1052                    let s = scales[v];
1053                    x.iter_mut()
1054                        .enumerate()
1055                        .for_each(|(i, a)| *a -= s * f32::from(code_at(c, bits, i)));
1056                    *sr = if residual_bits == 8 {
1057                        quantize(x, r)
1058                    } else {
1059                        quantize4(x, r)
1060                    };
1061                });
1062        }
1063        let mut screen_rows = vec![0u64; if screen { n * (screen_words + 1) } else { 0 }];
1064        if let Some(rot) = &screen_rot {
1065            screen_rows
1066                .par_chunks_exact_mut(screen_words + 1)
1067                .enumerate()
1068                .for_each_init(
1069                    || (Vec::<f32>::new(), vec![0f32; rot.padded()]),
1070                    |(x, xr), (v, srow)| {
1071                        centred(v, x);
1072                        rot.apply(x, xr);
1073                        let (mut norm2, mut abs) = (0f32, 0f32);
1074                        let mut pop = 0u32;
1075                        for (i, &a) in xr.iter().enumerate() {
1076                            norm2 += a * a;
1077                            abs += a.abs();
1078                            if a >= 0.0 {
1079                                srow[i / 64] |= 1u64 << (i % 64);
1080                                pop += 1;
1081                            }
1082                        }
1083                        // x_r ~= f * sign(x_r) with f = |x_r|^2 / sum |x_r,i| (the RaBitQ factor).
1084                        let f = if abs > 0.0 { norm2 / abs } else { 0.0 };
1085                        srow[screen_words] = u64::from(f.to_bits()) | (u64::from(pop) << 32);
1086                    },
1087                );
1088        }
1089
1090        let r = g.degree;
1091        // Stored vertex `v`'s list: the build's list of the vertex, renumbered, at most `r`.
1092        let list = |v: usize| g.neighbors(old(v)).iter().take(r).map(|&u| new(u));
1093        let (mut adj, mut adj3) = (Vec::new(), Vec::new());
1094        let (mut adjv, mut adjv_off) = (Vec::new(), Vec::new());
1095        if id_bytes == 3 {
1096            adj3 = Vec::<u8>::with_capacity(3 * n * r);
1097            huge_capacity(adj3.as_ptr(), adj3.capacity());
1098            for v in 0..n {
1099                let mut count = 0;
1100                for u in list(v) {
1101                    adj3.extend_from_slice(&u.min(NONE3).to_le_bytes()[..3]);
1102                    count += 1;
1103                }
1104                for _ in count..r {
1105                    adj3.extend_from_slice(&NONE3.to_le_bytes()[..3]);
1106                }
1107            }
1108        } else if id_bytes == 4 {
1109            adj = vec![NONE; n * r];
1110            for (v, dst) in adj.chunks_exact_mut(r).enumerate() {
1111                for (d, u) in dst.iter_mut().zip(list(v)) {
1112                    *d = u;
1113                }
1114            }
1115        } else {
1116            adjv_off.reserve(n + 1);
1117            let mut sorted: Vec<u32> = Vec::with_capacity(r);
1118            for v in 0..n {
1119                adjv_off.push(adjv.len() as u32);
1120                sorted.clear();
1121                sorted.extend(list(v).take_while(|&u| u != NONE));
1122                sorted.sort_unstable();
1123                let list = &sorted;
1124                let mut prev = v as u32;
1125                for (i, &u) in list.iter().enumerate() {
1126                    let delta = if i == 0 {
1127                        zigzag(i64::from(u) - i64::from(prev))
1128                    } else {
1129                        u64::from(u - prev - 1)
1130                    };
1131                    put_varint(&mut adjv, delta);
1132                    prev = u;
1133                }
1134            }
1135            adjv_off.push(adjv.len() as u32);
1136        }
1137        let entry = new(g.entry);
1138        let upper_ids = g.upper_ids.iter().map(|&u| new(u)).collect();
1139        let mut out = Self {
1140            n,
1141            dim,
1142            metric,
1143            padded,
1144            bits,
1145            cbytes,
1146            stride,
1147            degree: r,
1148            mean,
1149            half_norms: half_norms.into(),
1150            adj: adj.into(),
1151            adj3: adj3.into(),
1152            adj_off: Vec::new(),
1153            adjv: adjv.into(),
1154            adjv_off: adjv_off.into(),
1155            codes,
1156            scales: scales.into(),
1157            res_bits: residual_bits,
1158            res_stride,
1159            residual: residual.into(),
1160            res_scales: res_scales.into(),
1161            entry,
1162            upper_ids,
1163            layers: g.layers.clone(),
1164            ids: ids.into(),
1165            screen: screen_rows.into(),
1166            screen_words,
1167            screen_rot,
1168            screen_sigma: 0.0,
1169            block_bits: 0,
1170            block_dim: 0,
1171            block_m: 0,
1172            blocks_per_vertex: 0,
1173            blocks: Lines::zeroed(0),
1174            block_rot: None,
1175            block_kappa: 0.0,
1176        };
1177        if screen {
1178            out.screen_sigma = out.measure_screen_sigma(g);
1179        }
1180        Ok(out)
1181    }
1182
1183    /// Adds neighbour-block codes: every vertex stores, next to its neighbour list, a
1184    /// `bits`-per-coordinate code of each neighbour's rotated centred vector in the fast-scan
1185    /// layout (32 codes per block), so one streamed read and one table scan estimate all
1186    /// neighbours of an expanded vertex instead of one random row fetch per neighbour
1187    /// (SymphonyQG's layout; Gou et al., SIGMOD 2025). `bits` is 1, 2 or 4.
1188    pub fn with_blocks(mut self, g: &Graph, bits: u8) -> Result<Self, WarmError> {
1189        if !matches!(bits, 1 | 2 | 4) {
1190            return Err(invalid(format!("block bits must be 1, 2 or 4, not {bits}")));
1191        }
1192        if self.metric != VMetric::Angular {
1193            return Err(invalid("neighbour blocks are angular only".into()));
1194        }
1195        let (n, dim, r) = (self.n, self.dim, self.degree);
1196        let rot = Rotation::new(dim, 0, BLOCK_SEED);
1197        let bd = rot.padded();
1198        let m = bd * usize::from(bits) / 4;
1199        let per_vertex = r.div_ceil(BLOCK);
1200        let record = block_record_bytes(m);
1201        let ids = &self.ids;
1202        let old = |v: usize| if ids.is_empty() { v } else { ids[v] as usize };
1203        let row = |v: usize| &g.vectors[old(v) * dim..(old(v) + 1) * dim];
1204        let mut blocks = Lines::zeroed(n * per_vertex * record);
1205        blocks
1206            .bytes_mut()
1207            .par_chunks_exact_mut(per_vertex * record)
1208            .enumerate()
1209            .for_each_init(
1210                || {
1211                    let list = Vec::<u32>::with_capacity(r);
1212                    (
1213                        vec![0f32; bd],
1214                        vec![0f32; dim],
1215                        vec![0f32; bd],
1216                        vec![0u8; BLOCK * m],
1217                        list,
1218                    )
1219                },
1220                |(xr, diff, rr, nibbles, list), (v, dst)| {
1221                    // Codes are of the residual `r = x_u - x_v` of each neighbour `u` to the
1222                    // vertex `v` (as in SymphonyQG), so the code error scales with `|r|`; `e`
1223                    // folds in the query-independent terms `<x_v, r> - s * <levels, rot(x_v)>`,
1224                    // which makes the estimate's error `<rot(q - x_v), code error>`, scaled by
1225                    // the query's distance to `v` as well:
1226                    // `<q, x_u - mean> = <q, x_v - mean> + s * <levels, rot(q)> + e`.
1227                    // Every buffer is overwritten before it is read: `apply` writes all of its
1228                    // output, `diff` is rewritten per neighbour, `nibbles` is cleared per block.
1229                    let xv = row(v);
1230                    rot.apply(xv, xr);
1231                    let mut scales = [0f32; BLOCK];
1232                    let mut offsets = [0f32; BLOCK];
1233                    list.clear();
1234                    self.for_each_neighbour(v, |u| list.push(u));
1235                    for (b, rec) in dst.chunks_exact_mut(record).enumerate() {
1236                        nibbles.fill(0);
1237                        scales.fill(0.0);
1238                        offsets.fill(0.0);
1239                        for i in 0..BLOCK {
1240                            let Some(&u) = list.get(b * BLOCK + i) else {
1241                                break;
1242                            };
1243                            let xu = row(u as usize);
1244                            for (d, (a, b)) in diff.iter_mut().zip(xu.iter().zip(xv)) {
1245                                *d = a - b;
1246                            }
1247                            rot.apply(diff, rr);
1248                            let nib = &mut nibbles[i * m..(i + 1) * m];
1249                            let s = encode_block_code(rr, bits, nib);
1250                            let dot_lv: f32 =
1251                                (0..bd).map(|j| block_level(nib, bits, j) * xr[j]).sum();
1252                            let dot_vr: f32 = xv.iter().zip(diff.iter()).map(|(a, b)| a * b).sum();
1253                            scales[i] = s;
1254                            offsets[i] = dot_vr - s * dot_lv;
1255                        }
1256                        let rec_u8: &mut [u8] = bytemuck::cast_slice_mut(rec);
1257                        pack32(nibbles, m, &mut rec_u8[..16 * m]);
1258                        for (i, (s, e)) in scales.iter().zip(&offsets).enumerate() {
1259                            let at = 16 * m + 4 * i;
1260                            rec_u8[at..at + 4].copy_from_slice(&s.to_le_bytes());
1261                            let at = 16 * m + 4 * BLOCK + 4 * i;
1262                            rec_u8[at..at + 4].copy_from_slice(&e.to_le_bytes());
1263                        }
1264                    }
1265                },
1266            );
1267        self.block_bits = bits;
1268        self.block_dim = bd;
1269        self.block_m = m;
1270        self.blocks_per_vertex = per_vertex;
1271        self.blocks = blocks;
1272        self.block_rot = Some(rot);
1273        self.block_kappa = self.measure_block_kappa(g);
1274        Ok(self)
1275    }
1276
1277    /// Unit row `v` (stored numbering) of the build data.
1278    fn unit_row<'g>(&self, g: &'g Graph, v: usize) -> &'g [f32] {
1279        let o = if self.ids.is_empty() {
1280            v
1281        } else {
1282            self.ids[v] as usize
1283        };
1284        &g.vectors[o * self.dim..(o + 1) * self.dim]
1285    }
1286
1287    /// Fits `block_kappa` by least squares on a deterministic sample: the squared error of
1288    /// the block estimate against the exact centred inner product, regressed on
1289    /// `s^2 * |q - x_v|^2` (index vectors as queries, random expanded vertices).
1290    fn measure_block_kappa(&self, g: &Graph) -> f32 {
1291        let (nq, nv) = (32.min(self.n), 64.min(self.n));
1292        let mut rng = SplitMix(BLOCK_SEED ^ 1);
1293        let mut s = self.searcher();
1294        let (mut num, mut den) = (0f64, 0f64);
1295        let (mut diag_sum2, mut diag_n) = (0f64, 0f64);
1296        let mut scan = [0u32; BLOCK];
1297        for _ in 0..nq {
1298            let qi = (rng.next_u64() % self.n as u64) as usize;
1299            let q = self.unit_row(g, qi).to_vec();
1300            s.scratch.qf[..self.dim].copy_from_slice(&q);
1301            let sq = self.encode_query(
1302                &s.scratch.qf,
1303                &mut s.scratch.qc,
1304                &mut s.scratch.qtmp,
1305                &mut s.scratch.qrot,
1306                &mut s.scratch.qs,
1307            );
1308            if sq == 0.0 {
1309                continue;
1310            }
1311            self.encode_block_query(&s.scratch.qf, sq, &mut s.scratch.bqrot, &mut s.scratch.bq);
1312            let centred = |x: &[f32]| -> f32 {
1313                q.iter()
1314                    .zip(x)
1315                    .zip(&self.mean)
1316                    .map(|((a, b), m)| a * (b - m))
1317                    .sum()
1318            };
1319            let mut list = Vec::with_capacity(self.degree);
1320            for _ in 0..nv {
1321                let v = (rng.next_u64() % self.n as u64) as usize;
1322                let xv = self.unit_row(g, v);
1323                let ip_v = centred(xv);
1324                let d2: f32 = q.iter().zip(xv).map(|(a, b)| (a - b) * (a - b)).sum();
1325                list.clear();
1326                self.for_each_neighbour(v, |u| list.push(u));
1327                for b in 0..self.blocks_per_vertex {
1328                    let (codes, scales, offsets) = self.block(v, b);
1329                    scan32(codes, &s.scratch.bq.lut, self.block_m, &mut scan);
1330                    for (i, &sum) in scan.iter().enumerate() {
1331                        let Some(&u) = list.get(b * BLOCK + i) else {
1332                            break;
1333                        };
1334                        let sc = le_f32(scales, i);
1335                        let est = ip_v
1336                            + sc * (s.scratch.bq.delta * sum as f32 + s.scratch.bq.min_sum)
1337                            + le_f32(offsets, i);
1338                        let exact = centred(self.unit_row(g, u as usize));
1339                        let err = f64::from(est - exact);
1340                        num += err * err;
1341                        den += f64::from(sc * sc * d2);
1342                        diag_sum2 += f64::from(exact - ip_v).powi(2);
1343                        diag_n += 1.0;
1344                    }
1345                }
1346            }
1347        }
1348        if crate::compatibility::var_os("WARM_BLOCK_DIAG").is_some() && diag_n > 0.0 {
1349            eprintln!(
1350                "block diag: bits {} kappa {:.4} rms_err {:.5} rms_residual_ip {:.5} pairs {}",
1351                self.block_bits,
1352                num / den.max(1e-30),
1353                (num / diag_n).sqrt(),
1354                (diag_sum2 / diag_n).sqrt(),
1355                diag_n
1356            );
1357        }
1358        if den > 0.0 {
1359            (num / den) as f32
1360        } else {
1361            0.0
1362        }
1363    }
1364
1365    /// True if the index carries neighbour-block codes.
1366    pub fn has_blocks(&self) -> bool {
1367        self.block_bits > 0
1368    }
1369
1370    /// Bytes of the neighbour-block codes and scales.
1371    pub fn block_bytes(&self) -> usize {
1372        self.blocks.len
1373    }
1374
1375    /// Block `b` of vertex `v`: its packed codes, its 32 scales and its 32 offsets (both as
1376    /// little-endian f32 bytes).
1377    #[inline]
1378    fn block(&self, v: usize, b: usize) -> (&[u8], &[u8], &[u8]) {
1379        let record = block_record_bytes(self.block_m);
1380        let at = (v * self.blocks_per_vertex + b) * record;
1381        let bytes = self.blocks.bytes();
1382        // SAFETY: `i8` and `u8` have the same size and alignment; the slice stays in bounds.
1383        let rec =
1384            unsafe { std::slice::from_raw_parts(bytes.as_ptr().add(at).cast::<u8>(), record) };
1385        let (codes, rest) = rec.split_at(16 * self.block_m);
1386        let (scales, offsets) = rest.split_at(4 * BLOCK);
1387        (codes, scales, offsets)
1388    }
1389
1390    #[inline]
1391    fn prefetch_blocks(&self, v: usize) {
1392        let record = block_record_bytes(self.block_m) * self.blocks_per_vertex;
1393        let at = v * record;
1394        prefetch_bytes(&self.blocks.bytes()[at..at + record]);
1395    }
1396
1397    /// Builds the block lookup table of a query (`qf` padded, `sq` its 8-bit code scale).
1398    fn encode_block_query(&self, qf: &[f32], sq: f32, qrot: &mut [f32], bq: &mut BlockQuery) {
1399        let rot = self.block_rot.as_ref().expect("blocks present");
1400        rot.apply(&qf[..self.dim], qrot);
1401        block_table(qrot, self.block_bits, self.block_m, &mut bq.table);
1402        bq.quantize(self.block_m);
1403        bq.inv_sq = if sq > 0.0 { 1.0 / sq } else { 1.0 };
1404    }
1405
1406    /// Standard deviation of the screen estimate's error against the 8-bit estimate over
1407    /// index vectors used as queries (deterministic sample), in estimate units.
1408    fn measure_screen_sigma(&self, g: &Graph) -> f32 {
1409        let (nq, nv) = (128.min(self.n), 1024.min(self.n));
1410        let mut rng = SplitMix(SCREEN_SEED ^ 1);
1411        let mut qc = vec![0i8; self.padded];
1412        let mut qtmp = vec![0i8; self.padded];
1413        let mut qf = vec![0f32; self.padded];
1414        let mut qrot = vec![0f32; self.screen_words * 64];
1415        let mut qs = QueryCode::new(self.screen_words);
1416        let (mut sum, mut sum2, mut count) = (0f64, 0f64, 0usize);
1417        for _ in 0..nq {
1418            let q = (rng.next_u64() % self.n as u64) as usize;
1419            qf[..self.dim].copy_from_slice(&g.vectors[q * self.dim..(q + 1) * self.dim]);
1420            let sq = self.encode_query(&qf, &mut qc, &mut qtmp, &mut qrot, &mut qs);
1421            if sq == 0.0 {
1422                continue;
1423            }
1424            for _ in 0..nv {
1425                let v = (rng.next_u64() % self.n as u64) as u32;
1426                let d = f64::from(self.screen_est(&qs, v) - self.est(&qc, query_sum(&qc), v, 1.0));
1427                sum += d;
1428                sum2 += d * d;
1429                count += 1;
1430            }
1431        }
1432        if count == 0 {
1433            return 0.0;
1434        }
1435        let mean = sum / count as f64;
1436        ((sum2 / count as f64) - mean * mean).max(0.0).sqrt() as f32
1437    }
1438
1439    /// Writes the traversal code of `qf` (padded) into `qc` and, with a screen, the rotated
1440    /// query scaled like the code into `qs`; returns the query's quantization scale.
1441    fn encode_query(
1442        &self,
1443        qf: &[f32],
1444        qc: &mut [i8],
1445        qtmp: &mut [i8],
1446        qrot: &mut [f32],
1447        qs: &mut QueryCode,
1448    ) -> f32 {
1449        let sq = if self.bits == 8 {
1450            quantize(qf, qc)
1451        } else {
1452            let s = quantize(qf, qtmp);
1453            interleave_for_i4(qtmp, qc);
1454            s
1455        };
1456        if let Some(rot) = &self.screen_rot {
1457            rot.apply(&qf[..self.dim], qrot);
1458            if sq > 0.0 {
1459                // The 8-bit estimate drops the query scale, so the screen works in the same
1460                // units: rotated query divided by s_q.
1461                qrot.iter_mut().for_each(|v| *v /= sq);
1462            }
1463            qs.encode(qrot);
1464        }
1465        sq
1466    }
1467
1468    /// True if the index carries the 1-bit screen.
1469    pub fn has_screen(&self) -> bool {
1470        !self.screen.is_empty()
1471    }
1472
1473    /// Bytes of the screen rows.
1474    pub fn screen_bytes(&self) -> usize {
1475        self.screen.len() * 8
1476    }
1477
1478    /// 1-bit estimate in the units of [`VGraph::est`]: `-f * sum_i sign_i * q_i`, with the
1479    /// signed sum from the query's bit planes as `2 * (lo * pop + delta * raw) - sum(q)`.
1480    #[inline]
1481    fn screen_est(&self, qs: &QueryCode, v: u32) -> f32 {
1482        let w = self.screen_words;
1483        let row = &self.screen[v as usize * (w + 1)..(v as usize + 1) * (w + 1)];
1484        let raw = raw_one(&row[..w], &qs.planes, w) as f32;
1485        let meta = row[w];
1486        let f = f32::from_bits(meta as u32);
1487        let pop = (meta >> 32) as f32;
1488        -(f * (2.0 * (qs.lo * pop + qs.delta * raw) - qs.sum))
1489    }
1490
1491    #[inline]
1492    fn prefetch_screen(&self, v: u32) {
1493        let w = self.screen_words;
1494        prefetch(&self.screen[v as usize * (w + 1)..(v as usize + 1) * (w + 1)]);
1495    }
1496
1497    pub fn len(&self) -> usize {
1498        self.n
1499    }
1500
1501    pub fn is_empty(&self) -> bool {
1502        self.n == 0
1503    }
1504
1505    /// Huge pages behind this index's code rows (an upper bound, see the `/proc/self/smaps`
1506    /// note on [`huge_page_status`]); `None` where `/proc/self/smaps` is unreadable.
1507    pub fn code_rows_huge_bytes(&self) -> Option<u64> {
1508        let rows = self.codes.bytes();
1509        huge_bytes_behind(rows.as_ptr().cast(), rows.len())
1510    }
1511
1512    pub fn has_residual(&self) -> bool {
1513        self.res_bits > 0
1514    }
1515
1516    pub fn code_bits(&self) -> u8 {
1517        self.bits
1518    }
1519
1520    pub fn residual_bits(&self) -> u8 {
1521        self.res_bits
1522    }
1523
1524    /// Logical payload bytes, residual and id map included. See [`VGraphMemoryBreakdown`]
1525    /// for the accounting boundary and exclusions.
1526    pub fn index_bytes(&self) -> usize {
1527        self.memory_breakdown().total_bytes()
1528    }
1529
1530    pub fn memory_breakdown(&self) -> VGraphMemoryBreakdown {
1531        VGraphMemoryBreakdown {
1532            adjacency_ids_bytes: self.adj.len() * 4 + self.adj3.len() + self.adjv.len(),
1533            adjacency_offsets_bytes: (self.adj_off.len() + self.adjv_off.len()) * 4,
1534            traversal_codes_bytes: self.n * self.stride,
1535            traversal_scales_bytes: self.scales.len() * 4,
1536            residual_codes_bytes: self.residual.len(),
1537            residual_scales_bytes: self.res_scales.len() * 4,
1538            mean_bytes: self.mean.len() * 4,
1539            norms_bytes: self.half_norms.len() * 4,
1540            original_ids_bytes: self.ids.len() * 4,
1541            screen_bytes: self.screen_bytes(),
1542            neighbor_blocks_bytes: self.block_bytes(),
1543            upper_layers_bytes: upper_bytes(&self.upper_ids, &self.layers),
1544        }
1545    }
1546
1547    /// Bytes of the residual codes and their scales, which only `rerank` reads.
1548    pub fn residual_bytes(&self) -> usize {
1549        self.residual.len() + self.res_scales.len() * 4
1550    }
1551
1552    #[inline]
1553    fn code(&self, v: u32) -> &[i8] {
1554        let at = v as usize * self.stride;
1555        &self.codes.bytes()[at..at + self.cbytes]
1556    }
1557
1558    /// Appends the bottom-layer neighbours of stored vertex `v`, in list order, to `out`.
1559    pub(super) fn neighbours_into(&self, v: usize, out: &mut Vec<u32>) {
1560        self.for_each_neighbour(v, |u| out.push(u));
1561    }
1562
1563    /// The stored id of the bottom-layer entry vertex.
1564    pub(super) fn entry_vertex(&self) -> u32 {
1565        self.entry
1566    }
1567
1568    /// Calls `f` on every neighbour of `v`, in list order.
1569    #[inline]
1570    fn for_each_neighbour(&self, v: usize, mut f: impl FnMut(u32)) {
1571        if !self.adjv_off.is_empty() {
1572            return self.for_each_delta_neighbour(v, f);
1573        }
1574        let (start, end) = self.adjacency_span(v);
1575        if self.adj3.is_empty() {
1576            for &u in &self.adj[start..end] {
1577                if u == NONE {
1578                    break;
1579                }
1580                f(u);
1581            }
1582        } else {
1583            for b in self.adj3[3 * start..3 * end].as_chunks::<3>().0 {
1584                let u = u32::from_le_bytes([b[0], b[1], b[2], 0]);
1585                if u == NONE3 {
1586                    break;
1587                }
1588                f(u);
1589            }
1590        }
1591    }
1592
1593    #[cfg(feature = "experimental-neighbor-admission")]
1594    #[inline(always)]
1595    fn admit_neighbour<const SCREEN: bool>(
1596        &self,
1597        u: u32,
1598        visited: &mut Marks,
1599        fresh: &mut Vec<u32>,
1600    ) {
1601        if visited.insert(u) {
1602            if SCREEN {
1603                self.prefetch_screen(u);
1604            } else {
1605                prefetch_bytes(self.code(u));
1606                self.prefetch_scale(u);
1607            }
1608            fresh.push(u);
1609        }
1610    }
1611
1612    #[cfg(feature = "experimental-neighbor-admission")]
1613    #[inline(always)]
1614    fn admit_neighbours<const PROFILE: bool, const SCREEN: bool>(
1615        &self,
1616        v: usize,
1617        cap: usize,
1618        visited: &mut Marks,
1619        fresh: &mut Vec<u32>,
1620    ) -> usize {
1621        let mut neighbors_read = 0;
1622        if cap == 0 {
1623            self.for_each_neighbour(v, |u| {
1624                if PROFILE {
1625                    neighbors_read += 1;
1626                }
1627                self.admit_neighbour::<SCREEN>(u, visited, fresh);
1628            });
1629        } else {
1630            let mut seen = 0usize;
1631            self.for_each_neighbour(v, |u| {
1632                seen += 1;
1633                if seen <= cap {
1634                    self.admit_neighbour::<SCREEN>(u, visited, fresh);
1635                }
1636            });
1637        }
1638        neighbors_read
1639    }
1640
1641    #[inline]
1642    fn adjacency_span(&self, v: usize) -> (usize, usize) {
1643        if self.adj_off.is_empty() {
1644            (v * self.degree, (v + 1) * self.degree)
1645        } else {
1646            (self.adj_off[v] as usize, self.adj_off[v + 1] as usize)
1647        }
1648    }
1649
1650    /// Removes unused adjacency slots without changing neighbors or their order.
1651    /// Saving expands the lists into the original persisted format; reopening uses that
1652    /// format until this method is called again. Delta-coded lists already have exact spans.
1653    pub fn compact_adjacency(&mut self) -> Result<(), WarmError> {
1654        if !self.adj_off.is_empty() || !self.adjv_off.is_empty() {
1655            return Ok(());
1656        }
1657        let mut offsets = Vec::new();
1658        offsets
1659            .try_reserve_exact(self.n + 1)
1660            .map_err(|_| invalid("cannot allocate adjacency offsets".into()))?;
1661        offsets.push(0u32);
1662        for v in 0..self.n {
1663            let mut count = 0u32;
1664            self.for_each_neighbour(v, |_| count += 1);
1665            let end = offsets[v]
1666                .checked_add(count)
1667                .ok_or_else(|| invalid("compact adjacency exceeds u32 offsets".into()))?;
1668            offsets.push(end);
1669        }
1670        let edges = offsets[self.n] as usize;
1671        let width = if self.adj3.is_empty() { 4 } else { 3 };
1672        if (self.n * self.degree - edges).saturating_mul(width) <= offsets.len() * 4 {
1673            return Ok(());
1674        }
1675        if self.adj3.is_empty() {
1676            let mut compact = Vec::new();
1677            compact
1678                .try_reserve_exact(edges)
1679                .map_err(|_| invalid("cannot allocate compact adjacency".into()))?;
1680            for v in 0..self.n {
1681                self.for_each_neighbour(v, |u| compact.push(u));
1682            }
1683            self.adj = compact.into();
1684        } else {
1685            let bytes = edges
1686                .checked_mul(3)
1687                .ok_or_else(|| invalid("compact adjacency size overflows".into()))?;
1688            let mut compact = Vec::new();
1689            compact
1690                .try_reserve_exact(bytes)
1691                .map_err(|_| invalid("cannot allocate compact adjacency".into()))?;
1692            for v in 0..self.n {
1693                self.for_each_neighbour(v, |u| compact.extend_from_slice(&u.to_le_bytes()[..3]));
1694            }
1695            self.adj3 = compact.into();
1696        }
1697        self.adj_off = offsets;
1698        Ok(())
1699    }
1700
1701    /// [`VGraph::for_each_neighbour`] over a delta-coded list.
1702    #[cold]
1703    #[inline(never)]
1704    fn for_each_delta_neighbour(&self, v: usize, mut f: impl FnMut(u32)) {
1705        let span = &self.adjv[self.adjv_off[v] as usize..self.adjv_off[v + 1] as usize];
1706        let mut at = 0;
1707        let mut prev = v as u32;
1708        let mut first = true;
1709        while at < span.len() {
1710            let (delta, used) = get_varint(&span[at..]);
1711            at += used;
1712            let u = if first {
1713                (i64::from(prev) + unzigzag(delta)) as u32
1714            } else {
1715                prev + delta as u32 + 1
1716            };
1717            first = false;
1718            prev = u;
1719            f(u);
1720        }
1721    }
1722
1723    #[inline]
1724    fn prefetch_neighbours(&self, v: usize) {
1725        let (start, end) = self.adjacency_span(v);
1726        if !self.adjv_off.is_empty() {
1727            prefetch_u8(&self.adjv[self.adjv_off[v] as usize..self.adjv_off[v + 1] as usize]);
1728        } else if self.adj3.is_empty() {
1729            prefetch_ids(&self.adj[start..end]);
1730        } else {
1731            prefetch_u8(&self.adj3[3 * start..3 * end]);
1732        }
1733    }
1734
1735    /// Prefetches the per-vertex terms of an estimate of `v`: its scale and, for Euclidean,
1736    /// its norm term (each its own array, one line per random vertex).
1737    #[inline]
1738    fn prefetch_scale(&self, v: u32) {
1739        prefetch_f32(&self.scales[v as usize]);
1740        if self.metric == VMetric::Euclidean {
1741            prefetch_f32(&self.half_norms[v as usize]);
1742        }
1743    }
1744
1745    /// Prefetches what [`VGraph::fine`] reads for `v`.
1746    #[inline]
1747    fn prefetch_rerank_rows(&self, v: u32) {
1748        prefetch_bytes(self.code(v));
1749        self.prefetch_scale(v);
1750        if self.has_residual() {
1751            let v = v as usize;
1752            prefetch_bytes(&self.residual[v * self.res_stride..(v + 1) * self.res_stride]);
1753            prefetch_f32(&self.res_scales[v]);
1754        }
1755    }
1756
1757    #[inline]
1758    fn original(&self, v: u32) -> u32 {
1759        if self.ids.is_empty() {
1760            v
1761        } else {
1762            self.ids[v as usize]
1763        }
1764    }
1765
1766    /// Estimated negative inner product with the query code `qc`, whose coordinate sum is
1767    /// `qsum` (up to the per-query terms
1768    /// `<q, mean>` and `s_q`, which do not change the order). Euclidean adds `w |x|^2 / 2`
1769    /// with `w = 1 / s_q`, the order of `|q - x|^2` in the same units.
1770    #[inline]
1771    fn est(&self, qc: &[i8], qsum: i32, v: u32, w: f32) -> f32 {
1772        let dot = if self.bits == 8 {
1773            dot_i8_query(qc, self.code(v), qsum)
1774        } else {
1775            dot_i4_query(qc, self.code(v), qsum)
1776        };
1777        self.order(v, self.scales[v as usize] * dot as f32, w)
1778    }
1779
1780    /// The distance order of `v` from its estimated (or re-ranked) inner product `ip` in
1781    /// units of `1 / w`: `-ip` for angular, `w |x|^2 / 2 - ip` for Euclidean.
1782    #[inline]
1783    fn order(&self, v: u32, ip: f32, w: f32) -> f32 {
1784        if self.metric == VMetric::Euclidean {
1785            self.half_norms[v as usize] * w - ip
1786        } else {
1787            -ip
1788        }
1789    }
1790
1791    /// [`VGraph::est`] of four vertices with shared query loads; the same integer sums, so
1792    /// the same values.
1793    #[inline]
1794    fn est4<const QUERY_X4: bool>(&self, qc: &[i8], qsum: i32, v: [u32; 4], w: f32) -> [f32; 4] {
1795        let dots = if self.bits == 8 {
1796            if QUERY_X4 {
1797                dot_i8_query_x4(qc, v.map(|u| self.code(u)), qsum)
1798            } else {
1799                v.map(|u| dot_i8_query(qc, self.code(u), qsum))
1800            }
1801        } else {
1802            dot_i4_query_x4(qc, v.map(|u| self.code(u)), qsum)
1803        };
1804        std::array::from_fn(|i| self.order(v[i], self.scales[v[i] as usize] * dots[i] as f32, w))
1805    }
1806
1807    /// Re-rank score with the float query: negative estimated inner product with the
1808    /// reconstructed centred vector (Euclidean: plus `|x|^2 / 2`).
1809    fn fine(&self, qf: &[f32], v: u32) -> f32 {
1810        let primary = if self.bits == 8 {
1811            dot_f32_i8(qf, self.code(v))
1812        } else {
1813            dot_f32_i4(qf, self.code(v))
1814        };
1815        let mut ip = self.scales[v as usize] * primary;
1816        if self.has_residual() {
1817            let v = v as usize;
1818            let r = &self.residual[v * self.res_stride..(v + 1) * self.res_stride];
1819            let d = if self.res_bits == 8 {
1820                dot_f32_i8(qf, r)
1821            } else {
1822                dot_f32_i4(qf, r)
1823            };
1824            ip += self.res_scales[v] * d;
1825        }
1826        self.order(v, ip, 1.0)
1827    }
1828
1829    /// Greedy descent through the upper layers on estimated distances; returns the bottom-layer
1830    /// entry vertex and the number of estimates.
1831    fn descend(&self, qc: &[i8], qsum: i32, w: f32, reservoir: &mut Reservoir) -> (u32, usize) {
1832        if self.layers.is_empty() {
1833            return (self.entry, 0);
1834        }
1835        let mut p = 0u32;
1836        let mut best = self.est(qc, qsum, self.upper_ids[0], w);
1837        reservoir.offer(key(best, self.upper_ids[0]));
1838        let mut evals = 1;
1839        for layer in self.layers.iter().rev() {
1840            loop {
1841                // The list's rows are fetched together: each step's estimates are independent
1842                // of one another and only the step's winner decides the next list.
1843                let list = &layer[p as usize];
1844                for &u in list {
1845                    let id = self.upper_ids[u as usize];
1846                    prefetch_bytes(self.code(id));
1847                    self.prefetch_scale(id);
1848                }
1849                let mut moved = false;
1850                for &u in list {
1851                    let id = self.upper_ids[u as usize];
1852                    let d = self.est(qc, qsum, id, w);
1853                    reservoir.offer(key(d, id));
1854                    evals += 1;
1855                    if d < best {
1856                        best = d;
1857                        p = u;
1858                        moved = true;
1859                    }
1860                }
1861                if !moved {
1862                    break;
1863                }
1864            }
1865        }
1866        (self.upper_ids[p as usize], evals)
1867    }
1868
1869    /// A searcher with the knobs of the environment ([`SearchKnobs::from_env`]).
1870    pub fn searcher(&self) -> VSearcher<'_> {
1871        self.searcher_with(SearchKnobs::from_env())
1872    }
1873
1874    /// A searcher with explicit [`SearchKnobs`].
1875    pub fn searcher_with(&self, knobs: SearchKnobs) -> VSearcher<'_> {
1876        VSearcher {
1877            index: self,
1878            scratch: SearchScratch {
1879                #[cfg(not(target_arch = "wasm32"))]
1880                dim: self.dim,
1881                knobs,
1882                visited: Marks::new(self.n, self.has_blocks()),
1883                qf: vec![0.0; self.padded],
1884                qc: vec![0; self.padded],
1885                qsum: 0,
1886                qtmp: vec![0; self.padded],
1887                qrot: vec![0.0; self.screen_words * 64],
1888                qs: QueryCode::new(self.screen_words.max(1)),
1889                bqrot: vec![0.0; self.block_dim],
1890                bq: BlockQuery::new(self.block_m),
1891                scan: [0; BLOCK],
1892                fresh: Vec::with_capacity(self.degree),
1893                survivors: Vec::with_capacity(self.degree),
1894                screened: 0,
1895                pool: Pool {
1896                    keys: Vec::new(),
1897                    expanded: Vec::new(),
1898                    #[cfg(feature = "experimental-batch-pool")]
1899                    merge_keys: Vec::new(),
1900                    #[cfg(feature = "experimental-batch-pool")]
1901                    merge_expanded: Vec::new(),
1902                    cap: 1,
1903                    cursor: 0,
1904                },
1905                expanded: BinaryHeap::new(),
1906                reservoir: Reservoir::default(),
1907                keys: Vec::new(),
1908                results: Vec::new(),
1909            },
1910        }
1911    }
1912
1913    #[cfg(not(target_arch = "wasm32"))]
1914    pub(super) fn update_encoding(&self) -> super::update::Encoding {
1915        super::update::Encoding {
1916            bits: self.bits,
1917            residual_bits: self.res_bits,
1918            align: self.stride != self.cbytes,
1919            reorder: !self.ids.is_empty(),
1920            id_bytes: if !self.adjv_off.is_empty() {
1921                0
1922            } else if !self.adj3.is_empty() {
1923                3
1924            } else {
1925                4
1926            },
1927            screen: self.has_screen(),
1928            block_bits: self.block_bits,
1929        }
1930    }
1931
1932    #[cfg(not(target_arch = "wasm32"))]
1933    pub(super) fn search_reusing(
1934        &self,
1935        cache: &mut Option<SearchScratch>,
1936        query: &[f32],
1937        k: usize,
1938        params: VSearchParams,
1939        out: &mut Vec<u32>,
1940    ) -> Result<usize, WarmError> {
1941        let compatible = cache.as_ref().is_some_and(|s| {
1942            s.dim == self.dim
1943                && s.visited.n == self.n
1944                && s.qf.len() == self.padded
1945                && s.qrot.len() == self.screen_words * 64
1946                && s.bqrot.len() == self.block_dim
1947                && s.bq.table.len() == 16 * self.block_m
1948                && s.visited.traces.len() == if self.has_blocks() { self.n } else { 0 }
1949        });
1950        if !compatible {
1951            *cache = Some(self.searcher().scratch);
1952        }
1953        cache
1954            .as_mut()
1955            .expect("initialized search scratch")
1956            .search_validated(self, query, k, params, out)
1957    }
1958}
1959
1960/// Visited marks of one search: a bit per vertex, with an 8-bit epoch for the trace records of
1961/// the neighbour-block search.
1962struct Marks {
1963    /// One bit per vertex. The words with a bit set since the last [`Marks::next`] are listed
1964    /// in `touched` and cleared then, so a search's visited footprint is `n / 8` bytes (a
1965    /// million vertices fit the second-level cache) and clearing costs the words it set.
1966    words: Vec<u64>,
1967    touched: Vec<u32>,
1968    n: usize,
1969    epoch: u8,
1970    /// With neighbour blocks, one record per vertex instead of the bit.
1971    traces: Vec<Trace>,
1972}
1973
1974/// What one search knows about a vertex met through neighbour blocks: the precision-weighted
1975/// sum of the block estimates it received (one per expanded vertex listing it), their
1976/// weight, and whether it was expanded itself.
1977#[derive(Clone, Copy, Default)]
1978struct Trace {
1979    epoch: u8,
1980    expanded: bool,
1981    w: f32,
1982    wsum: f32,
1983    /// The vertex's key in the pool, or `u64::MAX` when it holds none.
1984    key: u64,
1985}
1986
1987impl Marks {
1988    /// Marks for `n` vertices, with trace records when `traces` is set.
1989    fn new(n: usize, traces: bool) -> Self {
1990        Self {
1991            words: vec![0; n.div_ceil(64)],
1992            touched: Vec::new(),
1993            n,
1994            epoch: 0,
1995            traces: vec![Trace::default(); if traces { n } else { 0 }],
1996        }
1997    }
1998
1999    fn next(&mut self) {
2000        for &w in &self.touched {
2001            self.words[w as usize] = 0;
2002        }
2003        self.touched.clear();
2004        self.epoch = self.epoch.wrapping_add(1);
2005        if self.epoch == 0 {
2006            self.traces.fill(Trace::default());
2007            self.epoch = 1;
2008        }
2009    }
2010
2011    /// The trace of `id` in this epoch, and whether it was just created (first encounter).
2012    #[inline]
2013    fn trace(&mut self, id: u32) -> (bool, &mut Trace) {
2014        let t = &mut self.traces[id as usize];
2015        let fresh = t.epoch != self.epoch;
2016        if fresh {
2017            *t = Trace {
2018                epoch: self.epoch,
2019                key: u64::MAX,
2020                ..Trace::default()
2021            };
2022        }
2023        (fresh, t)
2024    }
2025
2026    /// True if `id` was visited in this search (no mark is set).
2027    #[inline]
2028    fn seen(&self, id: u32) -> bool {
2029        self.words[(id >> 6) as usize] >> (id & 63) & 1 == 1
2030    }
2031
2032    /// True if `id` was not yet visited in this search.
2033    #[inline]
2034    fn insert(&mut self, id: u32) -> bool {
2035        let w = &mut self.words[(id >> 6) as usize];
2036        let bit = 1u64 << (id & 63);
2037        if *w & bit != 0 {
2038            return false;
2039        }
2040        if *w == 0 {
2041            self.touched.push(id >> 6);
2042        }
2043        *w |= bit;
2044        true
2045    }
2046}
2047
2048/// Best-first candidate pool as one sorted array: the `ef` best keys seen, nearest first,
2049/// each with an expanded flag, and a cursor at the nearest unexpanded entry. One insert is a
2050/// binary search and a short move; one expansion step is a scan from the cursor; the pool
2051/// rule ("the nearest unexpanded candidate is among the `ef` best seen") holds by
2052/// construction, so the search ends when no unexpanded entry is left. Two binary heaps
2053/// (queue and bounded pool) cost two sifts per admitted candidate and a pop per expansion,
2054/// and their results came back unsorted; the array is cache-resident at any practical `ef`.
2055struct Pool {
2056    keys: Vec<u64>,
2057    /// Expanded flags, bit `i` for `keys[i]`, in words; bits at `keys.len()` and above are 0.
2058    expanded: Vec<u64>,
2059    #[cfg(feature = "experimental-batch-pool")]
2060    merge_keys: Vec<u64>,
2061    #[cfg(feature = "experimental-batch-pool")]
2062    merge_expanded: Vec<u64>,
2063    cap: usize,
2064    cursor: usize,
2065}
2066
2067impl Pool {
2068    fn reset(&mut self, cap: usize) {
2069        self.keys.clear();
2070        self.expanded.clear();
2071        self.cap = cap.max(1);
2072        self.cursor = 0;
2073    }
2074
2075    fn is_full(&self) -> bool {
2076        self.keys.len() >= self.cap
2077    }
2078
2079    /// The worst key held, once the pool is full.
2080    fn worst(&self) -> Option<u64> {
2081        self.is_full().then(|| self.keys[self.keys.len() - 1])
2082    }
2083
2084    #[cfg(test)]
2085    fn is_expanded(&self, i: usize) -> bool {
2086        self.expanded[i / 64] >> (i % 64) & 1 == 1
2087    }
2088
2089    /// Opens an unexpanded flag at `pos` for `len` flags, shifting those at `pos..len` up one.
2090    #[inline]
2091    fn insert_flag(&mut self, pos: usize, len: usize) {
2092        let words = (len + 1).div_ceil(64);
2093        if self.expanded.len() < words {
2094            self.expanded.resize(words, 0);
2095        }
2096        let w = pos / 64;
2097        for i in (w + 1..words).rev() {
2098            self.expanded[i] = self.expanded[i] << 1 | self.expanded[i - 1] >> 63;
2099        }
2100        let low = (1u64 << (pos % 64)) - 1;
2101        let x = self.expanded[w];
2102        self.expanded[w] = x & low | (x & !low) << 1;
2103    }
2104
2105    /// Closes the flag at `pos` of `len` flags, shifting those above it down one.
2106    #[inline]
2107    fn remove_flag(&mut self, pos: usize, len: usize) {
2108        let words = len.div_ceil(64);
2109        let w = pos / 64;
2110        let low = (1u64 << (pos % 64)) - 1;
2111        let carry = |e: &[u64], i: usize| if i + 1 < words { e[i + 1] << 63 } else { 0 };
2112        let x = self.expanded[w];
2113        self.expanded[w] = x & low | (x >> 1) & !low | carry(&self.expanded, w);
2114        for i in w + 1..words {
2115            self.expanded[i] = self.expanded[i] >> 1 | carry(&self.expanded, i);
2116        }
2117    }
2118
2119    /// Admits `key` unless the pool is full of better ones; returns its position and the
2120    /// key evicted to make room.
2121    #[inline]
2122    fn insert(&mut self, key: u64) -> Option<(usize, Option<u64>)> {
2123        let n = self.keys.len();
2124        if n >= self.cap && key >= self.keys[n - 1] {
2125            return None;
2126        }
2127        let pos = self.keys.partition_point(|&t| t < key);
2128        let mut evicted = None;
2129        let mut len = n;
2130        if n >= self.cap {
2131            evicted = self.keys.pop();
2132            len -= 1;
2133            self.expanded[len / 64] &= !(1u64 << (len % 64));
2134        }
2135        self.keys.insert(pos, key);
2136        self.insert_flag(pos, len);
2137        if pos < self.cursor {
2138            self.cursor = pos;
2139        }
2140        Some((pos, evicted))
2141    }
2142
2143    /// `pending` has distinct vertex IDs absent from this pool, enforced by visited marks
2144    /// before row-code scoring. Re-admission through the fused path uses `insert` instead.
2145    #[cfg(feature = "experimental-batch-pool")]
2146    fn insert_batch(&mut self, pending: &mut Vec<u64>, k: usize) -> bool {
2147        debug_assert!(pending.iter().enumerate().all(|(position, &candidate)| {
2148            let id = key_id(candidate);
2149            pending[..position].iter().all(|&other| key_id(other) != id)
2150                && self.keys.iter().all(|&other| key_id(other) != id)
2151        }));
2152        if let Some(worst) = self.worst() {
2153            pending.retain(|&key| key < worst);
2154        }
2155        if pending.len() < 8 {
2156            let mut improved = false;
2157            for &key in pending.iter() {
2158                if let Some((pos, _)) = self.insert(key) {
2159                    improved |= pos < k;
2160                }
2161            }
2162            return improved;
2163        }
2164        pending.sort_unstable();
2165        // Sequential insertions reset patience if any new key enters the initial top-k.
2166        let improved = k > 0 && (self.keys.len() < k || pending[0] <= self.keys[k - 1]);
2167        let len = self.cap.min(self.keys.len().saturating_add(pending.len()));
2168        self.merge_keys.clear();
2169        self.merge_keys.reserve(len);
2170        self.merge_expanded.clear();
2171        self.merge_expanded.resize(len.div_ceil(64), 0);
2172        let (mut old, mut new) = (0, 0);
2173        for pos in 0..len {
2174            if new < pending.len() && (old == self.keys.len() || pending[new] <= self.keys[old]) {
2175                self.merge_keys.push(pending[new]);
2176                self.cursor = self.cursor.min(pos);
2177                new += 1;
2178            } else {
2179                self.merge_keys.push(self.keys[old]);
2180                let flag = self.expanded[old / 64] >> (old % 64) & 1;
2181                self.merge_expanded[pos / 64] |= flag << (pos % 64);
2182                old += 1;
2183            }
2184        }
2185        std::mem::swap(&mut self.keys, &mut self.merge_keys);
2186        std::mem::swap(&mut self.expanded, &mut self.merge_expanded);
2187        improved
2188    }
2189
2190    /// Drops `key` if present (a re-admitted vertex's older entry).
2191    fn remove(&mut self, key: u64) {
2192        let pos = self.keys.partition_point(|&t| t < key);
2193        if pos < self.keys.len() && self.keys[pos] == key {
2194            let len = self.keys.len();
2195            self.keys.remove(pos);
2196            self.remove_flag(pos, len);
2197            if pos < self.cursor {
2198                self.cursor -= 1;
2199            }
2200        }
2201    }
2202
2203    /// Index of the `n`-th (from 0) unexpanded key at or after the cursor.
2204    #[inline]
2205    fn unexpanded(&self, mut n: usize) -> Option<usize> {
2206        let len = self.keys.len();
2207        let mut i = self.cursor;
2208        while i < len {
2209            let free = !self.expanded[i / 64] >> (i % 64);
2210            if free == 0 {
2211                i = (i / 64 + 1) * 64;
2212                continue;
2213            }
2214            i += free.trailing_zeros() as usize;
2215            if i >= len {
2216                break;
2217            }
2218            if n == 0 {
2219                return Some(i);
2220            }
2221            n -= 1;
2222            i += 1;
2223        }
2224        None
2225    }
2226
2227    /// The nearest unexpanded key, now marked expanded.
2228    #[inline]
2229    fn next(&mut self) -> Option<u64> {
2230        match self.unexpanded(0) {
2231            Some(i) => {
2232                self.expanded[i / 64] |= 1u64 << (i % 64);
2233                self.cursor = i;
2234                Some(self.keys[i])
2235            }
2236            None => {
2237                self.cursor = self.keys.len();
2238                None
2239            }
2240        }
2241    }
2242
2243    /// The key [`Pool::next`] would return.
2244    #[inline]
2245    fn peek_next(&self) -> Option<u64> {
2246        self.peek_unexpanded(0)
2247    }
2248
2249    /// The `n`-th (from 0) unexpanded key in pool order.
2250    #[inline]
2251    fn peek_unexpanded(&self, n: usize) -> Option<u64> {
2252        self.unexpanded(n).map(|i| self.keys[i])
2253    }
2254}
2255
2256#[derive(Default)]
2257struct Reservoir {
2258    keys: Vec<u64>,
2259    cap: usize,
2260}
2261
2262impl Reservoir {
2263    fn reset(&mut self, cap: usize) {
2264        self.keys.clear();
2265        self.cap = cap;
2266    }
2267
2268    #[inline]
2269    fn offer(&mut self, key: u64) {
2270        if self.cap == 0 || (self.keys.len() == self.cap && key >= self.keys[self.keys.len() - 1]) {
2271            return;
2272        }
2273        let pos = match self.keys.binary_search(&key) {
2274            Ok(_) => return,
2275            Err(pos) => pos,
2276        };
2277        if self.keys.len() == self.cap {
2278            self.keys.pop();
2279        }
2280        self.keys.insert(pos, key);
2281    }
2282}
2283
2284/// Per-thread search state for one [`VGraph`].
2285pub struct VSearcher<'a> {
2286    index: &'a VGraph,
2287    scratch: SearchScratch,
2288}
2289
2290pub(super) struct SearchScratch {
2291    #[cfg(not(target_arch = "wasm32"))]
2292    dim: usize,
2293    visited: Marks,
2294    qf: Vec<f32>,
2295    /// Query code in the traversal kernel's layout, and its coordinate sum.
2296    qc: Vec<i8>,
2297    qsum: i32,
2298    qtmp: Vec<i8>,
2299    /// Rotated, scaled query and its bit planes for the screen.
2300    qrot: Vec<f32>,
2301    qs: QueryCode,
2302    /// Query rotated for the neighbour blocks, its lookup table, and one block's sums.
2303    bqrot: Vec<f32>,
2304    bq: BlockQuery,
2305    scan: [u32; BLOCK],
2306    fresh: Vec<u32>,
2307    /// Fresh neighbours that passed the screen.
2308    survivors: Vec<u32>,
2309    /// Neighbours the screen skipped in the last search.
2310    screened: usize,
2311    /// The `ef` best candidates seen, with their expansion state.
2312    pool: Pool,
2313    /// With blocks: the best `max(k, rerank)` expanded vertices by their own-code estimate,
2314    /// farthest on top. Block estimates only order the exploration; results come from here.
2315    expanded: BinaryHeap<u64>,
2316    reservoir: Reservoir,
2317    keys: Vec<u64>,
2318    results: Vec<(f32, u32)>,
2319    knobs: SearchKnobs,
2320}
2321
2322impl VSearcher<'_> {
2323    /// Neighbours the screen skipped (no 8-bit estimate) in the last search.
2324    pub fn last_screened(&self) -> usize {
2325        self.scratch.screened
2326    }
2327
2328    /// The knobs this searcher runs with.
2329    pub fn knobs(&self) -> SearchKnobs {
2330        self.scratch.knobs
2331    }
2332
2333    /// Searches this index with the searcher's reusable buffers.
2334    /// Returns `min(k, n)` distinct neighbors, restarting from an unseen vertex only when
2335    /// traversal would otherwise return too few candidates.
2336    #[inline]
2337    pub fn search(
2338        &mut self,
2339        query: &[f32],
2340        k: usize,
2341        params: VSearchParams,
2342        out: &mut Vec<u32>,
2343    ) -> Result<usize, WarmError> {
2344        self.scratch.search(self.index, query, k, params, out)
2345    }
2346}
2347
2348impl SearchScratch {
2349    /// Row estimates of the admitted neighbours in `fresh`, inserted into the pool in list
2350    /// order (four rows at a time with `batch_est`; the same values either way); returns the
2351    /// number of estimates. `stale` resets on an insert among the `k` best when `patience`
2352    /// is on.
2353    #[inline]
2354    fn score_fresh<const BATCH_POOL: bool, const QUERY_X4: bool>(
2355        &mut self,
2356        g: &VGraph,
2357        k: usize,
2358        patience: usize,
2359        stale: &mut usize,
2360        w: f32,
2361    ) -> usize {
2362        #[cfg(feature = "experimental-batch-pool")]
2363        if BATCH_POOL {
2364            return self.score_fresh_batch::<QUERY_X4>(g, k, patience, stale, w);
2365        }
2366        let qc = &self.qc[..];
2367        let n = self.fresh.len();
2368        let mut i = 0;
2369        if self.knobs.batch_est || (QUERY_X4 && g.bits == 8) {
2370            while i + 4 <= n {
2371                let ids = [
2372                    self.fresh[i],
2373                    self.fresh[i + 1],
2374                    self.fresh[i + 2],
2375                    self.fresh[i + 3],
2376                ];
2377                for (e, u) in g
2378                    .est4::<QUERY_X4>(qc, self.qsum, ids, w)
2379                    .into_iter()
2380                    .zip(ids)
2381                {
2382                    let kk = key(e, u);
2383                    self.reservoir.offer(kk);
2384                    if let Some((pos, _)) = self.pool.insert(kk) {
2385                        if patience > 0 && pos < k {
2386                            *stale = 0;
2387                        }
2388                    }
2389                }
2390                i += 4;
2391            }
2392        }
2393        for &u in &self.fresh[i..] {
2394            let kk = key(g.est(qc, self.qsum, u, w), u);
2395            self.reservoir.offer(kk);
2396            if let Some((pos, _)) = self.pool.insert(kk) {
2397                if patience > 0 && pos < k {
2398                    *stale = 0;
2399                }
2400            }
2401        }
2402        n
2403    }
2404
2405    #[cfg(feature = "experimental-batch-pool")]
2406    fn score_fresh_batch<const QUERY_X4: bool>(
2407        &mut self,
2408        g: &VGraph,
2409        k: usize,
2410        patience: usize,
2411        stale: &mut usize,
2412        w: f32,
2413    ) -> usize {
2414        let qc = &self.qc[..];
2415        let n = self.fresh.len();
2416        let (reservoir, pending) = (&mut self.reservoir, &mut self.keys);
2417        pending.clear();
2418        let mut offer = |estimate, id| {
2419            let key = key(estimate, id);
2420            reservoir.offer(key);
2421            pending.push(key);
2422        };
2423        let mut i = 0;
2424        if self.knobs.batch_est || (QUERY_X4 && g.bits == 8) {
2425            while i + 4 <= n {
2426                let ids = [
2427                    self.fresh[i],
2428                    self.fresh[i + 1],
2429                    self.fresh[i + 2],
2430                    self.fresh[i + 3],
2431                ];
2432                for (estimate, id) in g
2433                    .est4::<QUERY_X4>(qc, self.qsum, ids, w)
2434                    .into_iter()
2435                    .zip(ids)
2436                {
2437                    offer(estimate, id);
2438                }
2439                i += 4;
2440            }
2441        }
2442        for &id in &self.fresh[i..] {
2443            offer(g.est(qc, self.qsum, id, w), id);
2444        }
2445        if self.pool.insert_batch(pending, k) && patience > 0 {
2446            *stale = 0;
2447        }
2448        n
2449    }
2450
2451    /// Writes the ids of the (approximate) `k` nearest vectors by the index's metric (cosine
2452    /// or Euclidean) into `out`, nearest first, and returns the number of distance estimates
2453    /// (descent, traversal and re-rank). A query of the wrong dimension, with a non-finite
2454    /// value, or all zero on an angular index (cosine is undefined) is an error; the pool and
2455    /// heaps are bounded by the vertex count whatever `ef` is.
2456    fn search(
2457        &mut self,
2458        g: &VGraph,
2459        query: &[f32],
2460        k: usize,
2461        params: VSearchParams,
2462        out: &mut Vec<u32>,
2463    ) -> Result<usize, WarmError> {
2464        Self::validate_query(g, query)?;
2465        self.search_validated(g, query, k, params, out)
2466    }
2467
2468    fn validate_query(g: &VGraph, query: &[f32]) -> Result<(), WarmError> {
2469        if query.len() != g.dim {
2470            return Err(invalid(format!(
2471                "query has {} dimensions, the index has {}",
2472                query.len(),
2473                g.dim
2474            )));
2475        }
2476        if let Some(at) = query.iter().position(|v| !v.is_finite()) {
2477            return Err(invalid(format!(
2478                "non-finite query value at coordinate {at}"
2479            )));
2480        }
2481        if g.metric == VMetric::Angular && query.iter().all(|&v| v == 0.0) {
2482            return Err(invalid("zero query: cosine similarity is undefined".into()));
2483        }
2484        Ok(())
2485    }
2486
2487    // The caller has checked dimension, finite coordinates, and nonzero angular norm.
2488    fn search_validated(
2489        &mut self,
2490        g: &VGraph,
2491        query: &[f32],
2492        k: usize,
2493        params: VSearchParams,
2494        out: &mut Vec<u32>,
2495    ) -> Result<usize, WarmError> {
2496        self.search_inner::<
2497            false,
2498            { cfg!(feature = "experimental-batch-pool") },
2499            { cfg!(feature = "experimental-query-x4") },
2500        >(
2501            g,
2502            query,
2503            k,
2504            params,
2505            out,
2506            &mut QueryProfile::default(),
2507        )
2508    }
2509
2510    fn search_inner<const PROFILE: bool, const BATCH_POOL: bool, const QUERY_X4: bool>(
2511        &mut self,
2512        g: &VGraph,
2513        query: &[f32],
2514        k: usize,
2515        params: VSearchParams,
2516        out: &mut Vec<u32>,
2517        profile: &mut QueryProfile,
2518    ) -> Result<usize, WarmError> {
2519        if params.reservoir > 0 && (g.has_blocks() || (g.has_screen() && params.screen)) {
2520            return Err(invalid(
2521                "independent retention requires unscreened row-code search".into(),
2522            ));
2523        }
2524        let started = PROFILE.then(Instant::now);
2525        let ef = params.ef.max(k).max(1);
2526        self.reservoir.reset(if params.reservoir == 0 {
2527            0
2528        } else {
2529            params.reservoir.max(k).min(g.n)
2530        });
2531        // Only the order matters, so the query is used as given (no centring needed: the
2532        // `<q, mean>` term is common to all candidates) and its scale is dropped.
2533        self.qf[..g.dim].copy_from_slice(query);
2534        let sq = g.encode_query(
2535            &self.qf,
2536            &mut self.qc,
2537            &mut self.qtmp,
2538            &mut self.qrot,
2539            &mut self.qs,
2540        );
2541        // Weight of the Euclidean norm term in the estimate's units (unused for angular); a
2542        // zero query has every inner product zero, so the norms alone order it.
2543        let w = if sq > 0.0 { 1.0 / sq } else { 1.0 };
2544        // Only the x86 VNNI kernels read the sum.
2545        self.qsum = if cfg!(target_arch = "x86_64") {
2546            query_sum(&self.qc)
2547        } else {
2548            0
2549        };
2550        if let Some(started) = started {
2551            profile.query_encoding_ns = started.elapsed().as_nanos() as u64;
2552        }
2553        let qc = &self.qc[..];
2554        let blocks = g.has_blocks();
2555        if blocks {
2556            g.encode_block_query(&self.qf, sq, &mut self.bqrot, &mut self.bq);
2557        }
2558        let knobs = self.knobs;
2559        let block_screen = blocks && knobs.block_screen;
2560        let screen = g.has_screen() && params.screen && !blocks;
2561        let margin = params.screen_sigmas * g.screen_sigma;
2562        self.screened = 0;
2563
2564        self.visited.next();
2565        let entry_started = PROFILE.then(Instant::now);
2566        let (entry, mut evals) = g.descend(qc, self.qsum, w, &mut self.reservoir);
2567        if let Some(started) = entry_started {
2568            profile.entry_ns = started.elapsed().as_nanos() as u64;
2569            profile.entry_evaluations = evals;
2570        }
2571        self.visited.insert(entry);
2572        let fuse = blocks && params.fuse && !block_screen;
2573        if fuse {
2574            self.visited.trace(entry);
2575        }
2576        // Query terms of the error model: `|q - x_v|^2 = |q|^2 + 1 - 2 (<q, x_v - mean> + <q, mean>)`
2577        // for unit index rows.
2578        let (qn2, qm) = if blocks {
2579            let q = &self.qf[..g.dim];
2580            (
2581                q.iter().map(|a| a * a).sum::<f32>(),
2582                q.iter().zip(&g.mean).map(|(a, m)| a * m).sum::<f32>(),
2583            )
2584        } else {
2585            (0.0, 0.0)
2586        };
2587        self.pool.reset(ef);
2588        self.expanded.clear();
2589        let take = params.rerank.max(k).min(g.n);
2590        let k0 = key(g.est(qc, self.qsum, entry, w), entry);
2591        self.reservoir.offer(k0);
2592        self.pool.insert(k0);
2593        if fuse {
2594            self.visited.trace(entry).1.key = k0;
2595        }
2596        evals += 1;
2597        // Best-first expansion of the nearest unexpanded candidate among the `ef` best seen;
2598        // the pool holds exactly those, so the search ends when none is left. With
2599        // `patience`, it also ends after that many expansions without an insert among the
2600        // `k` best.
2601        let patience = params.patience;
2602        let frontier = if knobs.frontier.is_finite() {
2603            knobs.frontier.clamp(0.0, 1.0)
2604        } else {
2605            1.0
2606        };
2607        let mut stale = 0usize;
2608        let target = k.min(g.n);
2609        let enough = |search: &Self| {
2610            let count = if params.reservoir > 0 {
2611                search.reservoir.keys.len()
2612            } else if blocks && !block_screen {
2613                search.expanded.len()
2614            } else {
2615                search.pool.keys.len()
2616            };
2617            count >= target
2618        };
2619        let mut seed_cursor = 0usize;
2620        loop {
2621            let c = match self.pool.next() {
2622                Some(c) => c,
2623                None => {
2624                    if enough(self) {
2625                        break;
2626                    }
2627                    // Exhausting a full pool would have retained at least `target` distinct
2628                    // results. A short result therefore leaves room for an unseen seed.
2629                    let seed = (seed_cursor..g.n)
2630                        .find(|&v| {
2631                            if fuse {
2632                                self.visited.traces[v].epoch != self.visited.epoch
2633                            } else {
2634                                !self.visited.seen(v as u32)
2635                            }
2636                        })
2637                        .ok_or_else(|| {
2638                            invalid(
2639                                "search exhausted vertices before retaining enough candidates"
2640                                    .into(),
2641                            )
2642                        })?;
2643                    seed_cursor = seed + 1;
2644                    let seed = seed as u32;
2645                    self.visited.insert(seed);
2646                    let candidate = key(g.est(&self.qc, self.qsum, seed, w), seed);
2647                    self.reservoir.offer(candidate);
2648                    let (_, evicted) = self
2649                        .pool
2650                        .insert(candidate)
2651                        .ok_or_else(|| invalid("search could not admit a recovery seed".into()))?;
2652                    if fuse {
2653                        self.visited.trace(seed).1.key = candidate;
2654                        if let Some(evicted) = evicted {
2655                            self.visited.trace(key_id(evicted)).1.key = u64::MAX;
2656                        }
2657                    }
2658                    evals += 1;
2659                    if PROFILE {
2660                        profile.traversal_evaluations += 1;
2661                    }
2662                    stale = 0;
2663                    continue;
2664                }
2665            };
2666            if patience > 0 && stale >= patience && enough(self) {
2667                break;
2668            }
2669            if frontier < 1.0 && k > 0 && self.pool.is_full() {
2670                // Frontier rule: the nearest unexpanded candidate already lies past the point
2671                // `frontier` of the way from the `k`-th best estimate to the `ef`-th. At 1.0
2672                // it never fires (every pool key is at most the `ef`-th).
2673                let len = self.pool.keys.len();
2674                let kth = key_dist(self.pool.keys[k.min(len) - 1]);
2675                let worst = key_dist(self.pool.keys[len - 1]);
2676                if key_dist(c) > kth + frontier * (worst - kth) && enough(self) {
2677                    break;
2678                }
2679            }
2680            if fuse {
2681                // A re-admitted vertex may sit in the pool twice; expand it once.
2682                let (_, t) = self.visited.trace(key_id(c));
2683                if t.expanded {
2684                    continue;
2685                }
2686                t.expanded = true;
2687            }
2688            stale += 1;
2689            if PROFILE {
2690                profile.expansions += 1;
2691            }
2692            let adjacency_started = PROFILE.then(Instant::now);
2693            if let Some(next) = self.pool.peek_next() {
2694                g.prefetch_neighbours(key_id(next) as usize);
2695                if blocks {
2696                    g.prefetch_blocks(key_id(next) as usize);
2697                    if !block_screen {
2698                        prefetch_bytes(g.code(key_id(next)));
2699                    }
2700                }
2701            }
2702            if knobs.prefetch >= 2 {
2703                if let Some(second) = self.pool.peek_unexpanded(1) {
2704                    g.prefetch_neighbours(key_id(second) as usize);
2705                    if blocks {
2706                        g.prefetch_blocks(key_id(second) as usize);
2707                    }
2708                }
2709            }
2710            if block_screen {
2711                // The neighbour block screens and the pool holds row estimates, so the popped
2712                // key carries the row estimate of `v` (`key` round-trips the f32 bits). A fresh
2713                // neighbour whose block estimate, improved by `screen_sigmas` deviations of
2714                // the calibrated edge error, could not enter the full pool is skipped (it
2715                // stays visited); the survivors' rows are prefetched together, then estimated.
2716                let v = key_id(c) as usize;
2717                let ip_v = -key_dist(c) * sq;
2718                let d2_v = (qn2 + 1.0 - 2.0 * (ip_v + qm)).max(1e-6);
2719                let (delta, min_sum, inv_sq) = (self.bq.delta, self.bq.min_sum, self.bq.inv_sq);
2720                let worst = self.pool.worst().map(key_dist);
2721                let sigmas = params.screen_sigmas;
2722                self.fresh.clear();
2723                let list = &mut self.fresh;
2724                g.for_each_neighbour(v, |u| list.push(u));
2725                self.survivors.clear();
2726                for (b, ids) in self
2727                    .fresh
2728                    .chunks(BLOCK)
2729                    .take(g.blocks_per_vertex)
2730                    .enumerate()
2731                {
2732                    let (codes, scales, offsets) = g.block(v, b);
2733                    scan32(codes, &self.bq.lut, g.block_m, &mut self.scan);
2734                    let (scales, offsets) = (scales.as_chunks::<4>().0, offsets.as_chunks::<4>().0);
2735                    for (((&u, &raw), s), o) in ids.iter().zip(&self.scan).zip(scales).zip(offsets)
2736                    {
2737                        if !self.visited.insert(u) {
2738                            continue;
2739                        }
2740                        let s = f32::from_le_bytes(*s);
2741                        let ip = ip_v + s * (delta * raw as f32 + min_sum) + f32::from_le_bytes(*o);
2742                        if let Some(w) = worst {
2743                            let sigma = (g.block_kappa * s * s * d2_v).sqrt();
2744                            if -(ip + sigmas * sigma) * inv_sq >= w {
2745                                self.screened += 1;
2746                                continue;
2747                            }
2748                        }
2749                        prefetch_bytes(g.code(u));
2750                        g.prefetch_scale(u);
2751                        self.survivors.push(u);
2752                    }
2753                }
2754                std::mem::swap(&mut self.fresh, &mut self.survivors);
2755                evals += self.score_fresh::<BATCH_POOL, QUERY_X4>(g, k, patience, &mut stale, w);
2756                continue;
2757            }
2758            if blocks {
2759                // The expanded vertex is scored with its own code (the per-row estimate the
2760                // plain search uses everywhere) into the result pool. One table scan per
2761                // block then estimates all its neighbours, ordering the exploration; list
2762                // padding is skipped, and a vertex met again is re-admitted only when its
2763                // fused estimate improves on the entry it holds.
2764                let v = key_id(c) as usize;
2765                let est_v = g.est(&self.qc, self.qsum, v as u32, w);
2766                let ek = key(est_v, v as u32);
2767                evals += 1;
2768                // `<q, x_v - mean>` in float units: the estimate dropped the query scale.
2769                let ip_v = -est_v * sq;
2770                let d2_v = (qn2 + 1.0 - 2.0 * (ip_v + qm)).max(1e-6);
2771                if self.expanded.len() < take {
2772                    self.expanded.push(ek);
2773                } else if let Some(mut worst) = self.expanded.peek_mut() {
2774                    if ek < *worst {
2775                        *worst = ek;
2776                    }
2777                }
2778                let (delta, min_sum, inv_sq) = (self.bq.delta, self.bq.min_sum, self.bq.inv_sq);
2779                self.fresh.clear();
2780                let list = &mut self.fresh;
2781                g.for_each_neighbour(v, |u| list.push(u));
2782                for (b, ids) in self
2783                    .fresh
2784                    .chunks(BLOCK)
2785                    .take(g.blocks_per_vertex)
2786                    .enumerate()
2787                {
2788                    let (codes, scales, offsets) = g.block(v, b);
2789                    scan32(codes, &self.bq.lut, g.block_m, &mut self.scan);
2790                    let (scales, offsets) = (scales.as_chunks::<4>().0, offsets.as_chunks::<4>().0);
2791                    for (((&u, &raw), s), o) in ids.iter().zip(&self.scan).zip(scales).zip(offsets)
2792                    {
2793                        let s = f32::from_le_bytes(*s);
2794                        let ip = ip_v + s * (delta * raw as f32 + min_sum) + f32::from_le_bytes(*o);
2795                        if !fuse {
2796                            // Without fusion the block path keeps the plain path's byte marks
2797                            // (the per-vertex trace records would not stay in cache).
2798                            if !self.visited.insert(u) {
2799                                continue;
2800                            }
2801                            evals += 1;
2802                            if let Some((pos, _)) = self.pool.insert(key(-ip * inv_sq, u)) {
2803                                if patience > 0 && pos < k {
2804                                    stale = 0;
2805                                }
2806                            }
2807                            continue;
2808                        }
2809                        // Precision of this estimate under the calibrated error model.
2810                        let w = 1.0 / (g.block_kappa * s * s * d2_v + 1e-12);
2811                        let (fresh, t) = self.visited.trace(u);
2812                        let ip = if fresh {
2813                            t.w = w;
2814                            t.wsum = w * ip;
2815                            ip
2816                        } else {
2817                            if t.expanded {
2818                                continue;
2819                            }
2820                            t.w += w;
2821                            t.wsum += w * ip;
2822                            t.wsum / t.w
2823                        };
2824                        let kk = key(-ip * inv_sq, u);
2825                        let old = t.key;
2826                        evals += 1;
2827                        if !fresh && kk >= old {
2828                            continue;
2829                        }
2830                        if old != u64::MAX {
2831                            self.pool.remove(old);
2832                        }
2833                        // `t` borrows `visited` only, so it stays live across the pool update.
2834                        let Some((pos, evicted)) = self.pool.insert(kk) else {
2835                            t.key = u64::MAX;
2836                            continue;
2837                        };
2838                        t.key = kk;
2839                        if let Some(e) = evicted {
2840                            self.visited.trace(key_id(e)).1.key = u64::MAX;
2841                        }
2842                        if patience > 0 && pos < k {
2843                            stale = 0;
2844                        }
2845                    }
2846                }
2847                continue;
2848            }
2849            self.fresh.clear();
2850            let (visited, fresh) = (&mut self.visited, &mut self.fresh);
2851            #[cfg(feature = "experimental-neighbor-admission")]
2852            let neighbors_read = if screen {
2853                g.admit_neighbours::<PROFILE, true>(
2854                    key_id(c) as usize,
2855                    params.neighbour_cap,
2856                    visited,
2857                    fresh,
2858                )
2859            } else {
2860                g.admit_neighbours::<PROFILE, false>(
2861                    key_id(c) as usize,
2862                    params.neighbour_cap,
2863                    visited,
2864                    fresh,
2865                )
2866            };
2867            #[cfg(not(feature = "experimental-neighbor-admission"))]
2868            let neighbors_read = {
2869                let mut neighbors_read = 0;
2870                let admit = |u: u32, visited: &mut Marks, fresh: &mut Vec<u32>| {
2871                    if visited.insert(u) {
2872                        if screen {
2873                            g.prefetch_screen(u);
2874                        } else {
2875                            prefetch_bytes(g.code(u));
2876                            g.prefetch_scale(u);
2877                        }
2878                        fresh.push(u);
2879                    }
2880                };
2881                if params.neighbour_cap == 0 {
2882                    g.for_each_neighbour(key_id(c) as usize, |u| {
2883                        if PROFILE {
2884                            neighbors_read += 1;
2885                        }
2886                        admit(u, visited, fresh);
2887                    });
2888                } else {
2889                    let mut seen = 0usize;
2890                    g.for_each_neighbour(key_id(c) as usize, |u| {
2891                        seen += 1;
2892                        if seen <= params.neighbour_cap {
2893                            admit(u, visited, fresh);
2894                        }
2895                    });
2896                }
2897                neighbors_read
2898            };
2899            if screen {
2900                // A neighbour whose optimistic 1-bit estimate cannot enter the full pool is
2901                // skipped (it stays visited); the rest get the 8-bit estimate.
2902                let worst = self.pool.worst().map(key_dist);
2903                self.survivors.clear();
2904                for &u in &self.fresh {
2905                    if let Some(w) = worst {
2906                        if g.screen_est(&self.qs, u) - margin >= w {
2907                            self.screened += 1;
2908                            continue;
2909                        }
2910                    }
2911                    prefetch_bytes(g.code(u));
2912                    g.prefetch_scale(u);
2913                    self.survivors.push(u);
2914                }
2915                std::mem::swap(&mut self.fresh, &mut self.survivors);
2916            }
2917            if let Some(started) = adjacency_started {
2918                profile.adjacency_visited_prefetch_ns += started.elapsed().as_nanos() as u64;
2919                profile.neighbors_read += neighbors_read;
2920            }
2921            let scoring_started = PROFILE.then(Instant::now);
2922            let count = self.score_fresh::<BATCH_POOL, QUERY_X4>(g, k, patience, &mut stale, w);
2923            evals += count;
2924            if let Some(started) = scoring_started {
2925                profile.scoring_pool_ns += started.elapsed().as_nanos() as u64;
2926                profile.traversal_evaluations += count;
2927            }
2928            if knobs.prefetch >= 3 {
2929                // Lookahead: what the next expansion's admission touches for its unvisited
2930                // neighbours, one expansion early (its list was prefetched above).
2931                if let Some(next) = self.pool.peek_next() {
2932                    let visited = &self.visited;
2933                    g.for_each_neighbour(key_id(next) as usize, |u| {
2934                        if !visited.seen(u) {
2935                            if screen {
2936                                g.prefetch_screen(u);
2937                            } else {
2938                                prefetch_bytes(g.code(u));
2939                                g.prefetch_scale(u);
2940                            }
2941                        }
2942                    });
2943                }
2944            }
2945        }
2946
2947        let rerank_started = PROFILE.then(Instant::now);
2948        if PROFILE {
2949            profile.beam_retained = self.pool.keys.len();
2950            profile.traversal_evaluations += 1;
2951        }
2952        self.keys.clear();
2953        if params.reservoir > 0 {
2954            self.keys.extend_from_slice(&self.reservoir.keys);
2955        } else if blocks && !block_screen {
2956            self.keys.extend(self.expanded.drain());
2957            self.keys.sort_unstable();
2958        } else {
2959            self.keys.extend_from_slice(&self.pool.keys);
2960        }
2961        if PROFILE {
2962            profile.retained = self.keys.len();
2963        }
2964        self.results.clear();
2965        let take = take.min(self.keys.len());
2966        if params.rerank > 0 && knobs.rerank_z > 0.0 && sq > 0.0 && k > 0 && take > k {
2967            evals += self.rerank_gated(g, k, take, sq, knobs.rerank_z);
2968            if PROFILE {
2969                profile.rerank_evaluations = self.results.len();
2970            }
2971        } else if params.rerank > 0 {
2972            if PROFILE {
2973                profile.rerank_evaluations = take;
2974            }
2975            for &c in &self.keys[..take] {
2976                g.prefetch_rerank_rows(key_id(c));
2977            }
2978            // Ties (duplicate vectors) break by the original id, independent of the layout.
2979            self.results.extend(self.keys[..take].iter().map(|&c| {
2980                let id = key_id(c);
2981                (g.fine(&self.qf, id), g.original(id))
2982            }));
2983            evals += take;
2984            self.results
2985                .sort_unstable_by(|a, b| a.0.total_cmp(&b.0).then(a.1.cmp(&b.1)));
2986        } else {
2987            self.results.extend(
2988                self.keys[..take]
2989                    .iter()
2990                    .map(|&c| (0.0, g.original(key_id(c)))),
2991            );
2992        }
2993        out.clear();
2994        out.extend(self.results.iter().take(k).map(|r| r.1));
2995        if let Some(started) = rerank_started {
2996            profile.rerank_sort_ns = started.elapsed().as_nanos() as u64;
2997        }
2998        if let Some(started) = started {
2999            profile.total_ns = started.elapsed().as_nanos() as u64;
3000        }
3001        Ok(evals)
3002    }
3003}
3004
3005#[cfg(feature = "experimental")]
3006impl VSearcher<'_> {
3007    /// Profiles the ordinary row-code path; timers perturb latency and are not QPS evidence.
3008    pub fn profile_search(
3009        &mut self,
3010        query: &[f32],
3011        k: usize,
3012        params: VSearchParams,
3013        out: &mut Vec<u32>,
3014    ) -> Result<QueryProfile, WarmError> {
3015        let g = self.index;
3016        if g.has_screen()
3017            || g.has_blocks()
3018            || self.knobs() != SearchKnobs::OFF
3019            || params.neighbour_cap != 0
3020        {
3021            return Err(invalid("profiling requires the plain row-code path".into()));
3022        }
3023        SearchScratch::validate_query(g, query)?;
3024        let mut profile = QueryProfile::default();
3025        self.scratch
3026            .search_inner::<
3027                true,
3028                { cfg!(feature = "experimental-batch-pool") },
3029                { cfg!(feature = "experimental-query-x4") },
3030            >(
3031                g,
3032                query,
3033                k,
3034                params,
3035                out,
3036                &mut profile,
3037            )?;
3038        Ok(profile)
3039    }
3040}
3041
3042#[cfg(feature = "experimental")]
3043impl VGraph {
3044    /// Storage geometry and reconstruction error; diagnostic work is never part of search.
3045    pub fn encoding_diagnostic(&self, source: &Graph) -> Result<serde_json::Value, WarmError> {
3046        if source.n != self.n || source.dim < self.dim {
3047            return Err(invalid("encoding diagnostic source shape mismatch".into()));
3048        }
3049        let (mut coarse_error, mut fine_error) = (0f64, 0f64);
3050        for v in 0..self.n {
3051            let original = self.original(v as u32) as usize;
3052            for d in 0..self.dim {
3053                let expected =
3054                    f64::from(source.vectors[original * source.dim + d]) - f64::from(self.mean[d]);
3055                let primary = f64::from(self.scales[v])
3056                    * f64::from(code_at(self.code(v as u32), self.bits, d));
3057                let residual = if self.has_residual() {
3058                    f64::from(self.res_scales[v])
3059                        * f64::from(code_at(
3060                            &self.residual[v * self.res_stride..(v + 1) * self.res_stride],
3061                            self.res_bits,
3062                            d,
3063                        ))
3064                } else {
3065                    0.0
3066                };
3067                coarse_error += (primary - expected).powi(2);
3068                fine_error += (primary + residual - expected).powi(2);
3069            }
3070        }
3071        let count = (self.n * self.dim) as f64;
3072        Ok(
3073            serde_json::json!({"metric": format!("{:?}", self.metric), "dimensions": self.dim,
3074            "padded": self.padded, "primary_stride": self.stride, "residual_stride": self.res_stride,
3075            "norm_bytes": self.half_norms.len() * 4, "index_bytes": self.index_bytes(),
3076            "primary_coordinate_rmse": (coarse_error / count).sqrt(),
3077            "refined_coordinate_rmse": (fine_error / count).sqrt()}),
3078        )
3079    }
3080
3081    /// Runs an untimed traversal control against the graph used to encode this index.
3082    /// `oracle` is an original row ID; `exact` replaces routing scores only. Residual
3083    /// reranking, beam ordering and tie handling remain those of the compact index.
3084    /// Screens, blocks and runtime search knobs are excluded from this control.
3085    pub fn diagnose(
3086        &mut self,
3087        source: &Graph,
3088        query: &[f32],
3089        k: usize,
3090        params: VSearchParams,
3091        oracle: Option<u32>,
3092        exact: bool,
3093    ) -> Result<super::diagnostics::TraversalDiagnostic, WarmError> {
3094        if source.n != self.n
3095            || source.dim != self.dim + usize::from(self.metric == VMetric::InnerProduct)
3096            || source.degree != self.degree
3097            || if self.metric == VMetric::InnerProduct {
3098                source.metric != Metric::Euclidean
3099            } else {
3100                self.metric != source.metric.into()
3101            }
3102        {
3103            return Err(invalid(
3104                "diagnostic source graph does not match index geometry and metric".into(),
3105            ));
3106        }
3107        if self.has_screen() || self.has_blocks() {
3108            return Err(invalid(
3109                "diagnostic traversal does not support screens or blocks".into(),
3110            ));
3111        }
3112        if query.len() != self.dim
3113            || query.iter().any(|x| !x.is_finite())
3114            || (self.metric == VMetric::Angular && query.iter().all(|&x| x == 0.0))
3115            || k == 0
3116            || k > self.n
3117        {
3118            return Err(invalid(
3119                "invalid diagnostic query dimension, coordinates or k".into(),
3120            ));
3121        }
3122        let stored_oracle = match oracle {
3123            Some(id) if id as usize >= self.n => {
3124                return Err(invalid("diagnostic oracle ID outside index".into()));
3125            }
3126            Some(id) if self.ids.is_empty() => Some(id),
3127            Some(id) => Some(
3128                self.ids
3129                    .iter()
3130                    .position(|&original| original == id)
3131                    .expect("encoded IDs are a permutation") as u32,
3132            ),
3133            None => None,
3134        };
3135        let saved_entry = self.entry;
3136        let saved_layers = stored_oracle.map(|id| {
3137            self.entry = id;
3138            std::mem::take(&mut self.layers)
3139        });
3140        let result = if exact {
3141            self.diagnose_exact(source, query, k, params)
3142        } else {
3143            self.diagnose_coded(query, k, params)
3144        };
3145        self.entry = saved_entry;
3146        if let Some(layers) = saved_layers {
3147            self.layers = layers;
3148        }
3149        result
3150    }
3151
3152    fn diagnostic_entry(&self, mut score: impl FnMut(u32) -> f32, evaluated: &mut Vec<u32>) -> u32 {
3153        if self.layers.is_empty() {
3154            return self.entry;
3155        }
3156        let mut p = 0u32;
3157        let mut best = score(self.upper_ids[0]);
3158        evaluated.push(self.original(self.upper_ids[0]));
3159        for layer in self.layers.iter().rev() {
3160            loop {
3161                let mut moved = false;
3162                for &u in &layer[p as usize] {
3163                    let id = self.upper_ids[u as usize];
3164                    let candidate = score(id);
3165                    evaluated.push(self.original(id));
3166                    if candidate < best {
3167                        best = candidate;
3168                        p = u;
3169                        moved = true;
3170                    }
3171                }
3172                if !moved {
3173                    break;
3174                }
3175            }
3176        }
3177        self.upper_ids[p as usize]
3178    }
3179
3180    fn diagnose_coded(
3181        &self,
3182        query: &[f32],
3183        k: usize,
3184        params: VSearchParams,
3185    ) -> Result<super::diagnostics::TraversalDiagnostic, WarmError> {
3186        let mut report = super::diagnostics::TraversalDiagnostic::default();
3187        let mut searcher = self.searcher_with(SearchKnobs::OFF);
3188        searcher.search(query, k, params, &mut report.results)?;
3189        let scratch = &mut searcher.scratch;
3190        let sq = self.encode_query(
3191            &scratch.qf,
3192            &mut scratch.qc,
3193            &mut scratch.qtmp,
3194            &mut scratch.qrot,
3195            &mut scratch.qs,
3196        );
3197        let w = if sq > 0.0 { 1.0 / sq } else { 1.0 };
3198        // Descent does not set the base-layer marks; replay it after the production search.
3199        let qsum = query_sum(&scratch.qc);
3200        self.diagnostic_entry(
3201            |id| self.est(&scratch.qc, qsum, id, w),
3202            &mut report.evaluated,
3203        );
3204        report.evaluated.extend(
3205            (0..self.n as u32)
3206                .filter(|&id| scratch.visited.seen(id))
3207                .map(|id| self.original(id)),
3208        );
3209        report.evaluated.sort_unstable();
3210        report.evaluated.dedup();
3211        report
3212            .beam
3213            .extend(scratch.pool.keys.iter().map(|&c| self.original(key_id(c))));
3214        report.retained.extend(
3215            scratch
3216                .keys
3217                .iter()
3218                .map(|&candidate| self.original(key_id(candidate))),
3219        );
3220        if params.rerank > 0 {
3221            let take = params.rerank.max(k).min(self.n).min(report.retained.len());
3222            report
3223                .rerank_candidates
3224                .extend_from_slice(&report.retained[..take]);
3225        }
3226        Ok(report)
3227    }
3228
3229    fn diagnose_exact(
3230        &self,
3231        source: &Graph,
3232        query: &[f32],
3233        k: usize,
3234        params: VSearchParams,
3235    ) -> Result<super::diagnostics::TraversalDiagnostic, WarmError> {
3236        let mut report = super::diagnostics::TraversalDiagnostic::default();
3237        let score = |id: u32| {
3238            let original = self.original(id) as usize;
3239            let row = &source.vectors[original * source.dim..original * source.dim + self.dim];
3240            let ip: f64 = row
3241                .iter()
3242                .zip(query)
3243                .map(|(&x, &q)| f64::from(x) * f64::from(q))
3244                .sum();
3245            let norm = if self.metric == VMetric::Euclidean {
3246                row.iter().map(|&x| f64::from(x).powi(2)).sum::<f64>() * 0.5
3247            } else {
3248                0.0
3249            };
3250            (norm - ip) as f32
3251        };
3252        let mut reservoir = Reservoir::default();
3253        reservoir.reset(if params.reservoir == 0 {
3254            0
3255        } else {
3256            params.reservoir.max(k).min(self.n)
3257        });
3258        let entry = self.diagnostic_entry(
3259            |id| {
3260                let distance = score(id);
3261                reservoir.offer(key(distance, id));
3262                distance
3263            },
3264            &mut report.evaluated,
3265        );
3266        let mut pool = Pool {
3267            keys: Vec::new(),
3268            expanded: Vec::new(),
3269            #[cfg(feature = "experimental-batch-pool")]
3270            merge_keys: Vec::new(),
3271            #[cfg(feature = "experimental-batch-pool")]
3272            merge_expanded: Vec::new(),
3273            cap: params.ef.max(k).max(1),
3274            cursor: 0,
3275        };
3276        let mut seen = vec![false; self.n];
3277        seen[entry as usize] = true;
3278        report.evaluated.push(self.original(entry));
3279        let entry_key = key(score(entry), entry);
3280        reservoir.offer(entry_key);
3281        pool.insert(entry_key);
3282        let mut stale = 0usize;
3283        while let Some(candidate) = pool.next() {
3284            if params.patience > 0 && stale >= params.patience {
3285                break;
3286            }
3287            stale += 1;
3288            let mut count = 0usize;
3289            self.for_each_neighbour(key_id(candidate) as usize, |id| {
3290                count += 1;
3291                if (params.neighbour_cap > 0 && count > params.neighbour_cap) || seen[id as usize] {
3292                    return;
3293                }
3294                seen[id as usize] = true;
3295                report.evaluated.push(self.original(id));
3296                let kk = key(score(id), id);
3297                reservoir.offer(kk);
3298                if let Some((pos, _)) = pool.insert(kk) {
3299                    if params.patience > 0 && pos < k {
3300                        stale = 0;
3301                    }
3302                }
3303            });
3304        }
3305        report.evaluated.sort_unstable();
3306        report.evaluated.dedup();
3307        report
3308            .beam
3309            .extend(pool.keys.iter().map(|&c| self.original(key_id(c))));
3310        let keys = if params.reservoir > 0 {
3311            &reservoir.keys
3312        } else {
3313            &pool.keys
3314        };
3315        report
3316            .retained
3317            .extend(keys.iter().map(|&c| self.original(key_id(c))));
3318        let take = params.rerank.max(k).min(self.n).min(keys.len());
3319        if params.rerank > 0 {
3320            report
3321                .rerank_candidates
3322                .extend_from_slice(&report.retained[..take]);
3323            let mut qf = vec![0.0; self.padded];
3324            qf[..self.dim].copy_from_slice(query);
3325            let mut results: Vec<_> = keys[..take]
3326                .iter()
3327                .map(|&c| (self.fine(&qf, key_id(c)), self.original(key_id(c))))
3328                .collect();
3329            results.sort_unstable_by(|a, b| a.0.total_cmp(&b.0).then(a.1.cmp(&b.1)));
3330            report
3331                .results
3332                .extend(results.iter().take(k).map(|&(_, id)| id));
3333        } else {
3334            report
3335                .results
3336                .extend_from_slice(&report.retained[..k.min(take)]);
3337        }
3338        Ok(report)
3339    }
3340}
3341
3342#[cfg(test)]
3343mod serving_tests {
3344    use super::*;
3345
3346    #[test]
3347    fn connected_search_fingerprints_and_evaluation_counts() {
3348        let rows: Vec<f32> = (0..64)
3349            .flat_map(|v| (0..16).map(move |d| if v & (1 << (d % 6)) == 0 { -1.0 } else { 1.0 }))
3350            .collect();
3351        for metric in [VMetric::Angular, VMetric::Euclidean, VMetric::InnerProduct] {
3352            let mut graph = if metric == VMetric::InnerProduct {
3353                lifted(&rows, 16)
3354            } else {
3355                Graph::build(
3356                    &rows,
3357                    16,
3358                    &BuildParams {
3359                        metric: if metric == VMetric::Angular {
3360                            Metric::Angular
3361                        } else {
3362                            Metric::Euclidean
3363                        },
3364                        degree: 8,
3365                        build_ef: 16,
3366                        ..Default::default()
3367                    },
3368                )
3369                .unwrap()
3370            };
3371            graph.entry = 0;
3372            graph.layers.clear();
3373            graph.upper_ids.clear();
3374            for v in 0..graph.n {
3375                let neighbors: Vec<_> = [1, 3]
3376                    .map(|offset| ((v + offset) % graph.n) as u32)
3377                    .to_vec();
3378                graph.adj.set(v, &neighbors);
3379            }
3380            for bits in [4, 8] {
3381                let index = if metric == VMetric::InnerProduct {
3382                    VGraph::from_lifted_ip_graph(&graph, bits, 8, true, true, 3).unwrap()
3383                } else {
3384                    VGraph::from_graph_with(&graph, bits, 8, true, true, 3, false).unwrap()
3385                };
3386                for (knob_id, knobs) in [
3387                    SearchKnobs::OFF,
3388                    SearchKnobs {
3389                        prefetch: 3,
3390                        batch_est: true,
3391                        frontier: 0.8,
3392                        ..SearchKnobs::OFF
3393                    },
3394                ]
3395                .into_iter()
3396                .enumerate()
3397                {
3398                    let mut searcher = index.searcher_with(knobs);
3399                    let mut out = Vec::new();
3400                    let mut fingerprint = 0xcbf29ce484222325u64;
3401                    let mut evaluations = 0usize;
3402                    for ef in [8, 16] {
3403                        for reservoir in [0, 16] {
3404                            for patience in [0, 1] {
3405                                for query in rows.as_chunks::<16>().0.iter().step_by(7) {
3406                                    let count = searcher
3407                                        .search(
3408                                            query,
3409                                            5,
3410                                            VSearchParams {
3411                                                ef,
3412                                                rerank: 16,
3413                                                reservoir,
3414                                                patience,
3415                                                ..Default::default()
3416                                            },
3417                                            &mut out,
3418                                        )
3419                                        .unwrap();
3420                                    assert_eq!(out.len(), 5);
3421                                    evaluations += count;
3422                                    for value in
3423                                        out.iter()
3424                                            .map(|&id| u64::from(id))
3425                                            .chain(searcher.scratch.results.iter().map(
3426                                                |&(distance, _)| u64::from(distance.to_bits()),
3427                                            ))
3428                                            .chain([count as u64])
3429                                    {
3430                                        fingerprint =
3431                                            (fingerprint ^ value).wrapping_mul(0x100000001b3);
3432                                    }
3433                                }
3434                            }
3435                        }
3436                    }
3437                    println!("CONNECTED {metric:?} {bits} {knob_id} {fingerprint} {evaluations}");
3438                }
3439            }
3440        }
3441    }
3442
3443    fn disconnect(graph: &mut Graph, component_size: usize) {
3444        graph.entry = 0;
3445        graph.layers.clear();
3446        graph.upper_ids.clear();
3447        for v in 0..graph.n {
3448            let first = v / component_size * component_size;
3449            let neighbors: Vec<_> = (first..(first + component_size).min(graph.n))
3450                .filter(|&u| u != v)
3451                .map(|u| u as u32)
3452                .collect();
3453            graph.adj.set(v, &neighbors);
3454        }
3455    }
3456
3457    fn assert_complete(searcher: &VSearcher<'_>, out: &[u32], k: usize, rerank: usize) {
3458        assert_eq!(out.len(), k.min(searcher.index.n));
3459        assert_eq!(
3460            out.iter()
3461                .copied()
3462                .collect::<std::collections::BTreeSet<_>>()
3463                .len(),
3464            out.len()
3465        );
3466        assert!(out.iter().all(|&id| (id as usize) < searcher.index.n));
3467        assert!(searcher.scratch.keys.windows(2).all(|w| w[0] <= w[1]));
3468        if rerank > 0 {
3469            assert!(searcher.scratch.results.iter().all(|r| r.0.is_finite()));
3470            assert!(searcher
3471                .scratch
3472                .results
3473                .windows(2)
3474                .all(|w| { w[0].0.total_cmp(&w[1].0).then(w[0].1.cmp(&w[1].1)).is_le() }));
3475        }
3476    }
3477
3478    #[test]
3479    fn disconnected_graph_returns_the_requested_number_of_distinct_neighbors() {
3480        let rows: Vec<f32> = (0..67 * 2).map(|i| (i * 7 % 31 + 1) as f32).collect();
3481        for metric in [VMetric::Angular, VMetric::Euclidean, VMetric::InnerProduct] {
3482            let mut graph = if metric == VMetric::InnerProduct {
3483                lifted(&rows, 2)
3484            } else {
3485                Graph::build(
3486                    &rows,
3487                    2,
3488                    &BuildParams {
3489                        metric: if metric == VMetric::Angular {
3490                            Metric::Angular
3491                        } else {
3492                            Metric::Euclidean
3493                        },
3494                        degree: 2,
3495                        build_ef: 16,
3496                        ..Default::default()
3497                    },
3498                )
3499                .unwrap()
3500            };
3501            disconnect(&mut graph, 3);
3502            for bits in [4, 8] {
3503                for id_bytes in [0, 3, 4] {
3504                    let index = if metric == VMetric::InnerProduct {
3505                        VGraph::from_lifted_ip_graph(&graph, bits, 8, true, true, id_bytes).unwrap()
3506                    } else {
3507                        VGraph::from_graph_with(
3508                            &graph,
3509                            bits,
3510                            8,
3511                            true,
3512                            true,
3513                            id_bytes,
3514                            metric == VMetric::Angular,
3515                        )
3516                        .unwrap()
3517                    };
3518                    for (ef, reservoir, rerank, patience, neighbour_cap) in [
3519                        (0, 0, 0, 0, 0),
3520                        (8, 0, 67, 1, 1),
3521                        (8, 1, 67, 1, 0),
3522                        (8, 68, 67, 1, 1),
3523                    ] {
3524                        let knobs = SearchKnobs {
3525                            prefetch: 3,
3526                            batch_est: true,
3527                            frontier: 0.0,
3528                            rerank_z: 1.0,
3529                            ..SearchKnobs::OFF
3530                        };
3531                        let mut searcher = index.searcher_with(knobs);
3532                        let params = VSearchParams {
3533                            ef,
3534                            reservoir,
3535                            rerank,
3536                            patience,
3537                            neighbour_cap,
3538                            screen: reservoir == 0,
3539                            screen_sigmas: 0.0,
3540                            ..Default::default()
3541                        };
3542                        let mut out = Vec::new();
3543                        for k in [0, 1, 2, 3, 4, 63, 64, 65, 66, 67, 68, usize::MAX] {
3544                            let evaluations =
3545                                searcher.search(&[1.0, 2.0], k, params, &mut out).unwrap();
3546                            assert_complete(&searcher, &out, k, rerank);
3547                            let first = out.clone();
3548                            assert_eq!(
3549                                searcher.search(&[1.0, 2.0], k, params, &mut out).unwrap(),
3550                                evaluations
3551                            );
3552                            assert_eq!(out, first);
3553                        }
3554                    }
3555                }
3556            }
3557        }
3558    }
3559
3560    #[test]
3561    fn recovery_expands_the_seed_and_counts_its_score() {
3562        let rows = [
3563            1.0, 0.0, 2.0, 0.0, 3.0, 0.0, -10.0, 0.0, 5.0, 0.0, 6.0, 0.0, 100.0, 0.0,
3564        ];
3565        let mut graph = lifted(&rows, 2);
3566        disconnect(&mut graph, 1);
3567        graph.adj.set(0, &[1]);
3568        graph.adj.set(1, &[2]);
3569        graph.adj.set(3, &[5]);
3570        graph.adj.set(5, &[4]);
3571        let index = VGraph::from_lifted_ip_graph(&graph, 8, 8, false, false, 4).unwrap();
3572        let mut searcher = index.searcher_with(SearchKnobs::OFF);
3573        for reservoir in [0, 4] {
3574            let params = VSearchParams {
3575                ef: 4,
3576                rerank: 4,
3577                reservoir,
3578                patience: 1,
3579                ..Default::default()
3580            };
3581            let mut out = Vec::new();
3582            let evaluations = searcher.search(&[1.0, 0.0], 4, params, &mut out).unwrap();
3583            assert_eq!(out, [5, 4, 2, 1]);
3584            assert_eq!(evaluations, 10);
3585            assert!(!searcher.scratch.visited.seen(6));
3586            #[cfg(feature = "experimental")]
3587            {
3588                let profile = searcher
3589                    .profile_search(&[1.0, 0.0], 4, params, &mut out)
3590                    .unwrap();
3591                assert_eq!(profile.traversal_evaluations, 6);
3592                assert_eq!(
3593                    profile.entry_evaluations
3594                        + profile.traversal_evaluations
3595                        + profile.rerank_evaluations,
3596                    evaluations
3597                );
3598            }
3599        }
3600        searcher
3601            .search(
3602                &[0.0, 0.0],
3603                7,
3604                VSearchParams {
3605                    rerank: 7,
3606                    ..Default::default()
3607                },
3608                &mut Vec::new(),
3609            )
3610            .unwrap();
3611        assert_eq!(
3612            searcher
3613                .scratch
3614                .results
3615                .iter()
3616                .map(|r| r.1)
3617                .collect::<Vec<_>>(),
3618            [0, 1, 2, 3, 4, 5, 6]
3619        );
3620    }
3621
3622    #[test]
3623    fn disconnected_blocks_retain_enough_expansions_with_stopping() {
3624        let rows = crate::core::qgraph::tests::clustered(67, 16, 77);
3625        let mut graph = Graph::build(
3626            &rows,
3627            16,
3628            &BuildParams {
3629                degree: 32,
3630                build_ef: 64,
3631                ..Default::default()
3632            },
3633        )
3634        .unwrap();
3635        for component_size in [3, 32, 33] {
3636            disconnect(&mut graph, component_size);
3637            for bits in [1, 2, 4] {
3638                for id_bytes in [0, 3, 4] {
3639                    let index = VGraph::from_graph_with(&graph, 4, 8, true, true, id_bytes, false)
3640                        .unwrap()
3641                        .with_blocks(&graph, bits)
3642                        .unwrap();
3643                    for (fuse, block_screen) in [(false, false), (true, false), (false, true)] {
3644                        let mut searcher = index.searcher_with(SearchKnobs {
3645                            frontier: 0.0,
3646                            block_screen,
3647                            ..SearchKnobs::OFF
3648                        });
3649                        for rerank in [0, 67] {
3650                            let params = VSearchParams {
3651                                ef: 1,
3652                                rerank,
3653                                patience: 1,
3654                                fuse,
3655                                screen_sigmas: 0.0,
3656                                ..Default::default()
3657                            };
3658                            let mut out = Vec::new();
3659                            for k in [
3660                                0,
3661                                component_size - 1,
3662                                component_size,
3663                                component_size + 1,
3664                                66,
3665                                67,
3666                                68,
3667                            ] {
3668                                searcher.search(&rows[..16], k, params, &mut out).unwrap();
3669                                assert_complete(&searcher, &out, k, rerank);
3670                            }
3671                        }
3672                    }
3673                }
3674            }
3675        }
3676    }
3677
3678    #[test]
3679    fn recovery_reuses_marks_and_fused_traces_across_epoch_rollover() {
3680        let rows = crate::core::qgraph::tests::clustered(67, 16, 79);
3681        let mut graph = Graph::build(
3682            &rows,
3683            16,
3684            &BuildParams {
3685                degree: 2,
3686                build_ef: 16,
3687                ..Default::default()
3688            },
3689        )
3690        .unwrap();
3691        disconnect(&mut graph, 3);
3692        let index = VGraph::from_graph(&graph, true)
3693            .unwrap()
3694            .with_blocks(&graph, 2)
3695            .unwrap();
3696        let mut searcher = index.searcher_with(SearchKnobs::OFF);
3697        let mut out = Vec::new();
3698        for iteration in 0..260 {
3699            for fuse in [false, true] {
3700                let k = if iteration % 2 == 0 { 4 } else { 65 };
3701                let query = &rows[(iteration % 67) * 16..][..16];
3702                let params = VSearchParams {
3703                    ef: 1,
3704                    rerank: 67,
3705                    patience: 1,
3706                    fuse,
3707                    ..Default::default()
3708                };
3709                let mut fresh = index.searcher_with(SearchKnobs::OFF);
3710                let mut expected = Vec::new();
3711                let evaluations = fresh.search(query, k, params, &mut expected).unwrap();
3712                assert_eq!(
3713                    searcher.search(query, k, params, &mut out).unwrap(),
3714                    evaluations
3715                );
3716                assert_eq!(out, expected);
3717                assert_complete(&searcher, &out, k, params.rerank);
3718            }
3719        }
3720    }
3721
3722    #[test]
3723    fn singleton_search_bounds_k_without_recovery() {
3724        let graph = lifted(&[2.0, 0.0], 2);
3725        let index = VGraph::from_lifted_ip_graph(&graph, 8, 8, false, false, 4).unwrap();
3726        let mut searcher = index.searcher_with(SearchKnobs::OFF);
3727        for k in [0, 1, 2, usize::MAX] {
3728            let mut out = Vec::new();
3729            searcher
3730                .search(
3731                    &[1.0, 0.0],
3732                    k,
3733                    VSearchParams {
3734                        ef: 0,
3735                        ..Default::default()
3736                    },
3737                    &mut out,
3738                )
3739                .unwrap();
3740            assert_complete(&searcher, &out, k, 0);
3741        }
3742    }
3743
3744    fn lifted(rows: &[f32], dim: usize) -> Graph {
3745        let norms: Vec<f64> = rows
3746            .chunks_exact(dim)
3747            .map(|row| row.iter().map(|&x| f64::from(x).powi(2)).sum())
3748            .collect();
3749        let max = norms.iter().copied().fold(0.0, f64::max);
3750        let data: Vec<f32> = rows
3751            .chunks_exact(dim)
3752            .zip(norms)
3753            .flat_map(|(row, norm)| {
3754                row.iter()
3755                    .copied()
3756                    .chain([(max - norm).max(0.0).sqrt() as f32])
3757            })
3758            .collect();
3759        Graph::build_owned(
3760            data,
3761            dim + 1,
3762            &BuildParams {
3763                metric: Metric::Euclidean,
3764                degree: 2,
3765                build_ef: 16,
3766                ..Default::default()
3767            },
3768        )
3769        .unwrap()
3770    }
3771
3772    #[test]
3773    fn reservoir_recovers_a_beam_rejection_without_changing_expansions() {
3774        let mut graph = lifted(&[0.0, 0.0, 0.8, 0.0, 1.0, 381.0], 2);
3775        graph.entry = 0;
3776        graph.layers.clear();
3777        graph.upper_ids.clear();
3778        graph.adj.set(0, &[1, 2]);
3779        graph.adj.set(1, &[]);
3780        graph.adj.set(2, &[]);
3781        let index = VGraph::from_lifted_ip_graph(&graph, 8, 8, false, false, 4).unwrap();
3782        let mut searcher = index.searcher_with(SearchKnobs::OFF);
3783        let mut out = Vec::new();
3784        let params = VSearchParams {
3785            ef: 1,
3786            rerank: 3,
3787            ..Default::default()
3788        };
3789        searcher.search(&[1.0, 0.0], 1, params, &mut out).unwrap();
3790        assert_eq!(out, [1]);
3791        let beam = searcher.scratch.pool.keys.clone();
3792        let marks = searcher.scratch.visited.words.clone();
3793        for capacity in [1, 2, 3, usize::MAX] {
3794            let params = VSearchParams {
3795                reservoir: capacity,
3796                ..params
3797            };
3798            searcher.search(&[1.0, 0.0], 1, params, &mut out).unwrap();
3799            assert_eq!(searcher.scratch.pool.keys, beam);
3800            assert_eq!(searcher.scratch.visited.words, marks);
3801            assert_eq!(out, if capacity == 1 { vec![1] } else { vec![2] });
3802            assert_eq!(searcher.scratch.results.len(), capacity.min(3));
3803            #[cfg(feature = "experimental")]
3804            {
3805                let profile = searcher
3806                    .profile_search(&[1.0, 0.0], 1, params, &mut out)
3807                    .unwrap();
3808                assert_eq!(profile.expansions, 2);
3809                assert_eq!(profile.rerank_evaluations, capacity.min(3));
3810                assert_eq!(profile.beam_retained, 1);
3811                assert_eq!(profile.retained, capacity.min(3));
3812            }
3813        }
3814    }
3815
3816    #[test]
3817    fn retention_preserves_traversal_with_stopping_and_batch_scoring() {
3818        let rows = crate::core::qgraph::tests::clustered(256, 32, 77);
3819        for metric in [Metric::Angular, Metric::Euclidean] {
3820            let graph = Graph::build(
3821                &rows,
3822                32,
3823                &BuildParams {
3824                    metric,
3825                    degree: 8,
3826                    build_ef: 32,
3827                    ..Default::default()
3828                },
3829            )
3830            .unwrap();
3831            let index = VGraph::from_graph_with(&graph, 4, 8, true, true, 3, false).unwrap();
3832            for patience in [0, 1, 4] {
3833                for knobs in [
3834                    SearchKnobs::OFF,
3835                    SearchKnobs {
3836                        prefetch: 3,
3837                        batch_est: true,
3838                        frontier: 0.8,
3839                        block_screen: false,
3840                        rerank_z: 0.0,
3841                    },
3842                ] {
3843                    let mut a = index.searcher_with(knobs);
3844                    let mut b = index.searcher_with(knobs);
3845                    for q in rows.as_chunks::<32>().0.iter().take(8) {
3846                        let mut out = Vec::new();
3847                        let p = VSearchParams {
3848                            ef: 12,
3849                            rerank: 40,
3850                            patience,
3851                            ..Default::default()
3852                        };
3853                        let ea = a.search(q, 5, p, &mut out).unwrap();
3854                        let eb = b
3855                            .search(q, 5, VSearchParams { reservoir: 40, ..p }, &mut out)
3856                            .unwrap();
3857                        assert_eq!(a.scratch.pool.keys, b.scratch.pool.keys);
3858                        assert_eq!(a.scratch.pool.expanded, b.scratch.pool.expanded);
3859                        assert_eq!(a.scratch.visited.words, b.scratch.visited.words);
3860                        assert_eq!(ea - a.scratch.results.len(), eb - b.scratch.results.len());
3861                    }
3862                }
3863            }
3864        }
3865    }
3866
3867    #[test]
3868    fn reservoir_includes_upper_layer_scores_once() {
3869        let mut graph = lifted(&[0.0, 0.0, 0.8, 0.0, 1.0, 381.0], 2);
3870        graph.upper_ids = vec![0, 1, 2];
3871        graph.layers = vec![vec![vec![1, 2], vec![0, 2], vec![0, 1]]];
3872        for v in 0..3 {
3873            graph.adj.set(v, &[]);
3874        }
3875        let index = VGraph::from_lifted_ip_graph(&graph, 8, 8, false, false, 4).unwrap();
3876        let mut searcher = index.searcher_with(SearchKnobs::OFF);
3877        let mut out = Vec::new();
3878        for _ in 0..260 {
3879            searcher
3880                .search(
3881                    &[1.0, 0.0],
3882                    1,
3883                    VSearchParams {
3884                        ef: 1,
3885                        rerank: 3,
3886                        reservoir: 3,
3887                        ..Default::default()
3888                    },
3889                    &mut out,
3890                )
3891                .unwrap();
3892            assert_eq!(out, [2]);
3893            assert_eq!(searcher.scratch.keys.len(), 3);
3894        }
3895    }
3896
3897    #[test]
3898    fn direct_ip_preserves_magnitudes_zero_queries_and_topology() {
3899        let rows = [1.0, 0.0, 0.0, 3.0, 0.0, 3.0, -2.0, 0.0];
3900        let mut graph = lifted(&rows, 2);
3901        for v in 0..graph.n {
3902            graph.adj.set(v, &[((v + 1) % graph.n) as u32]);
3903        }
3904        for bits in [4, 8] {
3905            for reorder in [false, true] {
3906                let index =
3907                    VGraph::from_lifted_ip_graph(&graph, bits, 8, true, reorder, 3).unwrap();
3908                assert_eq!(index.metric(), VMetric::InnerProduct);
3909                assert!(index.half_norms.is_empty());
3910                for v in 0..index.n {
3911                    let mut neighbors = Vec::new();
3912                    index.for_each_neighbour(v, |u| neighbors.push(index.original(u)));
3913                    assert_eq!(
3914                        neighbors,
3915                        graph.neighbors(index.original(v as u32) as usize)
3916                    );
3917                }
3918                let mut searcher = index.searcher_with(SearchKnobs::OFF);
3919                let mut out = Vec::new();
3920                let params = VSearchParams {
3921                    ef: 4,
3922                    rerank: 4,
3923                    reservoir: 4,
3924                    ..Default::default()
3925                };
3926                for (query, expected) in [
3927                    ([1.0, 1.0], vec![1, 2, 0, 3]),
3928                    ([0.0, 0.0], vec![0, 1, 2, 3]),
3929                    ([-1.0, 0.0], vec![3, 1, 2, 0]),
3930                ] {
3931                    searcher.search(&query, 4, params, &mut out).unwrap();
3932                    assert_eq!(out, expected);
3933                }
3934                assert!(searcher
3935                    .search(&[f32::NAN, 0.0], 1, params, &mut out)
3936                    .is_err());
3937                assert!(searcher.search(&[0.0], 1, params, &mut out).is_err());
3938                assert!(index.with_blocks(&graph, 2).is_err());
3939            }
3940        }
3941    }
3942
3943    #[test]
3944    fn direct_ip_removes_the_128_dimension_stride_cliff() {
3945        let mut rows = vec![0.0; 3 * 128];
3946        rows[0] = 1.0;
3947        rows[128] = 2.0;
3948        rows[256] = -1.0;
3949        let graph = lifted(&rows, 128);
3950        let old = VGraph::from_graph_with(&graph, 8, 8, true, true, 3, false).unwrap();
3951        let new = VGraph::from_lifted_ip_graph(&graph, 8, 8, true, true, 3).unwrap();
3952        assert_eq!(
3953            (old.dim, old.padded, old.stride, old.res_stride),
3954            (129, 144, 256, 144)
3955        );
3956        assert_eq!(
3957            (new.dim, new.padded, new.stride, new.res_stride),
3958            (128, 128, 128, 128)
3959        );
3960        assert_eq!(
3961            old.index_bytes() - new.index_bytes(),
3962            148 * graph.n + 16 * 4
3963        );
3964        assert_eq!(old.entry, new.entry);
3965        assert_eq!(old.upper_ids, new.upper_ids);
3966        assert_eq!(old.layers, new.layers);
3967        assert_eq!(&*old.ids, &*new.ids);
3968        assert_eq!(&*old.adj3, &*new.adj3);
3969        for dim in [127, 129] {
3970            let mut rows = vec![0.0; 3 * dim];
3971            rows[0] = 1.0;
3972            rows[dim] = 2.0;
3973            rows[2 * dim] = -1.0;
3974            let graph = lifted(&rows, dim);
3975            let old = VGraph::from_graph_with(&graph, 8, 8, true, true, 3, false).unwrap();
3976            let new = VGraph::from_lifted_ip_graph(&graph, 8, 8, true, true, 3).unwrap();
3977            assert_eq!((new.stride, new.res_stride), (old.stride, old.res_stride));
3978            assert_eq!(old.index_bytes() - new.index_bytes(), 4 * graph.n);
3979        }
3980    }
3981}
3982
3983/// Format 02 adds the build parameters and a fingerprint of the vectors to the header, so a
3984/// cache written for other data or parameters is rejected instead of silently reused; 03 the
3985/// quantized-search flag; 04 the metric.
3986const CACHE_MAGIC: &[u8; 8] = b"HMSGRF04";
3987
3988/// Fingerprint of the prepared vectors (normalized for angular) a graph was built over.
3989fn data_fingerprint(vectors: &[f32]) -> u64 {
3990    let mut h = super::Fnv::new();
3991    vectors.iter().for_each(|x| h.word(x.to_bits()));
3992    h.finish()
3993}
3994
3995fn put_u32s(w: &mut impl std::io::Write, xs: &[u32]) -> std::io::Result<()> {
3996    w.write_all(&(xs.len() as u64).to_le_bytes())?;
3997    for x in xs {
3998        w.write_all(&x.to_le_bytes())?;
3999    }
4000    Ok(())
4001}
4002
4003fn get_u64(r: &mut impl std::io::Read) -> std::io::Result<u64> {
4004    let mut b = [0u8; 8];
4005    r.read_exact(&mut b)?;
4006    Ok(u64::from_le_bytes(b))
4007}
4008
4009fn get_u32s(r: &mut impl std::io::Read, max: u64) -> std::io::Result<Vec<u32>> {
4010    let len = get_u64(r)?;
4011    if len > max {
4012        return Err(std::io::Error::new(
4013            std::io::ErrorKind::InvalidData,
4014            "graph cache: list longer than the graph",
4015        ));
4016    }
4017    // Read through `take`, so the buffer grows with the bytes the file holds rather than
4018    // being allocated at the size its lengths claim.
4019    let want = len.saturating_mul(4);
4020    let mut bytes = Vec::new();
4021    std::io::Read::read_to_end(&mut std::io::Read::take(&mut *r, want), &mut bytes)?;
4022    if bytes.len() as u64 != want {
4023        return Err(std::io::Error::new(
4024            std::io::ErrorKind::UnexpectedEof,
4025            "graph cache: list shorter than its lengths",
4026        ));
4027    }
4028    Ok(bytes
4029        .as_chunks::<4>()
4030        .0
4031        .iter()
4032        .map(|c| u32::from_le_bytes(*c))
4033        .collect())
4034}
4035
4036fn put_lists<'a>(
4037    w: &mut impl std::io::Write,
4038    lists: impl Iterator<Item = &'a [u32]> + Clone,
4039) -> std::io::Result<()> {
4040    let lens: Vec<u32> = lists.clone().map(|l| l.len() as u32).collect();
4041    put_u32s(w, &lens)?;
4042    put_u32s(w, &lists.flatten().copied().collect::<Vec<u32>>())
4043}
4044
4045fn get_lists(r: &mut impl std::io::Read, n: u64) -> std::io::Result<Vec<Vec<u32>>> {
4046    let lens = get_u32s(r, n)?;
4047    let flat = get_u32s(r, lens.iter().map(|&l| u64::from(l)).sum())?;
4048    let mut out = Vec::with_capacity(lens.len());
4049    let mut at = 0;
4050    for l in lens {
4051        out.push(flat[at..at + l as usize].to_vec());
4052        at += l as usize;
4053    }
4054    Ok(out)
4055}
4056
4057/// Benchmark support: the graph structure (not the vectors) on disk, so that search
4058/// experiments on one graph need not rebuild it.
4059impl Graph {
4060    /// Writes the header (n, dim, degree, build_ef, alpha, seed, entry, data fingerprint,
4061    /// quantized_search, metric), the adjacency and the upper layers; the vectors are not
4062    /// stored.
4063    pub fn save(&self, path: &std::path::Path) -> std::io::Result<()> {
4064        if self.degree != self.build.degree {
4065            return Err(std::io::Error::new(
4066                std::io::ErrorKind::InvalidInput,
4067                "query-trained graph cannot use an untrained graph cache key; encode and save VGraph instead",
4068            ));
4069        }
4070        let mut w = std::io::BufWriter::new(std::fs::File::create(path)?);
4071        use std::io::Write;
4072        w.write_all(CACHE_MAGIC)?;
4073        let header = [
4074            self.n as u64,
4075            self.dim as u64,
4076            self.degree as u64,
4077            self.build.build_ef as u64,
4078            u64::from(self.build.alpha.to_bits()),
4079            self.build.seed,
4080            u64::from(self.entry),
4081            data_fingerprint(&self.vectors),
4082            u64::from(self.build.quantized_search),
4083            self.metric.code(),
4084        ];
4085        for x in header {
4086            w.write_all(&x.to_le_bytes())?;
4087        }
4088        put_lists(&mut w, self.adj.iter())?;
4089        put_u32s(&mut w, &self.upper_ids)?;
4090        w.write_all(&(self.layers.len() as u64).to_le_bytes())?;
4091        for l in &self.layers {
4092            put_lists(&mut w, l.iter().map(Vec::as_slice))?;
4093        }
4094        w.flush()
4095    }
4096
4097    /// Reads a graph written by [`Graph::save`] for the same `data` (prepared here exactly
4098    /// as [`Graph::build`] does) and the same `params`. Any mismatch in the vector count,
4099    /// dimension, degree, build_ef, alpha, seed, quantized search, metric or the data
4100    /// fingerprint is an error, so a stale cache never stands in for a build; delete the file
4101    /// to rebuild.
4102    pub fn load(
4103        data: &[f32],
4104        dim: usize,
4105        path: &std::path::Path,
4106        params: &BuildParams,
4107    ) -> std::io::Result<Self> {
4108        let bad = |m: String| std::io::Error::new(std::io::ErrorKind::InvalidData, m);
4109        let mut r = std::io::BufReader::new(std::fs::File::open(path)?);
4110        let mut magic = [0u8; 8];
4111        std::io::Read::read_exact(&mut r, &mut magic)?;
4112        if &magic != CACHE_MAGIC {
4113            return Err(bad(
4114                "graph cache: bad magic (an older format or not a cache file); delete it to rebuild"
4115                    .into(),
4116            ));
4117        }
4118        let mut header = [0u64; 10];
4119        for x in &mut header {
4120            *x = get_u64(&mut r)?;
4121        }
4122        let [n, d, degree, build_ef, alpha_bits, seed, entry, fingerprint, quantized, metric] =
4123            header;
4124        if dim == 0 || !data.len().is_multiple_of(dim) {
4125            return Err(bad("graph cache: ragged data".into()));
4126        }
4127        let mut vectors = data.to_vec();
4128        let norms = prepare_rows(&mut vectors, dim, params.metric);
4129        let checks: [(&str, u64, u64); 9] = [
4130            ("metric", metric, params.metric.code()),
4131            ("n", n, (data.len() / dim) as u64),
4132            ("dim", d, dim as u64),
4133            ("degree", degree, params.degree as u64),
4134            ("build_ef", build_ef, params.build_ef as u64),
4135            ("alpha", alpha_bits, u64::from(params.alpha.to_bits())),
4136            ("seed", seed, params.seed),
4137            ("data fingerprint", fingerprint, data_fingerprint(&vectors)),
4138            (
4139                "quantized_search",
4140                quantized,
4141                u64::from(params.quantized_search),
4142            ),
4143        ];
4144        if let Some((what, file, want)) = checks.iter().find(|(_, a, b)| a != b) {
4145            let show = |x: u64| {
4146                if *what == "alpha" {
4147                    f32::from_bits(x as u32).to_string()
4148                } else {
4149                    format!("{x:#x}")
4150                }
4151            };
4152            return Err(bad(format!(
4153                "graph cache: {what} {} in the file, {} requested; delete the file to rebuild",
4154                show(*file),
4155                show(*want)
4156            )));
4157        }
4158        let n = n as usize;
4159        if entry >= n as u64 {
4160            return Err(bad("graph cache: entry out of range".into()));
4161        }
4162        let lens = get_u32s(&mut r, n as u64)?;
4163        let flat = get_u32s(&mut r, lens.iter().map(|&l| u64::from(l)).sum())?;
4164        let adj = Adjacency::from_flat(lens, &flat, degree as usize)
4165            .ok_or_else(|| bad("graph cache: list longer than the degree".into()))?;
4166        drop(flat);
4167        let upper_ids = get_u32s(&mut r, n as u64)?;
4168        let nl = get_u64(&mut r)?;
4169        if nl > 64 {
4170            return Err(bad("graph cache: too many layers".into()));
4171        }
4172        let layers = (0..nl)
4173            .map(|_| get_lists(&mut r, n as u64))
4174            .collect::<std::io::Result<Vec<_>>>()?;
4175        let ok = adj.len() == n
4176            && adj.iter().flatten().all(|&u| (u as usize) < n)
4177            && upper_ids.iter().all(|&u| (u as usize) < n)
4178            && layers.iter().all(|l| {
4179                l.len() <= upper_ids.len() && l.iter().flatten().all(|&u| (u as usize) < l.len())
4180            });
4181        if !ok {
4182            return Err(bad("graph cache: id out of range".into()));
4183        }
4184        Ok(Self {
4185            n,
4186            dim,
4187            degree: degree as usize,
4188            metric: params.metric,
4189            vectors,
4190            norms,
4191            adj,
4192            entry: entry as u32,
4193            upper_ids,
4194            layers,
4195            build: params.clone(),
4196        })
4197    }
4198}
4199
4200#[cfg(test)]
4201mod tests {
4202    use super::super::kernels::SplitMix;
4203    use super::super::tests::{brute_force, clustered, gaussian};
4204    use super::*;
4205
4206    #[test]
4207    #[cfg(feature = "experimental-query-x4")]
4208    fn query_x4_estimates_match_serial_at_padding_boundaries() {
4209        for dim in [
4210            15, 16, 17, 31, 32, 33, 63, 64, 65, 111, 112, 113, 127, 128, 129,
4211        ] {
4212            let data = clustered(40, dim, 173);
4213            let queries = clustered(2, dim, 174);
4214            for metric in [Metric::Angular, Metric::Euclidean] {
4215                let graph = Graph::build(
4216                    &data,
4217                    dim,
4218                    &BuildParams {
4219                        metric,
4220                        degree: 8,
4221                        build_ef: 16,
4222                        quantized_search: true,
4223                        ..Default::default()
4224                    },
4225                )
4226                .unwrap();
4227                for bits in [4, 8] {
4228                    for align in [false, true] {
4229                        let index = VGraph::from_graph_with(&graph, bits, 8, align, true, 3, false)
4230                            .unwrap();
4231                        let mut searcher = index.searcher_with(SearchKnobs::OFF);
4232                        let scratch = &mut searcher.scratch;
4233                        for query in queries.chunks_exact(dim) {
4234                            scratch.qf[..dim].copy_from_slice(query);
4235                            index.encode_query(
4236                                &scratch.qf,
4237                                &mut scratch.qc,
4238                                &mut scratch.qtmp,
4239                                &mut scratch.qrot,
4240                                &mut scratch.qs,
4241                            );
4242                            let qsum = query_sum(&scratch.qc);
4243                            for ids in [[0, 1, 2, 3], [39, 16, 31, 0], [5, 5, 5, 5]] {
4244                                for w in [0.0, 1.0, 3.25] {
4245                                    let expected =
4246                                        ids.map(|u| index.est(&scratch.qc, qsum, u, w).to_bits());
4247                                    let actual = index
4248                                        .est4::<true>(&scratch.qc, qsum, ids, w)
4249                                        .map(f32::to_bits);
4250                                    assert_eq!(actual, expected, "dim {dim}, bits {bits}");
4251                                }
4252                            }
4253                        }
4254                    }
4255                }
4256            }
4257        }
4258    }
4259
4260    #[test]
4261    #[cfg(feature = "experimental-query-x4")]
4262    fn query_x4_preserves_admission_state_at_batch_boundaries() {
4263        let dim = 17;
4264        let data = clustered(48, dim, 175);
4265        let query = clustered(1, dim, 176);
4266        for metric in [Metric::Angular, Metric::Euclidean] {
4267            let graph = Graph::build(
4268                &data,
4269                dim,
4270                &BuildParams {
4271                    metric,
4272                    degree: 8,
4273                    build_ef: 16,
4274                    quantized_search: true,
4275                    ..Default::default()
4276                },
4277            )
4278            .unwrap();
4279            for bits in [4, 8] {
4280                let index = VGraph::from_graph_with(&graph, bits, 8, true, true, 3, false).unwrap();
4281                for count in (0..=9).chain(15..=17) {
4282                    for cap in [3, 4, 5, 8] {
4283                        for batch_est in [false, true] {
4284                            for patience in [0, 4] {
4285                                for reservoir in [0, 12] {
4286                                    let knobs = SearchKnobs {
4287                                        batch_est,
4288                                        ..SearchKnobs::OFF
4289                                    };
4290                                    let mut serial = index.searcher_with(knobs);
4291                                    let mut fused = index.searcher_with(knobs);
4292                                    for s in [&mut serial.scratch, &mut fused.scratch] {
4293                                        s.qf[..dim].copy_from_slice(&query);
4294                                        index.encode_query(
4295                                            &s.qf,
4296                                            &mut s.qc,
4297                                            &mut s.qtmp,
4298                                            &mut s.qrot,
4299                                            &mut s.qs,
4300                                        );
4301                                        s.qsum = query_sum(&s.qc);
4302                                        s.pool.reset(cap);
4303                                        s.reservoir.reset(reservoir);
4304                                        for id in 32..48 {
4305                                            let kk = key(index.est(&s.qc, s.qsum, id, 1.0), id);
4306                                            s.pool.insert(kk);
4307                                            s.reservoir.offer(kk);
4308                                        }
4309                                        s.pool.next();
4310                                        s.fresh.extend(0..count);
4311                                    }
4312                                    let (mut a, mut b) = (patience, patience);
4313                                    let expected = serial.scratch.score_fresh::<{
4314                                        cfg!(feature = "experimental-batch-pool")
4315                                    }, false>(
4316                                        &index,
4317                                        cap.min(3),
4318                                        patience,
4319                                        &mut a,
4320                                        1.0,
4321                                    );
4322                                    let actual = fused.scratch.score_fresh::<{
4323                                        cfg!(feature = "experimental-batch-pool")
4324                                    }, true>(
4325                                        &index,
4326                                        cap.min(3),
4327                                        patience,
4328                                        &mut b,
4329                                        1.0,
4330                                    );
4331                                    assert_eq!(actual, count as usize);
4332                                    assert_eq!(actual, expected);
4333                                    assert_eq!(a, b);
4334                                    assert_eq!(serial.scratch.pool.keys, fused.scratch.pool.keys);
4335                                    assert_eq!(
4336                                        serial.scratch.pool.expanded,
4337                                        fused.scratch.pool.expanded
4338                                    );
4339                                    assert_eq!(
4340                                        serial.scratch.pool.cursor,
4341                                        fused.scratch.pool.cursor
4342                                    );
4343                                    assert_eq!(
4344                                        serial.scratch.reservoir.keys,
4345                                        fused.scratch.reservoir.keys
4346                                    );
4347                                }
4348                            }
4349                        }
4350                    }
4351                }
4352            }
4353        }
4354    }
4355
4356    #[test]
4357    #[cfg(feature = "experimental-query-x4")]
4358    fn query_x4_preserves_search_results_work_and_reopening() {
4359        let directory = tempfile::tempdir().unwrap();
4360        for dim in [17, 33] {
4361            let mut data = clustered(80, dim, 177);
4362            let first = data[..dim].to_vec();
4363            data[7 * dim..8 * dim].copy_from_slice(&first);
4364            data[19 * dim..20 * dim].copy_from_slice(&first);
4365            let queries = clustered(2, dim, 178);
4366            for metric in [Metric::Angular, Metric::Euclidean] {
4367                let graph = Graph::build(
4368                    &data,
4369                    dim,
4370                    &BuildParams {
4371                        metric,
4372                        degree: 16,
4373                        build_ef: 32,
4374                        quantized_search: true,
4375                        ..Default::default()
4376                    },
4377                )
4378                .unwrap();
4379                for bits in [4, 8] {
4380                    for mode in 0..3 {
4381                        if mode > 0 && metric != Metric::Angular {
4382                            continue;
4383                        }
4384                        let mut index =
4385                            VGraph::from_graph_with(&graph, bits, 8, true, true, 3, mode == 1)
4386                                .unwrap();
4387                        if mode == 2 {
4388                            index = index.with_blocks(&graph, 4).unwrap();
4389                        }
4390                        let path = directory.path().join("query-x4.warm");
4391                        index.save(&path).unwrap();
4392                        let reopened = VGraph::open(&path, Residency::Mapped).unwrap();
4393                        let knobs = SearchKnobs {
4394                            block_screen: mode == 2,
4395                            ..SearchKnobs::OFF
4396                        };
4397                        let mut serial = index.searcher_with(knobs);
4398                        let mut fused = reopened.searcher_with(knobs);
4399                        for (ef, neighbour_cap, patience, reservoir) in
4400                            [(3, 0, 0, 0), (4, 3, 1, 0), (5, 4, 4, 0), (17, 5, 0, 24)]
4401                        {
4402                            if mode > 0 && reservoir > 0 {
4403                                continue;
4404                            }
4405                            let params = VSearchParams {
4406                                ef,
4407                                rerank: 24,
4408                                neighbour_cap,
4409                                patience,
4410                                reservoir,
4411                                ..Default::default()
4412                            };
4413                            for k in [0, 1, 5] {
4414                                for query in queries.chunks_exact(dim) {
4415                                    let (mut expected, mut actual) = (Vec::new(), Vec::new());
4416                                    let (mut a, mut b) =
4417                                        (QueryProfile::default(), QueryProfile::default());
4418                                    let ea = serial
4419                                        .scratch
4420                                        .search_inner::<
4421                                            true,
4422                                            { cfg!(feature = "experimental-batch-pool") },
4423                                            false,
4424                                        >(&index, query, k, params, &mut expected, &mut a)
4425                                        .unwrap();
4426                                    let eb = fused
4427                                        .scratch
4428                                        .search_inner::<
4429                                            true,
4430                                            { cfg!(feature = "experimental-batch-pool") },
4431                                            true,
4432                                        >(&reopened, query, k, params, &mut actual, &mut b)
4433                                        .unwrap();
4434                                    assert_eq!(expected, actual);
4435                                    assert_eq!(ea, eb);
4436                                    let work = |p: QueryProfile| {
4437                                        (
4438                                            p.expansions,
4439                                            p.neighbors_read,
4440                                            p.entry_evaluations,
4441                                            p.traversal_evaluations,
4442                                            p.rerank_evaluations,
4443                                            p.beam_retained,
4444                                            p.retained,
4445                                        )
4446                                    };
4447                                    assert_eq!(work(a), work(b));
4448                                    assert_eq!(serial.scratch.pool.keys, fused.scratch.pool.keys);
4449                                    assert_eq!(
4450                                        serial.scratch.pool.expanded,
4451                                        fused.scratch.pool.expanded
4452                                    );
4453                                    assert_eq!(
4454                                        serial.scratch.pool.cursor,
4455                                        fused.scratch.pool.cursor
4456                                    );
4457                                    assert_eq!(
4458                                        serial.scratch.visited.words,
4459                                        fused.scratch.visited.words
4460                                    );
4461                                    assert_eq!(
4462                                        serial.scratch.reservoir.keys,
4463                                        fused.scratch.reservoir.keys
4464                                    );
4465                                    assert_eq!(serial.last_screened(), fused.last_screened());
4466                                    let mut ordinary = Vec::new();
4467                                    assert_eq!(
4468                                        fused.search(query, k, params, &mut ordinary).unwrap(),
4469                                        eb
4470                                    );
4471                                    assert_eq!(ordinary, actual);
4472                                }
4473                            }
4474                        }
4475                    }
4476                }
4477            }
4478        }
4479    }
4480
4481    /// A refused huge-page advice falls back to base pages with identical arrays and results,
4482    /// and is counted. Runs the build twice in one process: advice failing, then as configured.
4483    #[test]
4484    fn refused_huge_page_advice_changes_nothing_but_the_counter() {
4485        // Large enough that the code rows span a 2 MiB-aligned interior.
4486        let (n, dim) = (40_000, 128);
4487        let data = clustered(n, dim, 3);
4488        let graph = Graph::build(
4489            &data,
4490            dim,
4491            &BuildParams {
4492                degree: 8,
4493                build_ef: 16,
4494                metric: Metric::Euclidean,
4495                ..BuildParams::default()
4496            },
4497        )
4498        .unwrap();
4499        let build = || VGraph::from_graph_with(&graph, 8, 8, true, true, 3, false).unwrap();
4500        FAIL_HUGE_ADVICE.store(true, Ordering::Relaxed);
4501        let failed_before = huge_page_status().failed_calls;
4502        let refused = build();
4503        FAIL_HUGE_ADVICE.store(false, Ordering::Relaxed);
4504        let advised = build();
4505        if cfg!(target_os = "linux") && huge_pages() {
4506            assert!(huge_page_status().failed_calls > failed_before);
4507        }
4508        assert_eq!(refused.codes.bytes(), advised.codes.bytes());
4509        assert_eq!(&refused.scales[..], &advised.scales[..]);
4510        assert_eq!(&refused.half_norms[..], &advised.half_norms[..]);
4511        let params = VSearchParams {
4512            ef: 40,
4513            rerank: 40,
4514            ..VSearchParams::default()
4515        };
4516        let (mut a, mut b) = (refused.searcher(), advised.searcher());
4517        let (mut ra, mut rb) = (Vec::new(), Vec::new());
4518        for q in data.chunks_exact(dim).step_by(997) {
4519            a.search(q, 10, params, &mut ra).unwrap();
4520            b.search(q, 10, params, &mut rb).unwrap();
4521            assert_eq!(ra, rb);
4522        }
4523    }
4524
4525    #[test]
4526    fn visited_bits_reset_between_searches() {
4527        // Ids at word edges and the last vertex of a count that is not a multiple of 64.
4528        let n = 64 * 5 + 3;
4529        let mut marks = Marks::new(n, false);
4530        let mut rng = SplitMix(5);
4531        for search in 0..300 {
4532            marks.next();
4533            let mut set = std::collections::HashSet::new();
4534            let fixed = [0, 63, 64, 127, n as u32 - 1];
4535            for i in 0..40 {
4536                let id = if i < fixed.len() && search % 2 == 0 {
4537                    fixed[i]
4538                } else {
4539                    (rng.next_u64() % n as u64) as u32
4540                };
4541                assert_eq!(marks.seen(id), set.contains(&id));
4542                assert_eq!(marks.insert(id), set.insert(id));
4543                assert!(marks.seen(id));
4544            }
4545            let ones: u32 = marks.words.iter().map(|w| w.count_ones()).sum();
4546            assert_eq!(ones as usize, set.len());
4547        }
4548    }
4549
4550    #[test]
4551    #[cfg(feature = "experimental-neighbor-admission")]
4552    fn neighbor_admission_matches_ordered_set_model_at_boundaries() {
4553        fn check<const PROFILE: bool, const SCREEN: bool>(
4554            index: &VGraph,
4555            vertex: usize,
4556            list: &[u32],
4557            cap: usize,
4558        ) {
4559            let mut visited = Marks::new(index.n, false);
4560            for initial in [&[][..], &[0, 63, 64, 127, 128, 130][..], list] {
4561                visited.next();
4562                let mut seen = std::collections::HashSet::new();
4563                for &u in initial {
4564                    visited.insert(u);
4565                    seen.insert(u);
4566                }
4567                let limit = if cap == 0 { list.len() } else { cap };
4568                let expected: Vec<_> = list
4569                    .iter()
4570                    .take(limit)
4571                    .copied()
4572                    .filter(|&u| seen.insert(u))
4573                    .collect();
4574                let mut fresh = Vec::new();
4575                let read = index.admit_neighbours::<PROFILE, SCREEN>(
4576                    vertex,
4577                    cap,
4578                    &mut visited,
4579                    &mut fresh,
4580                );
4581                assert_eq!(fresh, expected, "cap={cap}, vertex={vertex}");
4582                // The existing profile counts all uncapped list entries, including duplicates.
4583                assert_eq!(read, if PROFILE && cap == 0 { list.len() } else { 0 });
4584                for u in 0..index.n as u32 {
4585                    assert_eq!(visited.seen(u), seen.contains(&u));
4586                }
4587                let mut words = std::collections::HashSet::new();
4588                let touched: Vec<_> = initial
4589                    .iter()
4590                    .chain(&expected)
4591                    .map(|&u| u >> 6)
4592                    .filter(|&word| words.insert(word))
4593                    .collect();
4594                assert_eq!(visited.touched, touched);
4595                visited.next();
4596                assert!(visited.words.iter().all(|&word| word == 0));
4597                assert!(visited.touched.is_empty());
4598            }
4599        }
4600
4601        let dim = 17;
4602        let data = clustered(131, dim, 91);
4603        let lists = [
4604            Vec::new(),
4605            vec![130],
4606            vec![63, 64, 65, 127, 128, 129, 130, 63, 64, 0, 0],
4607            (0..63).rev().collect(),
4608            (0..64).rev().collect(),
4609            (0..65).rev().collect(),
4610        ];
4611        for metric in [Metric::Angular, Metric::Euclidean] {
4612            let mut graph = Graph::build(
4613                &data,
4614                dim,
4615                &BuildParams {
4616                    metric,
4617                    degree: 66,
4618                    build_ef: 80,
4619                    ..BuildParams::default()
4620                },
4621            )
4622            .unwrap();
4623            for bytes in [3, 4, 0] {
4624                let ordered: Vec<Vec<u32>> = lists
4625                    .iter()
4626                    .map(|list| {
4627                        let mut list = list.clone();
4628                        if bytes == 0 {
4629                            list.sort_unstable();
4630                            list.dedup();
4631                        }
4632                        list
4633                    })
4634                    .collect();
4635                for v in 0..graph.n {
4636                    graph.adj.set(v, &ordered[v % ordered.len()]);
4637                }
4638                let mut index = VGraph::from_graph_with(
4639                    &graph,
4640                    8,
4641                    8,
4642                    true,
4643                    false,
4644                    bytes,
4645                    metric == Metric::Angular,
4646                )
4647                .unwrap();
4648                for compact in [false, true] {
4649                    if compact {
4650                        index.compact_adjacency().unwrap();
4651                        assert!(bytes == 0 || !index.adj_off.is_empty());
4652                    }
4653                    for (vertex, list) in ordered.iter().enumerate() {
4654                        for cap in [0, 1, 2, 63, 64, 65, 66, usize::MAX] {
4655                            check::<false, false>(&index, vertex, list, cap);
4656                            check::<true, false>(&index, vertex, list, cap);
4657                            if metric == Metric::Angular {
4658                                check::<false, true>(&index, vertex, list, cap);
4659                                check::<true, true>(&index, vertex, list, cap);
4660                            }
4661                        }
4662                    }
4663                }
4664            }
4665        }
4666    }
4667
4668    /// The pool before its flags became bits: one `bool` per key, moved with the keys.
4669    struct ReferencePool {
4670        keys: Vec<u64>,
4671        expanded: Vec<bool>,
4672        cap: usize,
4673        cursor: usize,
4674    }
4675
4676    impl ReferencePool {
4677        fn insert(&mut self, key: u64) -> Option<(usize, Option<u64>)> {
4678            let n = self.keys.len();
4679            if n >= self.cap && key >= self.keys[n - 1] {
4680                return None;
4681            }
4682            let pos = self.keys.partition_point(|&t| t < key);
4683            let mut evicted = None;
4684            if n >= self.cap {
4685                evicted = self.keys.pop();
4686                self.expanded.pop();
4687            }
4688            self.keys.insert(pos, key);
4689            self.expanded.insert(pos, false);
4690            if pos < self.cursor {
4691                self.cursor = pos;
4692            }
4693            Some((pos, evicted))
4694        }
4695
4696        fn remove(&mut self, key: u64) {
4697            let pos = self.keys.partition_point(|&t| t < key);
4698            if pos < self.keys.len() && self.keys[pos] == key {
4699                self.keys.remove(pos);
4700                self.expanded.remove(pos);
4701                if pos < self.cursor {
4702                    self.cursor -= 1;
4703                }
4704            }
4705        }
4706
4707        fn next(&mut self) -> Option<u64> {
4708            while self.cursor < self.keys.len() {
4709                if !self.expanded[self.cursor] {
4710                    self.expanded[self.cursor] = true;
4711                    return Some(self.keys[self.cursor]);
4712                }
4713                self.cursor += 1;
4714            }
4715            None
4716        }
4717
4718        fn peek_unexpanded(&self, n: usize) -> Option<u64> {
4719            (self.cursor..self.keys.len())
4720                .filter(|&i| !self.expanded[i])
4721                .nth(n)
4722                .map(|i| self.keys[i])
4723        }
4724    }
4725
4726    #[test]
4727    fn pool_bits_follow_the_bool_pool() {
4728        let mut rng = SplitMix(9);
4729        let mut pool = Pool {
4730            keys: Vec::new(),
4731            expanded: Vec::new(),
4732            #[cfg(feature = "experimental-batch-pool")]
4733            merge_keys: Vec::new(),
4734            #[cfg(feature = "experimental-batch-pool")]
4735            merge_expanded: Vec::new(),
4736            cap: 1,
4737            cursor: 0,
4738        };
4739        // Capacities on both sides of each word boundary, and one wider than any list here.
4740        for cap in [1usize, 2, 10, 63, 64, 65, 127, 128, 129, 300, 10_000] {
4741            for _ in 0..20 {
4742                pool.reset(cap);
4743                let mut model = ReferencePool {
4744                    keys: Vec::new(),
4745                    expanded: Vec::new(),
4746                    cap,
4747                    cursor: 0,
4748                };
4749                for _ in 0..600 {
4750                    match rng.next_u64() % 8 {
4751                        0 | 1 => assert_eq!(pool.next(), model.next()),
4752                        2 if !model.keys.is_empty() => {
4753                            let k = model.keys[(rng.next_u64() % model.keys.len() as u64) as usize];
4754                            pool.remove(k);
4755                            model.remove(k);
4756                        }
4757                        3 => {
4758                            let n = (rng.next_u64() % 3) as usize;
4759                            assert_eq!(pool.peek_unexpanded(n), model.peek_unexpanded(n));
4760                        }
4761                        _ => {
4762                            let k = rng.next_u64() % 5000;
4763                            assert_eq!(pool.insert(k), model.insert(k));
4764                        }
4765                    }
4766                    assert_eq!(pool.keys, model.keys);
4767                    assert_eq!(pool.cursor, model.cursor);
4768                    let flags: Vec<bool> =
4769                        (0..pool.keys.len()).map(|i| pool.is_expanded(i)).collect();
4770                    assert_eq!(flags, model.expanded);
4771                    let set: u32 = pool.expanded.iter().map(|w| w.count_ones()).sum();
4772                    assert_eq!(set as usize, flags.iter().filter(|&&f| f).count());
4773                }
4774            }
4775        }
4776    }
4777
4778    #[test]
4779    #[cfg(feature = "experimental-batch-pool")]
4780    fn batch_pool_matches_sequential_insertion() {
4781        fn check(pool: &Pool, model: &ReferencePool) {
4782            assert_eq!(pool.keys, model.keys);
4783            assert_eq!(pool.cursor, model.cursor);
4784            let flags: Vec<_> = (0..pool.keys.len()).map(|i| pool.is_expanded(i)).collect();
4785            assert_eq!(flags, model.expanded);
4786            assert_eq!(
4787                pool.expanded
4788                    .iter()
4789                    .map(|word| word.count_ones())
4790                    .sum::<u32>() as usize,
4791                flags.iter().filter(|&&flag| flag).count()
4792            );
4793            for next in [0, 1, 2] {
4794                assert_eq!(pool.peek_unexpanded(next), model.peek_unexpanded(next));
4795            }
4796        }
4797
4798        let mut rng = SplitMix(539365412);
4799        for cap in [1usize, 2, 7, 8, 9, 63, 64, 65, 127, 128, 129] {
4800            for k in [0, 1, cap - 1, cap, cap + 1] {
4801                for count in [0, 1, 7, 8, 9, cap + 1] {
4802                    for fill in [0, cap / 2, cap] {
4803                        let mut pool = Pool {
4804                            keys: Vec::new(),
4805                            expanded: Vec::new(),
4806                            #[cfg(feature = "experimental-batch-pool")]
4807                            merge_keys: Vec::new(),
4808                            #[cfg(feature = "experimental-batch-pool")]
4809                            merge_expanded: Vec::new(),
4810                            cap,
4811                            cursor: 0,
4812                        };
4813                        let mut model = ReferencePool {
4814                            keys: Vec::new(),
4815                            expanded: Vec::new(),
4816                            cap,
4817                            cursor: 0,
4818                        };
4819                        let mut id = 0;
4820                        for _ in 0..fill {
4821                            let candidate = key(1000.0 + (rng.next_u64() % 16) as f32, id);
4822                            id += 1;
4823                            assert_eq!(pool.insert(candidate), model.insert(candidate));
4824                        }
4825                        for _ in 0..=fill {
4826                            assert_eq!(pool.next(), model.next());
4827                        }
4828                        for round in 0..8 {
4829                            let size = if round == 0 {
4830                                count
4831                            } else {
4832                                (rng.next_u64() % 17) as usize
4833                            };
4834                            let mut pending = Vec::new();
4835                            for _ in 0..size {
4836                                let distance = if round == 0 {
4837                                    (rng.next_u64() % 16) as f32
4838                                } else {
4839                                    (rng.next_u64() % 32) as f32 - 8.0
4840                                };
4841                                pending.push(key(distance, id));
4842                                id += 1;
4843                            }
4844                            let mut improved = false;
4845                            for &candidate in &pending {
4846                                if let Some((position, _)) = model.insert(candidate) {
4847                                    improved |= position < k;
4848                                }
4849                            }
4850                            assert_eq!(pool.insert_batch(&mut pending, k), improved);
4851                            check(&pool, &model);
4852                            for _ in 0..3 {
4853                                assert_eq!(pool.next(), model.next());
4854                            }
4855                            if round % 2 == 1 && !model.keys.is_empty() {
4856                                let removed =
4857                                    model.keys[(rng.next_u64() % model.keys.len() as u64) as usize];
4858                                pool.remove(removed);
4859                                model.remove(removed);
4860                            }
4861                            check(&pool, &model);
4862                        }
4863                    }
4864                }
4865            }
4866        }
4867    }
4868
4869    #[test]
4870    #[cfg(feature = "experimental-batch-pool")]
4871    fn batch_pool_equal_distances_preserve_distinct_ids_and_flags() {
4872        for cap in [1usize, 2, 7, 8, 9, 63, 64, 65, 127, 128, 129] {
4873            for fill in [cap - 1, cap] {
4874                for first_id in [10, 1000] {
4875                    let mut pool = Pool {
4876                        keys: Vec::new(),
4877                        expanded: Vec::new(),
4878                        #[cfg(feature = "experimental-batch-pool")]
4879                        merge_keys: Vec::new(),
4880                        #[cfg(feature = "experimental-batch-pool")]
4881                        merge_expanded: Vec::new(),
4882                        cap,
4883                        cursor: 0,
4884                    };
4885                    let mut model = ReferencePool {
4886                        keys: Vec::new(),
4887                        expanded: Vec::new(),
4888                        cap,
4889                        cursor: 0,
4890                    };
4891                    for id in first_id..first_id + fill as u32 {
4892                        assert_eq!(pool.insert(key(1.0, id)), model.insert(key(1.0, id)));
4893                    }
4894                    for _ in 0..=fill {
4895                        assert_eq!(pool.next(), model.next());
4896                    }
4897                    let mut pending: Vec<_> = [0, 200, 201, 202, 203, 204, 205, 206, 207]
4898                        .map(|id| key(1.0, id))
4899                        .to_vec();
4900                    let mut improved = false;
4901                    for &candidate in &pending {
4902                        if let Some((position, _)) = model.insert(candidate) {
4903                            improved |= position < 1;
4904                        }
4905                    }
4906                    assert_eq!(pool.insert_batch(&mut pending, 1), improved);
4907                    assert_eq!(pool.keys, model.keys);
4908                    assert_eq!(pool.cursor, model.cursor);
4909                    for (position, &expanded) in model.expanded.iter().enumerate() {
4910                        assert_eq!(pool.is_expanded(position), expanded);
4911                    }
4912                    while let Some(expected) = model.next() {
4913                        assert_eq!(pool.next(), Some(expected));
4914                    }
4915                    assert_eq!(pool.next(), None);
4916                }
4917            }
4918        }
4919    }
4920
4921    #[test]
4922    #[cfg(feature = "experimental-batch-pool")]
4923    fn batch_pool_preserves_search_results_work_and_reopening() {
4924        let dim = 16;
4925        let data = clustered(160, dim, 53);
4926        let queries = clustered(3, dim, 54);
4927        let directory = tempfile::tempdir().unwrap();
4928        for metric in [Metric::Angular, Metric::Euclidean] {
4929            let graph = Graph::build(
4930                &data,
4931                dim,
4932                &BuildParams {
4933                    metric,
4934                    degree: 32,
4935                    build_ef: 64,
4936                    quantized_search: true,
4937                    ..Default::default()
4938                },
4939            )
4940            .unwrap();
4941            for bits in [4, 8] {
4942                let index = VGraph::from_graph_with(&graph, bits, 8, true, true, 3, false).unwrap();
4943                let path = directory.path().join("batch-pool.warm");
4944                index.save(&path).unwrap();
4945                let reopened = VGraph::open(&path, Residency::Mapped).unwrap();
4946                for batch_est in [false, true] {
4947                    let knobs = SearchKnobs {
4948                        batch_est,
4949                        ..SearchKnobs::OFF
4950                    };
4951                    let mut plain = index.searcher_with(knobs);
4952                    let mut batched = reopened.searcher_with(knobs);
4953                    for ef in [7, 8, 9, 64, 129] {
4954                        for k in [0, 1, 7, 8, 9, 64] {
4955                            for patience in [0, 1, 4] {
4956                                for reservoir in [0, 80] {
4957                                    let params = VSearchParams {
4958                                        ef,
4959                                        rerank: 80,
4960                                        patience,
4961                                        reservoir,
4962                                        ..Default::default()
4963                                    };
4964                                    for query in queries.chunks_exact(dim) {
4965                                        let (mut expected, mut actual) = (Vec::new(), Vec::new());
4966                                        let (mut a, mut b) =
4967                                            (QueryProfile::default(), QueryProfile::default());
4968                                        let ea = plain
4969                                            .scratch
4970                                            .search_inner::<true, false, { cfg!(feature = "experimental-query-x4") }>(
4971                                                &index,
4972                                                query,
4973                                                k,
4974                                                params,
4975                                                &mut expected,
4976                                                &mut a,
4977                                            )
4978                                            .unwrap();
4979                                        let eb = batched
4980                                            .scratch
4981                                            .search_inner::<true, true, { cfg!(feature = "experimental-query-x4") }>(
4982                                                &reopened,
4983                                                query,
4984                                                k,
4985                                                params,
4986                                                &mut actual,
4987                                                &mut b,
4988                                            )
4989                                            .unwrap();
4990                                        assert_eq!(expected, actual);
4991                                        assert_eq!(ea, eb);
4992                                        let work = |p: QueryProfile| {
4993                                            (
4994                                                p.expansions,
4995                                                p.neighbors_read,
4996                                                p.entry_evaluations,
4997                                                p.traversal_evaluations,
4998                                                p.rerank_evaluations,
4999                                                p.beam_retained,
5000                                                p.retained,
5001                                            )
5002                                        };
5003                                        assert_eq!(work(a), work(b));
5004                                        assert_eq!(
5005                                            plain.scratch.pool.keys,
5006                                            batched.scratch.pool.keys
5007                                        );
5008                                        assert_eq!(
5009                                            plain.scratch.pool.expanded,
5010                                            batched.scratch.pool.expanded
5011                                        );
5012                                        assert_eq!(
5013                                            plain.scratch.pool.cursor,
5014                                            batched.scratch.pool.cursor
5015                                        );
5016                                        assert_eq!(
5017                                            plain.scratch.visited.words,
5018                                            batched.scratch.visited.words
5019                                        );
5020                                        assert_eq!(
5021                                            plain.scratch.reservoir.keys,
5022                                            batched.scratch.reservoir.keys
5023                                        );
5024                                        let mut unprofiled = Vec::new();
5025                                        assert_eq!(
5026                                            plain
5027                                                .search(query, k, params, &mut unprofiled)
5028                                                .unwrap(),
5029                                            ea
5030                                        );
5031                                        assert_eq!(unprofiled, expected);
5032                                        assert_eq!(
5033                                            batched
5034                                                .search(query, k, params, &mut unprofiled)
5035                                                .unwrap(),
5036                                            eb
5037                                        );
5038                                        assert_eq!(unprofiled, actual);
5039                                    }
5040                                }
5041                            }
5042                        }
5043                    }
5044                }
5045            }
5046        }
5047    }
5048
5049    /// Development microbenchmark (`--ignored`): a search-shaped stream (one expansion, then
5050    /// 20 inserts of keys drawn near the pool's range) on the bit pool and the bool pool.
5051    #[test]
5052    #[ignore]
5053    fn microbench_pool() {
5054        use std::hint::black_box;
5055        use std::time::Instant;
5056        for cap in [10usize, 20, 40, 80, 160, 320] {
5057            let stream: Vec<u64> = {
5058                let mut rng = SplitMix(cap as u64);
5059                (0..200_000).map(|_| rng.next_u64() % 100_000).collect()
5060            };
5061            let mut best = [f64::INFINITY; 2];
5062            for _ in 0..7 {
5063                for (which, b) in best.iter_mut().enumerate() {
5064                    let t = Instant::now();
5065                    let mut pool = Pool {
5066                        keys: Vec::new(),
5067                        expanded: Vec::new(),
5068                        #[cfg(feature = "experimental-batch-pool")]
5069                        merge_keys: Vec::new(),
5070                        #[cfg(feature = "experimental-batch-pool")]
5071                        merge_expanded: Vec::new(),
5072                        cap: 1,
5073                        cursor: 0,
5074                    };
5075                    let mut model = ReferencePool {
5076                        keys: Vec::new(),
5077                        expanded: Vec::new(),
5078                        cap,
5079                        cursor: 0,
5080                    };
5081                    pool.reset(cap);
5082                    let mut acc = 0u64;
5083                    for (i, &k) in stream.iter().enumerate() {
5084                        if which == 0 {
5085                            if i % 21 == 0 {
5086                                if pool.next().is_none() {
5087                                    pool.reset(cap);
5088                                }
5089                            } else if let Some((p, _)) = pool.insert(black_box(k)) {
5090                                acc += p as u64;
5091                            }
5092                        } else if i % 21 == 0 {
5093                            if model.next().is_none() {
5094                                model = ReferencePool {
5095                                    keys: Vec::new(),
5096                                    expanded: Vec::new(),
5097                                    cap,
5098                                    cursor: 0,
5099                                };
5100                            }
5101                        } else if let Some((p, _)) = model.insert(black_box(k)) {
5102                            acc += p as u64;
5103                        }
5104                    }
5105                    black_box(acc);
5106                    *b = b.min(t.elapsed().as_nanos() as f64 / stream.len() as f64);
5107                }
5108            }
5109            println!(
5110                "{{\"cap\": {cap}, \"bits_ns_per_op\": {:.2}, \"bools_ns_per_op\": {:.2}}}",
5111                best[0], best[1]
5112            );
5113        }
5114    }
5115
5116    #[test]
5117    fn compact_lists_preserve_order_search_and_persisted_bytes() {
5118        let dim = 16;
5119        let data = clustered(64, dim, 7);
5120        let queries = clustered(8, dim, 8);
5121        let mut graph = Graph::build(
5122            &data,
5123            dim,
5124            &BuildParams {
5125                degree: 16,
5126                build_ef: 32,
5127                ..BuildParams::default()
5128            },
5129        )
5130        .unwrap();
5131        for v in 0..graph.n {
5132            let count = [0, 1, 15, 16][v % 4];
5133            let ids: Vec<u32> = (1..=count).map(|j| ((v + j) % graph.n) as u32).collect();
5134            graph.adj.set(v, &ids);
5135        }
5136        let directory = tempfile::tempdir().unwrap();
5137        for bytes in [3, 4] {
5138            for bits in [4, 8] {
5139                let mut index =
5140                    VGraph::from_graph_with(&graph, bits, 8, true, true, bytes, false).unwrap();
5141                let lists: Vec<Vec<u32>> = (0..index.n)
5142                    .map(|v| {
5143                        let mut ids = Vec::new();
5144                        index.for_each_neighbour(v, |u| ids.push(u));
5145                        ids
5146                    })
5147                    .collect();
5148                let before = directory.path().join("before");
5149                let after = directory.path().join("after");
5150                index.save(&before).unwrap();
5151                let old_bytes = index.index_bytes();
5152                let old_memory = index.memory_breakdown();
5153                let params = VSearchParams {
5154                    ef: 32,
5155                    rerank: 16,
5156                    ..VSearchParams::default()
5157                };
5158                let expected: Vec<Vec<u32>> = {
5159                    let mut searcher = index.searcher_with(SearchKnobs::OFF);
5160                    queries
5161                        .chunks_exact(dim)
5162                        .map(|q| {
5163                            let mut out = Vec::new();
5164                            searcher.search(q, 10, params, &mut out).unwrap();
5165                            out
5166                        })
5167                        .collect()
5168                };
5169                index.compact_adjacency().unwrap();
5170                assert!(index.index_bytes() < old_bytes);
5171                let memory = index.memory_breakdown();
5172                assert_eq!(
5173                    old_memory.adjacency_ids_bytes,
5174                    graph.n * graph.degree * bytes as usize
5175                );
5176                assert_eq!(old_memory.adjacency_offsets_bytes, 0);
5177                assert_eq!(
5178                    memory.adjacency_ids_bytes,
5179                    lists.iter().map(Vec::len).sum::<usize>() * bytes as usize
5180                );
5181                assert_eq!(memory.adjacency_offsets_bytes, (graph.n + 1) * 4);
5182                assert_eq!(
5183                    old_bytes - index.index_bytes(),
5184                    old_memory.adjacency_ids_bytes
5185                        - memory.adjacency_ids_bytes
5186                        - memory.adjacency_offsets_bytes
5187                );
5188                for (v, want) in lists.iter().enumerate() {
5189                    let mut actual = Vec::new();
5190                    index.for_each_neighbour(v, |u| actual.push(u));
5191                    assert_eq!(&actual, want);
5192                }
5193                let compact_bytes = index.index_bytes();
5194                index.compact_adjacency().unwrap();
5195                assert_eq!(index.index_bytes(), compact_bytes);
5196                index.save(&after).unwrap();
5197                assert_eq!(
5198                    std::fs::read(&before).unwrap(),
5199                    std::fs::read(&after).unwrap()
5200                );
5201                for residency in [Residency::Mapped, Residency::Prefaulted, Residency::Heap] {
5202                    let reopened = VGraph::open(&after, residency).unwrap();
5203                    let (mut a, mut b) = (
5204                        index.searcher_with(SearchKnobs::OFF),
5205                        reopened.searcher_with(SearchKnobs::OFF),
5206                    );
5207                    for i in 0..260 {
5208                        let qi = i % expected.len();
5209                        let q = &queries[qi * dim..(qi + 1) * dim];
5210                        let (mut x, mut y) = (Vec::new(), Vec::new());
5211                        let evals = a.search(q, 10, params, &mut x).unwrap();
5212                        b.search(q, 10, params, &mut y).unwrap();
5213                        assert_eq!(x, expected[qi]);
5214                        assert_eq!(x, y);
5215                        #[cfg(feature = "experimental")]
5216                        {
5217                            let profile = a.profile_search(q, 10, params, &mut y).unwrap();
5218                            assert_eq!(x, y);
5219                            assert_eq!(
5220                                evals,
5221                                profile.entry_evaluations
5222                                    + profile.traversal_evaluations
5223                                    + profile.rerank_evaluations
5224                            );
5225                            assert!(profile.neighbors_read >= profile.traversal_evaluations - 1);
5226                        }
5227                        #[cfg(not(feature = "experimental"))]
5228                        let _ = evals;
5229                    }
5230                }
5231            }
5232        }
5233    }
5234
5235    /// A cache whose list lengths claim more data than the file holds is refused by reading,
5236    /// not by first allocating what the lengths claim.
5237    #[test]
5238    fn graph_cache_lengths_are_bounded_by_the_file() {
5239        let d = 16;
5240        let n = 64usize;
5241        let data = clustered(n, d, 4);
5242        let params = BuildParams {
5243            build_ef: 16,
5244            degree: 8,
5245            seed: 3,
5246            ..BuildParams::default()
5247        };
5248        let graph = Graph::build(&data, d, &params).expect("build");
5249        let dir = tempfile::tempdir().unwrap();
5250        let path = dir.path().join("graph.cache");
5251        graph.save(&path).unwrap();
5252        let mut bytes = std::fs::read(&path).unwrap();
5253        // The adjacency lengths list starts with `n` as a u64, the second such word after
5254        // the magic (the first is the header's own `n`).
5255        let lens = (8..bytes.len() - 8)
5256            .step_by(8)
5257            .filter(|&at| u64::from_le_bytes(bytes[at..at + 8].try_into().unwrap()) == n as u64)
5258            .nth(1)
5259            .expect("lengths list");
5260        for i in 0..n {
5261            let at = lens + 8 + 4 * i;
5262            bytes[at..at + 4].copy_from_slice(&u32::MAX.to_le_bytes());
5263        }
5264        let flat = lens + 8 + 4 * n;
5265        bytes[flat..flat + 8].copy_from_slice(&(n as u64 * u64::from(u32::MAX)).to_le_bytes());
5266        std::fs::write(&path, &bytes).unwrap();
5267        let error = Graph::load(&data, d, &path, &params)
5268            .err()
5269            .unwrap()
5270            .to_string();
5271        assert!(error.contains("shorter than its lengths"), "{error}");
5272    }
5273
5274    /// The cache round-trips the structure for the same data and parameters, and rejects a
5275    /// file written for other data or for any other build parameter.
5276    #[test]
5277    fn graph_cache_round_trips_and_rejects_mismatches() {
5278        let d = 16;
5279        let data = clustered(600, d, 4);
5280        let params = BuildParams {
5281            build_ef: 32,
5282            degree: 8,
5283            alpha: 1.1,
5284            seed: 9,
5285            ..BuildParams::default()
5286        };
5287        let graph = Graph::build(&data, d, &params).expect("build");
5288        let path = std::env::temp_dir().join(format!("warm-qgraph-cache-{}", std::process::id()));
5289        graph.save(&path).unwrap();
5290        let back = Graph::load(&data, d, &path, &params).unwrap();
5291        assert_eq!(back.structure_hash(), graph.structure_hash());
5292        assert!(back.adj == graph.adj && back.layers == graph.layers);
5293        assert_eq!(
5294            (back.entry, &back.upper_ids),
5295            (graph.entry, &graph.upper_ids)
5296        );
5297        assert_eq!(back.vectors, graph.vectors);
5298        let rejects = |data: &[f32], p: &BuildParams, what: &str| {
5299            let err = Graph::load(data, d, &path, p)
5300                .err()
5301                .expect("a mismatched cache was accepted")
5302                .to_string();
5303            assert!(err.contains(what), "{what}: {err}");
5304        };
5305        for (what, p) in [
5306            (
5307                "degree",
5308                BuildParams {
5309                    degree: 10,
5310                    ..params.clone()
5311                },
5312            ),
5313            (
5314                "build_ef",
5315                BuildParams {
5316                    build_ef: 33,
5317                    ..params.clone()
5318                },
5319            ),
5320            (
5321                "alpha",
5322                BuildParams {
5323                    alpha: 1.0,
5324                    ..params.clone()
5325                },
5326            ),
5327            (
5328                "seed",
5329                BuildParams {
5330                    seed: 10,
5331                    ..params.clone()
5332                },
5333            ),
5334            (
5335                "metric",
5336                BuildParams {
5337                    metric: Metric::Euclidean,
5338                    ..params.clone()
5339                },
5340            ),
5341        ] {
5342            rejects(&data, &p, what);
5343        }
5344        let mut other = data.clone();
5345        other[d * 7 + 3] += 0.5;
5346        rejects(&other, &params, "data fingerprint");
5347        rejects(&data[..d * 599], &params, "n");
5348        std::fs::write(&path, b"HMSGRF01").unwrap();
5349        rejects(&data, &params, "bad magic");
5350        std::fs::remove_file(&path).unwrap();
5351    }
5352
5353    /// On clustered rows whose norms spread over [0.1, 10], a Euclidean index (4-bit codes,
5354    /// 8-bit residual re-rank, reordered 3-byte ids: the benchmark layout) finds the exact
5355    /// Euclidean neighbours, while an angular index over the same rows, which ranks by
5356    /// direction alone, clearly does not.
5357    #[test]
5358    fn euclidean_index_finds_euclidean_neighbours() {
5359        let (d, n, k) = (32, 4000, 10);
5360        let spread = |rows: &mut [f32], salt: usize| {
5361            let count = rows.len() / d;
5362            for (i, x) in rows.chunks_exact_mut(d).enumerate() {
5363                let s = 0.1 * 100f32.powf(((i * 7919 + salt) % count) as f32 / count as f32);
5364                x.iter_mut().for_each(|v| *v *= s);
5365            }
5366        };
5367        // Queries are held-out rows of the same clusters.
5368        let mut data = clustered(n + 100, d, 7);
5369        spread(&mut data, 0);
5370        let queries = data.split_off(n * d);
5371        let l2 = |a: &[f32], b: &[f32]| -> f64 {
5372            a.iter()
5373                .zip(b)
5374                .map(|(x, y)| (f64::from(*x) - f64::from(*y)).powi(2))
5375                .sum()
5376        };
5377        let recall = |metric: Metric| -> f64 {
5378            let params = BuildParams {
5379                metric,
5380                degree: 32,
5381                build_ef: 100,
5382                ..BuildParams::default()
5383            };
5384            let g = Graph::build(&data, d, &params).expect("build");
5385            let index = VGraph::from_graph_with(&g, 4, 8, true, true, 3, false).expect("encode");
5386            let mut s = index.searcher_with(SearchKnobs::OFF);
5387            let mut out = Vec::new();
5388            let mut hits = 0;
5389            for q in queries.chunks_exact(d) {
5390                let mut all: Vec<(f64, u32)> = data
5391                    .chunks_exact(d)
5392                    .enumerate()
5393                    .map(|(i, x)| (l2(q, x), i as u32))
5394                    .collect();
5395                all.sort_by(|a, b| a.0.total_cmp(&b.0));
5396                let truth: Vec<u32> = all[..k].iter().map(|t| t.1).collect();
5397                let params = VSearchParams {
5398                    ef: 64,
5399                    rerank: 32,
5400                    ..VSearchParams::default()
5401                };
5402                s.search(q, k, params, &mut out).expect("search");
5403                hits += out.iter().filter(|x| truth.contains(x)).count();
5404            }
5405            hits as f64 / (k * 100) as f64
5406        };
5407        let (euclidean, angular) = (recall(Metric::Euclidean), recall(Metric::Angular));
5408        assert!(euclidean >= 0.95, "Euclidean recall {euclidean}");
5409        assert!(
5410            angular < euclidean - 0.3,
5411            "angular recall {angular} against Euclidean {euclidean}"
5412        );
5413    }
5414
5415    /// A Euclidean index refuses the angular-only encodings and accepts a zero query (its
5416    /// nearest rows are the ones of smallest norm).
5417    #[test]
5418    fn euclidean_index_limits() {
5419        let d = 16;
5420        let mut data = clustered(500, d, 2);
5421        for (i, x) in data.chunks_exact_mut(d).enumerate() {
5422            x.iter_mut().for_each(|v| *v *= 1.0 + i as f32 / 50.0);
5423        }
5424        let params = BuildParams {
5425            metric: Metric::Euclidean,
5426            degree: 16,
5427            build_ef: 64,
5428            ..BuildParams::default()
5429        };
5430        let g = Graph::build(&data, d, &params).expect("build");
5431        assert!(VGraph::from_graph_with(&g, 8, 0, false, false, 4, true).is_err());
5432        assert!(super::super::QGraph::from_graph(&g, &params).is_err());
5433        let blocks = VGraph::from_graph(&g, false)
5434            .expect("encode")
5435            .with_blocks(&g, 4);
5436        assert!(blocks.is_err());
5437        let index = VGraph::from_graph(&g, true).expect("encode");
5438        let mut out = Vec::new();
5439        let params = VSearchParams {
5440            ef: 500,
5441            rerank: 500,
5442            ..VSearchParams::default()
5443        };
5444        index
5445            .searcher_with(SearchKnobs::OFF)
5446            .search(&[0.0; 16], 1, params, &mut out)
5447            .expect("zero query");
5448        let smallest = data
5449            .chunks_exact(d)
5450            .enumerate()
5451            .min_by(|a, b| {
5452                let n2 = |x: &[f32]| x.iter().map(|v| v * v).sum::<f32>();
5453                n2(a.1).total_cmp(&n2(b.1))
5454            })
5455            .map(|m| m.0 as u32);
5456        assert_eq!(out.first().copied(), smallest);
5457    }
5458
5459    /// Per coordinate the 8-bit code is off by at most half a step, and the residual code
5460    /// shrinks that error by about another factor of 127.
5461    #[test]
5462    fn quantization_error_is_bounded() {
5463        let mut rng = SplitMix(12);
5464        let x: Vec<f32> = (0..100).map(|_| gaussian(&mut rng)).collect();
5465        let mut c = vec![0i8; 112];
5466        let s = quantize(&x, &mut c);
5467        let mut r: Vec<f32> = x
5468            .iter()
5469            .zip(&c)
5470            .map(|(a, &q)| a - s * f32::from(q))
5471            .collect();
5472        assert!(r.iter().all(|e| e.abs() <= s / 2.0 + 1e-6));
5473        assert!(c[100..].iter().all(|&q| q == 0));
5474        let mut rc = vec![0i8; 112];
5475        let sr = quantize(&r, &mut rc);
5476        r.iter_mut()
5477            .zip(&rc)
5478            .for_each(|(a, &q)| *a -= sr * f32::from(q));
5479        assert!(r.iter().all(|e| e.abs() <= s / 254.0 + 1e-6));
5480        assert_eq!(quantize(&[0.0; 4], &mut c[..16]), 0.0);
5481    }
5482
5483    fn recall(index: &VGraph, data: &[f32], queries: &[f32], d: usize, p: VSearchParams) -> f64 {
5484        let mut s = index.searcher();
5485        let mut out = Vec::new();
5486        let mut hits = 0;
5487        for q in queries.chunks_exact(d) {
5488            let truth = brute_force(data, d, q, 10);
5489            s.search(q, 10, p, &mut out).expect("search");
5490            assert_eq!(out.len(), 10);
5491            hits += out.iter().filter(|x| truth.contains(x)).count();
5492        }
5493        hits as f64 / (queries.len() / d * 10) as f64
5494    }
5495
5496    #[test]
5497    fn recall_on_clustered_data() {
5498        let d = 24;
5499        let data = clustered(4000, d, 1);
5500        let queries = clustered(100, d, 2);
5501        let graph = Graph::build(&data, d, &BuildParams::default()).expect("build");
5502        let plain = VGraph::from_graph(&graph, false).expect("encode");
5503        let fine = VGraph::from_graph(&graph, true).expect("encode");
5504        assert!(plain.index_bytes() < fine.index_bytes());
5505        let pool = VSearchParams {
5506            ef: 64,
5507            ..VSearchParams::default()
5508        };
5509        let rr = VSearchParams {
5510            ef: 64,
5511            rerank: 32,
5512            ..VSearchParams::default()
5513        };
5514        let r0 = recall(&plain, &data, &queries, d, pool);
5515        let r1 = recall(&fine, &data, &queries, d, rr);
5516        // Same regression guard as the edge-code index on this deliberately hard set.
5517        assert!(r0 >= 0.9, "recall {r0}");
5518        assert!(r1 >= r0, "rerank {r1} below pool order {r0}");
5519        let fine4 = VGraph::from_graph_with(&graph, 8, 4, false, false, 4, false).expect("encode");
5520        assert!(plain.index_bytes() < fine4.index_bytes());
5521        assert!(fine4.index_bytes() < fine.index_bytes());
5522        let r2 = recall(&fine4, &data, &queries, d, rr);
5523        assert!(r2 >= r0 - 0.02, "4-bit rerank {r2} below pool order {r0}");
5524        let lvq48 = VGraph::from_graph_with(&graph, 4, 8, true, false, 4, false).expect("encode");
5525        assert!(lvq48.index_bytes() < fine.index_bytes());
5526        // A sanity bound: 4-bit codes of 24 coordinates are coarse, so a wider pool and re-rank.
5527        let wide = VSearchParams {
5528            ef: 96,
5529            rerank: 64,
5530            ..VSearchParams::default()
5531        };
5532        let r3 = recall(&lvq48, &data, &queries, d, wide);
5533        assert!(
5534            r3 >= r0 - 0.05,
5535            "4-bit traversal with 8-bit residual {r3} vs {r0}"
5536        );
5537    }
5538
5539    /// Neighbour-block codes estimate every neighbour from the expanded vertex's record; the
5540    /// traversal must reach the same recall band as the per-row 8-bit estimates after the
5541    /// residual re-rank, at every block width, with and without renumbering or 3-byte ids.
5542    #[test]
5543    fn block_codes_keep_recall() {
5544        let d = 24;
5545        let data = clustered(4000, d, 1);
5546        let queries = clustered(100, d, 2);
5547        let graph = Graph::build(&data, d, &BuildParams::default()).expect("build");
5548        let plain = VGraph::from_graph(&graph, true).expect("encode");
5549        let rr = VSearchParams {
5550            ef: 96,
5551            rerank: 64,
5552            ..VSearchParams::default()
5553        };
5554        let r0 = recall(&plain, &data, &queries, d, rr);
5555        for bits in [4u8, 2, 1] {
5556            for (reorder, id_bytes) in [(false, 4u8), (true, 3)] {
5557                let b = VGraph::from_graph_with(&graph, 8, 8, true, reorder, id_bytes, false)
5558                    .expect("encode")
5559                    .with_blocks(&graph, bits)
5560                    .expect("blocks");
5561                assert!(b.has_blocks());
5562                assert!(b.index_bytes() > plain.index_bytes());
5563                let r = recall(&b, &data, &queries, d, rr);
5564                // 24 coordinates are few for 1-bit codes (a 1/sqrt(dim) relative error on
5565                // this deliberately hard set), so the band widens with coarseness; the
5566                // estimator itself is checked against exact inner products below.
5567                let slack = match bits {
5568                    4 => 0.03,
5569                    2 => 0.06,
5570                    _ => 0.30,
5571                };
5572                assert!(
5573                    r >= r0 - slack,
5574                    "{bits}-bit blocks (reorder {reorder}, ids {id_bytes}) recall {r} vs {r0}"
5575                );
5576            }
5577        }
5578        let e = match VGraph::from_graph(&graph, true)
5579            .expect("encode")
5580            .with_blocks(&graph, 3)
5581        {
5582            Ok(_) => panic!("3 bits accepted"),
5583            Err(e) => e.to_string(),
5584        };
5585        assert!(e.contains("block bits"), "{e}");
5586    }
5587
5588    /// The block estimate of a neighbour, in `est` units, must track the exact centred inner
5589    /// product: Pearson correlation over the neighbours of 64 vertices above a per-width bar.
5590    #[test]
5591    fn block_estimates_track_exact_inner_products() {
5592        let d = 64;
5593        let data = clustered(3000, d, 5);
5594        let graph = Graph::build(&data, d, &BuildParams::default()).expect("build");
5595        // The index works on unit rows; so must the exact products.
5596        let unit = &graph.vectors;
5597        for (bits, bar) in [(4u8, 0.998), (2, 0.98), (1, 0.90)] {
5598            let g = VGraph::from_graph(&graph, false)
5599                .expect("encode")
5600                .with_blocks(&graph, bits)
5601                .expect("blocks");
5602            let mut s = g.searcher();
5603            let q: Vec<f32> = unit[7 * d..8 * d].to_vec();
5604            s.scratch.qf[..d].copy_from_slice(&q);
5605            let sq = g.encode_query(
5606                &s.scratch.qf,
5607                &mut s.scratch.qc,
5608                &mut s.scratch.qtmp,
5609                &mut s.scratch.qrot,
5610                &mut s.scratch.qs,
5611            );
5612            g.encode_block_query(&s.scratch.qf, sq, &mut s.scratch.bqrot, &mut s.scratch.bq);
5613            let centred =
5614                |u: usize| -> f32 { (0..d).map(|j| q[j] * (unit[u * d + j] - g.mean[j])).sum() };
5615            let (mut xs, mut ys) = (Vec::new(), Vec::new());
5616            for v in 0..64usize {
5617                // The exact `<q, x_v - mean>` stands in for the expanded vertex's estimate.
5618                let ip_v = centred(v);
5619                let mut list = Vec::new();
5620                g.for_each_neighbour(v, |u| list.push(u));
5621                for b in 0..g.blocks_per_vertex {
5622                    let (codes, scales, offsets) = g.block(v, b);
5623                    scan32(codes, &s.scratch.bq.lut, g.block_m, &mut s.scratch.scan);
5624                    for i in 0..BLOCK {
5625                        let Some(&u) = list.get(b * BLOCK + i) else {
5626                            break;
5627                        };
5628                        let ip = ip_v
5629                            + le_f32(scales, i)
5630                                * (s.scratch.bq.delta * s.scratch.scan[i] as f32
5631                                    + s.scratch.bq.min_sum)
5632                            + le_f32(offsets, i);
5633                        xs.push(f64::from(ip));
5634                        ys.push(f64::from(centred(u as usize)));
5635                    }
5636                }
5637            }
5638            let n = xs.len() as f64;
5639            let (mx, my) = (xs.iter().sum::<f64>() / n, ys.iter().sum::<f64>() / n);
5640            let (mut sxy, mut sxx, mut syy) = (0.0, 0.0, 0.0);
5641            for (x, y) in xs.iter().zip(&ys) {
5642                sxy += (x - mx) * (y - my);
5643                sxx += (x - mx) * (x - mx);
5644                syy += (y - my) * (y - my);
5645            }
5646            let r = sxy / (sxx * syy).sqrt();
5647            eprintln!(
5648                "{bits}-bit block estimates: correlation {r:.4} over {} pairs",
5649                xs.len()
5650            );
5651            assert!(
5652                r > bar,
5653                "{bits}-bit block estimates: correlation {r} below {bar}"
5654            );
5655        }
5656    }
5657
5658    /// Delta-coded lists hold the same neighbour sets as the 4-byte ones, in fewer bytes,
5659    /// and the search returns the same ids; the block index decodes them in the same order
5660    /// it was built in.
5661    #[test]
5662    fn delta_coded_adjacency_round_trips() {
5663        let d = 24;
5664        let data = clustered(4000, d, 1);
5665        let queries = clustered(50, d, 2);
5666        let graph = Graph::build(&data, d, &BuildParams::default()).expect("build");
5667        let wide = VGraph::from_graph_with(&graph, 4, 8, true, true, 4, false).expect("encode");
5668        let delta = VGraph::from_graph_with(&graph, 4, 8, true, true, 0, false).expect("encode");
5669        assert!(
5670            delta.index_bytes() < wide.index_bytes(),
5671            "{} vs {}",
5672            delta.index_bytes(),
5673            wide.index_bytes()
5674        );
5675        for v in 0..wide.n {
5676            let (mut a, mut b) = (Vec::new(), Vec::new());
5677            wide.for_each_neighbour(v, |u| a.push(u));
5678            delta.for_each_neighbour(v, |u| b.push(u));
5679            a.sort_unstable();
5680            assert_eq!(a, b, "vertex {v}");
5681        }
5682        let p = VSearchParams {
5683            ef: 48,
5684            rerank: 16,
5685            ..VSearchParams::default()
5686        };
5687        let (mut s1, mut s2) = (wide.searcher(), delta.searcher());
5688        let (mut o1, mut o2) = (Vec::new(), Vec::new());
5689        for q in queries.chunks_exact(d) {
5690            s1.search(q, 10, p, &mut o1).expect("search");
5691            s2.search(q, 10, p, &mut o2).expect("search");
5692            assert_eq!(o1, o2);
5693        }
5694        let blocks = VGraph::from_graph_with(&graph, 8, 8, true, true, 0, false)
5695            .expect("encode")
5696            .with_blocks(&graph, 2)
5697            .expect("blocks");
5698        assert!(recall(&blocks, &data, &queries, d, p) > 0.8);
5699        assert_eq!(zigzag(-1), 1);
5700        assert_eq!(unzigzag(zigzag(-123_456)), -123_456);
5701        let mut buf = Vec::new();
5702        put_varint(&mut buf, 300);
5703        assert_eq!(get_varint(&buf), (300, 2));
5704    }
5705
5706    /// Renumbering and aligning the rows change the layout, not the search: the same ids come
5707    /// back (up to ties between equal estimates, which break by the stored id).
5708    #[test]
5709    fn reordered_aligned_index_returns_the_same_ids() {
5710        let d = 24;
5711        let data = clustered(3000, d, 5);
5712        let graph = Graph::build(&data, d, &BuildParams::default()).expect("build");
5713        let a = VGraph::from_graph(&graph, true).expect("encode");
5714        let b = VGraph::from_graph_with(&graph, 8, 8, true, true, 4, false).expect("encode");
5715        // 32-byte rows already sit inside one line; the id map is the only extra.
5716        assert_eq!(b.index_bytes(), a.index_bytes() + b.n * 4);
5717        let (mut sa, mut sb) = (a.searcher(), b.searcher());
5718        let (mut oa, mut ob) = (Vec::new(), Vec::new());
5719        let c = VGraph::from_graph_with(&graph, 8, 8, true, true, 3, false).expect("encode");
5720        assert_eq!(c.index_bytes(), b.index_bytes() - c.n * c.degree);
5721        let mut sc = c.searcher();
5722        let mut oc = Vec::new();
5723        let mut same = 0;
5724        for q in data.chunks_exact(d).take(200) {
5725            for p in [
5726                VSearchParams {
5727                    ef: 32,
5728                    ..VSearchParams::default()
5729                },
5730                VSearchParams {
5731                    ef: 48,
5732                    rerank: 16,
5733                    ..VSearchParams::default()
5734                },
5735            ] {
5736                sa.search(q, 10, p, &mut oa).expect("search");
5737                sb.search(q, 10, p, &mut ob).expect("search");
5738                sc.search(q, 10, p, &mut oc).expect("search");
5739                assert_eq!(ob, oc, "3-byte ids changed the search");
5740                same += usize::from(oa == ob);
5741            }
5742        }
5743        assert!(same >= 396, "{same} of 400 searches agree");
5744    }
5745
5746    #[test]
5747    fn build_and_search_are_deterministic() {
5748        let d = 20;
5749        let data = clustered(3000, d, 3);
5750        let params = BuildParams {
5751            seed: 9,
5752            ..BuildParams::default()
5753        };
5754        let a = VGraph::from_graph(&Graph::build(&data, d, &params).expect("build"), true)
5755            .expect("encode");
5756        let b = VGraph::from_graph(&Graph::build(&data, d, &params).expect("build"), true)
5757            .expect("encode");
5758        assert!(
5759            a.codes.bytes() == b.codes.bytes()
5760                && a.residual[..] == b.residual[..]
5761                && a.adj[..] == b.adj[..]
5762        );
5763        assert!(a.scales[..] == b.scales[..] && a.res_scales[..] == b.res_scales[..]);
5764        assert_eq!((a.entry, &a.mean), (b.entry, &b.mean));
5765        let (mut sa, mut sb) = (a.searcher(), b.searcher());
5766        let (mut oa, mut ob) = (Vec::new(), Vec::new());
5767        let p = VSearchParams {
5768            ef: 32,
5769            rerank: 16,
5770            ..VSearchParams::default()
5771        };
5772        for q in data.chunks_exact(d).take(50) {
5773            sa.search(q, 10, p, &mut oa).expect("search");
5774            sb.search(q, 10, p, &mut ob).expect("search");
5775            assert_eq!(oa, ob);
5776        }
5777    }
5778
5779    /// Every failure mode is an error that names its cause; no input panics.
5780    #[test]
5781    fn invalid_inputs_are_errors_that_name_the_cause() {
5782        let d = 16;
5783        let data = clustered(500, d, 4);
5784        let err = |r: Result<Graph, WarmError>| match r {
5785            Err(e) => e.to_string(),
5786            Ok(_) => panic!("an invalid input was accepted"),
5787        };
5788        assert!(err(Graph::build(&[], d, &BuildParams::default())).contains("empty"));
5789        assert!(err(Graph::build(&data[..d + 1], d, &BuildParams::default())).contains("ragged"));
5790        assert!(err(Graph::build(&data, 0, &BuildParams::default())).contains("ragged"));
5791        let mut bad = data.clone();
5792        bad[3 * d + 2] = f32::NAN;
5793        let e = err(Graph::build(&bad, d, &BuildParams::default()));
5794        assert!(e.contains("non-finite") && e.contains("row 3"), "{e}");
5795        let odd = BuildParams {
5796            degree: 7,
5797            ..BuildParams::default()
5798        };
5799        assert!(err(Graph::build(&data, d, &odd)).contains("even"));
5800        let graph = Graph::build(&data, d, &BuildParams::default()).expect("build");
5801        let e = match VGraph::from_graph_with(&graph, 5, 8, false, false, 4, false) {
5802            Err(e) => e.to_string(),
5803            Ok(_) => panic!("bits 5 was accepted"),
5804        };
5805        assert!(e.contains("bits"), "{e}");
5806        let index = VGraph::from_graph(&graph, true).expect("encode");
5807        let mut s = index.searcher();
5808        let mut out = Vec::new();
5809        let p = VSearchParams::default();
5810        let e = s
5811            .search(&data[..d - 1], 10, p, &mut out)
5812            .expect_err("short query")
5813            .to_string();
5814        assert!(e.contains("dimension"), "{e}");
5815        let mut q = data[..d].to_vec();
5816        q[5] = f32::INFINITY;
5817        let e = s.search(&q, 10, p, &mut out).expect_err("inf").to_string();
5818        assert!(e.contains("non-finite") && e.contains('5'), "{e}");
5819        let e = s
5820            .search(&vec![0.0; d], 10, p, &mut out)
5821            .expect_err("zero")
5822            .to_string();
5823        assert!(e.contains("zero"), "{e}");
5824        // Adversarial sizes stay bounded by the vertex count: a huge ef and k finish and
5825        // return at most n ids.
5826        let huge = VSearchParams {
5827            ef: usize::MAX,
5828            rerank: usize::MAX,
5829            ..VSearchParams::default()
5830        };
5831        let evals = s.search(&data[..d], 10, huge, &mut out).expect("search");
5832        assert_eq!(out.len(), 10);
5833        assert!(evals <= 2 * index.len() + 10);
5834        s.search(&data[..d], 5000, huge, &mut out).expect("search");
5835        assert!(out.len() <= index.len());
5836        // A zero data row is allowed and never returned ahead of real neighbours by a NaN.
5837        let mut with_zero = data.clone();
5838        with_zero[..d].fill(0.0);
5839        let g0 = Graph::build(&with_zero, d, &BuildParams::default()).expect("build");
5840        let i0 = VGraph::from_graph(&g0, true).expect("encode");
5841        let mut s0 = i0.searcher();
5842        s0.search(&data[d..2 * d], 10, p, &mut out).expect("search");
5843        assert_eq!(out.len(), 10);
5844        assert!(out.iter().all(|&v| (v as usize) < i0.len()));
5845    }
5846
5847    /// The screen only decides which neighbours get the 8-bit estimate: with a wide margin
5848    /// it skips nothing and the ids are identical; at one sigma recall stays close.
5849    #[test]
5850    fn screen_is_transparent_at_a_wide_margin() {
5851        let d = 24;
5852        let data = clustered(3000, d, 5);
5853        let queries = clustered(100, d, 2);
5854        let graph = Graph::build(&data, d, &BuildParams::default()).expect("build");
5855        let plain = VGraph::from_graph(&graph, true).expect("encode");
5856        let screened =
5857            VGraph::from_graph_with(&graph, 8, 8, false, false, 4, true).expect("encode");
5858        assert!(screened.has_screen() && screened.index_bytes() > plain.index_bytes());
5859        let (mut sa, mut sb) = (plain.searcher(), screened.searcher());
5860        let (mut oa, mut ob) = (Vec::new(), Vec::new());
5861        let wide = VSearchParams {
5862            ef: 64,
5863            rerank: 16,
5864            screen_sigmas: 1e6,
5865            ..VSearchParams::default()
5866        };
5867        for q in queries.chunks_exact(d) {
5868            let ea = sa.search(q, 10, wide, &mut oa).expect("search");
5869            let eb = sb.search(q, 10, wide, &mut ob).expect("search");
5870            assert_eq!((ea, &oa), (eb, &ob));
5871            assert_eq!(sb.last_screened(), 0);
5872        }
5873        let one = VSearchParams {
5874            ef: 64,
5875            rerank: 16,
5876            ..VSearchParams::default()
5877        };
5878        let r0 = recall(&plain, &data, &queries, d, one);
5879        let r1 = recall(&screened, &data, &queries, d, one);
5880        assert!(r1 >= r0 - 0.02, "screened {r1} vs {r0}");
5881        let mut skipped = 0;
5882        for q in queries.chunks_exact(d) {
5883            sb.search(q, 10, one, &mut ob).expect("search");
5884            skipped += sb.last_screened();
5885        }
5886        assert!(skipped > 0, "the screen never skipped a neighbour");
5887    }
5888}