use alloc::vec::Vec;
#[cfg(feature = "counters")]
use core::cell::Cell;
use crate::error::Error;
use crate::id::FactId;
const FACT_BYTES: usize = core::mem::size_of::<u32>();
const SCALE_BYTES: usize = core::mem::size_of::<f32>();
const HEAD: usize = FACT_BYTES + SCALE_BYTES;
const SIG_WORD_BYTES: usize = core::mem::size_of::<u64>();
#[derive(Debug, Default)]
pub struct VecScratch {
cand: Vec<(u32, u32)>,
top: Vec<(f32, u32)>,
query: Vec<u8>,
}
impl VecScratch {
pub fn new() -> Self {
Self::default()
}
}
#[derive(Debug)]
pub struct VecPool<'a> {
base: &'a [u8],
tail: Vec<u8>,
dim: usize,
max_bytes: usize,
#[cfg(feature = "counters")]
dots: Cell<u64>,
}
impl<'a> VecPool<'a> {
#[inline]
fn words(dim: usize) -> usize {
dim.div_ceil(64)
}
pub fn new(dim: usize, max_bytes: usize) -> Self {
Self {
base: &[],
tail: Vec::new(),
dim,
max_bytes,
#[cfg(feature = "counters")]
dots: Cell::new(0),
}
}
#[inline]
pub fn stride(&self) -> usize {
HEAD + Self::words(self.dim) * SIG_WORD_BYTES + self.dim
}
#[inline]
fn pool_len(&self) -> usize {
self.base.len() + self.tail.len()
}
#[inline]
pub(crate) fn slot_bytes(&self, i: usize) -> &[u8] {
let stride = self.stride();
let start = i * stride;
let base_len = self.base.len();
if start < base_len {
&self.base[start..start + stride]
} else {
let at = start - base_len;
&self.tail[at..at + stride]
}
}
#[inline]
pub fn len(&self) -> usize {
let pool_len = self.pool_len();
if pool_len == 0 {
0
} else {
pool_len / self.stride()
}
}
pub fn is_empty(&self) -> bool {
self.pool_len() == 0
}
pub fn pool_bytes(&self) -> usize {
self.pool_len()
}
#[inline]
pub fn slot_fact(&self, i: usize) -> u32 {
let slot = self.slot_bytes(i);
u32::from_le_bytes(slot[..FACT_BYTES].try_into().unwrap())
}
#[inline]
fn slot_scale(&self, i: usize) -> f32 {
let slot = self.slot_bytes(i);
f32::from_le_bytes(slot[FACT_BYTES..HEAD].try_into().unwrap())
}
#[inline]
pub(crate) fn quant(&self, i: usize) -> (f32, &[u8]) {
let stride = self.stride();
let q_off = HEAD + Self::words(self.dim) * SIG_WORD_BYTES;
let slot = self.slot_bytes(i);
let scale = f32::from_le_bytes(slot[FACT_BYTES..HEAD].try_into().unwrap());
(scale, &slot[q_off..stride])
}
#[inline]
pub(crate) fn sim(&self, a: u32, b: u32) -> f32 {
self.cosine_at(a as usize, b as usize)
}
fn encode_slot(&self, fact: u32, v: &[f32], out: &mut [u8]) -> Result<(), Error> {
if v.len() != self.dim {
return Err(Error::DimMismatch {
got: v.len(),
want: self.dim,
});
}
let mut norm_sq = 0.0f32;
for &x in v {
if !x.is_finite() {
return Err(Error::Invalid("vector must be finite"));
}
norm_sq += x * x;
}
let norm = libm::sqrtf(norm_sq);
if norm <= 0.0 {
return Err(Error::Invalid("vector must be nonzero"));
}
let inv_norm = 1.0 / norm;
let mut max_abs = 0.0f32;
for &x in v {
max_abs = max_abs.max(libm::fabsf(x * inv_norm));
}
let scale = max_abs / 127.0;
out[..FACT_BYTES].copy_from_slice(&fact.to_le_bytes());
out[FACT_BYTES..HEAD].copy_from_slice(&scale.to_le_bytes());
let words = Self::words(self.dim);
let q_off = HEAD + words * SIG_WORD_BYTES;
for (i, &x) in v.iter().enumerate() {
let qf = libm::roundf((x * inv_norm) / scale);
let qi = qf.clamp(-127.0, 127.0) as i32 as i8;
out[q_off + i] = qi as u8;
}
for w in 0..words {
let mut word = 0u64;
for b in 0..64 {
let i = w * 64 + b;
if i >= self.dim {
break;
}
if out[q_off + i] as i8 >= 0 {
word |= 1 << b;
}
}
out[HEAD + w * SIG_WORD_BYTES..HEAD + w * SIG_WORD_BYTES + SIG_WORD_BYTES]
.copy_from_slice(&word.to_le_bytes());
}
Ok(())
}
pub fn push(&mut self, fact: FactId, v: &[f32]) -> Result<u32, Error> {
let stride = self.stride();
let pool_len = self.pool_len();
if pool_len + stride > self.max_bytes {
return Err(Error::CapacityExceeded { what: "vectors" });
}
let index = u32::try_from(pool_len / stride).map_err(|_| Error::CapacityExceeded {
what: "vector slots",
})?;
let mut tail = core::mem::take(&mut self.tail);
let at = tail.len();
tail.resize(at + stride, 0);
let res = match self.encode_slot(fact.0, v, &mut tail[at..]) {
Ok(()) => Ok(index),
Err(e) => {
tail.truncate(at);
Err(e)
}
};
self.tail = tail;
res
}
pub(crate) fn quantized<'s>(&self, scratch: &'s VecScratch) -> (f32, &'s [u8]) {
let stride = self.stride();
let q_off = HEAD + Self::words(self.dim) * SIG_WORD_BYTES;
debug_assert_eq!(scratch.query.len(), stride);
(
f32::from_le_bytes(scratch.query[FACT_BYTES..HEAD].try_into().unwrap()),
&scratch.query[q_off..stride],
)
}
pub fn quantize_query(&self, v: &[f32], scratch: &mut VecScratch) -> Result<(), Error> {
let stride = self.stride();
scratch.query.clear();
scratch.query.resize(stride, 0);
let mut buf = core::mem::take(&mut scratch.query);
let res = self.encode_slot(0, v, &mut buf);
scratch.query = buf;
res
}
pub(crate) fn copy_slot(&mut self, src: &VecPool<'_>, i: u32) -> u32 {
debug_assert_eq!(self.dim, src.dim, "copy_slot across differing dims");
let stride = self.stride();
let index = (self.pool_len() / stride) as u32;
self.tail.extend_from_slice(src.slot_bytes(i as usize));
index
}
fn cosine_at(&self, a: usize, b: usize) -> f32 {
let stride = self.stride();
let q_off = HEAD + Self::words(self.dim) * SIG_WORD_BYTES;
let (sa, sb) = (self.slot_bytes(a), self.slot_bytes(b));
let dot = dot_i8(&sa[q_off..stride], &sb[q_off..stride]);
self.slot_scale(a) * self.slot_scale(b) * dot as f32
}
pub fn cosine_slots(&self, a: u32, b: u32) -> f32 {
let n = self.len();
if a as usize >= n || b as usize >= n {
return 0.0;
}
self.cosine_at(a as usize, b as usize)
}
pub fn search(
&self,
query: &[f32],
k: usize,
admit: &mut dyn FnMut(FactId) -> bool,
scratch: &mut VecScratch,
out: &mut Vec<(FactId, f32)>,
) -> Result<(), Error> {
out.clear();
let n = self.len();
if n == 0 || k == 0 {
return Ok(());
}
self.quantize_query(query, scratch)?;
let stride = self.stride();
let words = Self::words(self.dim);
let q_off = HEAD + words * SIG_WORD_BYTES;
let VecScratch {
cand, top, query, ..
} = scratch;
let q_sig = &query[HEAD..HEAD + words * SIG_WORD_BYTES];
cand.clear();
cand.reserve(n);
for i in 0..n {
let slot = self.slot_bytes(i);
let s_sig = &slot[HEAD..HEAD + words * SIG_WORD_BYTES];
let mut ham = 0u32;
for w in 0..words {
let a = u64::from_le_bytes(
q_sig[w * SIG_WORD_BYTES..w * SIG_WORD_BYTES + SIG_WORD_BYTES]
.try_into()
.unwrap(),
);
let b = u64::from_le_bytes(
s_sig[w * SIG_WORD_BYTES..w * SIG_WORD_BYTES + SIG_WORD_BYTES]
.try_into()
.unwrap(),
);
ham += (a ^ b).count_ones();
}
cand.push((ham, i as u32));
}
let c = (4 * k).max(64).min(n);
if cand.len() > c {
cand.select_nth_unstable(c - 1);
}
let q_scale = f32::from_le_bytes(query[FACT_BYTES..HEAD].try_into().unwrap());
let q_q = &query[q_off..q_off + self.dim];
top.clear();
#[cfg(feature = "counters")]
let mut dots = 0u64;
for &(_, slot) in cand[..c].iter() {
let sb = self.slot_bytes(slot as usize);
let fact = FactId(u32::from_le_bytes(sb[..FACT_BYTES].try_into().unwrap()));
if !admit(fact) {
continue;
}
let s_scale = f32::from_le_bytes(sb[FACT_BYTES..HEAD].try_into().unwrap());
let dot = dot_i8(q_q, &sb[q_off..stride]);
top.push((q_scale * s_scale * dot as f32, fact.0));
#[cfg(feature = "counters")]
{
dots += 1;
}
}
#[cfg(feature = "counters")]
self.dots.set(self.dots.get() + dots);
top.sort_unstable_by(|a, b| b.0.total_cmp(&a.0).then(a.1.cmp(&b.1)));
for &(score, id) in top.iter().take(k) {
out.push((FactId(id), score));
}
Ok(())
}
#[cfg(test)]
pub(crate) fn dump(&self) -> Vec<u8> {
let mut out = Vec::with_capacity(self.pool_len());
out.extend_from_slice(self.base);
out.extend_from_slice(&self.tail);
out
}
pub(crate) fn pieces(&self) -> [&[u8]; 2] {
[self.base, &self.tail]
}
pub(crate) fn from_parts(dim: usize, max_bytes: usize, bytes: &[u8]) -> Result<Self, Error> {
Self::frame_check(dim, max_bytes, bytes.len())?;
let mut pool = Self::new(dim, max_bytes);
pool.tail = bytes.to_vec();
Ok(pool)
}
pub(crate) fn from_parts_borrowed(
dim: usize,
max_bytes: usize,
bytes: &'a [u8],
) -> Result<Self, Error> {
Self::frame_check(dim, max_bytes, bytes.len())?;
let mut pool = Self::new(dim, max_bytes);
pool.base = bytes;
Ok(pool)
}
fn frame_check(dim: usize, max_bytes: usize, len: usize) -> Result<(), Error> {
if len > max_bytes {
return Err(Error::Corrupt("vector pool exceeds the configured ceiling"));
}
if dim == 0 {
if len != 0 {
return Err(Error::Corrupt("vector pool present with dim 0"));
}
return Ok(());
}
let stride = HEAD + Self::words(dim) * SIG_WORD_BYTES + dim;
if !len.is_multiple_of(stride) {
return Err(Error::Corrupt("vector pool is not a whole number of slots"));
}
Ok(())
}
pub(crate) fn validate(&self) -> Result<(), Error> {
if self.dim == 0 {
return Ok(());
}
let words = Self::words(self.dim);
let q_off = HEAD + words * SIG_WORD_BYTES;
for i in 0..self.len() {
let slot = self.slot_bytes(i);
let scale = f32::from_le_bytes(slot[FACT_BYTES..HEAD].try_into().unwrap());
if !scale.is_finite() || scale < 0.0 {
return Err(Error::Corrupt(
"vector slot scale is not finite and non-negative",
));
}
for w in 0..words {
let stored = u64::from_le_bytes(
slot[HEAD + w * SIG_WORD_BYTES..HEAD + w * SIG_WORD_BYTES + SIG_WORD_BYTES]
.try_into()
.unwrap(),
);
let mut expect = 0u64;
for b in 0..64 {
let j = w * 64 + b;
if j >= self.dim {
break;
}
if slot[q_off + j] as i8 >= 0 {
expect |= 1 << b;
}
}
if stored != expect {
return Err(Error::Corrupt(
"vector slot signature disagrees with its components",
));
}
}
}
Ok(())
}
#[cfg(feature = "counters")]
pub fn dots(&self) -> u64 {
self.dots.get()
}
#[cfg(feature = "counters")]
pub fn reset_dots(&self) {
self.dots.set(0);
}
}
#[inline]
pub(crate) fn dot_i8(a: &[u8], b: &[u8]) -> i32 {
let mut acc = 0i32;
for (&x, &y) in a.iter().zip(b.iter()) {
acc += i32::from(x as i8) * i32::from(y as i8);
}
acc
}
#[cfg(test)]
mod tests {
use super::*;
use alloc::vec;
struct Lcg(u64);
impl Lcg {
fn next(&mut self) -> f32 {
self.0 = self
.0
.wrapping_mul(6_364_136_223_846_793_005)
.wrapping_add(1_442_695_040_888_963_407);
((self.0 >> 40) as f32 / (1u64 << 24) as f32) * 2.0 - 1.0
}
fn vector(&mut self, dim: usize) -> Vec<f32> {
(0..dim).map(|_| self.next()).collect()
}
}
fn cosine_f32(a: &[f32], b: &[f32]) -> f32 {
let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
let na: f32 = libm::sqrtf(a.iter().map(|x| x * x).sum());
let nb: f32 = libm::sqrtf(b.iter().map(|x| x * x).sum());
dot / (na * nb)
}
#[test]
fn quantized_cosine_tracks_f32() {
let dim = 384;
let mut rng = Lcg(0x1234_5678);
let mut worst = 0.0f32;
for i in 0..200u32 {
let a = rng.vector(dim);
let b = rng.vector(dim);
let mut pool = VecPool::new(dim, usize::MAX);
pool.push(FactId(2 * i), &a).unwrap();
pool.push(FactId(2 * i + 1), &b).unwrap();
let q = pool.cosine_slots(0, 1);
let t = cosine_f32(&a, &b);
worst = worst.max(libm::fabsf(q - t));
}
assert!(
worst < 0.05,
"worst quantization error {worst} exceeds 0.05"
);
}
#[test]
fn golden_dim4() {
let dim = 4;
let mut pool = VecPool::new(dim, usize::MAX);
pool.push(FactId(0), &[1.0, 1.0, 0.0, 0.0]).unwrap();
pool.push(FactId(1), &[2.0, 2.0, 0.0, 0.0]).unwrap(); pool.push(FactId(2), &[0.0, 0.0, 1.0, 1.0]).unwrap(); assert!((pool.cosine_slots(0, 1) - 1.0).abs() < 1e-3);
assert!(pool.cosine_slots(0, 2).abs() < 1e-3);
pool.validate().unwrap();
let stride = pool.stride();
let sig = u64::from_le_bytes(pool.dump()[HEAD..HEAD + SIG_WORD_BYTES].try_into().unwrap());
assert_eq!(sig & 0b1111, 0b1111);
assert_eq!(pool.len(), 3);
assert_eq!(stride, HEAD + SIG_WORD_BYTES + 4);
}
#[test]
fn search_surfaces_the_nearest() {
let dim = 64;
let mut rng = Lcg(0xdead_beef);
let mut pool = VecPool::new(dim, usize::MAX);
let target = rng.vector(dim);
for i in 0..200u32 {
pool.push(FactId(i), &rng.vector(dim)).unwrap();
}
pool.push(FactId(500), &target).unwrap();
let mut scratch = VecScratch::new();
let mut out = Vec::new();
pool.search(&target, 5, &mut |_| true, &mut scratch, &mut out)
.unwrap();
assert_eq!(out[0].0, FactId(500), "exact match must rank first");
assert!(out[0].1 > 0.99, "self-cosine ≈ 1, got {}", out[0].1);
}
#[test]
fn degenerate_vectors_are_invalid() {
let mut pool = VecPool::new(3, usize::MAX);
assert_eq!(
pool.push(FactId(0), &[0.0, 0.0, 0.0]).unwrap_err(),
Error::Invalid("vector must be nonzero")
);
assert_eq!(
pool.push(FactId(0), &[1.0, f32::NAN, 0.0]).unwrap_err(),
Error::Invalid("vector must be finite")
);
assert!(matches!(
pool.push(FactId(0), &[1.0, 2.0]).unwrap_err(),
Error::DimMismatch { got: 2, want: 3 }
));
assert_eq!(pool.len(), 0);
assert!(pool.is_empty());
}
#[test]
fn accessors_and_edges() {
let dim = 4;
let mut pool = VecPool::new(dim, usize::MAX);
assert!(pool.is_empty());
assert_eq!(pool.pool_bytes(), 0);
let mut scratch = VecScratch::new();
let mut out = vec![(FactId(9), 1.0)];
pool.search(&[1.0; 4], 5, &mut |_| true, &mut scratch, &mut out)
.unwrap();
assert!(out.is_empty());
pool.push(FactId(0), &[1.0, 0.0, 0.0, 0.0]).unwrap();
assert!(!pool.is_empty());
assert_eq!(pool.pool_bytes(), pool.stride());
pool.search(&[1.0; 4], 0, &mut |_| true, &mut scratch, &mut out)
.unwrap();
assert!(out.is_empty());
assert_eq!(pool.cosine_slots(0, 9), 0.0);
let mut tight = VecPool::new(dim, 4);
assert_eq!(
tight.push(FactId(0), &[1.0, 0.0, 0.0, 0.0]).unwrap_err(),
Error::CapacityExceeded { what: "vectors" }
);
}
#[test]
fn from_parts_frames_slots() {
let dim = 8;
let mut pool = VecPool::new(dim, usize::MAX);
pool.push(FactId(0), &vec![0.5; dim]).unwrap();
pool.push(FactId(1), &vec![-0.5; dim]).unwrap();
let bytes = pool.dump();
let rebuilt = VecPool::from_parts(dim, usize::MAX, &bytes).unwrap();
assert_eq!(rebuilt.len(), 2);
rebuilt.validate().unwrap();
assert!(VecPool::from_parts(dim, usize::MAX, &bytes[..bytes.len() - 1]).is_err());
assert!(VecPool::from_parts(0, usize::MAX, &bytes).is_err());
assert!(VecPool::from_parts(dim, bytes.len() - 1, &bytes).is_err());
}
#[test]
fn validate_rejects_malformed_slots() {
let dim = 8;
let mut pool = VecPool::new(dim, usize::MAX);
pool.push(FactId(0), &vec![0.5; dim]).unwrap();
let good = pool.dump();
let mut bad = good.clone();
bad[FACT_BYTES..HEAD].copy_from_slice(&f32::NAN.to_le_bytes());
assert!(
VecPool::from_parts(dim, usize::MAX, &bad)
.unwrap()
.validate()
.is_err()
);
let mut bad = good.clone();
let q_off = HEAD + VecPool::words(dim) * SIG_WORD_BYTES;
bad[q_off] = (-1i8) as u8; assert!(
VecPool::from_parts(dim, usize::MAX, &bad)
.unwrap()
.validate()
.is_err()
);
}
#[test]
fn overlay_appends_to_tail_and_reads_span_the_boundary() {
let dim = 16;
let mut rng = Lcg(0x0ace_1a75);
let (va, vb, vc) = (rng.vector(dim), rng.vector(dim), rng.vector(dim));
let mut owned = VecPool::new(dim, usize::MAX);
owned.push(FactId(10), &va).unwrap();
owned.push(FactId(11), &vb).unwrap();
owned.push(FactId(12), &vc).unwrap();
let mut seed = VecPool::new(dim, usize::MAX);
seed.push(FactId(10), &va).unwrap();
seed.push(FactId(11), &vb).unwrap();
let base = seed.dump();
let base_snapshot = base.clone();
let mut pool = VecPool::from_parts_borrowed(dim, usize::MAX, &base).unwrap();
assert_eq!(pool.len(), 2);
let idx = pool.push(FactId(12), &vc).unwrap();
assert_eq!(idx, 2);
assert_eq!(pool.len(), 3);
assert_eq!(pool.slot_fact(0), 10); assert_eq!(pool.slot_fact(2), 12); assert!((pool.cosine_slots(0, 2) - owned.cosine_slots(0, 2)).abs() < 1e-6);
pool.validate().unwrap();
let mut scratch = VecScratch::new();
let mut out = Vec::new();
pool.search(&vc, 1, &mut |_| true, &mut scratch, &mut out)
.unwrap();
assert_eq!(out[0].0, FactId(12));
assert_eq!(pool.dump(), owned.dump());
assert_eq!(base, base_snapshot);
}
}