use crate::{
get_f32, get_u8, get_u16, get_u32, get_u64, put_f32, put_u8, put_u16, put_u32, put_u64,
};
use yo_common::{Code, Error, Result};
pub const IMAGE_TAG: u32 = u32::from_le_bytes(*b"YOIX");
pub const IMAGE_HEADER_LEN: usize = 100;
pub const PARTITION_ENTRY_LEN: usize = 16;
pub const POSTING_HEADER_LEN: usize = 16;
pub const META_LEN: usize = 16;
pub mod image_kind {
pub const VECTOR: u8 = 1;
}
pub mod metric {
pub const L2: u8 = 0;
pub const COSINE: u8 = 1;
pub const IP: u8 = 2;
pub const HAMMING: u8 = 3;
#[must_use]
pub const fn is_known(b: u8) -> bool {
matches!(b, L2 | COSINE | IP | HAMMING)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Default)]
pub struct ImageHeader {
pub kind: u8,
pub bits: u8,
pub metric: u8,
pub dim: u32,
pub partitions: u32,
pub seed: u64,
pub members: u64,
pub slots: u32,
pub posting: u32,
pub probe: u32,
pub rerank: u32,
pub sweep: u32,
pub widen: u32,
pub spill: u32,
pub slack: f32,
pub patience: u32,
pub centroids: Chain,
pub keys: Chain,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct Chain {
pub at: u64,
pub len: u64,
}
pub fn image_len(partitions: u32) -> Result<usize> {
(partitions as usize)
.checked_mul(PARTITION_ENTRY_LEN)
.and_then(|n| n.checked_add(IMAGE_HEADER_LEN))
.ok_or_else(|| {
Error::new(Code::Invalid, "that many partitions do not fit in an image")
.with_detail(format!("partitions={partitions}"))
})
}
impl ImageHeader {
pub fn encode(&self, into: &mut [u8]) -> Result<usize> {
let need = image_len(self.partitions)?;
if into.len() < need {
return Err(
Error::new(Code::Invalid, "buffer is shorter than the image root")
.with_detail(format!("have={} need={need}", into.len())),
);
}
put_u32(into, 0, IMAGE_TAG);
put_u8(into, 4, self.kind);
put_u8(into, 5, self.bits);
put_u8(into, 6, self.metric);
put_u8(into, 7, 0);
put_u32(into, 8, self.dim);
put_u32(into, 12, self.partitions);
put_u64(into, 16, self.seed);
put_u64(into, 24, self.members);
put_u32(into, 32, self.posting);
put_u32(into, 36, self.probe);
put_u32(into, 40, self.rerank);
put_u32(into, 44, self.sweep);
put_u32(into, 48, self.widen);
put_u32(into, 52, self.slots);
put_u64(into, 56, self.centroids.at);
put_u64(into, 64, self.centroids.len);
put_u64(into, 72, self.keys.at);
put_u64(into, 80, self.keys.len);
put_u32(into, 88, self.spill);
put_f32(into, 92, self.slack);
put_u32(into, 96, self.patience);
Ok(need)
}
pub fn decode(bytes: &[u8]) -> Result<ImageHeader> {
if bytes.len() < IMAGE_HEADER_LEN {
return Err(Error::new(Code::Corrupt, "shorter than an image header")
.with_detail(format!("len={}", bytes.len())));
}
let tag = get_u32(bytes, 0);
if tag != IMAGE_TAG {
return Err(Error::new(Code::Corrupt, "not an index image")
.with_detail(format!("tag={tag:#010x}")));
}
if get_u8(bytes, 7) != 0 {
return Err(Error::new(
Code::Corrupt,
"reserved image header bytes are set",
));
}
let h = ImageHeader {
kind: get_u8(bytes, 4),
bits: get_u8(bytes, 5),
metric: get_u8(bytes, 6),
dim: get_u32(bytes, 8),
partitions: get_u32(bytes, 12),
seed: get_u64(bytes, 16),
members: get_u64(bytes, 24),
posting: get_u32(bytes, 32),
probe: get_u32(bytes, 36),
rerank: get_u32(bytes, 40),
sweep: get_u32(bytes, 44),
widen: get_u32(bytes, 48),
slots: get_u32(bytes, 52),
spill: get_u32(bytes, 88),
slack: get_f32(bytes, 92),
patience: get_u32(bytes, 96),
centroids: Chain {
at: get_u64(bytes, 56),
len: get_u64(bytes, 64),
},
keys: Chain {
at: get_u64(bytes, 72),
len: get_u64(bytes, 80),
},
};
if h.dim == 0 || h.dim as usize > crate::vector::MAX_DIM {
return Err(Error::new(Code::Corrupt, "image dimension out of range")
.with_detail(format!("dim={}", h.dim)));
}
if h.bits != 1 && h.bits != 4 {
return Err(Error::new(Code::Corrupt, "unknown code width")
.with_detail(format!("bits={}", h.bits)));
}
if h.members > u64::from(h.slots) {
return Err(
Error::new(Code::Corrupt, "more members than the table has slots")
.with_detail(format!("members={} slots={}", h.members, h.slots)),
);
}
if !metric::is_known(h.metric) {
return Err(Error::new(Code::Corrupt, "unknown metric")
.with_detail(format!("metric={}", h.metric)));
}
let need = image_len(h.partitions)?;
if bytes.len() != need {
return Err(
Error::new(Code::Corrupt, "the image root is not the length it says")
.with_detail(format!("len={} need={need}", bytes.len())),
);
}
let want = u64::from(h.partitions) * u64::from(h.dim) * 4;
if h.centroids.len != want {
return Err(
Error::new(Code::Corrupt, "the centroid section is the wrong size")
.with_detail(format!("len={} want={want}", h.centroids.len)),
);
}
Ok(h)
}
}
pub fn put_partition(root: &mut [u8], i: u32, chain: Chain) -> Result<()> {
let at =
partition_offset(root.len(), i).ok_or_else(|| missing(i, root.len(), Code::Invalid))?;
put_u64(root, at, chain.at);
put_u64(root, at + 8, chain.len);
Ok(())
}
pub fn get_partition(root: &[u8], i: u32) -> Result<Chain> {
let at =
partition_offset(root.len(), i).ok_or_else(|| missing(i, root.len(), Code::Corrupt))?;
Ok(Chain {
at: get_u64(root, at),
len: get_u64(root, at + 8),
})
}
fn partition_offset(root_len: usize, i: u32) -> Option<usize> {
let at = IMAGE_HEADER_LEN + (i as usize) * PARTITION_ENTRY_LEN;
(at + PARTITION_ENTRY_LEN <= root_len).then_some(at)
}
fn missing(i: u32, root_len: usize, code: Code) -> Error {
Error::new(code, "no such partition in the image")
.with_detail(format!("partition={i} root={root_len}"))
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct PostingHeader {
pub count: u32,
pub code_bytes: u32,
pub stuck: u32,
}
pub fn posting_len(count: u32, code_bytes: u32) -> Result<usize> {
let count = count as usize;
let per = 8usize
.checked_add(8)
.and_then(|n| n.checked_add(code_bytes as usize))
.and_then(|n| n.checked_add(META_LEN))
.ok_or_else(|| Error::new(Code::Invalid, "a member of that size does not fit"))?;
count
.checked_mul(per)
.and_then(|n| n.checked_add(POSTING_HEADER_LEN))
.ok_or_else(|| {
Error::new(Code::Invalid, "that many members do not fit in a partition")
.with_detail(format!("count={count} code_bytes={code_bytes}"))
})
}
impl PostingHeader {
pub fn encode(&self, into: &mut [u8]) -> Result<usize> {
let need = posting_len(self.count, self.code_bytes)?;
if into.len() < need {
return Err(
Error::new(Code::Invalid, "buffer is shorter than the posting")
.with_detail(format!("have={} need={need}", into.len())),
);
}
put_u32(into, 0, self.count);
put_u32(into, 4, self.code_bytes);
put_u32(into, 8, self.stuck);
put_u32(into, 12, 0);
Ok(need)
}
pub fn decode(bytes: &[u8]) -> Result<PostingHeader> {
if bytes.len() < POSTING_HEADER_LEN {
return Err(Error::new(Code::Corrupt, "shorter than a posting header")
.with_detail(format!("len={}", bytes.len())));
}
if get_u32(bytes, 12) != 0 {
return Err(Error::new(
Code::Corrupt,
"reserved posting header bytes are set",
));
}
let h = PostingHeader {
count: get_u32(bytes, 0),
code_bytes: get_u32(bytes, 4),
stuck: get_u32(bytes, 8),
};
let Ok(need) = posting_len(h.count, h.code_bytes) else {
return Err(
Error::new(Code::Corrupt, "that many members do not fit in a partition")
.with_detail(format!("count={} code_bytes={}", h.count, h.code_bytes)),
);
};
if bytes.len() != need {
return Err(
Error::new(Code::Corrupt, "the posting is not the length it says")
.with_detail(format!("len={} need={need}", bytes.len())),
);
}
Ok(h)
}
#[must_use]
pub const fn ids_at(&self) -> usize {
POSTING_HEADER_LEN
}
#[must_use]
pub const fn tags_at(&self) -> usize {
self.ids_at() + self.count as usize * 8
}
#[must_use]
pub const fn codes_at(&self) -> usize {
self.tags_at() + self.count as usize * 8
}
#[must_use]
pub const fn meta_at(&self) -> usize {
self.codes_at() + self.count as usize * self.code_bytes as usize
}
}
pub fn put_floats(into: &mut [u8], values: &[f32]) -> Result<usize> {
let need = values.len() * 4;
if into.len() < need {
return Err(
Error::new(Code::Invalid, "buffer is shorter than the floats")
.with_detail(format!("have={} need={need}", into.len())),
);
}
for (i, v) in values.iter().enumerate() {
crate::put_f32(into, i * 4, *v);
}
Ok(need)
}
pub fn get_floats(bytes: &[u8], out: &mut [f32]) -> Result<()> {
let need = out.len() * 4;
if bytes.len() < need {
return Err(
Error::new(Code::Corrupt, "the section is shorter than its floats")
.with_detail(format!("len={} need={need}", bytes.len())),
);
}
for (i, slot) in out.iter_mut().enumerate() {
*slot = crate::get_f32(bytes, i * 4);
}
Ok(())
}
pub fn key_entry_len(klen: usize) -> Result<usize> {
if klen > crate::record::MAX_KEY_LEN {
return Err(
Error::new(Code::Invalid, "the key is longer than 65535 bytes")
.with_detail(format!("klen={klen}")),
);
}
Ok(10 + klen)
}
pub fn put_key(into: &mut [u8], id: u64, key: &[u8]) -> Result<usize> {
let need = key_entry_len(key.len())?;
if into.len() < need {
return Err(
Error::new(Code::Invalid, "buffer is shorter than the key entry")
.with_detail(format!("have={} need={need}", into.len())),
);
}
put_u64(into, 0, id);
put_u16(into, 8, key.len() as u16);
into[10..need].copy_from_slice(key);
Ok(need)
}
#[derive(Debug, Clone)]
pub struct Keys<'a> {
rest: &'a [u8],
}
impl<'a> Keys<'a> {
#[must_use]
pub const fn new(bytes: &'a [u8]) -> Keys<'a> {
Keys { rest: bytes }
}
#[must_use]
pub const fn done(&self) -> bool {
self.rest.is_empty()
}
}
impl<'a> Iterator for Keys<'a> {
type Item = (u64, &'a [u8]);
fn next(&mut self) -> Option<(u64, &'a [u8])> {
if self.rest.len() < 10 {
return None;
}
let id = get_u64(self.rest, 0);
let klen = get_u16(self.rest, 8) as usize;
let end = 10 + klen;
if self.rest.len() < end {
return None;
}
let key = &self.rest[10..end];
self.rest = &self.rest[end..];
Some((id, key))
}
}
#[cfg(test)]
mod tests {
use super::*;
fn header(partitions: u32, dim: u32) -> ImageHeader {
ImageHeader {
kind: image_kind::VECTOR,
bits: 1,
metric: metric::COSINE,
dim,
partitions,
seed: 0x0102_0304_0506_0708,
members: 9999,
slots: 10_000,
posting: 256,
probe: 8,
rerank: 4,
sweep: 4,
widen: 8,
spill: 4,
slack: 0.10,
patience: 3,
centroids: Chain {
at: 4096,
len: u64::from(partitions) * u64::from(dim) * 4,
},
keys: Chain { at: 8192, len: 123 },
}
}
#[test]
fn a_root_comes_back_field_for_field() {
let h = header(3, 128);
let mut buf = vec![0u8; image_len(3).unwrap()];
let wrote = h.encode(&mut buf).unwrap();
assert_eq!(wrote, buf.len());
assert_eq!(ImageHeader::decode(&buf).unwrap(), h);
}
#[test]
fn the_partition_table_is_addressed_and_not_walked() {
let h = header(4, 8);
let mut buf = vec![0u8; image_len(4).unwrap()];
h.encode(&mut buf).unwrap();
for i in 0..4 {
let chain = Chain {
at: 1000 + u64::from(i),
len: 64 + u64::from(i),
};
put_partition(&mut buf, i, chain).unwrap();
}
for i in 0..4 {
assert_eq!(
get_partition(&buf, i).unwrap(),
Chain {
at: 1000 + u64::from(i),
len: 64 + u64::from(i)
}
);
}
assert!(get_partition(&buf, 4).is_err(), "there is no fifth");
assert!(put_partition(&mut buf, 9, Chain::default()).is_err());
}
#[test]
fn a_root_that_is_not_a_root_is_refused() {
let h = header(2, 16);
let mut buf = vec![0u8; image_len(2).unwrap()];
h.encode(&mut buf).unwrap();
assert!(ImageHeader::decode(&buf).is_ok());
let mut wrong = buf.clone();
put_u32(&mut wrong, 0, 0xdead_beef);
assert_eq!(
ImageHeader::decode(&wrong).unwrap_err().code(),
Code::Corrupt,
"a chunk that is not an image was read as one"
);
let mut set = buf.clone();
set[7] = 1;
assert!(
ImageHeader::decode(&set).is_err(),
"byte 7 is reserved and a writer that set it disagreed with this layout"
);
let mut short = buf.clone();
put_u32(&mut short, 52, 3);
assert!(
ImageHeader::decode(&short).is_err(),
"a table with fewer slots than members cannot hold them"
);
let mut bits = buf.clone();
put_u8(&mut bits, 5, 2);
assert!(ImageHeader::decode(&bits).is_err(), "no two bit codes");
let mut met = buf.clone();
put_u8(&mut met, 6, 9);
assert!(ImageHeader::decode(&met).is_err(), "no ninth metric");
let mut dim = buf.clone();
put_u32(&mut dim, 8, 0);
assert!(ImageHeader::decode(&dim).is_err(), "no zero dimension");
let mut count = buf.clone();
put_u32(&mut count, 12, 99);
assert!(ImageHeader::decode(&count).is_err());
let mut cent = buf.clone();
put_u64(&mut cent, 64, 7);
assert!(
ImageHeader::decode(¢).is_err(),
"the centroid section has to be partitions times dim floats"
);
for len in 0..buf.len() {
assert!(
ImageHeader::decode(&buf[..len]).is_err(),
"{len} bytes decoded as a two partition image"
);
}
}
#[test]
fn a_partition_is_four_runs_that_do_not_overlap() {
let h = PostingHeader {
count: 5,
code_bytes: 16,
stuck: 12,
};
let mut buf = vec![0u8; posting_len(5, 16).unwrap()];
let wrote = h.encode(&mut buf).unwrap();
assert_eq!(wrote, buf.len());
assert_eq!(PostingHeader::decode(&buf).unwrap(), h);
assert_eq!(h.ids_at(), POSTING_HEADER_LEN);
assert_eq!(h.tags_at(), h.ids_at() + 40);
assert_eq!(h.codes_at(), h.tags_at() + 40);
assert_eq!(h.meta_at(), h.codes_at() + 80);
assert_eq!(h.meta_at() + 5 * META_LEN, buf.len());
}
#[test]
fn an_empty_partition_is_a_header_and_nothing_else() {
let h = PostingHeader {
count: 0,
code_bytes: 96,
stuck: 0,
};
let mut buf = vec![0u8; POSTING_HEADER_LEN];
h.encode(&mut buf).unwrap();
assert_eq!(PostingHeader::decode(&buf).unwrap(), h);
assert_eq!(h.meta_at(), buf.len());
}
#[test]
fn a_posting_that_is_not_the_length_it_claims_is_refused() {
let h = PostingHeader {
count: 3,
code_bytes: 8,
stuck: 0,
};
let mut buf = vec![0u8; posting_len(3, 8).unwrap()];
h.encode(&mut buf).unwrap();
for len in 0..buf.len() {
assert!(
PostingHeader::decode(&buf[..len]).is_err(),
"{len} bytes decoded as three members"
);
}
let mut set = buf.clone();
set[12] = 1;
assert!(PostingHeader::decode(&set).is_err(), "reserved");
let mut huge = buf.clone();
put_u32(&mut huge, 0, u32::MAX);
put_u32(&mut huge, 4, u32::MAX);
assert!(PostingHeader::decode(&huge).is_err());
}
#[test]
fn floats_go_down_and_come_back_bit_for_bit() {
let values = [0.0f32, -0.0, 1.5, -2.25, 1e-38, 3.4e38];
let mut buf = vec![0u8; values.len() * 4];
assert_eq!(put_floats(&mut buf, &values).unwrap(), buf.len());
let mut back = vec![0f32; values.len()];
get_floats(&buf, &mut back).unwrap();
for (a, b) in values.iter().zip(&back) {
assert_eq!(a.to_bits(), b.to_bits(), "{a} came back as {b}");
}
assert!(get_floats(&buf[..4], &mut back).is_err(), "not that many");
assert!(put_floats(&mut buf[..4], &values).is_err());
}
#[test]
fn the_key_table_walks_back_in_order() {
let entries: Vec<(u64, &[u8])> = vec![
(0, b"a".as_slice()),
(7, b"".as_slice()),
(3, b"a rather longer key than the first one".as_slice()),
];
let mut buf = Vec::new();
for (id, key) in &entries {
let mut one = vec![0u8; key_entry_len(key.len()).unwrap()];
let wrote = put_key(&mut one, *id, key).unwrap();
assert_eq!(wrote, one.len());
buf.extend_from_slice(&one);
}
let mut walk = Keys::new(&buf);
let got: Vec<(u64, &[u8])> = walk.by_ref().collect();
assert_eq!(got, entries);
assert!(walk.done(), "the walk left bytes behind");
}
#[test]
fn a_truncated_key_table_stops_rather_than_reading_past_it() {
let mut buf = vec![0u8; key_entry_len(4).unwrap()];
put_key(&mut buf, 1, b"abcd").unwrap();
assert!(Keys::new(&[]).done(), "no bytes is an empty table");
for len in 1..buf.len() {
let mut walk = Keys::new(&buf[..len]);
assert_eq!(walk.by_ref().count(), 0, "{len} bytes gave a whole key");
assert!(!walk.done(), "a short entry is not a finished table");
}
assert!(key_entry_len(70_000).is_err());
}
#[test]
fn the_layout_is_the_one_written_down() {
assert_eq!(IMAGE_HEADER_LEN, 100);
assert_eq!(PARTITION_ENTRY_LEN, 16);
assert_eq!(POSTING_HEADER_LEN, 16);
assert_eq!(META_LEN, 16);
assert_eq!(IMAGE_TAG, u32::from_le_bytes(*b"YOIX"));
assert_eq!(image_len(0).unwrap(), IMAGE_HEADER_LEN);
assert_eq!(image_len(1).unwrap(), IMAGE_HEADER_LEN + 16);
assert!(
posting_len(256, 96).unwrap() < 64 * 1024,
"a partition should be one chunk at the sizes it is tuned for"
);
}
}