use std::fs::File;
use std::io::{self, Read, Seek, SeekFrom, Write};
use std::path::Path;
pub(crate) const V7_MAGIC: &[u8; 4] = b"TV7\0";
pub(crate) const V7_VERSION: u8 = 2;
const BLOCK: usize = 32;
pub(crate) const MAX_OPS: usize = 1024;
static CRC_TABLE: std::sync::OnceLock<[u32; 256]> = std::sync::OnceLock::new();
pub(crate) fn crc32(data: &[u8]) -> u32 {
if data.len() < 4096 {
return crc32_one(data);
}
let third = data.len() / 3;
let (a, rest) = data.split_at(third);
let (b, c) = rest.split_at(third);
let (ca, cb, cc) = crc32_three(a, b, c);
let mut digest = [0u8; 12];
digest[..4].copy_from_slice(&ca.to_le_bytes());
digest[4..8].copy_from_slice(&cb.to_le_bytes());
digest[8..].copy_from_slice(&cc.to_le_bytes());
crc32_one(&digest)
}
fn crc32_one(data: &[u8]) -> u32 {
!crc32c_update(0xFFFF_FFFF, data)
}
fn crc32c_update(crc: u32, data: &[u8]) -> u32 {
#[cfg(target_arch = "aarch64")]
if std::arch::is_aarch64_feature_detected!("crc") {
return unsafe { crc32c_upd_hw_aarch64(crc, data) };
}
#[cfg(target_arch = "x86_64")]
if is_x86_feature_detected!("sse4.2") {
return unsafe { crc32c_upd_hw_x86(crc, data) };
}
crc32c_upd_soft(crc, data)
}
fn crc32_three(a: &[u8], b: &[u8], c: &[u8]) -> (u32, u32, u32) {
#[cfg(target_arch = "aarch64")]
if std::arch::is_aarch64_feature_detected!("crc") {
return unsafe { crc32c_three_hw_aarch64(a, b, c) };
}
#[cfg(target_arch = "x86_64")]
if is_x86_feature_detected!("sse4.2") {
return unsafe { crc32c_three_hw_x86(a, b, c) };
}
(crc32c_soft(a), crc32c_soft(b), crc32c_soft(c))
}
#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "crc")]
unsafe fn crc32c_upd_hw_aarch64(mut crc: u32, data: &[u8]) -> u32 {
use std::arch::aarch64::{__crc32cb, __crc32cd};
let (chunks, tail) = data.as_chunks::<8>();
for c in chunks {
crc = __crc32cd(crc, u64::from_le_bytes(*c));
}
for &b in tail {
crc = __crc32cb(crc, b);
}
crc
}
#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "crc")]
unsafe fn crc32c_three_hw_aarch64(a: &[u8], b: &[u8], c: &[u8]) -> (u32, u32, u32) {
use std::arch::aarch64::{__crc32cb, __crc32cd};
let n = a.len().min(b.len()).min(c.len()) / 8;
let (mut x, mut y, mut z) = (0xFFFF_FFFFu32, 0xFFFF_FFFFu32, 0xFFFF_FFFFu32);
for i in 0..n {
x = __crc32cd(x, u64::from_le_bytes(a[i * 8..i * 8 + 8].try_into().unwrap()));
y = __crc32cd(y, u64::from_le_bytes(b[i * 8..i * 8 + 8].try_into().unwrap()));
z = __crc32cd(z, u64::from_le_bytes(c[i * 8..i * 8 + 8].try_into().unwrap()));
}
let fin = |mut crc: u32, tail: &[u8]| {
for &v in tail {
crc = __crc32cb(crc, v);
}
!crc
};
(fin(x, &a[n * 8..]), fin(y, &b[n * 8..]), fin(z, &c[n * 8..]))
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "sse4.2")]
unsafe fn crc32c_upd_hw_x86(crc: u32, data: &[u8]) -> u32 {
use std::arch::x86_64::{_mm_crc32_u64, _mm_crc32_u8};
let mut wide = crc as u64;
let (chunks, tail) = data.as_chunks::<8>();
for c in chunks {
wide = _mm_crc32_u64(wide, u64::from_le_bytes(*c));
}
let mut crc = wide as u32;
for &b in tail {
crc = _mm_crc32_u8(crc, b);
}
crc
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "sse4.2")]
unsafe fn crc32c_three_hw_x86(a: &[u8], b: &[u8], c: &[u8]) -> (u32, u32, u32) {
use std::arch::x86_64::{_mm_crc32_u64, _mm_crc32_u8};
let n = a.len().min(b.len()).min(c.len()) / 8;
let (mut x, mut y, mut z) = (0xFFFF_FFFFu64, 0xFFFF_FFFFu64, 0xFFFF_FFFFu64);
for i in 0..n {
x = _mm_crc32_u64(x, u64::from_le_bytes(a[i * 8..i * 8 + 8].try_into().unwrap()));
y = _mm_crc32_u64(y, u64::from_le_bytes(b[i * 8..i * 8 + 8].try_into().unwrap()));
z = _mm_crc32_u64(z, u64::from_le_bytes(c[i * 8..i * 8 + 8].try_into().unwrap()));
}
let fin = |mut crc: u32, tail: &[u8]| {
for &v in tail {
crc = _mm_crc32_u8(crc, v);
}
!crc
};
(
fin(x as u32, &a[n * 8..]),
fin(y as u32, &b[n * 8..]),
fin(z as u32, &c[n * 8..]),
)
}
fn crc32c_soft(data: &[u8]) -> u32 {
!crc32c_upd_soft(0xFFFF_FFFF, data)
}
fn crc32c_upd_soft(mut crc: u32, data: &[u8]) -> u32 {
let table = CRC_TABLE.get_or_init(|| {
let mut t = [0u32; 256];
for (i, e) in t.iter_mut().enumerate() {
let mut c = i as u32;
for _ in 0..8 {
let mask = (c & 1).wrapping_neg();
c = (c >> 1) ^ (0x82F6_3B78 & mask);
}
*e = c;
}
t
});
for &b in data {
crc = (crc >> 8) ^ table[usize::from((crc as u8) ^ b)];
}
crc
}
struct DeltaDigest {
at: usize,
split: Option<(usize, usize)>,
parts: [u32; 3],
}
impl DeltaDigest {
fn new(total: usize) -> Self {
let split = (total >= 4096).then(|| {
let third = total / 3;
(third, third * 2)
});
Self {
at: 0,
split,
parts: [0xFFFF_FFFF; 3],
}
}
fn push(&mut self, mut data: &[u8]) {
let Some((s1, s2)) = self.split else {
self.parts[0] = crc32c_update(self.parts[0], data);
self.at += data.len();
return;
};
while !data.is_empty() {
let (part, end) = match self.at {
a if a < s1 => (0, s1),
a if a < s2 => (1, s2),
_ => (2, usize::MAX),
};
let take = data.len().min(end - self.at);
self.parts[part] = crc32c_update(self.parts[part], &data[..take]);
self.at += take;
data = &data[take..];
}
}
fn finish(self) -> u32 {
if self.split.is_none() {
return !self.parts[0];
}
let mut digest = [0u8; 12];
for (i, p) in self.parts.iter().enumerate() {
digest[i * 4..i * 4 + 4].copy_from_slice(&(!p).to_le_bytes());
}
crc32_one(&digest)
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct Geo {
pub kind: u8,
pub dim: usize,
pub bit_width: usize,
pub n_calib: usize,
}
impl Geo {
pub fn row_bytes(&self) -> usize {
let (_, n_byte_groups, _) = crate::pack::blocked_geometry(1, self.bit_width, self.dim);
n_byte_groups
}
fn n_levels(&self) -> usize {
1 << self.bit_width
}
fn id_bytes(&self, rows: usize) -> usize {
if self.kind == 1 {
rows * 8
} else {
0
}
}
pub fn sb_len(&self) -> usize {
23 + (self.n_levels() - 1) * 4 + self.n_levels() * 4 + 4 + self.n_calib * 8 + 4
}
fn op_size(&self) -> usize {
1 + self.row_bytes() + 4 + self.id_bytes(1)
}
pub fn hdr_probe_len(&self) -> usize {
16 + 31 * (self.row_bytes() + 4 + self.id_bytes(1))
+ 4
+ 4
+ MAX_OPS * 4
+ 12
+ 4
}
pub fn hdr_len(&self) -> usize {
16 + 31 * (self.row_bytes() + 4 + self.id_bytes(1))
+ 4
+ MAX_OPS * (5 + self.op_size())
+ 4
+ MAX_OPS * 4
+ 12
+ 4
}
fn hdr_at(&self, slot: usize) -> usize {
self.sb_len() + slot * self.hdr_len()
}
fn hdr_at_for_len(&self, slot: usize) -> usize {
self.hdr_at(slot)
}
#[cfg(test)]
pub fn hdr_at_for_test(&self, slot: usize) -> usize {
self.hdr_at(slot)
}
#[cfg(test)]
pub fn unit_at_for_test(&self, block: usize) -> usize {
self.unit_at(block)
}
pub fn unit_len(&self) -> usize {
BLOCK * self.row_bytes() + BLOCK * 4 + self.id_bytes(BLOCK)
}
pub fn unit_at(&self, block: usize) -> usize {
self.hdr_at(2) + block * self.unit_len()
}
pub fn full_len(&self, n_vectors: usize) -> usize {
self.unit_at(n_vectors / BLOCK)
}
}
#[derive(Clone, Copy, Debug)]
pub(crate) struct SyncCursor {
pub gen: u64,
pub n_synced: u64,
pub calib_gen: u64,
pub nonce: u64,
}
pub(crate) struct SyncSource<'a> {
pub kind: u8,
pub dim: usize,
pub bit_width: usize,
pub n_vectors: usize,
pub seq_blocks: &'a dyn Fn(usize, usize) -> Vec<u8>,
pub row_codes: &'a dyn Fn(usize, &mut Vec<u8>),
pub scales: &'a [f32],
pub ids: Option<&'a [u64]>,
pub tqplus_shift: &'a [f32],
pub tqplus_scale: &'a [f32],
pub boundaries: &'a [f32],
pub centroids: &'a [f32],
}
impl SyncSource<'_> {
fn geo(&self) -> Geo {
Geo {
kind: self.kind,
dim: self.dim,
bit_width: self.bit_width,
n_calib: self.tqplus_shift.len(),
}
}
}
fn superblock(src: &SyncSource<'_>, nonce: u64) -> Vec<u8> {
let mut sb = Vec::new();
sb.extend_from_slice(V7_MAGIC);
sb.push(V7_VERSION);
sb.push(src.bit_width as u8);
sb.push(src.kind);
sb.extend_from_slice(&(src.dim as u32).to_le_bytes());
sb.extend_from_slice(&nonce.to_le_bytes());
sb.extend_from_slice(&(MAX_OPS as u32).to_le_bytes());
for v in src.boundaries {
sb.extend_from_slice(&v.to_le_bytes());
}
for v in src.centroids {
sb.extend_from_slice(&v.to_le_bytes());
}
sb.extend_from_slice(&(src.tqplus_shift.len() as u32).to_le_bytes());
for v in src.tqplus_shift {
sb.extend_from_slice(&v.to_le_bytes());
}
for v in src.tqplus_scale {
sb.extend_from_slice(&v.to_le_bytes());
}
let c = crc32(&sb);
sb.extend_from_slice(&c.to_le_bytes());
debug_assert_eq!(sb.len(), src.geo().sb_len());
sb
}
fn header_slot(
src: &SyncSource<'_>,
gen: u64,
n: usize,
ops: &[(usize, Vec<usize>)],
delta: (&[usize], std::ops::Range<usize>, u32),
) -> Vec<u8> {
let geo = src.geo();
let n_tail = n % BLOCK;
let first_tail = n - n_tail;
let mut h = Vec::with_capacity(geo.hdr_len());
h.extend_from_slice(&gen.to_le_bytes());
h.extend_from_slice(&(n as u64).to_le_bytes());
for k in 0..n_tail {
let r = first_tail + k;
(src.row_codes)(r, &mut h);
h.extend_from_slice(&src.scales[r].to_le_bytes());
if let Some(ids) = src.ids {
h.extend_from_slice(&ids[r].to_le_bytes());
}
}
h.extend_from_slice(&(ops.len() as u32).to_le_bytes());
for (block, slots) in ops {
h.extend_from_slice(&(*block as u32).to_le_bytes());
h.push(slots.len() as u8);
for &s in slots {
h.push((s % BLOCK) as u8);
(src.row_codes)(s, &mut h);
h.extend_from_slice(&src.scales[s].to_le_bytes());
if let Some(ids) = src.ids {
h.extend_from_slice(&ids[s].to_le_bytes());
}
}
}
let (mat, app, delta_crc) = delta;
h.extend_from_slice(&(mat.len() as u32).to_le_bytes());
for &b in mat {
h.extend_from_slice(&(b as u32).to_le_bytes());
}
h.extend_from_slice(&(app.start as u32).to_le_bytes());
h.extend_from_slice(&(app.end as u32).to_le_bytes());
h.extend_from_slice(&delta_crc.to_le_bytes());
let c = crc32(&h);
h.extend_from_slice(&c.to_le_bytes());
debug_assert!(h.len() <= geo.hdr_len());
h
}
fn unit_bytes(src: &SyncSource<'_>, block: usize) -> Vec<u8> {
let geo = src.geo();
let from = block * BLOCK;
let codes = (src.seq_blocks)(from, from + BLOCK);
debug_assert_eq!(codes.len(), BLOCK * geo.row_bytes());
let mut u = Vec::with_capacity(geo.unit_len());
u.extend_from_slice(&codes);
drop(codes);
for lane in 0..BLOCK {
let r = from + lane;
let v = if r < src.n_vectors { src.scales[r] } else { 0.0 };
u.extend_from_slice(&v.to_le_bytes());
}
if let Some(ids) = src.ids {
for lane in 0..BLOCK {
let r = from + lane;
let v = if r < src.n_vectors { ids[r] } else { 0 };
u.extend_from_slice(&v.to_le_bytes());
}
}
debug_assert_eq!(u.len(), geo.unit_len());
u
}
#[derive(Debug, Default)]
pub(crate) struct Batch {
pub ops: Vec<(u64, Vec<u8>)>,
}
pub(crate) struct SyncPlan {
pub batches: Vec<Batch>,
pub new_cursor: SyncCursor,
pub carried: Vec<usize>,
}
pub(crate) fn plan_incremental(
src: &SyncSource<'_>,
cursor: SyncCursor,
pending: &std::collections::HashSet<usize>,
fresh: &std::collections::HashSet<usize>,
clear_target: Option<usize>,
) -> Option<SyncPlan> {
let geo = src.geo();
let old_blocks = (cursor.n_synced as usize) / BLOCK;
let new_blocks = src.n_vectors / BLOCK;
let live_blocks = old_blocks.min(new_blocks);
let gen = cursor.gen + 1;
let live = |s: &&usize| **s < live_blocks * BLOCK;
let fresh_units: std::collections::HashSet<usize> =
fresh.iter().filter(live).map(|&s| s / BLOCK).collect();
let mut materialize: Vec<usize> = pending
.iter()
.filter(live)
.map(|&s| s / BLOCK)
.filter(|b| !fresh_units.contains(b))
.collect();
materialize.sort_unstable();
materialize.dedup();
let mut carried: Vec<usize> = pending
.iter()
.chain(fresh.iter())
.filter(live)
.filter(|&&s| fresh_units.contains(&(s / BLOCK)))
.copied()
.collect();
carried.sort_unstable();
carried.dedup();
if carried.len() > MAX_OPS {
return None;
}
let mut groups: Vec<(usize, Vec<usize>)> = Vec::new();
for &s in &carried {
match groups.last_mut() {
Some((b, slots)) if *b == s / BLOCK => slots.push(s),
_ => groups.push((s / BLOCK, vec![s])),
}
}
let n_written = materialize.len() + new_blocks.saturating_sub(old_blocks);
let mut digest = DeltaDigest::new(8 + n_written * (4 + geo.unit_len()));
digest.push(&gen.to_le_bytes());
let mut batch = Batch::default();
for b in materialize.iter().copied().chain(old_blocks..new_blocks) {
let bytes = unit_bytes(src, b);
digest.push(&(b as u32).to_le_bytes());
digest.push(&bytes);
batch.ops.push((geo.unit_at(b) as u64, bytes));
}
let delta_crc = digest.finish();
let slot = (gen % 2) as usize;
batch.ops.push((
geo.hdr_at(slot) as u64,
header_slot(
src,
gen,
src.n_vectors,
&groups,
(&materialize, old_blocks..new_blocks, delta_crc),
),
));
let batches = if let Some(used) = clear_target {
let mut bytes = vec![0u8; used.clamp(8, geo.hdr_len())];
bytes[..8].copy_from_slice(&((gen % 2) ^ 1).to_le_bytes());
vec![Batch { ops: vec![(geo.hdr_at(slot) as u64, bytes)] }, batch]
} else {
vec![batch]
};
Some(SyncPlan {
batches,
new_cursor: SyncCursor {
gen,
n_synced: src.n_vectors as u64,
calib_gen: cursor.calib_gen,
nonce: cursor.nonce,
},
carried,
})
}
fn delta_digest(gen: u64, units: &[(usize, &[u8])]) -> u32 {
let total = units
.iter()
.fold(8usize, |n, (_, body)| n.saturating_add(4 + body.len()));
let mut d = DeltaDigest::new(total);
d.push(&gen.to_le_bytes());
for (b, body) in units {
d.push(&(*b as u32).to_le_bytes());
d.push(body);
}
d.finish()
}
fn delta_verified(
h: &ParsedHdr,
unit_len: usize,
mut feed_unit: impl FnMut(usize, &mut DeltaDigest) -> bool,
) -> bool {
let count = h.delta_mat.len().saturating_add(h.delta_app.len());
let mut d = DeltaDigest::new(8usize.saturating_add(count.saturating_mul(4 + unit_len)));
d.push(&h.gen.to_le_bytes());
for b in h.delta_mat.iter().copied().chain(h.delta_app.clone()) {
d.push(&(b as u32).to_le_bytes());
if !feed_unit(b, &mut d) {
return false;
}
}
d.finish() == h.delta_crc
}
fn fsync_commit(f: &File) -> io::Result<()> {
f.sync_all()
}
pub(crate) fn run_sync(path: &Path, plan: &SyncPlan) -> io::Result<SyncCursor> {
let mut f = std::fs::OpenOptions::new().read(true).write(true).open(path)?;
for batch in &plan.batches {
for (off, bytes) in &batch.ops {
f.seek(SeekFrom::Start(*off))?;
f.write_all(bytes)?;
}
f.flush()?;
fsync_commit(&f)?;
}
Ok(plan.new_cursor)
}
fn write_image<W: Write>(w: &mut W, src: &SyncSource<'_>, gen: u64, nonce: u64) -> io::Result<()> {
let geo = src.geo();
let n_blocks = src.n_vectors / BLOCK;
let h = header_slot(
src,
gen,
src.n_vectors,
&[],
(&[], 0..0, delta_digest(gen, &[])),
);
w.write_all(&superblock(src, nonce))?;
w.write_all(&h)?;
w.write_all(&vec![0u8; geo.hdr_len() - h.len()])?;
w.write_all(&vec![0u8; geo.hdr_len()])?;
for b in 0..n_blocks {
w.write_all(&unit_bytes(src, b))?;
}
Ok(())
}
pub(crate) const UNCLAIMED_NONCE: u64 = 0;
pub(crate) fn stream_image<W: Write>(w: &mut W, src: &SyncSource<'_>) -> io::Result<()> {
write_image(w, src, 0, UNCLAIMED_NONCE)
}
pub(crate) fn image_bytes(src: &SyncSource<'_>) -> Vec<u8> {
let geo = src.geo();
let mut buf = Vec::with_capacity(geo.full_len(src.n_vectors));
write_image(&mut buf, src, 0, UNCLAIMED_NONCE).expect("writing to a Vec<u8> cannot fail");
buf
}
pub(crate) fn image_len(src: &SyncSource<'_>) -> usize {
src.geo().full_len(src.n_vectors)
}
pub(crate) fn write_full(
path: &Path,
src: &SyncSource<'_>,
calib_gen: u64,
) -> io::Result<SyncCursor> {
write_full_with_durability(path, src, calib_gen, crate::io::Durability::Durable)
}
pub(crate) fn write_full_with_durability(
path: &Path,
src: &SyncSource<'_>,
calib_gen: u64,
durability: crate::io::Durability,
) -> io::Result<SyncCursor> {
write_full_inner(path, src, calib_gen, durability, crate::io::file_nonce())
}
pub(crate) fn write_snapshot(
path: &Path,
src: &SyncSource<'_>,
durability: crate::io::Durability,
) -> io::Result<()> {
write_full_inner(path, src, 0, durability, UNCLAIMED_NONCE).map(|_| ())
}
fn write_full_inner(
path: &Path,
src: &SyncSource<'_>,
calib_gen: u64,
durability: crate::io::Durability,
nonce: u64,
) -> io::Result<SyncCursor> {
let gen = 0u64;
crate::io::sweep_stale_tmps(path);
let (f, tmp) = crate::io::create_tmp(path)?;
let result = (|| {
let mut w = std::io::BufWriter::with_capacity(1 << 20, &f);
write_image(&mut w, src, gen, nonce)?;
w.flush()?;
drop(w);
match durability {
crate::io::Durability::Durable => fsync_commit(&f),
crate::io::Durability::Fast => Ok(()),
}
})();
let result = result.and_then(|()| {
drop(f);
crate::io::rename_atomic(&tmp, path)
});
if let Err(e) = result {
let _ = std::fs::remove_file(&tmp);
return Err(e);
}
if durability == crate::io::Durability::Durable {
crate::io::sync_parent_dir_after_commit(path);
}
Ok(SyncCursor {
gen,
n_synced: src.n_vectors as u64,
calib_gen,
nonce,
})
}
pub(crate) struct V7Load {
pub dim: usize,
pub bit_width: usize,
pub n_vectors: usize,
pub seq_blocked: Vec<u8>,
pub scales: Vec<f32>,
pub ids: Vec<u64>,
pub tqplus_shift: Vec<f32>,
pub tqplus_scale: Vec<f32>,
pub cursor: SyncCursor,
pub pending_slots: Vec<usize>,
}
fn bad(msg: impl Into<String>) -> io::Error {
io::Error::new(io::ErrorKind::InvalidData, msg.into())
}
fn read_u32(raw: &[u8], at: usize) -> io::Result<u32> {
raw.get(at..at + 4)
.map(|b| u32::from_le_bytes(b.try_into().unwrap()))
.ok_or_else(|| bad("unexpected end of file"))
}
fn read_u64_at(raw: &[u8], at: usize) -> io::Result<u64> {
raw.get(at..at + 8)
.map(|b| u64::from_le_bytes(b.try_into().unwrap()))
.ok_or_else(|| bad("unexpected end of file"))
}
fn read_f32(raw: &[u8], at: usize) -> io::Result<f32> {
raw.get(at..at + 4)
.map(|b| f32::from_le_bytes(b.try_into().unwrap()))
.ok_or_else(|| bad("unexpected end of file"))
}
type OpGroup = (usize, Vec<(usize, usize)>);
struct ParsedHdr {
gen: u64,
n: usize,
tail_at: usize,
groups: Vec<OpGroup>,
delta_mat: Vec<usize>,
delta_app: std::ops::Range<usize>,
delta_crc: u32,
used: usize,
}
fn parse_header_slot(raw: &[u8], geo: &Geo, slot: usize, file_len: usize) -> Option<ParsedHdr> {
let at = geo.hdr_at(slot);
let bytes = raw.get(at..at + geo.hdr_len())?;
parse_header_at(bytes, at, geo, slot, file_len)
}
fn parse_header_at(
bytes: &[u8],
at: usize,
geo: &Geo,
slot: usize,
file_len: usize,
) -> Option<ParsedHdr> {
let tail_row = geo.row_bytes() + 4 + geo.id_bytes(1);
let op_size = geo.op_size();
let gen = u64::from_le_bytes(bytes.get(..8)?.try_into().unwrap());
if (gen % 2) as usize != slot {
return None;
}
let n64 = u64::from_le_bytes(bytes.get(8..16)?.try_into().unwrap());
let n = usize::try_from(n64).ok()?;
let units_end = (n / BLOCK)
.checked_mul(geo.unit_len())
.and_then(|u| u.checked_add(geo.unit_at(0)))?;
if units_end > file_len {
return None;
}
let mut p = 16 + (n % BLOCK) * tail_row;
let n_units = u32::from_le_bytes(bytes.get(p..p + 4)?.try_into().unwrap()) as usize;
if n_units > MAX_OPS {
return None;
}
p += 4;
let mut groups = Vec::with_capacity(n_units);
for _ in 0..n_units {
let b = u32::from_le_bytes(bytes.get(p..p + 4)?.try_into().unwrap()) as usize;
if b >= n / BLOCK {
return None;
}
let n_ops = *bytes.get(p + 4)? as usize;
p += 5;
let mut ops = Vec::with_capacity(n_ops);
for _ in 0..n_ops {
let lane = *bytes.get(p)? as usize;
if lane >= BLOCK {
return None;
}
ops.push((b * BLOCK + lane, at + p + 1));
p += op_size;
}
groups.push((b, ops));
}
let n_mat = u32::from_le_bytes(bytes.get(p..p + 4)?.try_into().unwrap()) as usize;
if n_mat > MAX_OPS {
return None;
}
p += 4;
let mut delta_mat = Vec::with_capacity(n_mat);
for _ in 0..n_mat {
delta_mat.push(u32::from_le_bytes(bytes.get(p..p + 4)?.try_into().unwrap()) as usize);
p += 4;
}
let app_from = u32::from_le_bytes(bytes.get(p..p + 4)?.try_into().unwrap()) as usize;
let app_to = u32::from_le_bytes(bytes.get(p + 4..p + 8)?.try_into().unwrap()) as usize;
let delta_crc = u32::from_le_bytes(bytes.get(p + 8..p + 12)?.try_into().unwrap());
p += 12;
let stored = u32::from_le_bytes(bytes.get(p..p + 4)?.try_into().unwrap());
if crc32(bytes.get(..p)?) != stored {
return None;
}
Some(ParsedHdr {
gen,
n,
tail_at: at + 16,
groups,
delta_mat,
delta_app: app_from..app_to,
delta_crc,
used: p + 4,
})
}
pub(crate) fn load(path: &Path, expect_calib_gen: u64, expect_kind: u8) -> io::Result<V7Load> {
let f = File::open(path)?;
let on_disk = f.metadata()?.len();
let want = declared_len(&f).unwrap_or(on_disk).min(on_disk);
let mut raw = vec![
0u8;
usize::try_from(want).map_err(|_| io::Error::new(
io::ErrorKind::InvalidData,
"file too large for this platform"
))?
];
crate::io::read_exact_at(&f, &mut raw, 0)?;
load_image(raw, expect_calib_gen, expect_kind, &path.display().to_string())
}
fn declared_len(f: &File) -> Option<u64> {
let mut sb = [0u8; 64];
crate::io::read_exact_at(f, &mut sb, 0).ok()?;
if &sb[0..4] != V7_MAGIC || sb[4] != V7_VERSION {
return None;
}
let bit_width = sb[5] as usize;
let kind = sb[6];
if !(2..=4).contains(&bit_width) || kind > 1 {
return None;
}
let dim = u32::from_le_bytes(sb[7..11].try_into().ok()?) as usize;
if dim != 0 && (!dim.is_multiple_of(8) || dim > crate::MAX_DIM) {
return None;
}
let n_levels = 1usize << bit_width;
let n_calib_at = 23 + (2 * n_levels - 1) * 4;
let mut n_calib_bytes = [0u8; 4];
crate::io::read_exact_at(f, &mut n_calib_bytes, n_calib_at as u64).ok()?;
let n_calib = u32::from_le_bytes(n_calib_bytes) as usize;
if n_calib != 0 && n_calib != dim {
return None;
}
let geo = Geo { kind, dim, bit_width, n_calib };
let mut n = 0u64;
for slot in 0..2 {
let mut buf = [0u8; 16];
if crate::io::read_exact_at(f, &mut buf, geo.hdr_at_for_len(slot) as u64).is_err() {
return None;
}
n = n.max(u64::from_le_bytes(buf[8..16].try_into().ok()?));
}
let blocks = n / BLOCK as u64;
let units = (geo.unit_len() as u64).checked_mul(blocks)?;
(geo.unit_at(0) as u64).checked_add(units)
}
pub(crate) fn load_image(
mut raw: Vec<u8>,
expect_calib_gen: u64,
expect_kind: u8,
src: &str,
) -> io::Result<V7Load> {
if raw.len() < 11 || &raw[..4] != V7_MAGIC {
return Err(bad("not a v7 file"));
}
if raw[4] != V7_VERSION {
return Err(bad(format!("unsupported v7 revision {}", raw[4])));
}
let bit_width = raw[5] as usize;
if !(2..=4).contains(&bit_width) {
return Err(bad(format!("bit_width {bit_width} out of range")));
}
let kind = raw[6];
if kind != expect_kind {
return Err(bad(match kind {
1 => "this v7 file holds an IdMapIndex; load it with IdMapIndex::load".to_string(),
0 => {
"this v7 file holds a TurboQuantIndex; load it with TurboQuantIndex::load"
.to_string()
}
k => format!("unknown v7 index kind {k}"),
}));
}
let dim = read_u32(&raw, 7)? as usize;
if dim != 0 && (!dim.is_multiple_of(8) || dim > crate::MAX_DIM) {
return Err(bad(format!("dim {dim} invalid")));
}
let min_header = MAX_OPS.saturating_mul(dim.saturating_mul(bit_width) / 8);
if raw.len() < min_header {
return Err(bad(format!(
"truncated file: a {dim}-dim {bit_width}-bit image reserves at least \
{min_header} bytes of header, but the file is {} bytes",
raw.len(),
)));
}
let nonce = read_u64_at(&raw, 11)?;
let file_max_ops = read_u32(&raw, 19)? as usize;
if file_max_ops != MAX_OPS {
return Err(bad(format!(
"unsupported header ops capacity {file_max_ops} (this build supports {MAX_OPS})"
)));
}
let n_levels = 1usize << bit_width;
let mut off = 23;
if dim == 0 {
off += (2 * n_levels - 1) * 4;
} else {
let (canon_b, canon_c) = crate::codebook::codebook(bit_width, dim);
for want in canon_b.iter().chain(canon_c.iter()) {
if read_f32(&raw, off)? != *want {
return Err(bad("embedded codebook drifted from the canonical one"));
}
off += 4;
}
}
let n_calib = read_u32(&raw, off)? as usize;
off += 4;
if n_calib != 0 && n_calib != dim {
return Err(bad(format!("calibration length {n_calib} != dim {dim}")));
}
let mut tqplus_shift = Vec::with_capacity(n_calib);
let mut tqplus_scale = Vec::with_capacity(n_calib);
for k in 0..n_calib {
tqplus_shift.push(read_f32(&raw, off + k * 4)?);
}
off += n_calib * 4;
for k in 0..n_calib {
tqplus_scale.push(read_f32(&raw, off + k * 4)?);
}
off += n_calib * 4;
crate::io::validate_calibration(&tqplus_shift, &tqplus_scale)?;
let stored = read_u32(&raw, off)?;
if crc32(&raw[..off]) != stored {
return Err(bad("corrupt superblock (crc mismatch)"));
}
let geo = Geo {
kind,
dim,
bit_width,
n_calib,
};
let row_bytes = geo.row_bytes();
let tail_row = row_bytes + 4 + geo.id_bytes(1);
let parse_hdr = |slot: usize| parse_header_slot(&raw, &geo, slot, raw.len());
let delta_ok = |h: &ParsedHdr| -> bool {
delta_verified(h, geo.unit_len(), |b, d| {
let at = geo.unit_at(b);
raw.get(at..at + geo.unit_len()).map(|u| d.push(u)).is_some()
})
};
let mut cands: Vec<ParsedHdr> =
[parse_hdr(0), parse_hdr(1)].into_iter().flatten().collect();
cands.sort_by_key(|h| std::cmp::Reverse(h.gen));
let newest = cands.first().map(|h| h.gen);
let Some(chosen) = cands.into_iter().find(delta_ok) else {
return Err(bad("no valid commit header — unrecoverable v7 file"));
};
if newest.is_some_and(|g| g != chosen.gen) {
crate::warning::warn(&format!(
"{}: the newest commit (generation {}) is incomplete — its sync did \
not finish — so generation {} was loaded instead; changes made after \
that commit are lost",
src,
newest.expect("checked"),
chosen.gen,
));
}
let gen = chosen.gen;
let n_vectors = chosen.n;
if dim == 0 && n_vectors != 0 {
return Err(bad(format!(
"dim 0 with {n_vectors} rows: no dimension committed"
)));
}
let n_blocks = n_vectors / BLOCK;
let total_blocks = n_vectors.div_ceil(BLOCK);
let block_bytes = BLOCK * row_bytes;
debug_assert!(geo.unit_at(0) >= block_bytes, "compaction dest must trail the sources");
let n_tail = n_vectors % BLOCK;
let tail_copy: Vec<u8> = raw
.get(chosen.tail_at..chosen.tail_at + n_tail * tail_row)
.ok_or_else(|| bad("truncated commit tail"))?
.to_vec();
let op_size = row_bytes + 4 + geo.id_bytes(1);
type OwnedGroup = (usize, Vec<(usize, Vec<u8>)>);
let mut ops_owned: Vec<OwnedGroup> = Vec::with_capacity(chosen.groups.len());
for (b, ops) in &chosen.groups {
let mut owned = Vec::with_capacity(ops.len());
for &(slot, payload_at) in ops {
let payload = raw
.get(payload_at..payload_at + op_size)
.ok_or_else(|| bad("truncated pending op"))?
.to_vec();
owned.push((slot, payload));
}
ops_owned.push((*b, owned));
}
let mut scales: Vec<f32> = Vec::with_capacity(n_vectors);
let mut ids: Vec<u64> = Vec::with_capacity(if kind == 1 { n_vectors } else { 0 });
for b in 0..n_blocks {
let at = geo.unit_at(b);
if raw.len() < at + geo.unit_len() {
return Err(bad("truncated block unit"));
}
for lane in 0..BLOCK {
let so = at + block_bytes + lane * 4;
let v = f32::from_le_bytes(raw[so..so + 4].try_into().unwrap());
if !v.is_finite() || !(0.0..=crate::io::MAX_VECTOR_SCALE).contains(&v) {
return Err(bad(format!("invalid per-vector scale in block {b}")));
}
scales.push(v);
}
if kind == 1 {
for lane in 0..BLOCK {
let io_ = at + block_bytes + BLOCK * 4 + lane * 8;
ids.push(u64::from_le_bytes(raw[io_..io_ + 8].try_into().unwrap()));
}
}
raw.copy_within(at..at + block_bytes, b * block_bytes);
}
raw.truncate(n_blocks * block_bytes);
raw.resize(total_blocks * block_bytes, 0);
raw.shrink_to_fit();
let mut seq_blocked = raw;
for (b, ops) in &ops_owned {
for (slot, payload) in ops {
let lane = slot % BLOCK;
for g in 0..row_bytes {
seq_blocked[b * block_bytes + g * BLOCK + lane] = payload[g];
}
let v = f32::from_le_bytes(payload[row_bytes..row_bytes + 4].try_into().unwrap());
if !v.is_finite() || !(0.0..=crate::io::MAX_VECTOR_SCALE).contains(&v) {
return Err(bad("invalid per-vector scale in a pending op"));
}
scales[*slot] = v;
if kind == 1 {
ids[*slot] =
u64::from_le_bytes(payload[row_bytes + 4..row_bytes + 12].try_into().unwrap());
}
}
}
for k in 0..n_tail {
let r = n_blocks * BLOCK + k;
let lane = r % BLOCK;
let row = &tail_copy[k * tail_row..(k + 1) * tail_row];
for g in 0..row_bytes {
seq_blocked[n_blocks * block_bytes + g * BLOCK + lane] = row[g];
}
let v = f32::from_le_bytes(row[row_bytes..row_bytes + 4].try_into().unwrap());
if !v.is_finite() || !(0.0..=crate::io::MAX_VECTOR_SCALE).contains(&v) {
return Err(bad("invalid per-vector scale in the commit tail"));
}
scales.push(v);
if kind == 1 {
ids.push(u64::from_le_bytes(row[row_bytes + 4..row_bytes + 12].try_into().unwrap()));
}
}
Ok(V7Load {
dim,
bit_width,
n_vectors,
seq_blocked,
scales,
ids,
tqplus_shift,
tqplus_scale,
cursor: SyncCursor {
gen,
n_synced: n_vectors as u64,
calib_gen: expect_calib_gen,
nonce,
},
pending_slots: chosen
.groups
.iter()
.flat_map(|(_, ops)| ops.iter().map(|&(s, _)| s))
.collect(),
})
}
pub(crate) enum CursorState {
Intact {
stale_ahead: Option<usize>,
},
Foreign,
Replaced,
}
pub(crate) fn cursor_state(
path: &Path,
cursor: &SyncCursor,
geo: &Geo,
) -> io::Result<CursorState> {
let mut f = match File::open(path) {
Ok(f) => f,
Err(e) if e.kind() == io::ErrorKind::NotFound => return Ok(CursorState::Replaced),
Err(e) => return Err(e),
};
let file_len =
usize::try_from(f.metadata()?.len()).map_err(|_| bad("file too large"))?;
let mut head = [0u8; 19];
match f.read_exact(&mut head) {
Ok(()) => {}
Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => {
return Ok(CursorState::Replaced)
}
Err(e) => return Err(e),
}
if &head[..4] != V7_MAGIC || head[4] != V7_VERSION {
return Ok(CursorState::Replaced);
}
let file_nonce = u64::from_le_bytes(head[11..19].try_into().unwrap());
if file_nonce != cursor.nonce {
if file_nonce == UNCLAIMED_NONCE {
return Ok(CursorState::Replaced);
}
return Ok(CursorState::Foreign);
}
let mut read_slot = |slot: usize| -> Option<ParsedHdr> {
let at = geo.hdr_at(slot);
let avail = file_len.checked_sub(at)?;
let mut want = geo.hdr_probe_len().min(avail);
loop {
let mut buf = vec![0u8; want];
f.seek(SeekFrom::Start(at as u64)).ok()?;
f.read_exact(&mut buf).ok()?;
if let Some(h) = parse_header_at(&buf, at, geo, slot, file_len) {
return Some(h);
}
let full = geo.hdr_len().min(avail);
if want >= full {
return None;
}
want = full;
}
};
let mut cands: Vec<ParsedHdr> = [read_slot(0), read_slot(1)]
.into_iter()
.flatten()
.collect();
cands.sort_by_key(|h| std::cmp::Reverse(h.gen));
if cands.first().is_some_and(|h| h.gen == cursor.gen) {
return Ok(CursorState::Intact { stale_ahead: None });
}
let stale_ahead = cands
.iter()
.filter(|h| h.gen > cursor.gen)
.map(|h| h.used)
.max();
let mut adoptable: Option<u64> = None;
let mut unit_buf = vec![0u8; geo.unit_len()];
for h in &cands {
let verified = delta_verified(h, geo.unit_len(), |b, d| {
let at = geo.unit_at(b);
if at + geo.unit_len() > file_len {
return false;
}
if f.seek(SeekFrom::Start(at as u64)).is_err() || f.read_exact(&mut unit_buf).is_err() {
return false;
}
d.push(&unit_buf);
true
});
if verified {
adoptable = Some(h.gen);
break;
}
}
match adoptable {
Some(g) if g == cursor.gen => Ok(CursorState::Intact { stale_ahead }),
Some(_) => Ok(CursorState::Foreign),
None => Ok(CursorState::Replaced),
}
}
#[cfg(test)]
pub(crate) fn hdr_used_for_test(raw: &[u8], geo: &Geo, slot: usize) -> usize {
let tail_row = geo.row_bytes() + 4 + geo.id_bytes(1);
let op_size = geo.op_size();
let at = geo.hdr_at(slot);
let Some(bytes) = raw.get(at..at + geo.hdr_len()) else {
return 16;
};
let walk = || -> Option<usize> {
let n = usize::try_from(u64::from_le_bytes(bytes[8..16].try_into().ok()?)).ok()?;
let mut p = 16 + (n % BLOCK) * tail_row;
let n_units = u32::from_le_bytes(bytes.get(p..p + 4)?.try_into().ok()?) as usize;
p += 4;
if n_units > MAX_OPS {
return None;
}
for _ in 0..n_units {
let n_ops = *bytes.get(p + 4)? as usize;
p += 5 + n_ops * op_size;
}
let n_mat = u32::from_le_bytes(bytes.get(p..p + 4)?.try_into().ok()?) as usize;
if n_mat > MAX_OPS {
return None;
}
p += 4 + n_mat * 4 + 12;
bytes.get(p..p + 4)?;
Some(p + 4)
};
walk().unwrap_or(16)
}
#[cfg(test)]
pub(crate) fn reseal_for_test(bytes: &mut [u8], geo: &Geo) {
let sb = geo.sb_len();
if bytes.len() >= sb {
let c = crc32(&bytes[..sb - 4]);
bytes[sb - 4..sb].copy_from_slice(&c.to_le_bytes());
}
let hdr_len = geo.hdr_len();
let tail_row = geo.row_bytes() + 4 + geo.id_bytes(1);
let op_size = geo.op_size();
for slot in 0..2 {
let at = geo.hdr_at(slot);
if bytes.len() < at + hdr_len {
continue;
}
let h = &bytes[at..at + hdr_len];
let n = u64::from_le_bytes(h[8..16].try_into().unwrap()) as usize;
let Some(mut p) = (n % BLOCK).checked_mul(tail_row).and_then(|t| t.checked_add(16))
else {
continue;
};
let read_u32_at = |h: &[u8], q: usize| -> Option<u32> {
h.get(q..q + 4)
.map(|b| u32::from_le_bytes(b.try_into().unwrap()))
};
let Some(n_units) = read_u32_at(h, p) else { continue };
p += 4;
let mut ok = true;
for _ in 0..n_units {
let Some(&n_ops) = h.get(p + 4) else {
ok = false;
break;
};
let Some(np) = (n_ops as usize)
.checked_mul(op_size)
.and_then(|v| v.checked_add(p + 5))
else {
ok = false;
break;
};
p = np;
}
if !ok {
continue;
}
let Some(n_mat) = read_u32_at(h, p) else { continue };
let Some(np) = (n_mat as usize)
.checked_mul(4)
.and_then(|v| v.checked_add(p + 4 + 12))
else {
continue;
};
p = np;
if p + 4 <= hdr_len {
let c = crc32(&bytes[at..at + p]);
bytes[at + p..at + p + 4].copy_from_slice(&c.to_le_bytes());
}
}
}
pub(crate) fn is_v7(path: &Path) -> bool {
let mut magic = [0u8; 4];
File::open(path)
.and_then(|mut f| f.read_exact(&mut magic))
.map(|_| &magic == V7_MAGIC)
.unwrap_or(false)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn interleaved_crc_matches_the_reference_composition() {
let mut data = vec![0u8; 100_000];
let mut s = 0x9E37_79B9_7F4A_7C15u64;
for b in data.iter_mut() {
s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
*b = (s >> 56) as u8;
}
let third = data.len() / 3;
let (a, rest) = data.split_at(third);
let (b, c) = rest.split_at(third);
let mut digest = [0u8; 12];
digest[..4].copy_from_slice(&crc32c_soft(a).to_le_bytes());
digest[4..8].copy_from_slice(&crc32c_soft(b).to_le_bytes());
digest[8..].copy_from_slice(&crc32c_soft(c).to_le_bytes());
assert_eq!(crc32(&data), crc32c_soft(&digest));
let reference = crc32(&data);
data[50_000] ^= 1;
assert_ne!(crc32(&data), reference);
}
#[test]
fn crc_split_threshold_is_exact() {
for len in [4095usize, 4096, 4097] {
let data: Vec<u8> = (0..len).map(|i| (i % 249) as u8).collect();
let expect = if len < 4096 {
crc32c_soft(&data)
} else {
let third = len / 3;
let (a, rest) = data.split_at(third);
let (b, c) = rest.split_at(third);
let mut d = [0u8; 12];
d[..4].copy_from_slice(&crc32c_soft(a).to_le_bytes());
d[4..8].copy_from_slice(&crc32c_soft(b).to_le_bytes());
d[8..].copy_from_slice(&crc32c_soft(c).to_le_bytes());
crc32c_soft(&d)
};
assert_eq!(crc32(&data), expect, "len {len}");
}
}
#[test]
fn the_delta_digest_depends_on_every_byte() {
let mut body_a = vec![7u8; 1000];
let body_b = vec![9u8; 1000];
let base = delta_digest(3, &[(0usize, body_a.as_slice()), (1, body_b.as_slice())]);
for i in [0usize, 1, 499, 998, 999] {
body_a[i] ^= 1;
let changed = delta_digest(3, &[(0usize, body_a.as_slice()), (1, body_b.as_slice())]);
assert_ne!(base, changed, "flip at byte {i} must change the digest");
body_a[i] ^= 1;
}
assert_ne!(
base,
delta_digest(3, &[(1usize, body_a.as_slice()), (0, body_b.as_slice())]),
"block indices must bind"
);
assert_ne!(
base,
delta_digest(4, &[(0usize, body_a.as_slice()), (1, body_b.as_slice())]),
"the generation must bind"
);
let mut u1 = vec![1u8; 512];
let c = crc32(&u1);
u1.extend_from_slice(&c.to_le_bytes());
let mut u2 = vec![2u8; 512];
let c = crc32(&u2);
u2.extend_from_slice(&c.to_le_bytes());
assert_eq!(
delta_digest(1, &[(0usize, u1.as_slice())]),
delta_digest(1, &[(0usize, u2.as_slice())]),
"codeword payloads DO collide — which is why unit bodies must never embed their own CRC"
);
}
#[test]
fn the_streaming_digest_equals_the_materialized_one() {
let mut data = vec![0u8; 20_000];
let mut s = 0x2545_F491_4F6C_DD1Du64;
for b in data.iter_mut() {
s ^= s << 13;
s ^= s >> 7;
s ^= s << 17;
*b = (s >> 32) as u8;
}
for total in [0usize, 1, 7, 8, 4094, 4095, 4096, 4097, 6143, 6144, 12_289, 20_000] {
let buf = &data[..total];
let want = crc32(buf);
for chunk in [1usize, 3, 8, 64, 1000, 2048, 4096, 65_536] {
let mut d = DeltaDigest::new(total);
for piece in buf.chunks(chunk.max(1)) {
d.push(piece);
}
assert_eq!(d.finish(), want, "total {total}, chunk {chunk}");
}
let mut d = DeltaDigest::new(total);
let (mut at, mut step) = (0usize, 1usize);
while at < total {
let take = step.min(total - at);
d.push(&buf[at..at + take]);
at += take;
step = step * 2 + 1;
}
assert_eq!(d.finish(), want, "total {total}, ragged");
}
}
#[test]
fn geometry_is_pinned() {
for (kind, dim, bit_width, n_calib) in
[(0u8, 64usize, 4usize, 64usize), (1, 128, 2, 0), (0, 64, 3, 64)]
{
let geo = Geo {
kind,
dim,
bit_width,
n_calib,
};
let row = dim / (8 / bit_width);
let id1 = if kind == 1 { 8 } else { 0 };
let nl = 1usize << bit_width;
let tail_row = row + 4 + id1;
let op = 1 + row + 4 + id1;
assert_eq!(geo.row_bytes(), row, "row stride");
assert_eq!(geo.op_size(), op, "op size");
assert_eq!(
geo.sb_len(),
23 + (nl - 1) * 4 + nl * 4 + 4 + n_calib * 8 + 4,
"superblock"
);
assert_eq!(
geo.hdr_len(),
16 + 31 * tail_row + 4 + MAX_OPS * (5 + op) + 4 + MAX_OPS * 4 + 12 + 4,
"header slot"
);
assert_eq!(geo.unit_len(), 32 * row + 128 + 32 * id1, "unit");
assert_eq!(
geo.unit_at(3) - geo.unit_at(2),
geo.unit_len(),
"unit stride"
);
assert_eq!(
geo.hdr_probe_len(),
16 + 31 * tail_row + 4 + 4 + MAX_OPS * 4 + 12 + 4,
"header probe"
);
assert_eq!(
geo.hdr_len() - geo.hdr_probe_len(),
MAX_OPS * (5 + geo.op_size()),
"the probe must give up exactly the op-group region",
);
assert!(
geo.hdr_probe_len() < geo.hdr_len(),
"a probe that is not smaller than the slot saves nothing",
);
}
}
}