Skip to main content

plugmem_core/index/
vecpool.rs

1//! The vector index: flat quantized-vector storage with a two-phase
2//! flat search.
3//!
4//! Vectors are the fourth recall source. Storage is deliberately *not* an
5//! arena or a blob heap but one contiguous `Vec<u8>` of fixed-stride
6//! slots: flat search reads every live slot, so perfect locality matters
7//! more than sorted lookup. Slots are append-only; dead ones are dropped
8//! by `maintain`, never reused in place.
9//!
10//! One slot, all little-endian, `stride = 8 + 8·words + dim` bytes where
11//! `words = ceil(dim / 64)`:
12//!
13//! | off | size | field |
14//! |---|---|---|
15//! | 0 | 4 | `fact` — owning [`FactId`] |
16//! | 4 | 4 | `scale` — f32 quantization scale |
17//! | 8 | 8·words | `sig` — sign signature, bit `i` set iff `q[i] >= 0` |
18//! | 8+8·words | dim | `q` — the i8 components |
19//!
20//! Quantization is symmetric i8 over the L2-normalized vector
21//! (`scale = max|x/‖x‖| / 127`, `q = round((x/‖x‖) / scale)`), so a
22//! quantized cosine is `scale_a · scale_b · Σ qa·qb`. It is a pure
23//! function of the input — journal replay re-quantizes and reproduces
24//! every slot byte for byte.
25//!
26//! Search is two-phase: the sign signatures give a cheap Hamming
27//! prefilter (popcount over u64 words, SIMD-friendly), then only the best
28//! `max(4·k, 64)` candidates pay an exact quantized-cosine rescore. The
29//! `counters` feature counts those exact dot products — the deterministic
30//! cost metric the perf gate holds to `min(candidates, live)`.
31
32use alloc::vec::Vec;
33
34#[cfg(feature = "counters")]
35use core::cell::Cell;
36
37use crate::error::Error;
38use crate::id::FactId;
39
40/// Serialized width of a slot's owning `fact` id (little-endian `u32`).
41const FACT_BYTES: usize = core::mem::size_of::<u32>();
42
43/// Serialized width of a slot's quantization `scale` (little-endian `f32`).
44const SCALE_BYTES: usize = core::mem::size_of::<f32>();
45
46/// Fixed slot header: `fact u32` + `scale f32`.
47const HEAD: usize = FACT_BYTES + SCALE_BYTES;
48
49/// Bytes of one signature word (a little-endian `u64`): the sign signature is
50/// packed into `words(dim)` of these.
51const SIG_WORD_BYTES: usize = core::mem::size_of::<u64>();
52
53/// Reusable vector-search scratch, owned by the engine (the zero-alloc
54/// recall invariant: after warm-up a search allocates nothing).
55#[derive(Debug, Default)]
56pub struct VecScratch {
57    /// `(hamming, slot)` prefilter buffer.
58    cand: Vec<(u32, u32)>,
59    /// `(cosine, fact)` rescore buffer.
60    top: Vec<(f32, u32)>,
61    /// The quantized query slot (`stride` bytes; `fact` unused).
62    query: Vec<u8>,
63}
64
65impl VecScratch {
66    /// Empty scratch buffers.
67    pub fn new() -> Self {
68        Self::default()
69    }
70}
71
72/// Flat store of quantized vectors.
73///
74/// The slot bytes are an **overlay** of a borrowed base and an owned tail:
75/// the owned path (`new`/`push`/`from_parts`) keeps `base = &[]` with every
76/// slot in `tail`, byte-for-byte unchanged; the borrowed/overlay path
77/// (`from_parts_borrowed`/`from_parts_overlay`) maps an mmap'd section as
78/// `base` and appends new slots to `tail`. Because dead slots are
79/// dropped by `maintain` and never rewritten in place, a slot is wholly in
80/// `base` or wholly in `tail`, so a heap opened over a multi-gigabyte mmap
81/// grows without cloning it — reads dispatch on one comparison per slot.
82#[derive(Debug)]
83pub struct VecPool<'a> {
84    /// Borrowed base slots (an mmap'd section) — empty on the owned path.
85    base: &'a [u8],
86    /// Owned append tail; `push`/`copy_slot` extend this, never `base`.
87    tail: Vec<u8>,
88    dim: usize,
89    max_bytes: usize,
90    /// Exact dot products computed by searches (feature `counters`).
91    #[cfg(feature = "counters")]
92    dots: Cell<u64>,
93}
94
95impl<'a> VecPool<'a> {
96    /// Signature words for a dimension.
97    #[inline]
98    fn words(dim: usize) -> usize {
99        dim.div_ceil(64)
100    }
101
102    /// Creates an empty pool for `dim`-dimensional vectors (`dim == 0`
103    /// leaves the layer inert — the pool stays empty).
104    pub fn new(dim: usize, max_bytes: usize) -> Self {
105        Self {
106            base: &[],
107            tail: Vec::new(),
108            dim,
109            max_bytes,
110            #[cfg(feature = "counters")]
111            dots: Cell::new(0),
112        }
113    }
114
115    /// Byte stride of one slot.
116    #[inline]
117    pub fn stride(&self) -> usize {
118        HEAD + Self::words(self.dim) * SIG_WORD_BYTES + self.dim
119    }
120
121    /// Total bytes of the logical pool (`base` + `tail`).
122    #[inline]
123    fn pool_len(&self) -> usize {
124        self.base.len() + self.tail.len()
125    }
126
127    /// The `stride` bytes of slot `i`, dispatched to `base` or `tail`. The
128    /// base is a whole number of slots (framed on load), so slot `i` never
129    /// straddles the boundary. `pub(crate)` so the disk-first rebuild can stream
130    /// a survivor slot straight into a `Scratch`.
131    #[inline]
132    pub(crate) fn slot_bytes(&self, i: usize) -> &[u8] {
133        let stride = self.stride();
134        let start = i * stride;
135        let base_len = self.base.len();
136        if start < base_len {
137            &self.base[start..start + stride]
138        } else {
139            let at = start - base_len;
140            &self.tail[at..at + stride]
141        }
142    }
143
144    /// Number of stored slots.
145    #[inline]
146    pub fn len(&self) -> usize {
147        let pool_len = self.pool_len();
148        if pool_len == 0 {
149            0
150        } else {
151            pool_len / self.stride()
152        }
153    }
154
155    /// `true` when no vector is stored.
156    pub fn is_empty(&self) -> bool {
157        self.pool_len() == 0
158    }
159
160    /// Total bytes held.
161    pub fn pool_bytes(&self) -> usize {
162        self.pool_len()
163    }
164
165    /// The owning fact of slot `i`.
166    #[inline]
167    pub fn slot_fact(&self, i: usize) -> u32 {
168        let slot = self.slot_bytes(i);
169        u32::from_le_bytes(slot[..FACT_BYTES].try_into().unwrap())
170    }
171
172    /// The scale of slot `i`.
173    #[inline]
174    fn slot_scale(&self, i: usize) -> f32 {
175        let slot = self.slot_bytes(i);
176        f32::from_le_bytes(slot[FACT_BYTES..HEAD].try_into().unwrap())
177    }
178
179    /// The `(scale, q)` pair of slot `i` — the quantized payload a
180    /// distance evaluation needs. Shared with the HNSW graph, which
181    /// stores only neighbor ids and reads vectors from here.
182    #[inline]
183    pub(crate) fn quant(&self, i: usize) -> (f32, &[u8]) {
184        let stride = self.stride();
185        let q_off = HEAD + Self::words(self.dim) * SIG_WORD_BYTES;
186        let slot = self.slot_bytes(i);
187        let scale = f32::from_le_bytes(slot[FACT_BYTES..HEAD].try_into().unwrap());
188        (scale, &slot[q_off..stride])
189    }
190
191    /// Quantized cosine of two slots by index — the graph's edge metric.
192    /// Both indices must be in range (graph neighbors are validated on
193    /// load).
194    #[inline]
195    pub(crate) fn sim(&self, a: u32, b: u32) -> f32 {
196        self.cosine_at(a as usize, b as usize)
197    }
198
199    /// Quantizes `v` into `out` (which is sized to `stride`), writing the
200    /// full slot for `fact`. A pure, deterministic function of `v`.
201    ///
202    /// # Errors
203    ///
204    /// [`Error::DimMismatch`] if `v.len() != dim`; [`Error::Invalid`] if
205    /// `v` is not finite or is the zero vector (no direction to encode).
206    fn encode_slot(&self, fact: u32, v: &[f32], out: &mut [u8]) -> Result<(), Error> {
207        if v.len() != self.dim {
208            return Err(Error::DimMismatch {
209                got: v.len(),
210                want: self.dim,
211            });
212        }
213        let mut norm_sq = 0.0f32;
214        for &x in v {
215            if !x.is_finite() {
216                return Err(Error::Invalid("vector must be finite"));
217            }
218            norm_sq += x * x;
219        }
220        // `norm` is finite and non-negative (inputs are finite), so this
221        // rejects exactly the zero vector.
222        let norm = libm::sqrtf(norm_sq);
223        if norm <= 0.0 {
224            return Err(Error::Invalid("vector must be nonzero"));
225        }
226        let inv_norm = 1.0 / norm;
227        let mut max_abs = 0.0f32;
228        for &x in v {
229            max_abs = max_abs.max(libm::fabsf(x * inv_norm));
230        }
231        // Nonzero norm guarantees a nonzero max component, so scale > 0.
232        let scale = max_abs / 127.0;
233        out[..FACT_BYTES].copy_from_slice(&fact.to_le_bytes());
234        out[FACT_BYTES..HEAD].copy_from_slice(&scale.to_le_bytes());
235        let words = Self::words(self.dim);
236        let q_off = HEAD + words * SIG_WORD_BYTES;
237        for (i, &x) in v.iter().enumerate() {
238            let qf = libm::roundf((x * inv_norm) / scale);
239            let qi = qf.clamp(-127.0, 127.0) as i32 as i8;
240            out[q_off + i] = qi as u8;
241        }
242        // Sign signature from the quantized components.
243        for w in 0..words {
244            let mut word = 0u64;
245            for b in 0..64 {
246                let i = w * 64 + b;
247                if i >= self.dim {
248                    break;
249                }
250                if out[q_off + i] as i8 >= 0 {
251                    word |= 1 << b;
252                }
253            }
254            out[HEAD + w * SIG_WORD_BYTES..HEAD + w * SIG_WORD_BYTES + SIG_WORD_BYTES]
255                .copy_from_slice(&word.to_le_bytes());
256        }
257        Ok(())
258    }
259
260    /// Encodes one complete slot into reusable caller-owned scratch without
261    /// appending it to this pool. The disk-first reembed path writes that slot
262    /// straight to a [`Scratch`](crate::Scratch), keeping the new vector pool
263    /// out of RAM while it is produced.
264    pub(crate) fn encode_slot_into(
265        &self,
266        fact: FactId,
267        v: &[f32],
268        out: &mut Vec<u8>,
269    ) -> Result<(), Error> {
270        out.clear();
271        out.resize(self.stride(), 0);
272        self.encode_slot(fact.0, v, out)
273    }
274
275    /// Quantized cosine between a complete encoded query slot and one stored
276    /// slot. Similarity guarding encodes the incoming vector once, then uses
277    /// this for its bounded candidate scan without appending to the pool.
278    pub(crate) fn cosine_encoded_slot(&self, encoded: &[u8], slot: u32) -> f32 {
279        let stride = self.stride();
280        if encoded.len() != stride || slot as usize >= self.len() {
281            return 0.0;
282        }
283        let q_off = HEAD + Self::words(self.dim) * SIG_WORD_BYTES;
284        let stored = self.slot_bytes(slot as usize);
285        let query_scale = f32::from_le_bytes(encoded[FACT_BYTES..HEAD].try_into().unwrap());
286        let dot = dot_i8(&encoded[q_off..stride], &stored[q_off..stride]);
287        query_scale * self.slot_scale(slot as usize) * dot as f32
288    }
289
290    /// Quantizes `v` and appends it as `fact`'s slot, returning the slot
291    /// index. Ids are not stored sorted — the fact record keeps the index.
292    ///
293    /// # Errors
294    ///
295    /// [`Error::DimMismatch`]/[`Error::Invalid`] from quantization, or
296    /// [`Error::CapacityExceeded`] at the byte ceiling.
297    pub fn push(&mut self, fact: FactId, v: &[f32]) -> Result<u32, Error> {
298        let stride = self.stride();
299        let pool_len = self.pool_len();
300        if pool_len + stride > self.max_bytes {
301            return Err(Error::CapacityExceeded { what: "vectors" });
302        }
303        let index = u32::try_from(pool_len / stride).map_err(|_| Error::CapacityExceeded {
304            what: "vector slots",
305        })?;
306        // Take the tail out so `encode_slot(&self, ...)` can borrow `self`
307        // while we write into the (now-detached) buffer; the new slot always
308        // lands in the tail, never the borrowed base.
309        let mut tail = core::mem::take(&mut self.tail);
310        let at = tail.len();
311        tail.resize(at + stride, 0);
312        let res = match self.encode_slot(fact.0, v, &mut tail[at..]) {
313            Ok(()) => Ok(index),
314            Err(e) => {
315                // Roll the failed append back so the pool stays canonical.
316                tail.truncate(at);
317                Err(e)
318            }
319        };
320        self.tail = tail;
321        res
322    }
323
324    /// The `(scale, q)` view of a query previously quantized into
325    /// `scratch` by [`VecPool::quantize_query`] — what the graph search
326    /// and the tail scan consume.
327    pub(crate) fn quantized<'s>(&self, scratch: &'s VecScratch) -> (f32, &'s [u8]) {
328        let stride = self.stride();
329        let q_off = HEAD + Self::words(self.dim) * SIG_WORD_BYTES;
330        debug_assert_eq!(scratch.query.len(), stride);
331        (
332            f32::from_le_bytes(scratch.query[FACT_BYTES..HEAD].try_into().unwrap()),
333            &scratch.query[q_off..stride],
334        )
335    }
336
337    /// Quantizes a query vector into `scratch.query` (sized to `stride`,
338    /// `fact` left as `0`). Reused across searches.
339    pub fn quantize_query(&self, v: &[f32], scratch: &mut VecScratch) -> Result<(), Error> {
340        let stride = self.stride();
341        scratch.query.clear();
342        scratch.query.resize(stride, 0);
343        let mut buf = core::mem::take(&mut scratch.query);
344        let res = self.encode_slot(0, v, &mut buf);
345        scratch.query = buf;
346        res
347    }
348
349    /// Copies slot `i` of `src` verbatim (already quantized) into `self`,
350    /// returning its new index. Used by `maintain` compaction — the
351    /// quantized bytes are reproduced exactly, so a compacted snapshot is
352    /// byte-identical to a replayed one. `src` must share this pool's `dim`.
353    pub(crate) fn copy_slot(&mut self, src: &VecPool<'_>, i: u32) -> u32 {
354        debug_assert_eq!(self.dim, src.dim, "copy_slot across differing dims");
355        let stride = self.stride();
356        let index = (self.pool_len() / stride) as u32;
357        self.tail.extend_from_slice(src.slot_bytes(i as usize));
358        index
359    }
360
361    /// Appends an exact copy of one of this pool's slots for a new owning
362    /// fact. Retag revisions preserve embeddings without a model round trip or
363    /// a lossy dequantize/requantize cycle.
364    pub(crate) fn clone_slot_for_fact(&mut self, fact: FactId, source: u32) -> Result<u32, Error> {
365        let stride = self.stride();
366        let pool_len = self.pool_len();
367        if source as usize >= self.len() {
368            return Err(Error::Corrupt("retag vector slot is out of range"));
369        }
370        if pool_len + stride > self.max_bytes {
371            return Err(Error::CapacityExceeded { what: "vectors" });
372        }
373        let index = u32::try_from(pool_len / stride).map_err(|_| Error::CapacityExceeded {
374            what: "vector slots",
375        })?;
376        let source_start = source as usize * stride;
377        if source_start < self.base.len() {
378            self.tail
379                .extend_from_slice(&self.base[source_start..source_start + stride]);
380        } else {
381            let at = source_start - self.base.len();
382            let dst = self.tail.len();
383            self.tail.resize(dst + stride, 0);
384            self.tail.copy_within(at..at + stride, dst);
385        }
386        let at = self.tail.len() - stride;
387        self.tail[at..at + FACT_BYTES].copy_from_slice(&fact.0.to_le_bytes());
388        Ok(index)
389    }
390
391    /// The exact quantized cosine of two slots by their scales and i8
392    /// components.
393    fn cosine_at(&self, a: usize, b: usize) -> f32 {
394        let stride = self.stride();
395        let q_off = HEAD + Self::words(self.dim) * SIG_WORD_BYTES;
396        let (sa, sb) = (self.slot_bytes(a), self.slot_bytes(b));
397        let dot = dot_i8(&sa[q_off..stride], &sb[q_off..stride]);
398        self.slot_scale(a) * self.slot_scale(b) * dot as f32
399    }
400
401    /// Quantized cosine of slots `a` and `b` (similar-detection uses it on
402    /// two stored facts). Returns `0.0` if either index is out of range.
403    pub fn cosine_slots(&self, a: u32, b: u32) -> f32 {
404        let n = self.len();
405        if a as usize >= n || b as usize >= n {
406            return 0.0;
407        }
408        self.cosine_at(a as usize, b as usize)
409    }
410
411    /// Flat two-phase search: Hamming prefilter on signatures, then exact
412    /// quantized-cosine rescore of the best `max(4·k, 64)` candidates.
413    /// `admit` filters by the shared recall rule; writes the top `k`
414    /// `(fact, cosine)` into `out`, descending.
415    pub fn search(
416        &self,
417        query: &[f32],
418        k: usize,
419        admit: &mut dyn FnMut(FactId) -> bool,
420        scratch: &mut VecScratch,
421        out: &mut Vec<(FactId, f32)>,
422    ) -> Result<(), Error> {
423        out.clear();
424        let n = self.len();
425        if n == 0 || k == 0 {
426            return Ok(());
427        }
428        self.quantize_query(query, scratch)?;
429        let stride = self.stride();
430        let words = Self::words(self.dim);
431        let q_off = HEAD + words * SIG_WORD_BYTES;
432
433        // Phase 1: Hamming distance of every slot's signature to the query.
434        let VecScratch {
435            cand, top, query, ..
436        } = scratch;
437        let q_sig = &query[HEAD..HEAD + words * SIG_WORD_BYTES];
438        cand.clear();
439        cand.reserve(n);
440        for i in 0..n {
441            let slot = self.slot_bytes(i);
442            let s_sig = &slot[HEAD..HEAD + words * SIG_WORD_BYTES];
443            let mut ham = 0u32;
444            for w in 0..words {
445                let a = u64::from_le_bytes(
446                    q_sig[w * SIG_WORD_BYTES..w * SIG_WORD_BYTES + SIG_WORD_BYTES]
447                        .try_into()
448                        .unwrap(),
449                );
450                let b = u64::from_le_bytes(
451                    s_sig[w * SIG_WORD_BYTES..w * SIG_WORD_BYTES + SIG_WORD_BYTES]
452                        .try_into()
453                        .unwrap(),
454                );
455                ham += (a ^ b).count_ones();
456            }
457            cand.push((ham, i as u32));
458        }
459        let c = (4 * k).max(64).min(n);
460        if cand.len() > c {
461            cand.select_nth_unstable(c - 1);
462        }
463
464        // Phase 2: exact quantized cosine on the survivors.
465        let q_scale = f32::from_le_bytes(query[FACT_BYTES..HEAD].try_into().unwrap());
466        let q_q = &query[q_off..q_off + self.dim];
467        top.clear();
468        #[cfg(feature = "counters")]
469        let mut dots = 0u64;
470        for &(_, slot) in cand[..c].iter() {
471            let sb = self.slot_bytes(slot as usize);
472            let fact = FactId(u32::from_le_bytes(sb[..FACT_BYTES].try_into().unwrap()));
473            if !admit(fact) {
474                continue;
475            }
476            let s_scale = f32::from_le_bytes(sb[FACT_BYTES..HEAD].try_into().unwrap());
477            let dot = dot_i8(q_q, &sb[q_off..stride]);
478            top.push((q_scale * s_scale * dot as f32, fact.0));
479            #[cfg(feature = "counters")]
480            {
481                dots += 1;
482            }
483        }
484        #[cfg(feature = "counters")]
485        self.dots.set(self.dots.get() + dots);
486        top.sort_unstable_by(|a, b| b.0.total_cmp(&a.0).then(a.1.cmp(&b.1)));
487        for &(score, id) in top.iter().take(k) {
488            out.push((FactId(id), score));
489        }
490        Ok(())
491    }
492
493    /// The vector section as one contiguous buffer (`base ++ tail`).
494    /// Byte-identical to the owned pool holding the same slots, so an overlay
495    /// snapshot is canonical. Test-only: production streams via
496    /// [`VecPool::pieces`], which needs no owned copy.
497    #[cfg(test)]
498    pub(crate) fn dump(&self) -> Vec<u8> {
499        let mut out = Vec::with_capacity(self.pool_len());
500        out.extend_from_slice(self.base);
501        out.extend_from_slice(&self.tail);
502        out
503    }
504
505    /// The vector section as its two contiguous pieces (`base`, `tail`)
506    /// without concatenating them — lets the streaming snapshot writer
507    /// emit the dominant pool with no owned full-section copy.
508    /// Concatenated, the pieces equal the section (`base ++ tail`).
509    pub(crate) fn pieces(&self) -> [&[u8]; 2] {
510        [self.base, &self.tail]
511    }
512
513    /// Rebuilds a pool from its dumped section, checking only the framing
514    /// (length is a whole number of slots, fits the ceiling). Slot content
515    /// — scales, signatures, the fact bijection — is validated by the
516    /// engine's `validate_references` (it needs the fact records too).
517    pub(crate) fn from_parts(dim: usize, max_bytes: usize, bytes: &[u8]) -> Result<Self, Error> {
518        Self::frame_check(dim, max_bytes, bytes.len())?;
519        let mut pool = Self::new(dim, max_bytes);
520        pool.tail = bytes.to_vec();
521        Ok(pool)
522    }
523
524    /// Zero-copy sibling of [`VecPool::from_parts`]: the pool borrows the
525    /// dumped section (an mmap'd byte range) as its base instead of copying
526    /// it. Same framing checks; the lifetime ties the pool to `bytes`
527    /// Under the overlay write path a later [`VecPool::push`]
528    /// appends to an owned tail without cloning the base.
529    pub(crate) fn from_parts_borrowed(
530        dim: usize,
531        max_bytes: usize,
532        bytes: &'a [u8],
533    ) -> Result<Self, Error> {
534        Self::frame_check(dim, max_bytes, bytes.len())?;
535        let mut pool = Self::new(dim, max_bytes);
536        pool.base = bytes;
537        Ok(pool)
538    }
539
540    /// Frames the dumped section: a length within the ceiling and, for a
541    /// live layer, a whole number of slots. Shared by both `from_parts`
542    /// constructors so the owned and borrowed paths validate identically.
543    fn frame_check(dim: usize, max_bytes: usize, len: usize) -> Result<(), Error> {
544        if len > max_bytes {
545            return Err(Error::Corrupt("vector pool exceeds the configured ceiling"));
546        }
547        if dim == 0 {
548            if len != 0 {
549                return Err(Error::Corrupt("vector pool present with dim 0"));
550            }
551            return Ok(());
552        }
553        let stride = HEAD + Self::words(dim) * SIG_WORD_BYTES + dim;
554        if !len.is_multiple_of(stride) {
555            return Err(Error::Corrupt("vector pool is not a whole number of slots"));
556        }
557        Ok(())
558    }
559
560    /// Structural self-check of every slot (A.4): each scale is
561    /// finite and non-negative, and each signature bit agrees with the
562    /// sign of its quantized component. Keeps the panic-free contract:
563    /// after this, a search over the pool cannot read a malformed slot
564    /// into a NaN or disagree with the prefilter.
565    pub(crate) fn validate(&self) -> Result<(), Error> {
566        if self.dim == 0 {
567            return Ok(());
568        }
569        let words = Self::words(self.dim);
570        let q_off = HEAD + words * SIG_WORD_BYTES;
571        for i in 0..self.len() {
572            let slot = self.slot_bytes(i);
573            let scale = f32::from_le_bytes(slot[FACT_BYTES..HEAD].try_into().unwrap());
574            if !scale.is_finite() || scale < 0.0 {
575                return Err(Error::Corrupt(
576                    "vector slot scale is not finite and non-negative",
577                ));
578            }
579            for w in 0..words {
580                let stored = u64::from_le_bytes(
581                    slot[HEAD + w * SIG_WORD_BYTES..HEAD + w * SIG_WORD_BYTES + SIG_WORD_BYTES]
582                        .try_into()
583                        .unwrap(),
584                );
585                let mut expect = 0u64;
586                for b in 0..64 {
587                    let j = w * 64 + b;
588                    if j >= self.dim {
589                        break;
590                    }
591                    if slot[q_off + j] as i8 >= 0 {
592                        expect |= 1 << b;
593                    }
594                }
595                if stored != expect {
596                    return Err(Error::Corrupt(
597                        "vector slot signature disagrees with its components",
598                    ));
599                }
600            }
601        }
602        Ok(())
603    }
604
605    /// Exact dot products computed so far (feature `counters`).
606    #[cfg(feature = "counters")]
607    pub fn dots(&self) -> u64 {
608        self.dots.get()
609    }
610
611    /// Resets the dot counter (feature `counters`).
612    #[cfg(feature = "counters")]
613    pub fn reset_dots(&self) {
614        self.dots.set(0);
615    }
616}
617
618/// Integer dot product of two equal-length i8 slices held as bytes.
619/// `dim ≤ 4096` and `|q| ≤ 127`, so the sum fits `i32`
620/// (`4096 · 127² < 2³¹`).
621#[inline]
622pub(crate) fn dot_i8(a: &[u8], b: &[u8]) -> i32 {
623    let mut acc = 0i32;
624    for (&x, &y) in a.iter().zip(b.iter()) {
625        acc += i32::from(x as i8) * i32::from(y as i8);
626    }
627    acc
628}
629
630#[cfg(test)]
631mod tests {
632    use super::*;
633    use alloc::vec;
634
635    /// A tiny deterministic LCG yielding `f32` in `[-1, 1)` — no rng crate
636    /// in the core's test surface, and determinism is the repo law.
637    struct Lcg(u64);
638    impl Lcg {
639        fn next(&mut self) -> f32 {
640            self.0 = self
641                .0
642                .wrapping_mul(6_364_136_223_846_793_005)
643                .wrapping_add(1_442_695_040_888_963_407);
644            ((self.0 >> 40) as f32 / (1u64 << 24) as f32) * 2.0 - 1.0
645        }
646        fn vector(&mut self, dim: usize) -> Vec<f32> {
647            (0..dim).map(|_| self.next()).collect()
648        }
649    }
650
651    /// True (unquantized) cosine of two vectors.
652    fn cosine_f32(a: &[f32], b: &[f32]) -> f32 {
653        let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
654        let na: f32 = libm::sqrtf(a.iter().map(|x| x * x).sum());
655        let nb: f32 = libm::sqrtf(b.iter().map(|x| x * x).sum());
656        dot / (na * nb)
657    }
658
659    /// The quantized cosine tracks the true cosine within the documented
660    /// error band across many random pairs.
661    #[test]
662    fn quantized_cosine_tracks_f32() {
663        let dim = 384;
664        let mut rng = Lcg(0x1234_5678);
665        let mut worst = 0.0f32;
666        for i in 0..200u32 {
667            let a = rng.vector(dim);
668            let b = rng.vector(dim);
669            let mut pool = VecPool::new(dim, usize::MAX);
670            pool.push(FactId(2 * i), &a).unwrap();
671            pool.push(FactId(2 * i + 1), &b).unwrap();
672            let q = pool.cosine_slots(0, 1);
673            let t = cosine_f32(&a, &b);
674            worst = worst.max(libm::fabsf(q - t));
675        }
676        assert!(
677            worst < 0.05,
678            "worst quantization error {worst} exceeds 0.05"
679        );
680    }
681
682    /// Golden: a hand-built dim-4 example. `a` and `b` share a direction,
683    /// `c` is orthogonal-ish; quantized cosine ranks them accordingly and
684    /// the sign signature matches the component signs.
685    #[test]
686    fn golden_dim4() {
687        let dim = 4;
688        let mut pool = VecPool::new(dim, usize::MAX);
689        pool.push(FactId(0), &[1.0, 1.0, 0.0, 0.0]).unwrap();
690        pool.push(FactId(1), &[2.0, 2.0, 0.0, 0.0]).unwrap(); // same direction
691        pool.push(FactId(2), &[0.0, 0.0, 1.0, 1.0]).unwrap(); // orthogonal
692        // Parallel vectors → cosine ≈ 1.
693        assert!((pool.cosine_slots(0, 1) - 1.0).abs() < 1e-3);
694        // Orthogonal vectors → cosine ≈ 0.
695        assert!(pool.cosine_slots(0, 2).abs() < 1e-3);
696        // Structural self-check passes and the signature is well-formed.
697        pool.validate().unwrap();
698        // Slot 0 signature: components (positive, positive, +0, +0) → all
699        // sign bits set for the first four bits.
700        let stride = pool.stride();
701        let sig = u64::from_le_bytes(pool.dump()[HEAD..HEAD + SIG_WORD_BYTES].try_into().unwrap());
702        assert_eq!(sig & 0b1111, 0b1111);
703        assert_eq!(pool.len(), 3);
704        // dim 4: one signature word + 4 i8 components.
705        assert_eq!(stride, HEAD + SIG_WORD_BYTES + 4);
706    }
707
708    /// The two-phase search returns the true nearest neighbor at the top.
709    #[test]
710    fn search_surfaces_the_nearest() {
711        let dim = 64;
712        let mut rng = Lcg(0xdead_beef);
713        let mut pool = VecPool::new(dim, usize::MAX);
714        let target = rng.vector(dim);
715        // 200 random vectors, then the target itself as fact 500.
716        for i in 0..200u32 {
717            pool.push(FactId(i), &rng.vector(dim)).unwrap();
718        }
719        pool.push(FactId(500), &target).unwrap();
720        let mut scratch = VecScratch::new();
721        let mut out = Vec::new();
722        pool.search(&target, 5, &mut |_| true, &mut scratch, &mut out)
723            .unwrap();
724        assert_eq!(out[0].0, FactId(500), "exact match must rank first");
725        assert!(out[0].1 > 0.99, "self-cosine ≈ 1, got {}", out[0].1);
726    }
727
728    /// A zero or non-finite vector has no direction to quantize.
729    #[test]
730    fn degenerate_vectors_are_invalid() {
731        let mut pool = VecPool::new(3, usize::MAX);
732        assert_eq!(
733            pool.push(FactId(0), &[0.0, 0.0, 0.0]).unwrap_err(),
734            Error::Invalid("vector must be nonzero")
735        );
736        assert_eq!(
737            pool.push(FactId(0), &[1.0, f32::NAN, 0.0]).unwrap_err(),
738            Error::Invalid("vector must be finite")
739        );
740        assert!(matches!(
741            pool.push(FactId(0), &[1.0, 2.0]).unwrap_err(),
742            Error::DimMismatch { got: 2, want: 3 }
743        ));
744        // A failed push leaves the pool canonical (nothing appended).
745        assert_eq!(pool.len(), 0);
746        assert!(pool.is_empty());
747    }
748
749    /// Accessors and the edge branches: empty pool, `k == 0`, an
750    /// out-of-range cosine, and the byte ceiling.
751    #[test]
752    fn accessors_and_edges() {
753        let dim = 4;
754        let mut pool = VecPool::new(dim, usize::MAX);
755        assert!(pool.is_empty());
756        assert_eq!(pool.pool_bytes(), 0);
757        let mut scratch = VecScratch::new();
758        let mut out = vec![(FactId(9), 1.0)];
759        // Empty pool: search clears the output and returns nothing.
760        pool.search(&[1.0; 4], 5, &mut |_| true, &mut scratch, &mut out)
761            .unwrap();
762        assert!(out.is_empty());
763
764        pool.push(FactId(0), &[1.0, 0.0, 0.0, 0.0]).unwrap();
765        assert!(!pool.is_empty());
766        assert_eq!(pool.pool_bytes(), pool.stride());
767        // k == 0 short-circuits.
768        pool.search(&[1.0; 4], 0, &mut |_| true, &mut scratch, &mut out)
769            .unwrap();
770        assert!(out.is_empty());
771        // An out-of-range slot index yields a zero cosine, never a panic.
772        assert_eq!(pool.cosine_slots(0, 9), 0.0);
773
774        // The byte ceiling: a push past `max_bytes` is a typed error.
775        let mut tight = VecPool::new(dim, 4);
776        assert_eq!(
777            tight.push(FactId(0), &[1.0, 0.0, 0.0, 0.0]).unwrap_err(),
778            Error::CapacityExceeded { what: "vectors" }
779        );
780    }
781
782    /// `from_parts` accepts a whole number of slots and rejects the rest.
783    #[test]
784    fn from_parts_frames_slots() {
785        let dim = 8;
786        let mut pool = VecPool::new(dim, usize::MAX);
787        pool.push(FactId(0), &vec![0.5; dim]).unwrap();
788        pool.push(FactId(1), &vec![-0.5; dim]).unwrap();
789        let bytes = pool.dump();
790        let rebuilt = VecPool::from_parts(dim, usize::MAX, &bytes).unwrap();
791        assert_eq!(rebuilt.len(), 2);
792        rebuilt.validate().unwrap();
793        // One byte short of a slot boundary is corrupt.
794        assert!(VecPool::from_parts(dim, usize::MAX, &bytes[..bytes.len() - 1]).is_err());
795        // A non-empty pool with dim 0 is corrupt.
796        assert!(VecPool::from_parts(0, usize::MAX, &bytes).is_err());
797        // Bytes past the configured ceiling are corrupt.
798        assert!(VecPool::from_parts(dim, bytes.len() - 1, &bytes).is_err());
799    }
800
801    /// The structural self-check rejects malformed slots — the panic-free
802    /// contract for the vector section on hostile input.
803    #[test]
804    fn validate_rejects_malformed_slots() {
805        let dim = 8;
806        let mut pool = VecPool::new(dim, usize::MAX);
807        pool.push(FactId(0), &vec![0.5; dim]).unwrap();
808        let good = pool.dump();
809
810        // A non-finite scale (bytes 4..8) is rejected.
811        let mut bad = good.clone();
812        bad[FACT_BYTES..HEAD].copy_from_slice(&f32::NAN.to_le_bytes());
813        assert!(
814            VecPool::from_parts(dim, usize::MAX, &bad)
815                .unwrap()
816                .validate()
817                .is_err()
818        );
819
820        // A signature bit that disagrees with its component's sign is
821        // rejected: flip one i8 component negative without touching sig.
822        let mut bad = good.clone();
823        let q_off = HEAD + VecPool::words(dim) * SIG_WORD_BYTES;
824        bad[q_off] = (-1i8) as u8; // was positive (sig bit 0 set)
825        assert!(
826            VecPool::from_parts(dim, usize::MAX, &bad)
827                .unwrap()
828                .validate()
829                .is_err()
830        );
831    }
832
833    /// Overlay open: a pool over a borrowed base grows through an owned tail
834    /// without touching the base, and every accessor (fact/scale/cosine/
835    /// search) spans the base/tail boundary. The dump is byte-identical to
836    /// the fully-owned pool holding the same slots.
837    #[test]
838    fn overlay_appends_to_tail_and_reads_span_the_boundary() {
839        let dim = 16;
840        let mut rng = Lcg(0x0ace_1a75);
841        let (va, vb, vc) = (rng.vector(dim), rng.vector(dim), rng.vector(dim));
842
843        // Fully-owned reference pool with all three vectors.
844        let mut owned = VecPool::new(dim, usize::MAX);
845        owned.push(FactId(10), &va).unwrap();
846        owned.push(FactId(11), &vb).unwrap();
847        owned.push(FactId(12), &vc).unwrap();
848
849        // Base = first two vectors, serialized as if from an mmap; the third
850        // is appended through the overlay open.
851        let mut seed = VecPool::new(dim, usize::MAX);
852        seed.push(FactId(10), &va).unwrap();
853        seed.push(FactId(11), &vb).unwrap();
854        let base = seed.dump();
855        let base_snapshot = base.clone();
856
857        // For this append-only store, the overlay open is `from_parts_borrowed`
858        // (a borrowed base that a later `push` extends via the owned tail).
859        let mut pool = VecPool::from_parts_borrowed(dim, usize::MAX, &base).unwrap();
860        assert_eq!(pool.len(), 2);
861        let idx = pool.push(FactId(12), &vc).unwrap();
862        assert_eq!(idx, 2);
863        assert_eq!(pool.len(), 3);
864
865        // Accessors read base slots (0,1) and the tail slot (2) alike.
866        assert_eq!(pool.slot_fact(0), 10); // base
867        assert_eq!(pool.slot_fact(2), 12); // tail
868        // A cosine between a base slot and the tail slot matches the owned
869        // pool's — the overlay changes representation, not values.
870        assert!((pool.cosine_slots(0, 2) - owned.cosine_slots(0, 2)).abs() < 1e-6);
871        pool.validate().unwrap();
872
873        // Search finds the appended (tail) vector by querying it exactly.
874        let mut scratch = VecScratch::new();
875        let mut out = Vec::new();
876        pool.search(&vc, 1, &mut |_| true, &mut scratch, &mut out)
877            .unwrap();
878        assert_eq!(out[0].0, FactId(12));
879
880        // The dump is canonical (== owned) and the borrowed base is untouched.
881        assert_eq!(pool.dump(), owned.dump());
882        assert_eq!(base, base_snapshot);
883    }
884}