#[cfg(not(target_pointer_width = "64"))]
compile_error!("turbovec requires a 64-bit target (target_pointer_width = \"64\")");
pub mod codebook;
pub mod encode;
pub mod error;
pub mod id_map;
pub mod convert;
pub mod io;
mod io_v7;
pub mod pack;
pub mod rotation;
pub mod search;
pub mod warning;
#[cfg(test)]
mod kernel_tests;
pub use error::{AddError, CalibrateError, ConstructError, FromPartsError, SearchError};
pub use id_map::{IdMapIndex, IdSearchResults};
pub use warning::{set_warning_hook, WarningHook};
use std::path::Path;
use std::sync::OnceLock;
const BLOCK: usize = 32;
pub const MAX_DIM: usize = 16384;
const FLUSH_EVERY: usize = 256;
const MAX_INPUT_MAGNITUDE: f32 = 1e16;
#[cfg(test)]
thread_local! {
static FORCE_REPACK_PANIC: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
}
#[cfg(test)]
thread_local! {
static FORCE_ENCODE_PANIC: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
}
#[cfg(test)]
thread_local! {
static FORCE_FIT_PANIC: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
}
#[cfg(test)]
thread_local! {
static FORCE_SWAP_REMOVE_PANIC: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
}
pub const MIN_INPUT_NORM: f32 = 1e-10;
pub const MIN_CALIBRATION_ROWS: usize = encode::MIN_CALIBRATION_ROWS;
pub const RECOMMENDED_CALIBRATION_ROWS: usize = encode::RECOMMENDED_CALIBRATION_ROWS;
pub fn expected_codebook(bit_width: usize, dim: usize) -> (Vec<f32>, Vec<f32>) {
assert!(
(2..=4).contains(&bit_width),
"bit_width must be 2, 3 or 4, got {bit_width}"
);
assert!(
dim >= 8 && dim % 8 == 0,
"dim must be a positive multiple of 8, got {dim}"
);
assert!(
dim <= MAX_DIM,
"dim must be <= {MAX_DIM} (MAX_DIM), got {dim}"
);
codebook::codebook(bit_width, dim)
}
pub fn first_invalid_coord(values: &[f32], dim: usize) -> Option<(usize, usize, f32)> {
encode::par_first_invalid_coord(values, dim, MAX_INPUT_MAGNITUDE)
}
pub fn validation_parallelizes(len: usize) -> bool {
len > encode::VALIDATE_CHUNK
}
#[derive(Debug)]
struct BlockedCache {
data: Vec<u8>,
n_blocks: usize,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum CalibrationState {
Uncalibrated,
Calibrated,
}
#[derive(Debug)]
pub struct TurboQuantIndex {
dim: Option<usize>,
bit_width: usize,
n_vectors: usize,
packed_codes: OnceLock<Vec<u8>>,
scales: Vec<f32>,
tqplus_shift: Vec<f32>,
tqplus_scale: Vec<f32>,
rotation: OnceLock<rotation::Rotation>,
boundaries: OnceLock<Vec<f32>>,
centroids: OnceLock<Vec<f32>>,
blocked: OnceLock<BlockedCache>,
encode_scratch: Vec<f32>,
encode_scratch_prev: usize,
sync_cursor: Option<io_v7::SyncCursor>,
sync_path: Option<std::path::PathBuf>,
sync_pending: std::collections::HashSet<usize>,
sync_capture_buf: Vec<u8>,
sync_capture_at: Vec<(u32, u32)>,
sync_fresh: std::collections::HashSet<usize>,
calib_gen: u64,
}
fn retain_scratch(scratch: &mut Vec<f32>, prev: usize, want: usize) -> usize {
let target = prev.saturating_add(prev / 2);
if scratch.capacity() > target.saturating_mul(2) {
scratch.truncate(target);
scratch.shrink_to(target);
}
want
}
pub(crate) fn reserve_mostly_exact<T>(v: &mut Vec<T>, additional: usize) {
let len = v.len();
if v.capacity() - len >= additional {
return;
}
let headroom = len / 8;
if additional <= headroom {
v.reserve_exact(additional + headroom);
} else {
v.reserve(additional);
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct SearchResults {
pub scores: Vec<f32>,
pub indices: Vec<i64>,
pub nq: usize,
pub k: usize,
}
impl SearchResults {
pub fn scores_for_query(&self, qi: usize) -> &[f32] {
&self.scores[qi * self.k..(qi + 1) * self.k]
}
pub fn indices_for_query(&self, qi: usize) -> &[i64] {
&self.indices[qi * self.k..(qi + 1) * self.k]
}
}
impl TurboQuantIndex {
fn packed(&self) -> &Vec<u8> {
self.packed_codes.get_or_init(|| {
let (Some(dim), Some(cache)) = (self.dim, self.blocked.get()) else {
debug_assert!(
self.n_vectors == 0,
"packed_codes unset with no blocked cache but n_vectors > 0"
);
return Vec::new();
};
if self.n_vectors == 0 {
return Vec::new();
}
let (_, nbg, _) = pack::blocked_geometry(self.n_vectors, self.bit_width, dim);
let seq = pack::native_to_seq(&cache.data, self.bit_width, nbg);
pack::seq_to_packed(&seq, self.n_vectors, self.bit_width, dim)
})
}
pub fn packed_ready(&self) -> bool {
self.packed_codes.get().is_some()
}
fn packed_mut(&mut self) -> &mut Vec<u8> {
self.packed();
self.packed_codes
.get_mut()
.expect("packed_codes just materialized")
}
pub fn new(dim: usize, bit_width: usize) -> Result<Self, ConstructError> {
if !(2..=4).contains(&bit_width) {
return Err(ConstructError::BitWidthOutOfRange(bit_width));
}
if dim == 0 || dim % 8 != 0 {
return Err(ConstructError::DimNotPositiveMultipleOf8(dim));
}
if dim > MAX_DIM {
return Err(ConstructError::DimTooLarge { dim, max: MAX_DIM });
}
Ok(Self {
dim: Some(dim),
bit_width,
n_vectors: 0,
packed_codes: OnceLock::from(Vec::new()),
scales: Vec::new(),
tqplus_shift: Vec::new(),
tqplus_scale: Vec::new(),
rotation: OnceLock::new(),
boundaries: OnceLock::new(),
centroids: OnceLock::new(),
blocked: OnceLock::new(),
encode_scratch: Vec::new(),
encode_scratch_prev: 0,
sync_cursor: None,
sync_path: None,
sync_pending: std::collections::HashSet::new(),
sync_fresh: std::collections::HashSet::new(),
sync_capture_buf: Vec::new(),
sync_capture_at: Vec::new(),
calib_gen: 0,
})
}
pub fn new_lazy(bit_width: usize) -> Result<Self, ConstructError> {
if !(2..=4).contains(&bit_width) {
return Err(ConstructError::BitWidthOutOfRange(bit_width));
}
Ok(Self {
dim: None,
bit_width,
n_vectors: 0,
packed_codes: OnceLock::from(Vec::new()),
scales: Vec::new(),
tqplus_shift: Vec::new(),
tqplus_scale: Vec::new(),
rotation: OnceLock::new(),
boundaries: OnceLock::new(),
centroids: OnceLock::new(),
blocked: OnceLock::new(),
encode_scratch: Vec::new(),
encode_scratch_prev: 0,
sync_cursor: None,
sync_path: None,
sync_pending: std::collections::HashSet::new(),
sync_fresh: std::collections::HashSet::new(),
sync_capture_buf: Vec::new(),
sync_capture_at: Vec::new(),
calib_gen: 0,
})
}
pub fn add(&mut self, vectors: &[f32]) {
let dim = self.dim.expect(
"TurboQuantIndex dim is not set; use add_2d(vectors, dim) on the \
first add or construct via TurboQuantIndex::new(dim, bit_width)",
);
let n = vectors.len() / dim;
assert_eq!(
vectors.len(),
n * dim,
"vectors length must be a multiple of dim"
);
if n == 0 {
return;
}
if let Some((vi, ci, v)) = first_invalid_coord(vectors, dim) {
panic!(
"invalid input value at vector {vi}, coord {ci}: {v} \
(must be finite and |value| < 1e16 to avoid f32 norm overflow)",
);
}
self.encode_and_append(vectors, n, dim);
}
#[cfg(test)]
pub(crate) fn force_encode_panic(on: bool) {
FORCE_ENCODE_PANIC.with(|f| f.set(on));
}
#[cfg(test)]
pub(crate) fn force_repack_panic(on: bool) {
FORCE_REPACK_PANIC.with(|f| f.set(on));
}
#[cfg(test)]
pub(crate) fn force_encode_panic_after_append(on: bool) {
encode::force_panic_after_append(on);
}
#[cfg(test)]
pub(crate) fn force_fit_panic(on: bool) {
FORCE_FIT_PANIC.with(|f| f.set(on));
}
#[cfg(test)]
pub(crate) fn force_swap_remove_panic(on: bool) {
FORCE_SWAP_REMOVE_PANIC.with(|f| f.set(on));
}
fn encode_and_append(&mut self, vectors: &[f32], n: usize, dim: usize) {
debug_assert!(
self.sync_capture_at.iter().all(|&(s, _)| {
let s = s as usize;
s < self.n_vectors || s >= self.n_vectors + n
}),
"a slot about to be written by add still holds a removal capture",
);
let rotation = self
.rotation
.get_or_init(|| rotation::Rotation::new(dim));
if self.boundaries.get().is_none() || self.centroids.get().is_none() {
let (boundaries, centroids) = codebook::codebook(self.bit_width, dim);
let _ = self.boundaries.set(boundaries);
let _ = self.centroids.set(centroids);
}
let boundaries = self
.boundaries
.get()
.expect("boundaries cache is initialized");
let centroids = self
.centroids
.get()
.expect("centroids cache is initialized");
let existing = if self.tqplus_shift.is_empty() {
None
} else {
Some((self.tqplus_shift.as_slice(), self.tqplus_scale.as_slice()))
};
let lazy_append = self.n_vectors > 0
&& self.packed_codes.get().is_none()
&& self.blocked.get().is_some();
if !lazy_append {
self.packed();
}
let mut scratch = std::mem::take(&mut self.encode_scratch);
let mut packed_codes = self.packed_codes.take().unwrap_or_default();
debug_assert!(
lazy_append || self.n_vectors == 0 || !packed_codes.is_empty(),
"eager add must start from materialized packed rows"
);
let mut scales_buf = std::mem::take(&mut self.scales);
let packed_len_before = packed_codes.len();
let scales_len_before = scales_buf.len();
let bit_width = self.bit_width;
let encode_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
#[cfg(test)]
if FORCE_ENCODE_PANIC.with(|f| f.replace(false)) {
panic!("forced encode panic (test)");
}
encode::encode(
vectors,
n,
dim,
rotation,
boundaries,
centroids,
bit_width,
existing,
&mut scratch,
&mut packed_codes,
&mut scales_buf,
)
}));
if let Err(panic) = encode_result {
{
scales_buf.truncate(scales_len_before);
if !lazy_append {
packed_codes.truncate(packed_len_before);
self.packed_codes = OnceLock::from(packed_codes);
}
self.scales = scales_buf;
self.encode_scratch = scratch;
std::panic::resume_unwind(panic);
}
}
self.encode_scratch_prev = retain_scratch(&mut scratch, self.encode_scratch_prev, n * dim);
self.encode_scratch = scratch;
let old_n = self.n_vectors;
let new_n = old_n + n;
if lazy_append {
let bit_width = self.bit_width;
let cache = self
.blocked
.get_mut()
.expect("lazy_append requires a blocked cache");
pack::append_lanes(&mut cache.data, &packed_codes, old_n, n, bit_width, dim);
let (new_n_blocks, _, _) = pack::blocked_geometry(new_n, bit_width, dim);
cache.n_blocks = new_n_blocks;
self.scales = scales_buf;
self.n_vectors = new_n;
return;
}
if self.blocked.get().is_none() {
let bit_width = self.bit_width;
let built = match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
#[cfg(test)]
if FORCE_REPACK_PANIC.with(|f| f.replace(false)) {
panic!("forced repack panic (test)");
}
pack::repack(&packed_codes, new_n, bit_width, dim)
})) {
Ok(built) => built,
Err(panic) => {
packed_codes.truncate(packed_len_before);
scales_buf.truncate(scales_len_before);
self.packed_codes = OnceLock::from(packed_codes);
self.scales = scales_buf;
std::panic::resume_unwind(panic);
}
};
let (data, n_blocks) = built;
let _ = self.blocked.set(BlockedCache { data, n_blocks });
} else {
let (new_n_blocks, n_byte_groups, _) =
pack::blocked_geometry(new_n, self.bit_width, dim);
let block_bytes = n_byte_groups * BLOCK;
let first_block = old_n / BLOCK;
let bit_width = self.bit_width;
let patch = match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
#[cfg(test)]
if FORCE_REPACK_PANIC.with(|f| f.replace(false)) {
panic!("forced repack panic (test)");
}
pack::repack_block_range(
&packed_codes,
new_n,
bit_width,
dim,
first_block,
new_n_blocks,
)
})) {
Ok(patch) => patch,
Err(panic) => {
packed_codes.truncate(packed_len_before);
scales_buf.truncate(scales_len_before);
self.packed_codes = OnceLock::from(packed_codes);
self.scales = scales_buf;
std::panic::resume_unwind(panic);
}
};
let cache = self.blocked.get_mut().expect("blocked present");
cache.data.truncate(first_block * block_bytes);
reserve_mostly_exact(&mut cache.data, patch.len());
cache.data.extend_from_slice(&patch);
cache.n_blocks = new_n_blocks;
}
drop(packed_codes);
self.scales = scales_buf;
self.n_vectors = new_n;
}
pub fn add_2d(&mut self, vectors: &[f32], dim: usize) -> Result<(), AddError> {
match self.dim {
Some(existing) if existing != dim => {
return Err(AddError::DimMismatch { existing, got: dim });
}
Some(_) => {}
None => {
if dim == 0 {
return Err(AddError::ZeroDim);
}
if dim % 8 != 0 {
return Err(AddError::DimNotMultipleOf8(dim));
}
if dim > MAX_DIM {
return Err(AddError::DimTooLarge { dim, max: MAX_DIM });
}
}
}
if let Some((vi, ci, v)) = first_invalid_coord(vectors, dim) {
return Err(AddError::InvalidInputValue {
vector_index: vi,
coord_index: ci,
value: v,
});
}
assert_eq!(
vectors.len() % dim,
0,
"vectors length must be a multiple of dim"
);
if vectors.is_empty() {
return Ok(());
}
if self.dim.is_none() {
self.dim = Some(dim);
if let Err(panic) =
std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| self.add(vectors)))
{
self.dim = None;
self.rotation = OnceLock::new();
self.boundaries = OnceLock::new();
self.centroids = OnceLock::new();
std::panic::resume_unwind(panic);
}
return Ok(());
}
self.add(vectors);
Ok(())
}
pub fn search(&self, queries: &[f32], k: usize) -> SearchResults {
self.search_with_mask(queries, k, None)
}
pub fn try_search(&self, queries: &[f32], k: usize) -> Result<SearchResults, SearchError> {
self.try_search_with_mask(queries, k, None)
}
pub fn search_with_mask(
&self,
queries: &[f32],
k: usize,
mask: Option<&[bool]>,
) -> SearchResults {
self.try_search_with_mask(queries, k, mask)
.unwrap_or_else(|e| panic!("{e}"))
}
pub fn try_search_with_mask(
&self,
queries: &[f32],
k: usize,
mask: Option<&[bool]>,
) -> Result<SearchResults, SearchError> {
let Some(dim) = self.dim else {
return Ok(SearchResults {
scores: Vec::new(),
indices: Vec::new(),
nq: 0,
k: 0,
});
};
let nq = queries.len() / dim;
if queries.len() != nq * dim {
return Err(SearchError::QueryBufferNotMultipleOfDim {
queries_len: queries.len(),
dim,
});
}
if let Some((vi, ci, v)) = first_invalid_coord(queries, dim) {
return Err(SearchError::InvalidQueryValue {
query_index: vi,
coord_index: ci,
value: v,
});
}
if self.n_vectors == 0 {
if let Some(m) = mask {
if !m.is_empty() {
return Err(SearchError::MaskLengthMismatch {
expected: 0,
got: m.len(),
});
}
}
return Ok(SearchResults {
scores: Vec::new(),
indices: Vec::new(),
nq,
k: 0,
});
}
let rotation = self
.rotation
.get_or_init(|| rotation::Rotation::new(dim));
let centroids = self.centroids.get_or_init(|| {
let (_, c) = codebook::codebook(self.bit_width, dim);
c
});
let blocked = self.blocked.get_or_init(|| {
let (data, n_blocks) =
pack::repack(self.packed(), self.n_vectors, self.bit_width, dim);
BlockedCache { data, n_blocks }
});
if let Some(m) = mask {
if m.len() != self.n_vectors {
return Err(SearchError::MaskLengthMismatch {
expected: self.n_vectors,
got: m.len(),
});
}
}
let packed_mask = mask.map(|m| {
let n_words = self.n_vectors.div_ceil(64);
let mut buf = Vec::with_capacity(n_words);
let mut allowed = 0usize;
for chunk in m.chunks(64) {
let mut word = 0u64;
for (bit, &b) in chunk.iter().enumerate() {
word |= (b as u64) << bit;
}
allowed += word.count_ones() as usize;
buf.push(word);
}
debug_assert_eq!(buf.len(), n_words);
(buf, allowed)
});
let n_allowed = packed_mask.as_ref().map_or(self.n_vectors, |p| p.1);
let packed_mask = packed_mask.map(|p| p.0);
let effective_k = k.min(self.n_vectors).min(n_allowed);
let (scores, indices) = search::search(
queries,
nq,
rotation,
&blocked.data,
centroids,
&self.scales,
&self.tqplus_shift,
&self.tqplus_scale,
self.bit_width,
dim,
self.n_vectors,
blocked.n_blocks,
k,
packed_mask.as_deref(),
);
Ok(SearchResults {
scores,
indices,
nq,
k: effective_k,
})
}
pub fn prepare(&self) {
let Some(dim) = self.dim else { return };
if self.n_vectors == 0 {
return;
}
self.rotation
.get_or_init(|| rotation::Rotation::new(dim));
self.centroids.get_or_init(|| {
let (_, c) = codebook::codebook(self.bit_width, dim);
c
});
self.blocked.get_or_init(|| {
let (data, n_blocks) =
pack::repack(self.packed(), self.n_vectors, self.bit_width, dim);
BlockedCache { data, n_blocks }
});
}
pub(crate) fn sync_watermark(&self) -> usize {
self.sync_cursor
.map(|c| (c.n_synced as usize) / BLOCK * BLOCK)
.unwrap_or(0)
}
#[cfg(all(test, target_arch = "x86_64"))]
pub(crate) fn captured_len_for_test(&self) -> usize {
self.sync_capture_at.len()
}
#[cfg(all(test, target_arch = "x86_64"))]
pub(crate) fn sync_capture_buf_len_for_test(&self) -> usize {
self.sync_capture_buf.len()
}
fn capture_lookup(&self) -> std::collections::HashMap<usize, (usize, usize)> {
if self.packed_codes.get().is_some() || self.sync_capture_at.is_empty() {
return std::collections::HashMap::new();
}
let dim = self.dim.expect("captures exist, so dim is committed");
let (_, n_byte_groups, _) =
pack::blocked_geometry(self.n_vectors, self.bit_width, dim);
let mut m = std::collections::HashMap::with_capacity(self.sync_capture_at.len());
for &(slot, off) in &self.sync_capture_at {
m.insert(slot as usize, (off as usize, n_byte_groups));
}
m
}
fn mark_dirty(&mut self, slot: usize) {
if slot < self.sync_watermark() {
self.sync_fresh.insert(slot);
}
}
#[cfg(test)]
pub(crate) fn plan_next_sync(&mut self, kind: u8, ids: Option<&[u64]>) -> io_v7::SyncPlan {
let dim = self.dim.expect("plan_next_sync on a lazy index");
if self.blocked.get().is_none() {
self.packed();
}
if self.boundaries.get().is_none() || self.centroids.get().is_none() {
let (b, c) = codebook::codebook(self.bit_width, dim);
let _ = self.boundaries.set(b);
let _ = self.centroids.set(c);
}
let seq_blocks = |from: usize, to: usize| self.seq_blocks_range(from, to);
let capture_lookup = self.capture_lookup();
let row_codes = |idx: usize, out: &mut Vec<u8>| {
match capture_lookup.get(&idx) {
Some(&(off, len)) => {
out.extend_from_slice(&self.sync_capture_buf[off..off + len])
}
None => out.extend_from_slice(&self.seq_row(idx)),
}
};
let source = io_v7::SyncSource {
kind,
dim,
bit_width: self.bit_width,
n_vectors: self.n_vectors,
seq_blocks: &seq_blocks,
row_codes: &row_codes,
scales: &self.scales,
ids,
tqplus_shift: &self.tqplus_shift,
tqplus_scale: &self.tqplus_scale,
boundaries: self.boundaries.get().expect("seeded above"),
centroids: self.centroids.get().expect("seeded above"),
};
let stale_ahead = match (&self.sync_cursor, &self.sync_path) {
(Some(c), Some(p)) => {
let geo = io_v7::Geo {
kind,
dim,
bit_width: self.bit_width,
n_calib: self.tqplus_shift.len(),
};
match io_v7::cursor_state(p, c, &geo) {
Ok(io_v7::CursorState::Intact { stale_ahead }) => stale_ahead,
_ => None,
}
}
_ => None,
};
io_v7::plan_incremental(
&source,
self.sync_cursor.expect("plan_next_sync on an unbound index"),
&self.sync_pending,
&self.sync_fresh,
stale_ahead,
)
.expect("plan_next_sync: ops exceed the header capacity")
}
fn seq_row(&self, idx: usize) -> Vec<u8> {
let dim = self.dim.expect("seq_row on a dim-less index");
let packed_row = dim * self.bit_width / 8;
let (_, row_bytes, _) = pack::blocked_geometry(1, self.bit_width, dim);
if let Some(packed) = self.packed_codes.get() {
return pack::extract_codes_flat(
&packed[idx * packed_row..(idx + 1) * packed_row],
1,
self.bit_width,
dim,
);
}
let cache = self.blocked.get().expect("no code layout materialized");
let b = idx / BLOCK;
let lane = idx % BLOCK;
(0..row_bytes)
.map(|g| pack::read_code(&cache.data, self.bit_width, row_bytes, b, g, lane))
.collect()
}
fn seq_blocks_range(&self, from: usize, to: usize) -> Vec<u8> {
debug_assert!(from.is_multiple_of(BLOCK) && to.is_multiple_of(BLOCK) && from <= to);
let dim = self.dim.expect("seq_blocks_range on a dim-less index");
let packed_row = dim * self.bit_width / 8;
let (_, row_bytes, _) = pack::blocked_geometry(1, self.bit_width, dim);
if let Some(packed) = self.packed_codes.get() {
let flat = pack::extract_codes_flat(
&packed[from * packed_row..to * packed_row],
to - from,
self.bit_width,
dim,
);
let n = to - from;
return pack::pack_blocked_sequential(
n,
n / BLOCK,
row_bytes,
n / BLOCK * row_bytes * BLOCK,
&flat,
);
}
let cache = self.blocked.get().expect("no code layout materialized");
let block_bytes = row_bytes * BLOCK;
pack::native_to_seq(
&cache.data[from / BLOCK * block_bytes..to / BLOCK * block_bytes],
self.bit_width,
row_bytes,
)
}
pub fn sync(&mut self, path: impl AsRef<Path>) -> std::io::Result<()> {
self.sync_v7_impl(path.as_ref(), 0, None)
}
fn with_sync_source<R>(
&self,
kind: u8,
ids_full: Option<&[u64]>,
f: impl FnOnce(&io_v7::SyncSource<'_>) -> R,
) -> std::io::Result<R> {
let dim = self.dim.unwrap_or(0);
if self.blocked.get().is_none() {
self.packed();
}
let n_levels = 1usize << self.bit_width;
let (lazy_b, lazy_c) = (vec![0.0f32; n_levels - 1], vec![0.0f32; n_levels]);
if dim != 0 && (self.boundaries.get().is_none() || self.centroids.get().is_none()) {
let (b, c) = codebook::codebook(self.bit_width, dim);
let _ = self.boundaries.set(b);
let _ = self.centroids.set(c);
}
let seq_blocks = |from: usize, to: usize| self.seq_blocks_range(from, to);
let capture_lookup = self.capture_lookup();
let row_codes = |idx: usize, out: &mut Vec<u8>| match capture_lookup.get(&idx) {
Some(&(off, len)) => out.extend_from_slice(&self.sync_capture_buf[off..off + len]),
None => out.extend_from_slice(&self.seq_row(idx)),
};
let source = io_v7::SyncSource {
kind,
dim,
bit_width: self.bit_width,
n_vectors: self.n_vectors,
seq_blocks: &seq_blocks,
row_codes: &row_codes,
scales: &self.scales,
ids: ids_full,
tqplus_shift: &self.tqplus_shift,
tqplus_scale: &self.tqplus_scale,
boundaries: if dim == 0 {
&lazy_b
} else {
self.boundaries.get().expect("seeded above")
},
centroids: if dim == 0 {
&lazy_c
} else {
self.centroids.get().expect("seeded above")
},
};
Ok(f(&source))
}
pub(crate) fn v7_image(&self, kind: u8, ids_full: Option<&[u64]>) -> std::io::Result<Vec<u8>> {
self.with_sync_source(kind, ids_full, io_v7::image_bytes)
}
pub(crate) fn v7_image_len(&self, kind: u8, ids_full: Option<&[u64]>) -> std::io::Result<usize> {
self.with_sync_source(kind, ids_full, io_v7::image_len)
}
pub(crate) fn sync_v7_impl(
&mut self,
path: &Path,
kind: u8,
ids_full: Option<&[u64]>,
) -> std::io::Result<()> {
let Some(dim) = self.dim else {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"cannot sync a lazy index that has never seen an add or calibrate",
));
};
if self.blocked.get().is_none() {
self.packed();
}
if self.boundaries.get().is_none() || self.centroids.get().is_none() {
let (b, c) = codebook::codebook(self.bit_width, dim);
let _ = self.boundaries.set(b);
let _ = self.centroids.set(c);
}
let geo = io_v7::Geo {
kind,
dim,
bit_width: self.bit_width,
n_calib: self.tqplus_shift.len(),
};
let bound = matches!(
(&self.sync_cursor, &self.sync_path),
(Some(_), Some(p)) if p == path
);
let state = if bound {
let c = self.sync_cursor.as_ref().expect("checked above");
io_v7::cursor_state(path, c, &geo)?
} else {
io_v7::CursorState::Replaced
};
let incremental = bound && {
let c = self.sync_cursor.as_ref().expect("checked above");
c.calib_gen == self.calib_gen
};
if matches!(state, io_v7::CursorState::Foreign) {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
format!(
"the v7 file at {} no longer matches this index's last sync \
(another writer advanced or replaced it); load() the file to \
adopt its state, or choose a different path",
path.display()
),
));
}
let result = {
let seq_blocks = |from: usize, to: usize| self.seq_blocks_range(from, to);
let capture_lookup = self.capture_lookup();
let row_codes = |idx: usize, out: &mut Vec<u8>| {
match capture_lookup.get(&idx) {
Some(&(off, len)) => {
out.extend_from_slice(&self.sync_capture_buf[off..off + len])
}
None => out.extend_from_slice(&self.seq_row(idx)),
}
};
let source = io_v7::SyncSource {
kind,
dim,
bit_width: self.bit_width,
n_vectors: self.n_vectors,
seq_blocks: &seq_blocks,
row_codes: &row_codes,
scales: &self.scales,
ids: ids_full,
tqplus_shift: &self.tqplus_shift,
tqplus_scale: &self.tqplus_scale,
boundaries: self.boundaries.get().expect("seeded above"),
centroids: self.centroids.get().expect("seeded above"),
};
match state {
io_v7::CursorState::Intact { stale_ahead } if incremental => {
let c = self.sync_cursor.expect("checked above");
match io_v7::plan_incremental(
&source,
c,
&self.sync_pending,
&self.sync_fresh,
stale_ahead,
) {
Some(plan) => {
io_v7::run_sync(path, &plan).map(|c| (c, plan.carried))
}
None => io_v7::write_full(path, &source, self.calib_gen)
.map(|c| (c, Vec::new())),
}
}
_ => io_v7::write_full(path, &source, self.calib_gen)
.map(|c| (c, Vec::new())),
}
};
match result {
Ok((cursor, carried)) => {
self.sync_pending = carried.into_iter().collect();
self.sync_fresh.clear();
self.sync_capture_buf.clear();
self.sync_capture_at.clear();
self.sync_cursor = Some(io_v7::SyncCursor {
calib_gen: self.calib_gen,
..cursor
});
self.sync_path = Some(path.to_path_buf());
Ok(())
}
Err(e) => {
self.sync_cursor = None;
self.sync_path = None;
self.sync_capture_buf.clear();
self.sync_capture_at.clear();
Err(e)
}
}
}
fn load_v7(path: &Path) -> std::io::Result<Self> {
let l = io_v7::load(path, 0, 0)?;
let bind = (l.cursor.nonce != io_v7::UNCLAIMED_NONCE).then_some(path);
Self::from_v7(l, bind)
}
pub(crate) fn from_v7(l: io_v7::V7Load, path: Option<&Path>) -> std::io::Result<Self> {
let n_blocks = l.n_vectors.div_ceil(BLOCK);
let (_, nbg, _) = pack::blocked_geometry(l.n_vectors, l.bit_width, l.dim);
let native = pack::seq_into_native(l.seq_blocked, l.bit_width, nbg);
let (tqplus_shift, tqplus_scale) =
Self::normalize_calibration(l.tqplus_shift, l.tqplus_scale);
let (boundaries, centroids) = if l.dim == 0 {
(Vec::new(), Vec::new())
} else {
codebook::codebook(l.bit_width, l.dim)
};
let blocked = OnceLock::new();
let boundaries_lock = OnceLock::new();
let centroids_lock = OnceLock::new();
let packed_codes = if l.n_vectors == 0 {
OnceLock::from(Vec::new())
} else {
let _ = blocked.set(BlockedCache {
data: native,
n_blocks,
});
let _ = boundaries_lock.set(boundaries);
let _ = centroids_lock.set(centroids);
OnceLock::new()
};
Ok(Self {
dim: (l.dim != 0).then_some(l.dim),
bit_width: l.bit_width,
n_vectors: l.n_vectors,
packed_codes,
scales: l.scales,
tqplus_shift,
tqplus_scale,
rotation: OnceLock::new(),
boundaries: boundaries_lock,
centroids: centroids_lock,
blocked,
encode_scratch: Vec::new(),
encode_scratch_prev: 0,
sync_cursor: path.map(|_| l.cursor),
sync_path: path.map(|p| p.to_path_buf()),
sync_pending: l.pending_slots.iter().copied().collect(),
sync_fresh: std::collections::HashSet::new(),
sync_capture_buf: Vec::new(),
sync_capture_at: Vec::new(),
calib_gen: 0,
})
}
pub fn write(&self, path: impl AsRef<Path>) -> std::io::Result<()> {
self.write_with_durability(path, io::Durability::Durable)
}
pub fn write_with_durability(
&self,
path: impl AsRef<Path>,
durability: io::Durability,
) -> std::io::Result<()> {
self.with_sync_source(0, None, |src| {
io_v7::write_snapshot(path.as_ref(), src, durability)
})?
}
pub fn codes_blocked_seq(&self) -> Vec<u8> {
let Some(dim) = self.dim else {
return Vec::new();
};
if self.n_vectors == 0 {
return Vec::new();
}
if let Some(cache) = self.blocked.get() {
let (_, nbg, _) = pack::blocked_geometry(self.n_vectors, self.bit_width, dim);
return pack::native_to_seq(&cache.data, self.bit_width, nbg);
}
pack::repack_seq(self.packed(), self.n_vectors, self.bit_width, dim)
}
pub fn codebook_for_write(&self) -> (Vec<f32>, Vec<f32>) {
let n_levels = 1usize << self.bit_width;
let Some(dim) = self.dim else {
return (vec![0.0; n_levels - 1], vec![0.0; n_levels]);
};
if self.n_vectors == 0 {
return (vec![0.0; n_levels - 1], vec![0.0; n_levels]);
}
if self.boundaries.get().is_none() || self.centroids.get().is_none() {
let (b, c) = codebook::codebook(self.bit_width, dim);
let _ = self.boundaries.set(b);
let _ = self.centroids.set(c);
}
let boundaries = self.boundaries.get().expect("boundaries just seeded");
let centroids = self.centroids.get().expect("centroids just seeded");
(boundaries.clone(), centroids.clone())
}
pub fn write_to_writer<W: std::io::Write>(&self, w: &mut W) -> std::io::Result<()> {
self.with_sync_source(0, None, |src| io_v7::stream_image(w, src))?
}
pub fn serialized_len(&self) -> usize {
self.v7_image_len(0, None)
.expect("with_sync_source handles the lazy sentinel, so this cannot fail")
}
pub fn to_bytes(&self) -> Vec<u8> {
self.v7_image(0, None)
.expect("with_sync_source handles the lazy sentinel, so this cannot fail")
}
pub fn load_from_reader<R: std::io::Read>(r: &mut R) -> std::io::Result<Self> {
let mut raw = Vec::new();
std::io::Read::read_to_end(r, &mut raw)?;
Self::from_v7_image(raw)
}
fn from_v7_image(raw: Vec<u8>) -> std::io::Result<Self> {
let l = io_v7::load_image(raw, 0, 0, "the byte image")?;
Self::from_v7(l, None)
}
pub fn from_bytes(bytes: &[u8]) -> std::io::Result<Self> {
Self::load_from_reader(&mut &bytes[..])
}
pub fn load(path: impl AsRef<Path>) -> std::io::Result<Self> {
if io_v7::is_v7(path.as_ref()) {
return Self::load_v7(path.as_ref());
}
Err(io::legacy_format_error(path.as_ref()))
}
fn normalize_calibration(
tqplus_shift: Vec<f32>,
tqplus_scale: Vec<f32>,
) -> (Vec<f32>, Vec<f32>) {
let declares_nothing = tqplus_shift.iter().all(|&x| x == 0.0)
&& tqplus_scale.iter().all(|&x| x == 1.0);
if declares_nothing {
return (Vec::new(), Vec::new());
}
(tqplus_shift, tqplus_scale)
}
pub fn from_parts(
dim: Option<usize>,
bit_width: usize,
n_vectors: usize,
packed_codes: Vec<u8>,
scales: Vec<f32>,
tqplus_shift: Vec<f32>,
tqplus_scale: Vec<f32>,
) -> Result<Self, FromPartsError> {
if !(2..=4).contains(&bit_width) {
return Err(FromPartsError::BitWidthOutOfRange(bit_width));
}
if tqplus_shift.len() != tqplus_scale.len() {
return Err(FromPartsError::TqplusLengthMismatch {
shift_len: tqplus_shift.len(),
scale_len: tqplus_scale.len(),
});
}
match dim {
Some(d) => {
if d == 0 || d % 8 != 0 {
return Err(FromPartsError::DimNotPositiveMultipleOf8(d));
}
if d > MAX_DIM {
return Err(FromPartsError::DimTooLarge { dim: d, max: MAX_DIM });
}
let expected_packed = (d / 8)
.checked_mul(bit_width)
.and_then(|x| x.checked_mul(n_vectors))
.ok_or(FromPartsError::PackedCodesSizeOverflow {
n_vectors,
dim: d,
bit_width,
})?;
if packed_codes.len() != expected_packed {
return Err(FromPartsError::PackedCodesLengthMismatch {
expected: expected_packed,
got: packed_codes.len(),
});
}
if scales.len() != n_vectors {
return Err(FromPartsError::ScalesLengthMismatch {
expected: n_vectors,
got: scales.len(),
});
}
if !tqplus_shift.is_empty() && tqplus_shift.len() != d {
return Err(FromPartsError::TqplusLengthNotDim {
got: tqplus_shift.len(),
dim: d,
});
}
}
None => {
if n_vectors != 0 {
return Err(FromPartsError::LazyMustHaveZeroVectors(n_vectors));
}
if !packed_codes.is_empty() {
return Err(FromPartsError::LazyMustHaveEmptyPackedCodes(
packed_codes.len(),
));
}
if !scales.is_empty() {
return Err(FromPartsError::LazyMustHaveEmptyScales(scales.len()));
}
if !tqplus_shift.is_empty() {
return Err(FromPartsError::LazyMustHaveEmptyTqplus(tqplus_shift.len()));
}
}
}
if let Some((i, &s)) = scales
.iter()
.enumerate()
.find(|(_, s)| {
!s.is_finite() || **s < 0.0 || **s > crate::io::MAX_VECTOR_SCALE
})
{
return Err(FromPartsError::InvalidScaleValue { slot: i, value: s });
}
if let Some((i, &v)) = tqplus_shift
.iter()
.enumerate()
.find(|(_, v)| {
!v.is_finite() || v.abs() > crate::io::max_tqplus_shift(tqplus_shift.len())
})
{
return Err(FromPartsError::InvalidTqplusShiftValue { coord: i, value: v });
}
if let Some((i, &v)) = tqplus_scale
.iter()
.enumerate()
.find(|(_, v)| {
!v.is_finite() || **v < crate::io::min_tqplus_scale(tqplus_scale.len())
})
{
return Err(FromPartsError::InvalidTqplusScaleValue { coord: i, value: v });
}
let (tqplus_shift, tqplus_scale) =
Self::normalize_calibration(tqplus_shift, tqplus_scale);
Ok(Self {
dim,
bit_width,
n_vectors,
packed_codes: OnceLock::from(packed_codes),
scales,
tqplus_shift,
tqplus_scale,
rotation: OnceLock::new(),
boundaries: OnceLock::new(),
centroids: OnceLock::new(),
blocked: OnceLock::new(),
encode_scratch: Vec::new(),
encode_scratch_prev: 0,
sync_cursor: None,
sync_path: None,
sync_pending: std::collections::HashSet::new(),
sync_fresh: std::collections::HashSet::new(),
sync_capture_buf: Vec::new(),
sync_capture_at: Vec::new(),
calib_gen: 0,
})
}
pub fn packed_codes(&self) -> &[u8] {
self.packed()
}
pub fn scales(&self) -> &[f32] {
&self.scales
}
pub fn tqplus_shift(&self) -> &[f32] {
&self.tqplus_shift
}
pub fn tqplus_scale(&self) -> &[f32] {
&self.tqplus_scale
}
pub fn calibrate_2d(&mut self, sample: &[f32], dim: usize) -> Result<(), CalibrateError> {
match self.dim {
Some(existing) if existing != dim => {
return Err(CalibrateError::DimMismatch { existing, got: dim });
}
Some(_) => {}
None => {
if dim == 0 {
return Err(CalibrateError::ZeroDim);
}
if dim % 8 != 0 {
return Err(CalibrateError::DimNotMultipleOf8(dim));
}
if dim > MAX_DIM {
return Err(CalibrateError::DimTooLarge { dim, max: MAX_DIM });
}
}
}
if sample.len() % dim != 0 {
return Err(CalibrateError::SampleBufferNotMultipleOfDim {
sample_len: sample.len(),
dim,
});
}
let n = sample.len() / dim;
if n < MIN_CALIBRATION_ROWS {
return Err(CalibrateError::SampleTooSmall {
rows: n,
min: MIN_CALIBRATION_ROWS,
});
}
if let Some((vi, ci, v)) = first_invalid_coord(sample, dim) {
return Err(CalibrateError::InvalidInputValue {
vector_index: vi,
coord_index: ci,
value: v,
});
}
let rotation = self.rotation.get_or_init(|| rotation::Rotation::new(dim));
if self.boundaries.get().is_none() || self.centroids.get().is_none() {
let (b, c) = codebook::codebook(self.bit_width, dim);
let _ = self.boundaries.set(b);
let _ = self.centroids.set(c);
}
let centroids = self.centroids.get().expect("centroids seeded above");
let mut scratch = std::mem::take(&mut self.encode_scratch);
let fitted = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
#[cfg(test)]
if FORCE_FIT_PANIC.with(|f| f.replace(false)) {
panic!("forced fit panic (test)");
}
encode::fit_calibration(sample, n, dim, rotation, centroids, &mut scratch)
}));
self.encode_scratch_prev = retain_scratch(&mut scratch, self.encode_scratch_prev, n * dim);
self.encode_scratch = scratch;
let (shift, scale_tq) = match fitted {
Ok(pair) => pair,
Err(panic) => {
if self.dim.is_none() {
self.rotation = OnceLock::new();
self.boundaries = OnceLock::new();
self.centroids = OnceLock::new();
}
std::panic::resume_unwind(panic);
}
};
debug_assert_eq!(shift.len(), dim, "fit returns a full-length pair");
if shift.iter().all(|&x| x == 0.0) && scale_tq.iter().all(|&x| x == 1.0) {
if self.dim.is_none() {
self.rotation = OnceLock::new();
self.boundaries = OnceLock::new();
self.centroids = OnceLock::new();
}
return Err(CalibrateError::DegenerateSample);
}
let reencoded = (self.n_vectors > 0).then(|| self.reencode_stored_rows(dim, (&shift, &scale_tq)));
self.dim = Some(dim);
if let Some((packed, scales)) = reencoded {
self.packed_codes = OnceLock::from(packed);
self.scales = scales;
self.blocked = OnceLock::new();
}
self.tqplus_shift = shift;
self.tqplus_scale = scale_tq;
self.calib_gen += 1;
self.sync_capture_buf.clear();
self.sync_capture_at.clear();
Ok(())
}
pub fn calibrate(&mut self, sample: &[f32]) -> Result<(), CalibrateError> {
let dim = self.dim.expect(
"TurboQuantIndex dim is not set; use calibrate_2d(sample, dim) on a \
lazy index or construct via TurboQuantIndex::new(dim, bit_width)",
);
self.calibrate_2d(sample, dim)
}
fn reencode_stored_rows(&self, dim: usize, new_pair: (&[f32], &[f32])) -> (Vec<u8>, Vec<f32>) {
let n = self.n_vectors;
let bits = self.bit_width;
let bytes_per_row = bits * (dim / 8);
let packed = self.packed();
let boundaries = self.boundaries.get().expect("populated index has a codebook");
let centroids = self.centroids.get().expect("populated index has a codebook");
let identity;
let (old_shift, old_inv): (&[f32], Vec<f32>) = if self.tqplus_shift.is_empty() {
identity = vec![0.0f32; dim];
(&identity, vec![1.0f32; dim])
} else {
(
self.tqplus_shift.as_slice(),
self.tqplus_scale.iter().map(|s| 1.0 / s).collect(),
)
};
let cpb = 8 / bits;
let field = if bits == 3 { 4 } else { bits };
let mask = (1u8 << bits) - 1;
let n_byte_groups = dim / cpb;
let mut new_packed = Vec::with_capacity(n * bytes_per_row);
let mut new_scales = Vec::with_capacity(n);
const REENCODE_CHUNK_ROWS: usize = 4096;
let chunk_rows = REENCODE_CHUNK_ROWS.min(n);
let mut recon = vec![0.0f32; chunk_rows * dim];
let mut norms = vec![0.0f32; chunk_rows];
let mut start = 0usize;
while start < n {
let rows = chunk_rows.min(n - start);
let codes_flat = pack::extract_codes_flat(
&packed[start * bytes_per_row..(start + rows) * bytes_per_row],
rows,
bits,
dim,
);
for i in 0..rows {
let row = &codes_flat[i * n_byte_groups..(i + 1) * n_byte_groups];
let out = &mut recon[i * dim..(i + 1) * dim];
let mut sumsq = 0.0f64;
for (d, slot) in out.iter_mut().enumerate() {
let code = ((row[d / cpb] >> ((cpb - 1 - d % cpb) * field)) & mask) as usize;
let x = centroids[code] * old_inv[d] - old_shift[d];
*slot = x;
sumsq += f64::from(x) * f64::from(x);
}
norms[i] = (f64::from(self.scales[start + i]) * sumsq) as f32;
}
encode::encode_prerotated(
&recon[..rows * dim],
&norms[..rows],
rows,
dim,
boundaries,
centroids,
bits,
Some(new_pair),
&mut new_packed,
&mut new_scales,
);
start += rows;
}
(new_packed, new_scales)
}
pub fn calibration_state(&self) -> CalibrationState {
if self.tqplus_shift.is_empty() {
CalibrationState::Uncalibrated
} else {
CalibrationState::Calibrated
}
}
pub fn swap_remove(&mut self, idx: usize) -> usize {
#[cfg(test)]
if FORCE_SWAP_REMOVE_PANIC.with(|f| f.replace(false)) {
panic!("forced swap_remove panic (test)");
}
assert!(
idx < self.n_vectors,
"index {idx} out of bounds (n_vectors = {})",
self.n_vectors
);
if self.sync_cursor.is_some() {
self.mark_dirty(idx);
let last = self.n_vectors - 1;
if last != idx {
self.mark_dirty(last);
}
}
let dim = self.dim.expect("n_vectors > 0 but dim is None");
let bytes_per_vec = dim * self.bit_width / 8;
let (_, capture_lane, _) = pack::blocked_geometry(self.n_vectors, self.bit_width, dim);
let capture_at = if cfg!(target_arch = "x86_64")
&& self.sync_cursor.is_some()
&& idx < self.sync_watermark()
&& self.packed_codes.get().is_none()
{
let capture_reuse = self
.sync_capture_at
.iter()
.find(|&&(s, _)| s as usize == idx)
.map(|&(_, off)| off as usize);
capture_reuse.or_else(|| {
let off = self.sync_capture_buf.len();
(off < io_v7::MAX_OPS * capture_lane).then_some(off)
})
} else {
None
};
let last = self.n_vectors - 1;
debug_assert!(
self.packed_codes.get().is_some() || self.blocked.get().is_some(),
"swap_remove: neither packed_codes nor the blocked cache is present"
);
if self.packed_codes.get().is_some() {
if idx != last {
let src = last * bytes_per_vec;
let dst = idx * bytes_per_vec;
self.packed_mut().copy_within(src..src + bytes_per_vec, dst);
}
self.packed_mut().truncate(last * bytes_per_vec);
}
if idx != last {
self.scales[idx] = self.scales[last];
}
self.scales.truncate(last);
self.n_vectors -= 1;
let capture_buf = &mut self.sync_capture_buf;
if let Some(cache) = self.blocked.get_mut() {
let (new_n_blocks, n_byte_groups, _) =
pack::blocked_geometry(self.n_vectors, self.bit_width, dim);
let block_bytes = n_byte_groups * BLOCK;
if idx != last {
let dst = capture_at.map(|off| {
if off == capture_buf.len() {
capture_buf.resize(off + n_byte_groups, 0);
}
&mut capture_buf[off..off + n_byte_groups]
});
pack::move_lane(&mut cache.data, self.bit_width, n_byte_groups, last, idx, dst);
}
pack::zero_lane(&mut cache.data, self.bit_width, n_byte_groups, last);
cache.data.truncate(new_n_blocks * block_bytes);
cache.n_blocks = new_n_blocks;
}
if !self.sync_capture_at.is_empty() {
self.sync_capture_at
.retain(|&(s, _)| s as usize != idx && s as usize != last);
}
if idx != last {
if let (Some(off), Ok(slot)) = (capture_at, u32::try_from(idx)) {
self.sync_capture_at.push((slot, off as u32));
}
}
last
}
pub fn len(&self) -> usize {
self.n_vectors
}
pub fn is_empty(&self) -> bool {
self.n_vectors == 0
}
#[deprecated(
since = "0.10.0",
note = "returns 0 for a lazy index, which is unsafe to do arithmetic with; use dim_opt()"
)]
pub fn dim(&self) -> usize {
self.dim.unwrap_or(0)
}
pub fn dim_opt(&self) -> Option<usize> {
self.dim
}
pub fn bit_width(&self) -> usize {
self.bit_width
}
}
#[cfg(test)]
mod scratch_retention_tests {
use super::TurboQuantIndex;
const DIM: usize = 256;
fn rows(n: usize, dim: usize) -> Vec<f32> {
(0..n * dim)
.map(|i| ((i % 97) as f32 / 97.0) - 0.5)
.collect()
}
#[test]
fn one_shot_bulk_add_releases_the_encode_scratch() {
let n = 24_000;
let mut idx = TurboQuantIndex::new(DIM, 4).unwrap();
idx.add_2d(&rows(n, DIM), DIM).unwrap();
assert_eq!(idx.len(), n);
assert!(
idx.encode_scratch.capacity() < n * DIM / 4,
"one-shot bulk add retained {} scratch elements (batch was {})",
idx.encode_scratch.capacity(),
n * DIM,
);
}
#[test]
fn repeated_same_size_adds_keep_the_scratch_warm() {
let n = 24_000;
let batch = rows(n, DIM);
let mut idx = TurboQuantIndex::new(DIM, 4).unwrap();
for _ in 0..3 {
idx.add_2d(&batch, DIM).unwrap();
}
assert_eq!(idx.len(), 3 * n);
assert!(
idx.encode_scratch.capacity() >= n * DIM,
"steady same-size adds dropped the warm scratch to {} elements (need {})",
idx.encode_scratch.capacity(),
n * DIM,
);
}
#[test]
fn growing_batch_sizes_keep_their_growth_headroom() {
let mut idx = TurboQuantIndex::new(DIM, 4).unwrap();
let mut n = 8_000;
let mut last = 0;
for _ in 0..6 {
idx.add_2d(&rows(n, DIM), DIM).unwrap();
last = n;
n += n / 20; }
assert!(
idx.encode_scratch.capacity() >= last * DIM,
"a growing batch size left only {} scratch elements after a \
{}-element add, so the next add must grow and be shrunk again",
idx.encode_scratch.capacity(),
last * DIM,
);
}
#[test]
fn jittering_batch_sizes_keep_their_growth_headroom() {
let mut idx = TurboQuantIndex::new(DIM, 4).unwrap();
let sizes = [9_000, 11_000, 9_500, 10_800, 9_200, 10_400];
for n in sizes {
idx.add_2d(&rows(n, DIM), DIM).unwrap();
}
let biggest_recent = 10_400;
assert!(
idx.encode_scratch.capacity() >= biggest_recent * DIM,
"a jittering batch size left only {} scratch elements, below \
the {} the last add needed",
idx.encode_scratch.capacity(),
biggest_recent * DIM,
);
}
#[test]
fn a_step_up_in_batch_size_is_not_shrunk_back() {
let small = 6_000;
let big = 3 * small;
let mut idx = TurboQuantIndex::new(DIM, 4).unwrap();
idx.add_2d(&rows(small, DIM), DIM).unwrap();
idx.add_2d(&rows(big, DIM), DIM).unwrap();
assert!(
idx.encode_scratch.capacity() >= big * DIM,
"a {small}->{big} step left only {} scratch elements, below the \
{} the larger batch needed",
idx.encode_scratch.capacity(),
big * DIM,
);
}
#[test]
fn a_one_off_spike_does_not_stay_pinned() {
let n = 6_000;
let mut idx = TurboQuantIndex::new(DIM, 4).unwrap();
idx.add_2d(&rows(n, DIM), DIM).unwrap();
idx.add_2d(&rows(4 * n, DIM), DIM).unwrap();
idx.add_2d(&rows(n, DIM), DIM).unwrap();
assert!(
idx.encode_scratch.capacity() < 4 * n * DIM,
"a 4x spike left {} scratch elements pinned after the batch \
size dropped back to {}",
idx.encode_scratch.capacity(),
n * DIM,
);
}
#[test]
fn retain_scratch_truncates_before_shrinking() {
let big = 8 << 20;
let mut scratch: Vec<f32> = vec![0.0; big];
let prev = super::retain_scratch(&mut scratch, 0, big);
assert_eq!(prev, big, "returns this call's demand");
assert_eq!(
scratch.capacity(),
0,
"a buffer no recent add needed was not released",
);
}
}
#[cfg(test)]
mod from_parts_tests {
use super::TurboQuantIndex;
use crate::FromPartsError;
#[test]
fn from_parts_rejects_packed_codes_length_mismatch() {
let err = TurboQuantIndex::from_parts(
Some(64),
4,
2,
vec![0u8; 32],
vec![1.0f32; 2],
Vec::new(),
Vec::new(),
)
.unwrap_err();
assert!(matches!(
err,
FromPartsError::PackedCodesLengthMismatch { expected: 64, got: 32 }
));
}
#[test]
fn from_parts_rejects_lazy_with_nonzero_n_vectors() {
let err = TurboQuantIndex::from_parts(
None,
4,
5,
Vec::new(),
Vec::new(),
Vec::new(),
Vec::new(),
)
.unwrap_err();
assert!(matches!(err, FromPartsError::LazyMustHaveZeroVectors(5)));
}
#[test]
fn from_parts_accepts_lazy_uncommitted() {
let idx = TurboQuantIndex::from_parts(
None,
4,
0,
Vec::new(),
Vec::new(),
Vec::new(),
Vec::new(),
)
.unwrap();
assert_eq!(idx.dim_opt(), None);
assert_eq!(idx.len(), 0);
}
#[test]
fn from_parts_accepts_eager_with_consistent_lengths() {
let idx = TurboQuantIndex::from_parts(
Some(64),
4,
2,
vec![0u8; 64],
vec![1.0f32; 2],
Vec::new(),
Vec::new(),
)
.unwrap();
assert_eq!(idx.dim_opt(), Some(64));
assert_eq!(idx.len(), 2);
assert!(idx.tqplus_shift().is_empty());
assert!(idx.tqplus_scale().is_empty());
}
}
#[cfg(all(test, target_arch = "x86_64"))]
mod x86_scalar_fallback_tests {
use super::TurboQuantIndex;
use crate::search::FORCE_SCALAR_FALLBACK;
use std::sync::atomic::Ordering;
fn unit_vectors(n: usize, dim: usize, seed: u64) -> Vec<f32> {
let mut s = seed.wrapping_add(0x9E3779B97F4A7C15);
let mut out = vec![0.0f32; n * dim];
for row in out.chunks_mut(dim) {
let mut norm = 0.0f64;
for x in row.iter_mut() {
s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
let v = ((s >> 33) as f64 / (1u64 << 31) as f64) - 1.0;
*x = v as f32;
norm += v * v;
}
let inv = 1.0 / (norm.sqrt() + 1e-9);
for x in row.iter_mut() {
*x = (*x as f64 * inv) as f32;
}
}
out
}
fn topk_sets(indices: &[i64], nq: usize, k: usize) -> Vec<std::collections::BTreeSet<i64>> {
(0..nq)
.map(|q| indices[q * k..(q + 1) * k].iter().copied().collect())
.collect()
}
fn exact_topk(
data: &[f32],
queries: &[f32],
n: usize,
dim: usize,
nq: usize,
k: usize,
) -> Vec<std::collections::BTreeSet<i64>> {
(0..nq)
.map(|q| {
let qr = &queries[q * dim..(q + 1) * dim];
let mut scored: Vec<(f32, i64)> = (0..n)
.map(|i| {
let row = &data[i * dim..(i + 1) * dim];
(qr.iter().zip(row).map(|(a, b)| a * b).sum::<f32>(), i as i64)
})
.collect();
scored.sort_by(|a, b| {
b.0.partial_cmp(&a.0).unwrap().then_with(|| a.1.cmp(&b.1))
});
scored[..k].iter().map(|p| p.1).collect()
})
.collect()
}
#[test]
fn scalar_fallback_matches_simd_topk() {
let dim = 64;
let n = 600;
let nq = 12;
let k = 16;
for &bits in &[2usize, 3, 4] {
let data = unit_vectors(n, dim, 11);
let mut idx = TurboQuantIndex::new(dim, bits).unwrap();
idx.add(&data);
let queries = unit_vectors(nq, dim, 22);
FORCE_SCALAR_FALLBACK.store(false, Ordering::Relaxed);
let simd = idx.search(&queries, k);
FORCE_SCALAR_FALLBACK.store(true, Ordering::Relaxed);
let scalar = idx.search(&queries, k);
FORCE_SCALAR_FALLBACK.store(false, Ordering::Relaxed);
assert_eq!(simd.k, scalar.k, "bits={bits}: differing result width");
let simd_sets = topk_sets(&simd.indices, nq, simd.k);
let scalar_sets = topk_sets(&scalar.indices, nq, scalar.k);
if bits == 4 && crate::pack::vector_major_for(bits, dim / (8 / bits)) {
let exact = exact_topk(&data, &queries, n, dim, nq, k);
let hits = |sets: &[std::collections::BTreeSet<i64>]| -> usize {
sets.iter().zip(exact.iter()).map(|(a, b)| a.intersection(b).count()).sum()
};
let (hs, hc) = (hits(&simd_sets), hits(&scalar_sets));
assert!(
hs + 2 >= hc,
"bits={bits}: scalar fallback ({hc}/{}) beat the SIMD kernel ({hs}/{}) \
against exact scores — the kernel is the more accurate of the two",
nq * k,
nq * k,
);
let agree: usize = simd_sets
.iter()
.zip(scalar_sets.iter())
.map(|(a, b)| a.intersection(b).count())
.sum();
assert!(
agree * 4 >= nq * k * 3,
"bits={bits}: scalar fallback agreed with SIMD on only {agree}/{} slots",
nq * k,
);
} else {
assert_eq!(
simd_sets, scalar_sets,
"bits={bits}: scalar fallback returned a different top-k than SIMD",
);
}
}
}
}
#[cfg(test)]
mod v7_crash_tests {
use super::*;
use std::path::{Path, PathBuf};
const DIM: usize = 64;
fn rows(n: usize, seed: u64) -> Vec<f32> {
let mut v = vec![0.0f32; n * DIM];
let mut s = seed | 1;
for x in v.iter_mut() {
s ^= s << 13;
s ^= s >> 7;
s ^= s << 17;
*x = ((s >> 40) as f32 / (1u64 << 23) as f32) - 0.5;
}
for row in v.chunks_mut(DIM) {
let norm: f32 = row.iter().map(|x| x * x).sum::<f32>().sqrt();
for x in row.iter_mut() {
*x /= norm;
}
}
v
}
fn temp(name: &str) -> PathBuf {
let mut p = std::env::temp_dir();
let nonce = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos();
p.push(format!("turbovec-v7crash-{nonce}-{name}"));
std::fs::create_dir(&p).unwrap();
p.push("index.tv");
p
}
fn apply(file: &mut Vec<u8>, off: u64, bytes: &[u8]) {
let end = off as usize + bytes.len();
if file.len() < end {
file.resize(end, 0);
}
file[off as usize..end].copy_from_slice(bytes);
}
fn state_of(scratch: &Path, bytes: &[u8]) -> Option<Vec<u8>> {
std::fs::write(scratch, bytes).unwrap();
TurboQuantIndex::load(scratch).ok().map(|i| i.to_bytes())
}
#[test]
fn a_reused_generation_cannot_resurrect_the_commit_it_overwrites() {
let path = temp("reuse");
let scratch = path.with_file_name("reuse-scratch.tv");
let mut idx = TurboQuantIndex::new(DIM, 4).unwrap();
idx.calibrate(&rows(1024, 70)).unwrap();
idx.add(&rows(200, 71));
idx.sync(&path).unwrap();
idx.swap_remove(10);
idx.sync(&path).unwrap();
let at_gen1 = std::fs::read(&path).unwrap();
let state_a = TurboQuantIndex::load(&path).unwrap().to_bytes();
let mut a = TurboQuantIndex::load(&path).unwrap();
a.add(&rows(1, 72));
let plan_a = a.plan_next_sync(0, None);
assert_eq!(plan_a.batches.len(), 1, "no stale header yet, so no repair batch");
let mut crashed = at_gen1.clone();
let ops = &plan_a.batches[0].ops;
let (hdr_off, hdr_bytes) = ops.last().expect("the header is the last op");
apply(&mut crashed, *hdr_off, hdr_bytes);
std::fs::write(&path, &crashed).unwrap();
assert_eq!(
state_of(&scratch, &crashed).expect("must load"),
state_a,
"attempt A names a unit that never landed, so it must be rejected"
);
let mut b = TurboQuantIndex::load(&path).unwrap();
b.add(&rows(3, 73));
let plan_b = b.plan_next_sync(0, None);
assert_eq!(
plan_b.batches.len(),
2,
"a sync landing on a rejected header of its own generation must \
invalidate it behind its own barrier first"
);
assert_eq!(
plan_b.batches[0].ops.len(),
1,
"the repair batch is one write"
);
let mut done = crashed.clone();
for batch in &plan_b.batches {
for (off, bytes) in &batch.ops {
apply(&mut done, *off, bytes);
}
}
let state_b = state_of(&scratch, &done).expect("full plan must load");
assert_eq!(state_b, b.to_bytes(), "the standard oracle");
assert_ne!(state_b, state_a);
for bi in 0..plan_b.batches.len() {
let ops = &plan_b.batches[bi].ops;
for (oj, (off, bytes)) in ops.iter().enumerate() {
for cut in 0..=bytes.len() {
let mut torn = crashed.clone();
for prev in &plan_b.batches[..bi] {
for (o, x) in &prev.ops {
apply(&mut torn, *o, x);
}
}
for (o, x) in &ops[..oj] {
apply(&mut torn, *o, x);
}
apply(&mut torn, *off, &bytes[..cut]);
let got = state_of(&scratch, &torn).unwrap_or_else(|| {
panic!("batch {bi} op {oj} cut {cut}: unloadable")
});
assert!(
got == state_a || got == state_b,
"batch {bi} op {oj} cut {cut}: resurrected an abandoned commit"
);
}
}
}
}
#[test]
fn a_sync_torn_at_any_byte_recovers_the_previous_commit() {
let path = temp("torn");
let scratch = path.with_file_name("torn-scratch.tv");
let mut idx = TurboQuantIndex::new(DIM, 4).unwrap();
idx.calibrate(&rows(1024, 15)).unwrap();
idx.add(&rows(100, 16));
idx.sync(&path).unwrap();
let base = std::fs::read(&path).unwrap();
let state_a = TurboQuantIndex::load(&path).unwrap().to_bytes();
idx.add(&rows(37, 17));
idx.swap_remove(5);
idx.swap_remove(60);
idx.swap_remove(idx.len() - 1);
let plan = idx.plan_next_sync(0, None);
assert_eq!(plan.batches.len(), 1, "single-fsync sync");
let mut done = base.clone();
for b in &plan.batches {
for (off, bytes) in &b.ops {
apply(&mut done, *off, bytes);
}
}
let state_b = state_of(&scratch, &done).expect("full plan must load");
assert_eq!(state_b, idx.to_bytes(), "the standard oracle");
let header_batch = plan.batches.len() - 1;
for bi in 0..plan.batches.len() {
let ops = &plan.batches[bi].ops;
for schedule in 0..3 {
for (oj, (off, bytes)) in ops.iter().enumerate() {
for cut in 0..=bytes.len() {
let mut torn = base.clone();
for prev in &plan.batches[..bi] {
for (o, b) in &prev.ops {
apply(&mut torn, *o, b);
}
}
match schedule {
0 => {
for (o, b) in &ops[..oj] {
apply(&mut torn, *o, b);
}
}
1 => {
for (o, b) in ops[oj + 1..].iter().rev() {
apply(&mut torn, *o, b);
}
}
_ => {}
}
apply(&mut torn, *off, &bytes[..cut]);
let got = state_of(&scratch, &torn).unwrap_or_else(|| {
panic!("batch {bi} op {oj} sched {schedule} cut {cut}: unloadable")
});
let complete = bi == header_batch
&& cut == bytes.len()
&& (schedule == 0 && oj == ops.len() - 1
|| schedule == 1 && oj == 0
|| ops.len() == 1);
let want = if complete { &state_b } else { &state_a };
assert_eq!(
&got, want,
"batch {bi} op {oj} sched {schedule} cut {cut}: wrong state"
);
}
}
}
}
}
#[test]
fn a_recovery_load_syncs_forward_and_survives_a_second_tear() {
let path = temp("double");
let scratch = path.with_file_name("double-scratch.tv");
let mut idx = TurboQuantIndex::new(DIM, 4).unwrap();
idx.calibrate(&rows(1024, 25)).unwrap();
idx.add(&rows(96, 26));
idx.sync(&path).unwrap();
idx.swap_remove(3);
idx.swap_remove(40);
assert_eq!(idx.plan_next_sync(0, None).batches.len(), 1);
idx.sync(&path).unwrap();
let base = std::fs::read(&path).unwrap();
let state_a = TurboQuantIndex::load(&path).unwrap().to_bytes();
assert_eq!(state_a, idx.to_bytes());
idx.add(&rows(40, 27));
let plan = idx.plan_next_sync(0, None);
assert_eq!(plan.batches.len(), 1, "single-fsync sync");
assert!(plan.carried.is_empty());
let mut crashed = base.clone();
let ops = &plan.batches[0].ops;
for (off, bytes) in &ops[..ops.len() - 1] {
apply(&mut crashed, *off, bytes);
}
std::fs::write(&path, &crashed).unwrap();
let mut rec = TurboQuantIndex::load(&path).unwrap();
assert_eq!(rec.to_bytes(), state_a, "recovery must be exact");
rec.add(&rows(5, 28));
rec.swap_remove(10);
rec.sync(&path).unwrap();
let loaded = TurboQuantIndex::load(&path).unwrap();
assert_eq!(loaded.to_bytes(), rec.to_bytes());
let plan2 = {
let mut again = TurboQuantIndex::load(scratch_write(&scratch, &crashed)).unwrap();
again.add(&rows(5, 28));
again.swap_remove(10);
again.plan_next_sync(0, None)
};
let state_b = {
let mut done = crashed.clone();
for b in &plan2.batches {
for (off, bytes) in &b.ops {
apply(&mut done, *off, bytes);
}
}
state_of(&scratch, &done).expect("full plan2 must load")
};
for bi in 0..plan2.batches.len() {
let ops = &plan2.batches[bi].ops;
for (oj, (off, bytes)) in ops.iter().enumerate() {
for cut in [0, bytes.len() / 2, bytes.len()] {
let mut torn = crashed.clone();
for prev in &plan2.batches[..bi] {
for (o, b) in &prev.ops {
apply(&mut torn, *o, b);
}
}
for (o, b) in &ops[..oj] {
apply(&mut torn, *o, b);
}
apply(&mut torn, *off, &bytes[..cut]);
let got = state_of(&scratch, &torn)
.unwrap_or_else(|| panic!("batch {bi} op {oj} cut {cut}: unloadable"));
assert!(
got == state_a || got == state_b,
"batch {bi} op {oj} cut {cut}: neither adjacent commit"
);
}
}
}
}
fn scratch_write<'a>(p: &'a Path, bytes: &[u8]) -> &'a Path {
std::fs::write(p, bytes).unwrap();
p
}
#[test]
fn barrier_counts_match_the_change() {
let path = temp("barriers");
let mut idx = TurboQuantIndex::new(DIM, 4).unwrap();
idx.calibrate(&rows(1024, 35)).unwrap();
idx.add(&rows(70, 36));
idx.sync(&path).unwrap();
idx.add(&rows(40, 37));
assert_eq!(
idx.plan_next_sync(0, None).batches.len(),
1,
"pure append: one fsync"
);
idx.sync(&path).unwrap();
idx.add(&rows(3, 38)); idx.swap_remove(idx.len() - 1); assert_eq!(
idx.plan_next_sync(0, None).batches.len(),
1,
"tail-only change: header alone"
);
idx.sync(&path).unwrap();
idx.swap_remove(0);
assert_eq!(
idx.plan_next_sync(0, None).batches.len(),
1,
"committed-unit removal: one fsync — the op rides the header"
);
idx.sync(&path).unwrap();
idx.add(&rows(1, 39));
let plan = idx.plan_next_sync(0, None);
assert_eq!(plan.batches.len(), 1, "materialization folds into the batch");
assert!(
plan.batches[0].ops.len() >= 2,
"the materialized unit write precedes the header op"
);
assert!(plan.carried.is_empty(), "no fresh dirt: nothing carried");
idx.sync(&path).unwrap();
let loaded = TurboQuantIndex::load(&path).unwrap();
assert_eq!(loaded.to_bytes(), idx.to_bytes());
}
#[test]
fn an_absurd_header_n_is_refused_not_allocated() {
let path = temp("hostilen");
let mut idx = TurboQuantIndex::new(DIM, 4).unwrap();
idx.calibrate(&rows(1024, 65)).unwrap();
idx.add(&rows(64, 66)); idx.sync(&path).unwrap();
let mut bytes = std::fs::read(&path).unwrap();
let geo = io_v7::Geo {
kind: 0,
dim: DIM,
bit_width: 4,
n_calib: DIM,
};
let at = geo.hdr_at_for_test(0);
bytes[at + 8..at + 16].copy_from_slice(&(u64::MAX / 2).to_le_bytes());
let used = 16 + 4; let c = io_v7::crc32(&bytes[at..at + used]);
bytes[at + used..at + used + 4].copy_from_slice(&c.to_le_bytes());
std::fs::write(&path, &bytes).unwrap();
assert!(
TurboQuantIndex::load(&path).is_err(),
"an absurd n must refuse, not allocate"
);
}
#[test]
fn bit_rot_in_any_byte_is_never_served_silently() {
let path = temp("rot");
let scratch = path.with_file_name("rot-scratch.tv");
let mut idx = TurboQuantIndex::new(DIM, 4).unwrap();
idx.calibrate(&rows(1024, 45)).unwrap();
idx.add(&rows(70, 46));
idx.sync(&path).unwrap();
let prev = TurboQuantIndex::load(&path).unwrap().to_bytes();
idx.add(&rows(20, 47));
idx.swap_remove(5);
idx.sync(&path).unwrap();
let cur = TurboQuantIndex::load(&path).unwrap().to_bytes();
let file = std::fs::read(&path).unwrap();
let geo = io_v7::Geo {
kind: 0,
dim: DIM,
bit_width: 4,
n_calib: DIM,
};
let newest_hdr = geo.hdr_at_for_test(1)..geo.hdr_at_for_test(1) + geo.hdr_len();
let structural_end = geo.unit_at_for_test(0).min(file.len());
for at in 0..structural_end {
let mut bytes = file.clone();
bytes[at] ^= 1 << (at % 8);
match state_of(&scratch, &bytes) {
None => {}
Some(got) if got == cur => {}
Some(got) if got == prev && newest_hdr.contains(&at) => {}
Some(_) => panic!("flip at byte {at} served a state it must not"),
}
}
}
}
#[cfg(test)]
mod v7_crash_blocked_tests {
use super::*;
use std::path::PathBuf;
const DIM: usize = 64;
fn rows(n: usize, seed: u64) -> Vec<f32> {
let mut v = vec![0.0f32; n * DIM];
let mut s = seed | 1;
for x in v.iter_mut() {
s ^= s << 13;
s ^= s >> 7;
s ^= s << 17;
*x = ((s >> 40) as f32 / (1u64 << 23) as f32) - 0.5;
}
for row in v.chunks_mut(DIM) {
let norm: f32 = row.iter().map(|x| x * x).sum::<f32>().sqrt();
for x in row.iter_mut() {
*x /= norm;
}
}
v
}
fn temp(name: &str) -> PathBuf {
let mut p = std::env::temp_dir();
let nonce = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos();
p.push(format!("turbovec-v7blk-{nonce}-{name}"));
std::fs::create_dir(&p).unwrap();
p.push("index.tv");
p
}
fn apply(file: &mut Vec<u8>, off: u64, bytes: &[u8]) {
let end = off as usize + bytes.len();
if file.len() < end {
file.resize(end, 0);
}
file[off as usize..end].copy_from_slice(bytes);
}
#[test]
fn blocked_only_capture_survives_a_torn_sync() {
let path = temp("blkcap");
let scratch = path.with_file_name("scratch.tv");
let mut idx = TurboQuantIndex::new(DIM, 4).unwrap();
idx.calibrate(&rows(1024, 55)).unwrap();
idx.add(&rows(100, 56));
idx.sync(&path).unwrap();
let mut idx = TurboQuantIndex::load(&path).unwrap();
assert!(!idx.packed_ready(), "reload must be blocked-only");
idx.swap_remove(1); idx.swap_remove(37); idx.sync(&path).unwrap();
let state_a = TurboQuantIndex::load(&path).unwrap().to_bytes();
assert_eq!(state_a, idx.to_bytes(), "gathered op bytes must be exact");
let base = std::fs::read(&path).unwrap();
idx.add(&rows(1, 57));
let plan = idx.plan_next_sync(0, None);
assert_eq!(plan.batches.len(), 1, "single-fsync sync");
let mut torn = base.clone();
let ops = &plan.batches[0].ops;
for (off, bytes) in &ops[..ops.len() - 1] {
apply(&mut torn, *off, bytes);
}
std::fs::write(&scratch, &torn).unwrap();
let recovered = TurboQuantIndex::load(&scratch).unwrap();
assert_eq!(recovered.to_bytes(), state_a, "recovery must be exact");
idx.sync(&path).unwrap();
assert_eq!(TurboQuantIndex::load(&path).unwrap().to_bytes(), idx.to_bytes());
}
}
#[cfg(test)]
mod v7_delta_tests {
use super::*;
use std::path::PathBuf;
const DIM: usize = 64;
fn rows(n: usize, seed: u64) -> Vec<f32> {
let mut v = vec![0.0f32; n * DIM];
let mut s = seed | 1;
for x in v.iter_mut() {
s ^= s << 13;
s ^= s >> 7;
s ^= s << 17;
*x = ((s >> 40) as f32 / (1u64 << 23) as f32) - 0.5;
}
for row in v.chunks_mut(DIM) {
let norm: f32 = row.iter().map(|x| x * x).sum::<f32>().sqrt();
for x in row.iter_mut() {
*x /= norm;
}
}
v
}
fn temp(name: &str) -> PathBuf {
let mut p = std::env::temp_dir();
let nonce = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos();
p.push(format!("turbovec-v7delta-{nonce}-{name}"));
std::fs::create_dir(&p).unwrap();
p.push("index.tv");
p
}
#[test]
#[cfg(target_arch = "x86_64")]
fn the_capture_arena_is_bounded_under_slot_churn() {
let path = temp("arena-bound");
let mut idx = TurboQuantIndex::new(DIM, 4).unwrap();
idx.calibrate(&rows(1024, 61)).unwrap();
idx.add(&rows(96, 62));
idx.sync(&path).unwrap();
let mut idx = TurboQuantIndex::load(&path).unwrap();
let mut eager = TurboQuantIndex::new(DIM, 4).unwrap();
eager.calibrate(&rows(1024, 61)).unwrap();
eager.add(&rows(96, 62));
let lane = DIM / 2; for i in 0..3 * io_v7::MAX_OPS {
let victim = if i % 2 == 0 { 5 } else { 7 };
idx.swap_remove(victim);
idx.add(&rows(1, 63 + i as u64));
eager.swap_remove(victim);
eager.add(&rows(1, 63 + i as u64));
assert!(
idx.sync_capture_buf_len_for_test() <= io_v7::MAX_OPS * lane,
"arena exceeded its bound at churn round {i}"
);
}
assert!(
idx.captured_len_for_test() > 0,
"the churn never engaged the capture path — vacuous test"
);
idx.sync(&path).unwrap();
let loaded = TurboQuantIndex::load(&path).unwrap();
assert_eq!(loaded.to_bytes(), idx.to_bytes());
assert_eq!(
loaded.to_bytes(),
eager.to_bytes(),
"capture-path result diverged from the capture-free oracle"
);
}
#[test]
#[cfg(target_arch = "x86_64")] fn the_removal_capture_engages_below_the_block_floor() {
let path = temp("capture-engages");
let mut idx = TurboQuantIndex::new(DIM, 4).unwrap();
idx.add(&rows(200, 5));
idx.sync(&path).unwrap();
let mut idx = TurboQuantIndex::load(&path).unwrap();
assert_eq!(idx.captured_len_for_test(), 0, "nothing captured yet");
idx.swap_remove(5);
assert_eq!(
idx.captured_len_for_test(),
1,
"a removal at slot 5, far below the committed floor of 192, must \
be captured — if only the boundary slot captures, the sync is \
back to re-reading every row it serializes",
);
idx.swap_remove(100);
assert_eq!(idx.captured_len_for_test(), 2, "and again mid-range");
let before = idx.captured_len_for_test();
idx.swap_remove(idx.len() - 1);
assert_eq!(
idx.captured_len_for_test(),
before,
"a pop above the committed floor has no redo op to feed",
);
idx.add(&rows(40, 6)); let before = idx.captured_len_for_test();
let floor = 192;
assert!(idx.len() > floor + 2, "need slack above the floor");
idx.swap_remove(floor + 1); assert_eq!(
idx.captured_len_for_test(),
before,
"a removal above the committed floor must not be captured: the \
file holds nothing for that slot, so no redo op will ever read \
it. Capturing anyway is work `remove()` pays for nothing — and \
`remove()` is the cell this optimization must not tax.",
);
idx.sync(&path).unwrap();
assert_eq!(idx.captured_len_for_test(), 0, "cleared by the commit");
}
#[test]
fn a_spliced_header_without_its_data_is_not_adopted() {
let path = temp("splice");
let mut idx = TurboQuantIndex::new(DIM, 4).unwrap();
idx.calibrate(&rows(1024, 11)).unwrap();
idx.add(&rows(96, 12));
idx.sync(&path).unwrap();
idx.swap_remove(3);
idx.add(&rows(1, 13));
idx.sync(&path).unwrap();
let pre = std::fs::read(&path).unwrap();
let state_pre = TurboQuantIndex::load(&path).unwrap().to_bytes();
idx.swap_remove(40);
idx.add(&rows(1, 14));
let plan = idx.plan_next_sync(0, None);
let (hdr_off, hdr_bytes) = plan.batches[0]
.ops
.last()
.expect("plan has a header op")
.clone();
assert!(
plan.batches[0].ops.len() > 1,
"the sync must materialize at least one unit for this test"
);
let mut spliced = pre.clone();
let end = hdr_off as usize + hdr_bytes.len();
if spliced.len() < end {
spliced.resize(end, 0);
}
spliced[hdr_off as usize..end].copy_from_slice(&hdr_bytes);
std::fs::write(&path, &spliced).unwrap();
let got = TurboQuantIndex::load(&path).unwrap();
assert_eq!(
got.to_bytes(),
state_pre,
"a header without its data must fall back, not resurrect stale blocks"
);
}
#[test]
fn an_out_of_range_op_block_is_refused_not_indexed() {
let path = temp("hostileb");
let mut idx = TurboQuantIndex::new(DIM, 4).unwrap();
idx.calibrate(&rows(1024, 31)).unwrap();
idx.add(&rows(65, 32)); idx.sync(&path).unwrap();
idx.swap_remove(3);
idx.sync(&path).unwrap();
let mut bytes = std::fs::read(&path).unwrap();
let geo = io_v7::Geo {
kind: 0,
dim: DIM,
bit_width: 4,
n_calib: DIM,
};
let at = geo.hdr_at_for_test(1);
let gb = at + 16 + 4;
bytes[gb..gb + 4].copy_from_slice(&u32::MAX.to_le_bytes());
let row = DIM / 2;
let op_size = 1 + row + 4;
let used = 16 + 4 + 5 + op_size + 4 + 12;
let c = io_v7::crc32(&bytes[at..at + used]);
bytes[at + used..at + used + 4].copy_from_slice(&c.to_le_bytes());
std::fs::write(&path, &bytes).unwrap();
let _ = TurboQuantIndex::load(&path);
}
#[test]
fn a_negative_tail_scale_is_refused() {
let path = temp("negscale");
let mut idx = TurboQuantIndex::new(DIM, 4).unwrap();
idx.calibrate(&rows(1024, 41)).unwrap();
idx.add(&rows(65, 42)); idx.sync(&path).unwrap();
let mut bytes = std::fs::read(&path).unwrap();
let geo = io_v7::Geo {
kind: 0,
dim: DIM,
bit_width: 4,
n_calib: DIM,
};
let row_bytes = DIM / 2;
let at = geo.hdr_at_for_test(0);
let sc = at + 16 + row_bytes;
let v = f32::from_le_bytes(bytes[sc..sc + 4].try_into().unwrap());
bytes[sc..sc + 4].copy_from_slice(&(-v.max(0.5)).to_le_bytes());
let used = 16 + (row_bytes + 4) + 4 + 16;
let c = io_v7::crc32(&bytes[at..at + used]);
bytes[at + used..at + used + 4].copy_from_slice(&c.to_le_bytes());
std::fs::write(&path, &bytes).unwrap();
assert!(
TurboQuantIndex::load(&path).is_err(),
"a negative tail scale must refuse the load"
);
}
#[test]
fn a_zero_tqplus_scale_is_refused() {
let path = temp("zeroscale");
let mut idx = TurboQuantIndex::new(DIM, 4).unwrap();
idx.calibrate(&rows(1024, 51)).unwrap();
idx.add(&rows(40, 52));
idx.sync(&path).unwrap();
let mut bytes = std::fs::read(&path).unwrap();
let nl = 16;
let scale5 = 23 + (nl - 1) * 4 + nl * 4 + 4 + DIM * 4 + 5 * 4;
bytes[scale5..scale5 + 4].copy_from_slice(&0.0f32.to_le_bytes());
let sb_end = 23 + (nl - 1) * 4 + nl * 4 + 4 + DIM * 8;
let c = io_v7::crc32(&bytes[..sb_end]);
bytes[sb_end..sb_end + 4].copy_from_slice(&c.to_le_bytes());
std::fs::write(&path, &bytes).unwrap();
let err = TurboQuantIndex::load(&path).unwrap_err();
assert!(err.to_string().contains("TQ+ scale"), "{err}");
}
#[test]
fn a_queued_compaction_still_respects_the_foreign_file_guard() {
let path = temp("calibforeign");
let mut a = TurboQuantIndex::new(DIM, 4).unwrap();
a.calibrate(&rows(1024, 53)).unwrap();
a.add(&rows(40, 54));
a.sync(&path).unwrap();
let mut b = TurboQuantIndex::new(DIM, 4).unwrap();
b.calibrate(&rows(1024, 55)).unwrap();
b.add(&rows(20, 56));
b.sync(&path).unwrap();
a.calibrate(&rows(1024, 57)).unwrap(); let err = a.sync(&path).unwrap_err();
assert_eq!(err.kind(), std::io::ErrorKind::InvalidData, "{err}");
let loaded = TurboQuantIndex::load(&path).unwrap();
assert_eq!(loaded.to_bytes(), b.to_bytes(), "B's file must be untouched");
}
#[test]
fn every_field_tamper_loads_politely() {
let path = temp("matrix");
let scratch = path.with_file_name("matrix-scratch.tv");
let mut idx = TurboQuantIndex::new(DIM, 4).unwrap();
idx.calibrate(&rows(1024, 61)).unwrap();
idx.add(&rows(70, 62));
idx.sync(&path).unwrap();
idx.swap_remove(3); idx.add(&rows(2, 63));
idx.sync(&path).unwrap();
let base = std::fs::read(&path).unwrap();
let geo = io_v7::Geo {
kind: 0,
dim: DIM,
bit_width: 4,
n_calib: DIM,
};
let try_load = |bytes: &[u8], what: &str| {
std::fs::write(&scratch, bytes).unwrap();
let r = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
let _ = TurboQuantIndex::load(&scratch);
}));
assert!(r.is_ok(), "loader panicked on tamper: {what}");
};
let hostile: [u8; 4] = [0xFF, 0x00, 0x80, 0x01];
let mut targets: Vec<usize> = Vec::new();
let sb_end = geo.hdr_at_for_test(0);
targets.extend(0..sb_end);
for slot in [0usize, 1] {
let at = geo.hdr_at_for_test(slot);
let used = io_v7::hdr_used_for_test(&base, &geo, slot) + 8;
let end = at + geo.hdr_len();
targets.extend(at..(at + used).min(end));
targets.extend(((at + used)..end).step_by(251));
}
let structural_end = geo.unit_at_for_test(0);
let _ = structural_end;
for &at in targets.iter().filter(|&&a| a < base.len()) {
for v in hostile {
if base[at] == v {
continue;
}
let mut bytes = base.clone();
bytes[at] = v;
io_v7::reseal_for_test(&mut bytes, &geo);
try_load(&bytes, &format!("byte {at} <- {v:#04x}"));
}
}
let mut s = 0x1234_5678_9ABC_DEF0u64;
for i in 0..400 {
let mut bytes = base.clone();
for _ in 0..1 + (i % 4) {
s ^= s << 13;
s ^= s >> 7;
s ^= s << 17;
let at = (s as usize) % bytes.len();
bytes[at] = (s >> 32) as u8;
}
io_v7::reseal_for_test(&mut bytes, &geo);
try_load(&bytes, &format!("random tamper {i}"));
}
for cut in [0usize, 4, 11, 19, structural_end / 2, structural_end] {
try_load(&base[..cut.min(base.len())], &format!("truncate {cut}"));
}
}
#[test]
fn a_negative_block_scale_is_refused() {
let path = temp("negblockscale");
let mut idx = TurboQuantIndex::new(DIM, 4).unwrap();
idx.calibrate(&rows(1024, 71)).unwrap();
idx.add(&rows(64, 72));
idx.sync(&path).unwrap();
let mut bytes = std::fs::read(&path).unwrap();
let geo = io_v7::Geo {
kind: 0,
dim: DIM,
bit_width: 4,
n_calib: DIM,
};
let row_bytes = DIM / 2;
let sc = geo.unit_at_for_test(0) + 32 * row_bytes + 7 * 4;
let v = f32::from_le_bytes(bytes[sc..sc + 4].try_into().unwrap());
bytes[sc..sc + 4].copy_from_slice(&(-v.max(0.5)).to_le_bytes());
std::fs::write(&path, &bytes).unwrap();
let err = TurboQuantIndex::load(&path).unwrap_err();
assert!(err.to_string().contains("scale"), "{err}");
}
#[test]
fn a_mass_removal_stays_incremental() {
let path = temp("massremove");
let mut idx = TurboQuantIndex::new(DIM, 4).unwrap();
idx.calibrate(&rows(1024, 73)).unwrap();
idx.add(&rows(200, 74));
idx.sync(&path).unwrap();
let nonce_before = std::fs::read(&path).unwrap()[11..19].to_vec();
for i in 0..80 {
idx.swap_remove(i);
}
idx.sync(&path).unwrap();
let nonce_after = std::fs::read(&path).unwrap()[11..19].to_vec();
assert_eq!(nonce_before, nonce_after, "sync degraded to a full rewrite");
let loaded = TurboQuantIndex::load(&path).unwrap();
assert_eq!(loaded.to_bytes(), idx.to_bytes());
idx.add(&rows(1, 75));
idx.sync(&path).unwrap();
let loaded = TurboQuantIndex::load(&path).unwrap();
assert_eq!(loaded.to_bytes(), idx.to_bytes());
}
#[test]
fn a_torn_materialize_of_a_delta_named_unit_recovers() {
let path = temp("mat-delta");
let mut idx = TurboQuantIndex::new(DIM, 4).unwrap();
idx.calibrate(&rows(1024, 41)).unwrap();
idx.add(&rows(32, 42));
idx.sync(&path).unwrap();
idx.add(&rows(64, 43));
idx.swap_remove(40);
idx.sync(&path).unwrap();
let committed = TurboQuantIndex::load(&path).unwrap().to_bytes();
let plan = idx.plan_next_sync(0, None);
let batch = &plan.batches[0];
let mut f = std::fs::OpenOptions::new().write(true).open(&path).unwrap();
use std::io::{Seek, SeekFrom, Write};
for (off, bytes) in &batch.ops[..batch.ops.len() - 1] {
f.seek(SeekFrom::Start(*off)).unwrap();
f.write_all(bytes).unwrap();
}
drop(f);
match TurboQuantIndex::load(&path) {
Ok(loaded) => assert_eq!(
loaded.to_bytes(),
committed,
"loaded state differs from the previous commit"
),
Err(e) => panic!("file refused to load: {e}"),
}
}
#[test]
fn an_op_overflow_falls_back_to_a_full_rewrite() {
let path = temp("ovf-full");
let mut idx = TurboQuantIndex::new(DIM, 4).unwrap();
idx.calibrate(&rows(1024, 51)).unwrap();
idx.add(&rows(2_112, 52));
idx.sync(&path).unwrap();
let nonce_before = std::fs::read(&path).unwrap()[11..19].to_vec();
for v in (5..5 + io_v7::MAX_OPS + 1).rev() {
idx.swap_remove(v);
}
idx.sync(&path).unwrap();
let nonce_after = std::fs::read(&path).unwrap()[11..19].to_vec();
assert_ne!(nonce_before, nonce_after, "overflow must full-rewrite");
let loaded = TurboQuantIndex::load(&path).unwrap();
assert_eq!(loaded.to_bytes(), idx.to_bytes());
idx.add(&rows(1, 53));
idx.sync(&path).unwrap();
assert_eq!(TurboQuantIndex::load(&path).unwrap().to_bytes(), idx.to_bytes());
}
#[test]
#[allow(clippy::permissions_set_readonly_false)]
fn a_failed_sync_recovers_via_full_write() {
let path = temp("failedsync");
let mut idx = TurboQuantIndex::new(DIM, 4).unwrap();
idx.calibrate(&rows(1024, 91)).unwrap();
idx.add(&rows(64, 92));
idx.sync(&path).unwrap();
idx.add(&rows(32, 93));
let mut perm = std::fs::metadata(&path).unwrap().permissions();
perm.set_readonly(true);
std::fs::set_permissions(&path, perm.clone()).unwrap();
assert!(idx.sync(&path).is_err(), "read-only sync must fail");
perm.set_readonly(false);
std::fs::set_permissions(&path, perm).unwrap();
idx.sync(&path).unwrap();
let loaded = TurboQuantIndex::load(&path).unwrap();
assert_eq!(loaded.to_bytes(), idx.to_bytes());
}
#[test]
fn sync_keeps_working_after_a_fallback_load() {
let path = temp("wedge");
let mut idx = TurboQuantIndex::new(DIM, 4).unwrap();
idx.calibrate(&rows(1024, 21)).unwrap();
idx.add(&rows(64, 22));
idx.sync(&path).unwrap();
let state_g0 = TurboQuantIndex::load(&path).unwrap().to_bytes();
idx.add(&rows(32, 23));
idx.sync(&path).unwrap();
let geo = io_v7::Geo {
kind: 0,
dim: DIM,
bit_width: 4,
n_calib: DIM,
};
let mut bytes = std::fs::read(&path).unwrap();
let at = geo.unit_at(2); for b in bytes[at..at + geo.unit_len()].iter_mut() {
*b = 0;
}
std::fs::write(&path, &bytes).unwrap();
let mut rec = TurboQuantIndex::load(&path).unwrap();
assert_eq!(rec.to_bytes(), state_g0, "fallback must be exact");
rec.add(&rows(5, 24));
rec.sync(&path).unwrap();
let after = TurboQuantIndex::load(&path).unwrap();
assert_eq!(after.to_bytes(), rec.to_bytes());
let path2 = temp("wedge-trunc");
let mut idx = TurboQuantIndex::new(DIM, 4).unwrap();
idx.calibrate(&rows(1024, 25)).unwrap();
idx.add(&rows(64, 26));
idx.sync(&path2).unwrap();
let pre_len = std::fs::metadata(&path2).unwrap().len();
let state_g0 = TurboQuantIndex::load(&path2).unwrap().to_bytes();
idx.add(&rows(32, 27));
idx.sync(&path2).unwrap();
let full = std::fs::read(&path2).unwrap();
std::fs::write(&path2, &full[..pre_len as usize]).unwrap();
let mut rec = TurboQuantIndex::load(&path2).unwrap();
assert_eq!(rec.to_bytes(), state_g0);
rec.swap_remove(0);
rec.sync(&path2).unwrap();
assert_eq!(
TurboQuantIndex::load(&path2).unwrap().to_bytes(),
rec.to_bytes()
);
}
}