Skip to main content

vecq_core/
store.rs

1//! Quantized store: nibble-packed 4-bit codes + per-vector norm correction,
2//! with asymmetric inner-product scoring (query f32, database 4-bit).
3
4use crate::lloyd;
5use crate::rhdh::{padded_dim, Rhdh};
6use std::collections::HashMap;
7
8/// A quantized vector database in memory.
9///
10/// Each vector is stored as `padded_dim / 2` bytes of 4-bit Lloyd-Max codes
11/// (computed after RHDH rotation) plus one f32 correction factor. The score
12/// against an f32 query is an unbiased estimate of the cosine similarity
13/// after undoing the per-vector quantization scale.
14///
15/// Vectors can be stored anonymously via [`VecqIndex::add`] or under a
16/// caller-chosen `u64` key via [`VecqIndex::add_keyed`]. Keyed vectors can be
17/// removed in place (tombstoned); tombstoned slots keep their storage but are
18/// skipped by searches and dropped by [`VecqIndex::compact`] and by
19/// [`VecqIndex::to_bytes`](crate::format). Slot indices stay stable until a
20/// compaction, so integrators can treat a slot as a transient handle while
21/// keys are the durable identity.
22pub struct VecqIndex {
23    pub(crate) dim: usize,
24    padded: usize,
25    pub(crate) seed: u64,
26    transform: Rhdh,
27    pub(crate) codes: Vec<u8>, // n * padded/2 nibbles, low nibble = dim i*2
28    pub(crate) scales: Vec<f32>, // per-vector dequantization scale
29    pub(crate) n: usize,       // total slots in use (live + tombstoned)
30    keys: Vec<Option<u64>>,    // slot -> caller key (keyed slots only)
31    key_to_slot: HashMap<u64, usize>,
32    alive: Vec<bool>, // slot -> not tombstoned
33    live: usize,      // number of non-tombstoned slots
34}
35
36impl VecqIndex {
37    /// Create an empty index for `dim`-dimensional unit vectors.
38    /// `seed` must be persisted with the index for cross-platform determinism.
39    pub fn new(dim: usize, seed: u64) -> Self {
40        let padded = padded_dim(dim);
41        Self {
42            dim,
43            padded,
44            seed,
45            transform: Rhdh::new(dim, seed),
46            codes: Vec::new(),
47            scales: Vec::new(),
48            n: 0,
49            keys: Vec::new(),
50            key_to_slot: HashMap::new(),
51            alive: Vec::new(),
52            live: 0,
53        }
54    }
55
56    /// Number of live (searchable) vectors.
57    pub fn len(&self) -> usize {
58        self.live
59    }
60
61    pub fn is_empty(&self) -> bool {
62        self.live == 0
63    }
64
65    /// Total slots in use, including tombstoned ones
66    /// (`slots() == len() + tombstones()`).
67    pub fn slots(&self) -> usize {
68        self.n
69    }
70
71    /// Number of tombstoned slots awaiting [`VecqIndex::compact`].
72    pub fn tombstones(&self) -> usize {
73        self.n - self.live
74    }
75
76    pub fn dim(&self) -> usize {
77        self.dim
78    }
79
80    pub fn seed(&self) -> u64 {
81        self.seed
82    }
83
84    #[cfg(test)]
85    pub(crate) fn padded(&self) -> usize {
86        self.padded
87    }
88
89    // -- crate-internal accessors used by the persistence format ---------
90
91    pub(crate) fn padded_dim(&self) -> usize {
92        self.padded
93    }
94
95    pub(crate) fn live_slots(&self) -> usize {
96        self.live
97    }
98
99    pub(crate) fn slot_alive(&self, slot: usize) -> bool {
100        self.alive[slot]
101    }
102
103    pub(crate) fn slot_scale(&self, slot: usize) -> f32 {
104        self.scales[slot]
105    }
106
107    pub(crate) fn slot_codes(&self, slot: usize, bpv: usize) -> &[u8] {
108        &self.codes[slot * bpv..(slot + 1) * bpv]
109    }
110
111    /// Mark the index as holding `count` dense (all-live, keyless) slots;
112    /// used after loading from the file format.
113    pub(crate) fn init_dense(&mut self, count: usize) {
114        self.keys = vec![None; count];
115        self.alive = vec![true; count];
116        self.key_to_slot.clear();
117        self.live = count;
118    }
119
120    /// Quantize and add one vector (any norm; normalized internally).
121    ///
122    /// Returns the slot index holding the vector (stable until compaction).
123    pub fn add(&mut self, v: &[f32]) -> usize {
124        self.append_slot(v, None)
125    }
126
127    /// Quantize and add one vector under a caller-chosen `u64` key.
128    ///
129    /// If `key` already exists, the vector is replaced in place (the slot
130    /// index is preserved, matching usearch's insert semantics). Otherwise a
131    /// new slot is appended. Returns the slot index holding the vector.
132    pub fn add_keyed(&mut self, key: u64, v: &[f32]) -> usize {
133        if let Some(&slot) = self.key_to_slot.get(&key) {
134            let scale = self.encode_into(slot * (self.padded / 2), v);
135            self.scales[slot] = scale;
136            return slot;
137        }
138        self.append_slot(v, Some(key))
139    }
140
141    /// Remove a keyed vector. The slot becomes a tombstone: its storage is
142    /// kept (slot indices stay stable) but searches skip it until
143    /// [`VecqIndex::compact`]. Returns `false` if the key is unknown.
144    pub fn remove_keyed(&mut self, key: u64) -> bool {
145        match self.key_to_slot.remove(&key) {
146            Some(slot) => {
147                self.alive[slot] = false;
148                self.keys[slot] = None;
149                self.live -= 1;
150                true
151            }
152            None => false,
153        }
154    }
155
156    /// Look up the key stored at `slot` (`None` for anonymous slots,
157    /// tombstones, or out-of-range indices).
158    pub fn key_of(&self, slot: usize) -> Option<u64> {
159        self.keys.get(slot).copied().flatten()
160    }
161
162    /// Whether `key` currently identifies a live vector.
163    pub fn contains_key(&self, key: u64) -> bool {
164        self.key_to_slot.contains_key(&key)
165    }
166
167    /// Rebuild the index in place, dropping tombstoned slots.
168    ///
169    /// All remaining vectors keep their keys; **slot indices shift** to become
170    /// dense (0..len). Search results are unchanged.
171    pub fn compact(&mut self) {
172        if self.live == self.n {
173            return;
174        }
175        let bpv = self.padded / 2;
176        let mut codes = Vec::with_capacity(self.live * bpv);
177        let mut scales = Vec::with_capacity(self.live);
178        let mut keys = Vec::with_capacity(self.live);
179        let mut alive = Vec::with_capacity(self.live);
180        self.key_to_slot.clear();
181        for slot in 0..self.n {
182            if self.alive[slot] {
183                codes.extend_from_slice(&self.codes[slot * bpv..(slot + 1) * bpv]);
184                scales.push(self.scales[slot]);
185                if let Some(key) = self.keys[slot] {
186                    self.key_to_slot.insert(key, keys.len());
187                }
188                keys.push(self.keys[slot]);
189                alive.push(true);
190            }
191        }
192        self.codes = codes;
193        self.scales = scales;
194        self.keys = keys;
195        self.alive = alive;
196        self.n = self.live;
197    }
198
199    /// Append one vector as a new slot; returns the slot index.
200    fn append_slot(&mut self, v: &[f32], key: Option<u64>) -> usize {
201        let slot = self.n;
202        let scale = self.encode_into(slot * (self.padded / 2), v);
203        self.scales.push(scale);
204        self.keys.push(key);
205        if let Some(key) = key {
206            self.key_to_slot.insert(key, slot);
207        }
208        self.alive.push(true);
209        self.n += 1;
210        self.live += 1;
211        slot
212    }
213
214    /// Quantize `v` into the code bytes starting at `base` (extending
215    /// `codes` when appending); returns the unit-norm correction scale.
216    fn encode_into(&mut self, base: usize, v: &[f32]) -> f32 {
217        assert_eq!(v.len(), self.dim, "vector dim mismatch");
218        let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
219        assert!(norm > 0.0, "zero vector");
220        let unit: Vec<f32> = v.iter().map(|x| x / norm).collect();
221
222        let mut rotated = Vec::with_capacity(self.padded);
223        self.transform.apply(&unit, &mut rotated);
224
225        // Quantize to 4-bit codes, nibble-packed.
226        let bytes_per_vec = self.padded / 2;
227        if self.codes.len() < base + bytes_per_vec {
228            self.codes.resize(base + bytes_per_vec, 0);
229        }
230        let mut sum_sq = 0f32;
231        for (i, &x) in rotated.iter().enumerate() {
232            let code = lloyd::quantize_4bit(x);
233            let b = base + i / 2;
234            let byte = if i % 2 == 0 {
235                (self.codes[b] & 0xF0) | code
236            } else {
237                (self.codes[b] & 0x0F) | (code << 4)
238            };
239            self.codes[b] = byte;
240            sum_sq += lloyd::dequantize_4bit(code).powi(2);
241        }
242
243        // Scale so that the stored vector is unit-norm: dequantized vector q
244        // has norm sqrt(sum_sq); asymmetric scoring multiplies by 1/sqrt(sum_sq).
245        1.0 / sum_sq.sqrt()
246    }
247
248    /// Prepare an f32 query in rotated space (call once per query).
249    pub fn prepare_query(&self, q: &[f32]) -> PreparedQuery {
250        assert_eq!(q.len(), self.dim);
251        let norm: f32 = q.iter().map(|x| x * x).sum::<f32>().sqrt();
252        let unit: Vec<f32> = q.iter().map(|x| x / norm).collect();
253        let mut rotated = Vec::with_capacity(self.padded);
254        self.transform.apply(&unit, &mut rotated);
255        // Normalize so ||rotated|| == 1 despite the unnormalized FWHT
256        // (which scales norms by sqrt(padded)).
257        let rnorm: f32 = rotated.iter().map(|x| x * x).sum::<f32>().sqrt();
258        for x in rotated.iter_mut() {
259            *x /= rnorm;
260        }
261        // Precompute dequantized lookup for fast scoring: for each of the 16
262        // codes, the contribution when multiplied with the query coordinate.
263        let mut lut = [0f32; 16];
264        for (c, slot) in lut.iter_mut().enumerate() {
265            *slot = lloyd::dequantize_4bit(c as u8);
266        }
267        PreparedQuery { rotated, lut, norm }
268    }
269
270    /// Asymmetric score of vector `idx` against a prepared query.
271    /// Returns estimated cosine similarity in [-1, 1].
272    ///
273    /// Dispatches to the explicit NEON path on aarch64, the explicit AVX2
274    /// path on x86_64 when the host supports it (runtime detection), and the
275    /// fixed 8-bucket scalar path otherwise. All use the identical
276    /// association order (per code byte: mul, mul, add, then add into bucket
277    /// j; final pairwise tree), so they produce the same f32 bits — guarded
278    /// by `neon_matches_scalar_bitwise` / `avx2_matches_scalar_bitwise` in
279    /// tests.
280    #[inline]
281    pub fn score(&self, pq: &PreparedQuery, idx: usize) -> f32 {
282        let base = idx * (self.padded / 2);
283        let codes = &self.codes[base..base + self.padded / 2];
284        let q = &pq.rotated[..self.padded];
285        #[cfg(target_arch = "aarch64")]
286        {
287            // NEON is baseline on aarch64.
288            let raw = unsafe { neon::score_neon(codes, q, &pq.lut) };
289            raw * self.scales[idx]
290        }
291        #[cfg(target_arch = "x86_64")]
292        {
293            if avx2::available() {
294                // SAFETY: feature availability checked immediately above.
295                let raw = unsafe { avx2::score_avx2(codes, q, &pq.lut) };
296                raw * self.scales[idx]
297            } else {
298                score_scalar(codes, q, &pq.lut) * self.scales[idx]
299            }
300        }
301        #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
302        {
303            score_scalar(codes, q, &pq.lut) * self.scales[idx]
304        }
305    }
306
307    /// Brute-force top-k search. Returns (slot index, score) sorted by score
308    /// desc. Tombstoned slots are skipped.
309    ///
310    /// Uses a bounded min-heap of size k (no O(n log n) sort, no O(n)
311    /// allocation per query): push while the heap is not full, then only
312    /// push-and-pop when the candidate beats the current k-th score.
313    pub fn search(&self, q: &[f32], k: usize) -> Vec<(usize, f32)> {
314        self.search_slots(q, k)
315    }
316
317    /// Keyed variant of [`VecqIndex::search`]: returns (key, score) sorted by
318    /// score desc, restricted to live keyed vectors.
319    pub fn search_keyed(&self, q: &[f32], k: usize) -> Vec<(u64, f32)> {
320        self.search_slots(q, k)
321            .into_iter()
322            .filter_map(|(slot, s)| self.key_of(slot).map(|key| (key, s)))
323            .collect()
324    }
325
326    fn search_slots(&self, q: &[f32], k: usize) -> Vec<(usize, f32)> {
327        use std::cmp::Reverse;
328        use std::collections::BinaryHeap;
329
330        let pq = self.prepare_query(q);
331        let k = k.min(self.live).max(1);
332        let bpv = self.padded / 2;
333        // f32 -> u32 monotonic key (NaN-safe, preserves total order):
334        // flip all bits for negatives, flip sign bit for positives.
335        let key = |s: f32| -> u32 {
336            let b = s.to_bits();
337            if b & 0x8000_0000 != 0 {
338                !b
339            } else {
340                b ^ 0x8000_0000
341            }
342        };
343        let mut heap: BinaryHeap<Reverse<(u32, usize)>> = BinaryHeap::with_capacity(k + 1);
344        let consider = |s: f32, idx: usize, heap: &mut BinaryHeap<Reverse<(u32, usize)>>| {
345            let ks = key(s);
346            if heap.len() < k {
347                heap.push(Reverse((ks, idx)));
348            } else if ks > heap.peek().map(|r| r.0 .0).unwrap_or(0) {
349                heap.push(Reverse((ks, idx)));
350                heap.pop();
351            }
352        };
353        #[cfg(target_arch = "aarch64")]
354        let q_rot = &pq.rotated[..self.padded];
355        #[cfg(target_arch = "x86_64")]
356        let use_avx2 = avx2::available();
357        let mut idx = 0;
358        #[cfg(target_arch = "aarch64")]
359        {
360            // Batch 4 vectors per pass: shared q loads + LUT setup. Tombstoned
361            // slots are still scored (keeping the batch dense) but filtered
362            // before entering the heap.
363            while idx + 4 <= self.n {
364                let codes4 = &self.codes[idx * bpv..(idx + 4) * bpv];
365                let raw = unsafe { neon::score_neon4(codes4, q_rot, &pq.lut) };
366                for (v, &r) in raw.iter().enumerate() {
367                    let si = idx + v;
368                    if self.alive[si] {
369                        consider(r * self.scales[si], si, &mut heap);
370                    }
371                }
372                idx += 4;
373            }
374        }
375        #[cfg(target_arch = "x86_64")]
376        {
377            if use_avx2 {
378                // Batch 4 vectors per pass: shared q deinterleave. Tombstoned
379                // slots are still scored (keeping the batch dense) but
380                // filtered before entering the heap.
381                while idx + 4 <= self.n {
382                    let codes4 = &self.codes[idx * bpv..(idx + 4) * bpv];
383                    // SAFETY: AVX2 availability checked via `use_avx2`.
384                    let raw =
385                        unsafe { avx2::score_avx24(codes4, &pq.rotated[..self.padded], &pq.lut) };
386                    for (v, &r) in raw.iter().enumerate() {
387                        let si = idx + v;
388                        if self.alive[si] {
389                            consider(r * self.scales[si], si, &mut heap);
390                        }
391                    }
392                    idx += 4;
393                }
394            }
395        }
396        while idx < self.n {
397            if self.alive[idx] {
398                consider(self.score(&pq, idx), idx, &mut heap);
399            }
400            idx += 1;
401        }
402        // Inverse of `key`: undo the sign flip to recover the exact f32 bits.
403        let key_undo = |k: u32| -> u32 {
404            if k & 0x8000_0000 != 0 {
405                k ^ 0x8000_0000 // was a positive float
406            } else {
407                !k // was a negative float
408            }
409        };
410        let mut out: Vec<(usize, f32)> = heap
411            .into_iter()
412            .map(|r| (r.0 .1, f32::from_bits(key_undo(r.0 .0))))
413            .collect();
414        out.sort_by(|a, b| b.1.partial_cmp(&a.1).expect("no NaN scores"));
415        out
416    }
417}
418
419/// Reference 8-bucket scalar scoring. Bucket j accumulates byte j, j+8, ...
420/// of every 8-byte block; final reduction is a fixed pairwise tree.
421#[cfg_attr(target_arch = "aarch64", cfg(test))]
422pub(crate) fn score_scalar(codes: &[u8], q: &[f32], lut: &[f32; 16]) -> f32 {
423    let nb = codes.len();
424    let mut acc = [0f32; 8];
425    let mut i = 0;
426    while i + 8 <= nb {
427        for j in 0..8 {
428            let b = codes[i + j];
429            let c = (i + j) * 2;
430            acc[j] += q[c] * lut[(b & 0x0F) as usize] + q[c + 1] * lut[(b >> 4) as usize];
431        }
432        i += 8;
433    }
434    let mut tail = 0f32;
435    while i < nb {
436        let b = codes[i];
437        let c = i * 2;
438        tail += q[c] * lut[(b & 0x0F) as usize] + q[c + 1] * lut[(b >> 4) as usize];
439        i += 1;
440    }
441    let s01 = acc[0] + acc[1];
442    let s23 = acc[2] + acc[3];
443    let s45 = acc[4] + acc[5];
444    let s67 = acc[6] + acc[7];
445    (s01 + s23) + (s45 + s67) + tail
446}
447
448#[cfg(target_arch = "aarch64")]
449mod neon {
450    //! Explicit NEON scoring path, bit-identical to [`score_scalar`].
451    //!
452    //! Why manual intrinsics: the scalar loop gathers `lut[nibble]` with a
453    //! data-dependent index, which LLVM's vectorizer refuses to
454    //! auto-vectorize (verified in disassembly: zero fmla in `search`).
455    //! The LUT gather maps naturally to `vqtbl4q_u8`.
456    //!
457    //! Bit-identity with the scalar path is structural: per 8-byte block,
458    //! byte j's term `q_even*lut[lo] + q_odd*lut[hi]` (vmul, vmul, vadd —
459    //! Rust never contracts into FMA) is added into accumulator lane j,
460    //! blocks in increasing order, and the final reduction uses the same
461    //! pairwise tree. Guarded by `neon_matches_scalar_bitwise`.
462    use std::arch::aarch64::*;
463
464    /// Gather 16 f32 from the 16-entry LUT given per-lane nibble indices.
465    ///
466    /// The 64-byte LUT (16 little-endian f32) is a `uint8x16x4_t` table.
467    /// Four `vqtbl4q_u8` gathers produce byte-plane k (k=0..3) of all 16
468    /// floats; a 4x16 byte transpose then rebuilds the 4 f32x4 registers.
469    #[inline]
470    unsafe fn gather16(tbl: uint8x16x4_t, nibbles: uint8x16_t) -> [float32x4_t; 4] {
471        let idx = vmulq_u8(nibbles, vdupq_n_u8(4)); // byte offset of each lane's float
472        let one = vdupq_n_u8(1);
473        let two = vdupq_n_u8(2);
474        let b0 = vqtbl4q_u8(tbl, idx);
475        let b1 = vqtbl4q_u8(tbl, vaddq_u8(idx, one));
476        let b2 = vqtbl4q_u8(tbl, vaddq_u8(idx, two));
477        let b3 = vqtbl4q_u8(tbl, vaddq_u8(idx, vdupq_n_u8(3)));
478        // Transpose: float j = (b0[j], b1[j], b2[j], b3[j]).
479        let z01 = vzip1q_u8(b0, b1); // u16 lanes (b0j, b1j)
480        let z23 = vzip1q_u8(b2, b3); // u16 lanes (b2j, b3j)
481        let z01b = vzip2q_u8(b0, b1);
482        let z23b = vzip2q_u8(b2, b3);
483        let lo16 = vreinterpretq_u16_u8(z01);
484        let hi16 = vreinterpretq_u16_u8(z23);
485        let lo16b = vreinterpretq_u16_u8(z01b);
486        let hi16b = vreinterpretq_u16_u8(z23b);
487        [
488            vreinterpretq_f32_u8(vreinterpretq_u8_u16(vzip1q_u16(lo16, hi16))),
489            vreinterpretq_f32_u8(vreinterpretq_u8_u16(vzip2q_u16(lo16, hi16))),
490            vreinterpretq_f32_u8(vreinterpretq_u8_u16(vzip1q_u16(lo16b, hi16b))),
491            vreinterpretq_f32_u8(vreinterpretq_u8_u16(vzip2q_u16(lo16b, hi16b))),
492        ]
493    }
494
495    /// NEON scoring over one vector's codes. See module docs.
496    #[inline]
497    pub unsafe fn score_neon(codes: &[u8], q: &[f32], lut: &[f32; 16]) -> f32 {
498        // Build the 64-byte LUT table for vqtbl4q_u8.
499        let mut bytes = [0u8; 64];
500        for (c, &v) in lut.iter().enumerate() {
501            bytes[c * 4..c * 4 + 4].copy_from_slice(&v.to_le_bytes());
502        }
503        let tbl = uint8x16x4_t(
504            vld1q_u8(bytes[0..16].as_ptr()),
505            vld1q_u8(bytes[16..32].as_ptr()),
506            vld1q_u8(bytes[32..48].as_ptr()),
507            vld1q_u8(bytes[48..64].as_ptr()),
508        );
509        // acc_lo lanes 0-3 = scalar buckets 0-3; acc_hi lanes 0-3 = 4-7.
510        let mut acc_lo = vdupq_n_f32(0.0);
511        let mut acc_hi = vdupq_n_f32(0.0);
512        let nb = codes.len();
513        let mut i = 0;
514        while i + 8 <= nb {
515            let b8 = vld1_u8(codes.as_ptr().add(i)); // 8 code bytes (safe load)
516                                                     // Nibble layout for the gather: lanes 0-7 = low nibbles (even
517                                                     // dims), lanes 8-15 = high nibbles (odd dims).
518            let lo = vand_u8(b8, vdup_n_u8(0x0F));
519            let hi = vshr_n_u8(b8, 4);
520            let nibbles = vcombine_u8(lo, hi); // [lo_0..lo_7, hi_0..hi_7]
521            let g = gather16(tbl, nibbles);
522            // g[0] = lut[lo_0..3], g[1] = lut[lo_4..7],
523            // g[2] = lut[hi_0..3], g[3] = lut[hi_4..7].
524            // Load q[2i .. 2i+16) and deinterleave even/odd dims.
525            let q0 = vld1q_f32(q.as_ptr().add(i * 2));
526            let q1 = vld1q_f32(q.as_ptr().add(i * 2 + 4));
527            let q2 = vld1q_f32(q.as_ptr().add(i * 2 + 8));
528            let q3 = vld1q_f32(q.as_ptr().add(i * 2 + 12));
529            let q_even_lo = vuzp1q_f32(q0, q1); // dims 2i, 2i+2, 2i+4, 2i+6
530            let q_even_hi = vuzp1q_f32(q2, q3);
531            let q_odd_lo = vuzp2q_f32(q0, q1);
532            let q_odd_hi = vuzp2q_f32(q2, q3);
533            // term = q_even*lut[lo] + q_odd*lut[hi]  (mul, mul, add — no FMA)
534            let t_lo = vaddq_f32(vmulq_f32(q_even_lo, g[0]), vmulq_f32(q_odd_lo, g[2]));
535            let t_hi = vaddq_f32(vmulq_f32(q_even_hi, g[1]), vmulq_f32(q_odd_hi, g[3]));
536            acc_lo = vaddq_f32(acc_lo, t_lo);
537            acc_hi = vaddq_f32(acc_hi, t_hi);
538            i += 8;
539        }
540        // Extract buckets and reduce with the scalar pairwise tree.
541        let mut acc = [0f32; 8];
542        acc[0] = vgetq_lane_f32(acc_lo, 0);
543        acc[1] = vgetq_lane_f32(acc_lo, 1);
544        acc[2] = vgetq_lane_f32(acc_lo, 2);
545        acc[3] = vgetq_lane_f32(acc_lo, 3);
546        acc[4] = vgetq_lane_f32(acc_hi, 0);
547        acc[5] = vgetq_lane_f32(acc_hi, 1);
548        acc[6] = vgetq_lane_f32(acc_hi, 2);
549        acc[7] = vgetq_lane_f32(acc_hi, 3);
550        // Scalar tail for the last (< 8) code bytes. padded is a multiple of
551        // 8 elements (padded/2 bytes multiple of 4), so nb % 8 is 0 or 4.
552        let mut tail = 0f32;
553        while i < nb {
554            let b = codes[i];
555            let c = i * 2;
556            tail += q[c] * lut[(b & 0x0F) as usize] + q[c + 1] * lut[(b >> 4) as usize];
557            i += 1;
558        }
559        let s01 = acc[0] + acc[1];
560        let s23 = acc[2] + acc[3];
561        let s45 = acc[4] + acc[5];
562        let s67 = acc[6] + acc[7];
563        (s01 + s23) + (s45 + s67) + tail
564    }
565
566    /// Score 4 consecutive vectors at once, amortizing the q loads and LUT
567    /// table setup across all 4. Each vector accumulates in the exact same
568    /// per-lane order as [`score_neon`], so results are bit-identical.
569    ///
570    /// Returns raw (pre-scale) scores; the caller multiplies by `scales`.
571    #[inline]
572    pub unsafe fn score_neon4(codes4: &[u8], q: &[f32], lut: &[f32; 16]) -> [f32; 4] {
573        let mut bytes = [0u8; 64];
574        for (c, &v) in lut.iter().enumerate() {
575            bytes[c * 4..c * 4 + 4].copy_from_slice(&v.to_le_bytes());
576        }
577        let tbl = uint8x16x4_t(
578            vld1q_u8(bytes[0..16].as_ptr()),
579            vld1q_u8(bytes[16..32].as_ptr()),
580            vld1q_u8(bytes[32..48].as_ptr()),
581            vld1q_u8(bytes[48..64].as_ptr()),
582        );
583        let nb = codes4.len() / 4; // bytes per vector
584        let mut acc_lo = [vdupq_n_f32(0.0); 4];
585        let mut acc_hi = [vdupq_n_f32(0.0); 4];
586        let mut i = 0;
587        while i + 8 <= nb {
588            // Shared q loads for this block: q[2i .. 2i+16), deinterleaved.
589            let q0 = vld1q_f32(q.as_ptr().add(i * 2));
590            let q1 = vld1q_f32(q.as_ptr().add(i * 2 + 4));
591            let q2 = vld1q_f32(q.as_ptr().add(i * 2 + 8));
592            let q3 = vld1q_f32(q.as_ptr().add(i * 2 + 12));
593            let q_even_lo = vuzp1q_f32(q0, q1);
594            let q_even_hi = vuzp1q_f32(q2, q3);
595            let q_odd_lo = vuzp2q_f32(q0, q1);
596            let q_odd_hi = vuzp2q_f32(q2, q3);
597            for v in 0..4 {
598                let b8 = vld1_u8(codes4.as_ptr().add(v * nb + i));
599                let lo = vand_u8(b8, vdup_n_u8(0x0F));
600                let hi = vshr_n_u8(b8, 4);
601                let nibbles = vcombine_u8(lo, hi);
602                let g = gather16(tbl, nibbles);
603                let t_lo = vaddq_f32(vmulq_f32(q_even_lo, g[0]), vmulq_f32(q_odd_lo, g[2]));
604                let t_hi = vaddq_f32(vmulq_f32(q_even_hi, g[1]), vmulq_f32(q_odd_hi, g[3]));
605                acc_lo[v] = vaddq_f32(acc_lo[v], t_lo);
606                acc_hi[v] = vaddq_f32(acc_hi[v], t_hi);
607            }
608            i += 8;
609        }
610        let mut out = [0f32; 4];
611        for v in 0..4 {
612            // Same lane extraction + pairwise reduction as score_neon.
613            let mut a = [0f32; 8];
614            a[0] = vgetq_lane_f32(acc_lo[v], 0);
615            a[1] = vgetq_lane_f32(acc_lo[v], 1);
616            a[2] = vgetq_lane_f32(acc_lo[v], 2);
617            a[3] = vgetq_lane_f32(acc_lo[v], 3);
618            a[4] = vgetq_lane_f32(acc_hi[v], 0);
619            a[5] = vgetq_lane_f32(acc_hi[v], 1);
620            a[6] = vgetq_lane_f32(acc_hi[v], 2);
621            a[7] = vgetq_lane_f32(acc_hi[v], 3);
622            // Scalar tail for the last (< 8) code bytes.
623            let mut tail = 0f32;
624            let mut j = i;
625            while j < nb {
626                let b = codes4[v * nb + j];
627                let c = j * 2;
628                tail += q[c] * lut[(b & 0x0F) as usize] + q[c + 1] * lut[(b >> 4) as usize];
629                j += 1;
630            }
631            let s01 = a[0] + a[1];
632            let s23 = a[2] + a[3];
633            let s45 = a[4] + a[5];
634            let s67 = a[6] + a[7];
635            out[v] = (s01 + s23) + (s45 + s67) + tail;
636        }
637        out
638    }
639}
640
641/// Explicit AVX2 scoring path (x86_64), bit-identical to [`score_scalar`].
642///
643/// Unlike NEON (baseline on aarch64), AVX2 is not universal on x86_64, so the
644/// path is selected at runtime with `is_x86_feature_detected!` and the
645/// kernels are `#[target_feature(enable = "avx2")]`.
646///
647/// Bit-identity with the scalar path is structural: lane j of the
648/// accumulator corresponds to scalar bucket j. Per 8-byte block, lane j
649/// computes `q[2b]*lut[lo_b] + q[2b+1]*lut[hi_b]` (b = block start + j;
650/// vmul, vmul, vadd — no FMA contraction) and adds it into lane j, blocks
651/// in increasing order — the same per-bucket term and accumulation order as
652/// the scalar loop. The LUT gather uses `vgatherdps` on the 16-entry table
653/// where NEON uses `vqtbl4q_u8`. Final reduction is the same pairwise tree.
654#[cfg(target_arch = "x86_64")]
655mod avx2 {
656    use std::arch::x86_64::*;
657
658    /// Whether the host CPU supports AVX2.
659    pub fn available() -> bool {
660        std::is_x86_feature_detected!("avx2")
661    }
662
663    /// Gather `lut[nibble]` for 8 nibbles into an 8-lane vector.
664    #[inline]
665    unsafe fn gather8(lut: &[f32; 16], nibbles: __m128i) -> __m256 {
666        // The gather's scale of 4 turns each nibble index into an f32 byte
667        // offset — no pre-shift needed.
668        let idx = _mm256_cvtepu8_epi32(nibbles);
669        _mm256_i32gather_ps(lut.as_ptr(), idx, 4)
670    }
671
672    /// Deinterleave the 16 f32 at `q` into even dims (8 lanes) and odd dims
673    /// (8 lanes): {d0,d2,..,d14} and {d1,d3,..,d15}.
674    #[inline]
675    unsafe fn deinterleave16(q: *const f32) -> (__m256, __m256) {
676        let qa = _mm256_loadu_ps(q);
677        let qb = _mm256_loadu_ps(q.add(8));
678        // shuffle_ps picks {a0,a2,b0,b2} (even) / {a1,a3,b1,b3} (odd) per
679        // 128-bit half; the vpermps index vector then interleaves the halves
680        // into contiguous even/odd streams {d0,d2,..,d14} / {d1,d3,..,d15}.
681        let fixup = _mm256_setr_epi32(0, 1, 4, 5, 2, 3, 6, 7);
682        let even = _mm256_permutevar8x32_ps(_mm256_shuffle_ps(qa, qb, 0x88), fixup);
683        let odd = _mm256_permutevar8x32_ps(_mm256_shuffle_ps(qa, qb, 0xDD), fixup);
684        (even, odd)
685    }
686
687    /// AVX2 scoring over one vector's codes. See module docs.
688    #[inline]
689    #[target_feature(enable = "avx2")]
690    pub unsafe fn score_avx2(codes: &[u8], q: &[f32], lut: &[f32; 16]) -> f32 {
691        let mut acc = _mm256_setzero_ps();
692        let nb = codes.len();
693        let mut i = 0;
694        while i + 8 <= nb {
695            let b8 = _mm_loadl_epi64(codes.as_ptr().add(i) as *const __m128i);
696            let g_lo = gather8(lut, _mm_and_si128(b8, _mm_set1_epi8(0x0F)));
697            let g_hi = gather8(
698                lut,
699                _mm_and_si128(_mm_srli_epi16(b8, 4), _mm_set1_epi8(0x0F)),
700            );
701            let (even, odd) = deinterleave16(q.as_ptr().add(i * 2));
702            // term = q_even*lut[lo] + q_odd*lut[hi]  (mul, mul, add — no FMA)
703            let term = _mm256_add_ps(_mm256_mul_ps(even, g_lo), _mm256_mul_ps(odd, g_hi));
704            acc = _mm256_add_ps(acc, term);
705            i += 8;
706        }
707        // Extract lanes and reduce with the scalar pairwise tree.
708        let mut a = [0f32; 8];
709        _mm256_storeu_ps(a.as_mut_ptr(), acc);
710        // Scalar tail for the last (< 8) code bytes. padded is a multiple of
711        // 8 elements (padded/2 bytes multiple of 4), so nb % 8 is 0 or 4.
712        let mut tail = 0f32;
713        while i < nb {
714            let b = codes[i];
715            let c = i * 2;
716            tail += q[c] * lut[(b & 0x0F) as usize] + q[c + 1] * lut[(b >> 4) as usize];
717            i += 1;
718        }
719        let s01 = a[0] + a[1];
720        let s23 = a[2] + a[3];
721        let s45 = a[4] + a[5];
722        let s67 = a[6] + a[7];
723        (s01 + s23) + (s45 + s67) + tail
724    }
725
726    /// Score 4 consecutive vectors at once, amortizing the q deinterleave
727    /// across all 4. Each vector accumulates in the exact same per-lane order
728    /// as [`score_avx2`], so results are bit-identical. Returns raw
729    /// (pre-scale) scores; the caller multiplies by `scales`.
730    #[inline]
731    #[target_feature(enable = "avx2")]
732    pub unsafe fn score_avx24(codes4: &[u8], q: &[f32], lut: &[f32; 16]) -> [f32; 4] {
733        let nb = codes4.len() / 4; // bytes per vector
734        let mut acc = [_mm256_setzero_ps(); 4];
735        let mut i = 0;
736        while i + 8 <= nb {
737            // Shared q loads + deinterleave for this block.
738            let (even, odd) = deinterleave16(q.as_ptr().add(i * 2));
739            for (v, acc_v) in acc.iter_mut().enumerate() {
740                let b8 = _mm_loadl_epi64(codes4.as_ptr().add(v * nb + i) as *const __m128i);
741                let g_lo = gather8(lut, _mm_and_si128(b8, _mm_set1_epi8(0x0F)));
742                let g_hi = gather8(
743                    lut,
744                    _mm_and_si128(_mm_srli_epi16(b8, 4), _mm_set1_epi8(0x0F)),
745                );
746                let term = _mm256_add_ps(_mm256_mul_ps(even, g_lo), _mm256_mul_ps(odd, g_hi));
747                *acc_v = _mm256_add_ps(*acc_v, term);
748            }
749            i += 8;
750        }
751        let mut out = [0f32; 4];
752        for v in 0..4 {
753            let mut a = [0f32; 8];
754            _mm256_storeu_ps(a.as_mut_ptr(), acc[v]);
755            // Scalar tail for the last (< 8) code bytes.
756            let mut tail = 0f32;
757            let mut j = i;
758            while j < nb {
759                let b = codes4[v * nb + j];
760                let c = j * 2;
761                tail += q[c] * lut[(b & 0x0F) as usize] + q[c + 1] * lut[(b >> 4) as usize];
762                j += 1;
763            }
764            let s01 = a[0] + a[1];
765            let s23 = a[2] + a[3];
766            let s45 = a[4] + a[5];
767            let s67 = a[6] + a[7];
768            out[v] = (s01 + s23) + (s45 + s67) + tail;
769        }
770        out
771    }
772}
773
774/// A query preprocessed in the quantized domain.
775pub struct PreparedQuery {
776    rotated: Vec<f32>,
777    lut: [f32; 16],
778    #[allow(dead_code)]
779    norm: f32,
780}
781
782/// Exact cosine similarity between two f32 vectors (ground truth helper).
783pub fn cosine_f32(a: &[f32], b: &[f32]) -> f32 {
784    let mut dot = 0f32;
785    let mut na = 0f32;
786    let mut nb = 0f32;
787    for i in 0..a.len() {
788        dot += a[i] * b[i];
789        na += a[i] * a[i];
790        nb += b[i] * b[i];
791    }
792    dot / (na.sqrt() * nb.sqrt())
793}
794
795#[cfg(test)]
796mod tests {
797    use super::*;
798
799    fn rand_unit(dim: usize, seed: u64) -> Vec<f32> {
800        // xorshift normals, then normalize
801        let mut x = seed | 1;
802        let mut v = Vec::with_capacity(dim);
803        for _ in 0..dim {
804            x ^= x << 13;
805            x ^= x >> 7;
806            x ^= x << 17;
807            let u1 = ((x >> 11) as f64 / (1u64 << 53) as f64).max(1e-12);
808            x ^= x << 16;
809            let u2 = (x >> 11) as f64 / (1u64 << 53) as f64;
810            v.push(((-2.0 * u1.ln()).sqrt() * (std::f64::consts::TAU * u2).cos()) as f32);
811        }
812        let norm: f32 = v.iter().map(|a| a * a).sum::<f32>().sqrt();
813        v.into_iter().map(|a| a / norm).collect()
814    }
815
816    #[test]
817    fn score_correlates_with_exact_cosine() {
818        let dim = 128;
819        let mut idx = VecqIndex::new(dim, 7);
820        let base: Vec<Vec<f32>> = (0..200).map(|i| rand_unit(dim, i + 1)).collect();
821        for v in &base {
822            idx.add(v);
823        }
824        let q = rand_unit(dim, 999);
825        let pq = idx.prepare_query(&q);
826        let exact: Vec<f32> = base.iter().map(|v| cosine_f32(&q, v)).collect();
827        let mut max_err = 0f32;
828        for (i, &e) in exact.iter().enumerate().take(200) {
829            let est = idx.score(&pq, i);
830            max_err = max_err.max((est - e).abs());
831        }
832        assert!(max_err < 0.2, "max score error {max_err}");
833    }
834
835    #[test]
836    fn score_reproducible_and_close_to_naive() {
837        let dim = 128;
838        let mut idx = VecqIndex::new(dim, 13);
839        for i in 0..50 {
840            idx.add(&rand_unit(dim, i + 21));
841        }
842        let q = rand_unit(dim, 321);
843        let pq = idx.prepare_query(&q);
844        for vi in 0..50 {
845            let base = vi * (idx.padded() / 2);
846            let mut naive = 0f32;
847            for i in 0..idx.padded() {
848                let b = idx.codes[base + i / 2];
849                let code = if i % 2 == 0 { b & 0x0F } else { b >> 4 };
850                naive += pq.rotated[i] * pq.lut[code as usize];
851            }
852            let s = idx.score(&pq, vi);
853            assert_eq!(s.to_bits(), idx.score(&pq, vi).to_bits());
854            assert!((s - naive * idx.scales[vi]).abs() < 1e-5, "vector {vi}");
855        }
856    }
857
858    #[test]
859    fn neon_matches_scalar_bitwise() {
860        let dim = 128;
861        let mut idx = VecqIndex::new(dim, 42);
862        for i in 0..30 {
863            idx.add(&rand_unit(dim, i + 500));
864        }
865        let q = rand_unit(dim, 777);
866        let pq = idx.prepare_query(&q);
867        for vi in 0..30 {
868            let base = vi * (idx.padded() / 2);
869            let codes = &idx.codes[base..base + idx.padded() / 2];
870            let qslice = &pq.rotated[..idx.padded()];
871            #[cfg(target_arch = "aarch64")]
872            {
873                let neon = unsafe { neon::score_neon(codes, qslice, &pq.lut) };
874                let scalar = score_scalar(codes, qslice, &pq.lut);
875                assert_eq!(
876                    neon.to_bits(),
877                    scalar.to_bits(),
878                    "vector {vi}: NEON and scalar diverged"
879                );
880            }
881            #[cfg(not(target_arch = "aarch64"))]
882            {
883                let _ = (base, codes, qslice);
884            }
885        }
886    }
887
888    #[cfg(target_arch = "aarch64")]
889    #[test]
890    fn neon4_matches_neon_bitwise() {
891        let dim = 128;
892        let mut idx = VecqIndex::new(dim, 91);
893        for i in 0..12 {
894            idx.add(&rand_unit(dim, i + 90));
895        }
896        let q = rand_unit(dim, 1234);
897        let pq = idx.prepare_query(&q);
898        let bpv = idx.padded() / 2;
899        for chunk_start in (0..12).step_by(4) {
900            let codes4 = &idx.codes[chunk_start * bpv..(chunk_start + 4) * bpv];
901            let qslice = &pq.rotated[..idx.padded()];
902            {
903                let batched = unsafe { neon::score_neon4(codes4, qslice, &pq.lut) };
904                for v in 0..4 {
905                    let single = unsafe {
906                        neon::score_neon(&codes4[v * bpv..(v + 1) * bpv], qslice, &pq.lut)
907                    };
908                    assert_eq!(
909                        batched[v].to_bits(),
910                        single.to_bits(),
911                        "chunk {chunk_start} vec {v}: neon4 diverged from neon"
912                    );
913                }
914            }
915        }
916    }
917
918    #[cfg(target_arch = "x86_64")]
919    #[test]
920    fn avx2_matches_scalar_bitwise() {
921        if !avx2::available() {
922            return; // host without AVX2: scalar path is the only path
923        }
924        // dim 128: padded 128 (bpv 32, 4 full blocks). dim 8: padded 8
925        // (bpv 4) — exercises the 4-byte scalar tail after the block loop.
926        for (dim, seed) in [(128, 42), (8, 43)] {
927            let mut idx = VecqIndex::new(dim, seed);
928            for i in 0..30 {
929                idx.add(&rand_unit(dim, i + 500));
930            }
931            let q = rand_unit(dim, 777);
932            let pq = idx.prepare_query(&q);
933            for vi in 0..30 {
934                let base = vi * (idx.padded() / 2);
935                let codes = &idx.codes[base..base + idx.padded() / 2];
936                let qslice = &pq.rotated[..idx.padded()];
937                let avx2raw = unsafe { avx2::score_avx2(codes, qslice, &pq.lut) };
938                let scalar = score_scalar(codes, qslice, &pq.lut);
939                assert_eq!(
940                    avx2raw.to_bits(),
941                    scalar.to_bits(),
942                    "dim {dim} vector {vi}: AVX2 and scalar diverged"
943                );
944            }
945        }
946    }
947
948    #[cfg(target_arch = "x86_64")]
949    #[test]
950    fn avx24_matches_avx2_bitwise() {
951        if !avx2::available() {
952            return;
953        }
954        let dim = 128;
955        let mut idx = VecqIndex::new(dim, 91);
956        for i in 0..12 {
957            idx.add(&rand_unit(dim, i + 90));
958        }
959        let q = rand_unit(dim, 1234);
960        let pq = idx.prepare_query(&q);
961        let bpv = idx.padded() / 2;
962        for chunk_start in (0..12).step_by(4) {
963            let codes4 = &idx.codes[chunk_start * bpv..(chunk_start + 4) * bpv];
964            let qslice = &pq.rotated[..idx.padded()];
965            let batched = unsafe { avx2::score_avx24(codes4, qslice, &pq.lut) };
966            for v in 0..4 {
967                let single =
968                    unsafe { avx2::score_avx2(&codes4[v * bpv..(v + 1) * bpv], qslice, &pq.lut) };
969                assert_eq!(
970                    batched[v].to_bits(),
971                    single.to_bits(),
972                    "chunk {chunk_start} vec {v}: avx24 diverged from avx2"
973                );
974            }
975        }
976    }
977
978    #[cfg(target_arch = "x86_64")]
979    #[test]
980    fn search_dispatch_matches_scalar_on_avx2_hosts() {
981        // Whatever path dispatch picks, results must equal the scalar
982        // reference bit for bit.
983        let dim = 128;
984        let mut idx = VecqIndex::new(dim, 55);
985        for i in 0..30 {
986            idx.add(&rand_unit(dim, i + 800));
987        }
988        let q = rand_unit(dim, 888);
989        let pq = idx.prepare_query(&q);
990        for vi in 0..30 {
991            let base = vi * (idx.padded() / 2);
992            let codes = &idx.codes[base..base + idx.padded() / 2];
993            let qslice = &pq.rotated[..idx.padded()];
994            let scalar = score_scalar(codes, qslice, &pq.lut) * idx.scales[vi];
995            assert_eq!(idx.score(&pq, vi).to_bits(), scalar.to_bits());
996        }
997    }
998
999    #[test]
1000    fn search_returns_sorted_results() {
1001        let dim = 64;
1002        let mut idx = VecqIndex::new(dim, 3);
1003        for i in 0..50 {
1004            idx.add(&rand_unit(dim, i * 31 + 5));
1005        }
1006        let q = rand_unit(dim, 77);
1007        let res = idx.search(&q, 5);
1008        assert_eq!(res.len(), 5);
1009        for w in res.windows(2) {
1010            assert!(w[0].1 >= w[1].1);
1011        }
1012    }
1013
1014    #[test]
1015    fn quantized_size_is_one_eighth() {
1016        let dim = 384;
1017        let mut idx = VecqIndex::new(dim, 1);
1018        idx.add(&rand_unit(dim, 11));
1019        assert_eq!(idx.codes.len(), 512 / 2);
1020    }
1021
1022    #[test]
1023    fn keyed_add_search_remove() {
1024        let dim = 64;
1025        let mut idx = VecqIndex::new(dim, 5);
1026        for i in 0..50u64 {
1027            idx.add_keyed(1000 + i, &rand_unit(dim, i * 17 + 3));
1028        }
1029        assert_eq!(idx.len(), 50);
1030        assert!(idx.contains_key(1000));
1031        assert!(!idx.contains_key(999));
1032
1033        let q = rand_unit(dim, 77);
1034        let keyed = idx.search_keyed(&q, 5);
1035        assert_eq!(keyed.len(), 5);
1036        for w in keyed.windows(2) {
1037            assert!(w[0].1 >= w[1].1);
1038        }
1039        // Keys from search_keyed must all exist and match positional results.
1040        let positional = idx.search(&q, 5);
1041        for ((key, ks), (slot, ps)) in keyed.iter().zip(positional.iter()) {
1042            assert_eq!(key, &idx.key_of(*slot).unwrap());
1043            assert_eq!(ks.to_bits(), ps.to_bits());
1044        }
1045
1046        // Remove the top hit: it must vanish from results, others keep scores.
1047        let top_key = keyed[0].0;
1048        assert!(idx.remove_keyed(top_key));
1049        assert!(!idx.remove_keyed(top_key), "second remove is a no-op");
1050        assert!(!idx.remove_keyed(12345), "unknown key returns false");
1051        assert_eq!(idx.len(), 49);
1052        assert_eq!(idx.tombstones(), 1);
1053        let keyed2 = idx.search_keyed(&q, 5);
1054        assert!(!keyed2.iter().any(|(k, _)| *k == top_key));
1055        for (k, s) in keyed2.iter() {
1056            let old = keyed.iter().find(|(ok, _)| ok == k).map(|(_, os)| *os);
1057            if let Some(os) = old {
1058                assert_eq!(s.to_bits(), os.to_bits(), "key {k} score changed");
1059            }
1060        }
1061    }
1062
1063    #[test]
1064    fn keyed_add_same_key_replaces() {
1065        let dim = 32;
1066        let mut idx = VecqIndex::new(dim, 9);
1067        idx.add_keyed(7, &rand_unit(dim, 101));
1068        idx.add_keyed(7, &rand_unit(dim, 202));
1069        assert_eq!(idx.len(), 1, "replace must not grow the index");
1070        assert_eq!(idx.tombstones(), 0);
1071        // The stored vector is the second one: query near it, key 7 wins.
1072        let q = rand_unit(dim, 202);
1073        let res = idx.search_keyed(&q, 1);
1074        assert_eq!(res[0].0, 7);
1075    }
1076
1077    #[test]
1078    fn keyed_slot_indices_stay_stable_across_remove_and_serialize() {
1079        let dim = 64;
1080        let mut idx = VecqIndex::new(dim, 15);
1081        for i in 0..20u64 {
1082            idx.add_keyed(i, &rand_unit(dim, i + 300));
1083        }
1084        let q = rand_unit(dim, 404);
1085        let before = idx.search(&q, 20);
1086        // Remove two vectors: remaining slot indices must not shift.
1087        idx.remove_keyed(idx.key_of(before[0].0).unwrap());
1088        idx.remove_keyed(idx.key_of(before[5].0).unwrap());
1089        let after = idx.search(&q, 20);
1090        assert_eq!(after.len(), 18);
1091        for (slot, s) in &after {
1092            let old = before.iter().find(|(os, _)| os == slot);
1093            assert!(old.is_some(), "slot {slot} moved after remove");
1094            assert_eq!(old.unwrap().1.to_bits(), s.to_bits());
1095        }
1096        // Serializing drops tombstones on disk but must not disturb memory.
1097        let bytes = idx.to_bytes();
1098        let disk = VecqIndex::from_bytes(&bytes).unwrap();
1099        assert_eq!(disk.len(), 18);
1100        assert_eq!(idx.search(&q, 20), after, "in-memory results unchanged");
1101    }
1102
1103    #[test]
1104    fn compact_drops_tombstones_and_preserves_results() {
1105        let dim = 64;
1106        let mut idx = VecqIndex::new(dim, 21);
1107        for i in 0..40u64 {
1108            idx.add_keyed(10 * i, &rand_unit(dim, i + 61));
1109        }
1110        for i in 0..20u64 {
1111            assert!(idx.remove_keyed(10 * i));
1112        }
1113        let q = rand_unit(dim, 123);
1114        let expected = idx.search_keyed(&q, 20);
1115        idx.compact();
1116        assert_eq!(idx.tombstones(), 0);
1117        assert_eq!(idx.len(), 20);
1118        assert_eq!(idx.search_keyed(&q, 20), expected);
1119        // Round-trip after compact: keys are not persisted by design, so the
1120        // reloaded index is searchable positionally. f16 scales perturb
1121        // scores by <1e-3, so compare order and approximate scores.
1122        let bytes = idx.to_bytes();
1123        let back = VecqIndex::from_bytes(&bytes).unwrap();
1124        let reloaded = back.search(&q, 20);
1125        assert_eq!(reloaded.len(), 20);
1126        for ((slot, s), (key, ks)) in reloaded.iter().zip(expected.iter()) {
1127            assert_eq!(idx.key_of(*slot), Some(*key));
1128            assert!(
1129                (s - ks).abs() < 1e-3,
1130                "key {key} score drifted: {s} vs {ks}"
1131            );
1132        }
1133    }
1134
1135    #[test]
1136    fn keyed_search_on_empty_and_drained_index() {
1137        let dim = 32;
1138        let mut idx = VecqIndex::new(dim, 31);
1139        assert!(idx.search_keyed(&rand_unit(dim, 1), 3).is_empty());
1140        idx.add_keyed(1, &rand_unit(dim, 2));
1141        idx.add_keyed(2, &rand_unit(dim, 3));
1142        assert!(idx.remove_keyed(1));
1143        assert!(idx.remove_keyed(2));
1144        assert!(idx.is_empty(), "drained index reports empty");
1145        assert_eq!(idx.tombstones(), 2);
1146        assert!(idx.search_keyed(&rand_unit(dim, 4), 3).is_empty());
1147    }
1148
1149    #[test]
1150    fn keyed_index_from_file_supports_keyed_adds() {
1151        let dim = 64;
1152        let mut idx = VecqIndex::new(dim, 41);
1153        for i in 0..10u64 {
1154            idx.add(&rand_unit(dim, i + 700));
1155        }
1156        let bytes = idx.to_bytes();
1157        let mut back = VecqIndex::from_bytes(&bytes).unwrap();
1158        back.add_keyed(555, &rand_unit(dim, 999));
1159        assert!(back.contains_key(555));
1160        assert_eq!(back.len(), 11);
1161        let q = rand_unit(dim, 999);
1162        assert_eq!(back.search_keyed(&q, 1)[0].0, 555);
1163    }
1164
1165    #[test]
1166    #[should_panic(expected = "vector dim mismatch")]
1167    fn keyed_add_dim_mismatch_panics() {
1168        let mut idx = VecqIndex::new(32, 3);
1169        idx.add_keyed(1, &[0.5; 64]);
1170    }
1171}