use std::path::Path;
use crate::index::snapshot::ArtifactData;
use crate::index::trigram::storage::format::LEXICON_MAGIC;
use crate::index::mmap::mmap_open;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LexiconEntry {
pub trigram: [u8; 3],
pub offset: u64,
pub len: u32,
}
#[derive(Debug)]
pub struct Lexicon {
data: ArtifactData,
count: usize,
}
impl Lexicon {
const ENTRY_SIZE: usize = 15;
pub fn encode(entries: &[LexiconEntry]) -> std::io::Result<Vec<u8>> {
let mut data =
Vec::with_capacity(LEXICON_MAGIC.len() + 4 + entries.len() * Self::ENTRY_SIZE);
data.extend_from_slice(&LEXICON_MAGIC);
let count = u32::try_from(entries.len()).map_err(|_| {
std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"lexicon entry count exceeds u32::MAX",
)
})?;
data.extend_from_slice(&count.to_le_bytes());
for e in entries {
data.extend_from_slice(&e.trigram);
data.extend_from_slice(&e.offset.to_le_bytes());
data.extend_from_slice(&e.len.to_le_bytes());
}
Ok(data)
}
fn bytes(&self) -> &[u8] {
self.data.as_ref()
}
pub fn from_artifact(data: ArtifactData) -> std::io::Result<Self> {
let bytes = data.as_ref();
let count = Self::validate(bytes)?;
Ok(Self { data, count })
}
pub fn create(path: &Path, entries: &[LexiconEntry]) -> std::io::Result<Self> {
let data = Self::encode(entries)?;
std::fs::write(path, &data)?;
Self::open(path)
}
pub fn open(path: &Path) -> std::io::Result<Self> {
let mmap = mmap_open(path)?;
Self::from_artifact(ArtifactData::Mmap(mmap))
}
fn validate(bytes: &[u8]) -> std::io::Result<usize> {
let magic_len = LEXICON_MAGIC.len();
if bytes.len() < magic_len + 4 {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"lexicon too short for magic+count",
));
}
if bytes[..magic_len] != LEXICON_MAGIC {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"unexpected lexicon magic",
));
}
let n = u32::from_le_bytes(bytes[magic_len..magic_len + 4].try_into().unwrap()) as usize;
let expected_bytes = n * Self::ENTRY_SIZE;
if bytes.len() < magic_len + 4 + expected_bytes {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"lexicon truncated",
));
}
let entries = &bytes[magic_len + 4..];
let mut prev: Option<([u8; 3], u64)> = None;
for chunk in entries.chunks_exact(Self::ENTRY_SIZE) {
let tri: [u8; 3] = chunk[..3].try_into().unwrap();
let posting_off = u64::from_le_bytes(chunk[3..11].try_into().unwrap());
if let Some((prev_tri, prev_off)) = prev {
if tri <= prev_tri {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
format!("lexicon trigram {tri:?} out of order (prev {prev_tri:?})"),
));
}
if posting_off < prev_off {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
format!(
"lexicon entry {tri:?} posting offset {posting_off} less than previous {prev_off}",
),
));
}
}
prev = Some((tri, posting_off));
}
Ok(n)
}
#[must_use]
pub const fn len(&self) -> usize {
self.count
}
#[must_use]
pub const fn is_empty(&self) -> bool {
self.count == 0
}
#[must_use]
pub fn get(&self, tri: [u8; 3]) -> Option<LexiconEntry> {
if self.count == 0 {
return None;
}
let bytes = self.bytes();
let magic_len = LEXICON_MAGIC.len();
let data_start = magic_len + 4;
let mut lo = 0;
let mut hi = self.count;
while lo < hi {
let mid = lo + (hi - lo) / 2;
let offset = data_start + mid * Self::ENTRY_SIZE;
let entry_tri: [u8; 3] = bytes[offset..offset + 3].try_into().unwrap();
match entry_tri.cmp(&tri) {
std::cmp::Ordering::Less => lo = mid + 1,
std::cmp::Ordering::Greater => hi = mid,
std::cmp::Ordering::Equal => {
let off =
u64::from_le_bytes(bytes[offset + 3..offset + 11].try_into().unwrap());
let len =
u32::from_le_bytes(bytes[offset + 11..offset + 15].try_into().unwrap());
return Some(LexiconEntry {
trigram: entry_tri,
offset: off,
len,
});
}
}
}
None
}
#[must_use]
pub fn posting_byte_end(&self, offset: u64, payload_len: usize) -> usize {
if self.count == 0 {
return payload_len;
}
let bytes = self.bytes();
let data_start = LEXICON_MAGIC.len() + 4;
let mut lo = 0usize;
let mut hi = self.count;
while lo < hi {
let mid = lo + (hi - lo) / 2;
let off = data_start + mid * Self::ENTRY_SIZE + 3;
let entry_off =
u64::from_le_bytes(bytes[off..off + 8].try_into().expect("entry offset"));
if entry_off <= offset {
lo = mid + 1;
} else {
hi = mid;
}
}
if lo < self.count {
let off = data_start + lo * Self::ENTRY_SIZE + 3;
usize::try_from(u64::from_le_bytes(
bytes[off..off + 8].try_into().expect("entry offset"),
))
.unwrap_or(payload_len)
} else {
payload_len
}
}
#[must_use]
pub const fn iter(&self) -> LexiconIter<'_> {
LexiconIter {
lexicon: self,
pos: 0,
}
}
}
impl<'a> IntoIterator for &'a Lexicon {
type Item = LexiconEntry;
type IntoIter = LexiconIter<'a>;
fn into_iter(self) -> Self::IntoIter {
LexiconIter {
lexicon: self,
pos: 0,
}
}
}
pub struct LexiconIter<'_a> {
lexicon: &'_a Lexicon,
pos: usize,
}
impl Iterator for LexiconIter<'_> {
type Item = LexiconEntry;
fn next(&mut self) -> Option<Self::Item> {
if self.pos >= self.lexicon.count {
return None;
}
let bytes = self.lexicon.bytes();
let magic_len = LEXICON_MAGIC.len();
let offset = magic_len + 4 + self.pos * Lexicon::ENTRY_SIZE;
let tri: [u8; 3] = bytes[offset..offset + 3].try_into().unwrap();
let off = u64::from_le_bytes(bytes[offset + 3..offset + 11].try_into().unwrap());
let len = u32::from_le_bytes(bytes[offset + 11..offset + 15].try_into().unwrap());
self.pos += 1;
Some(LexiconEntry {
trigram: tri,
offset: off,
len,
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Write;
use tempfile::TempDir;
fn make_entry(tri: [u8; 3], offset: u64, len: u32) -> LexiconEntry {
LexiconEntry {
trigram: tri,
offset,
len,
}
}
fn create_lexicon(entries: &[LexiconEntry]) -> Lexicon {
let tmp = TempDir::new().expect("create temp dir");
let path = tmp.path().join("lexicon.bin");
Lexicon::create(&path, entries).expect("create lexicon")
}
#[test]
fn create_sets_len() {
let entries = vec![make_entry(*b"abc", 0, 4), make_entry(*b"def", 4, 8)];
let lexicon = create_lexicon(&entries);
assert_eq!(lexicon.len(), 2);
}
#[test]
fn empty_lexicon_reports_is_empty() {
let lexicon = create_lexicon(&[]);
assert!(lexicon.is_empty());
assert_eq!(lexicon.len(), 0);
}
#[test]
fn get_finds_first_middle_and_last() {
let entries = vec![
make_entry(*b"aaa", 0, 4),
make_entry(*b"bbb", 4, 4),
make_entry(*b"ccc", 8, 4),
];
let lexicon = create_lexicon(&entries);
assert!(lexicon.get(*b"aaa").is_some());
assert!(lexicon.get(*b"bbb").is_some());
assert!(lexicon.get(*b"ccc").is_some());
}
#[test]
fn get_returns_none_for_absent_trigram() {
let entries = vec![make_entry(*b"abc", 0, 4)];
let lexicon = create_lexicon(&entries);
assert!(lexicon.get(*b"xyz").is_none());
}
#[test]
fn get_returns_none_for_empty_lexicon() {
let lexicon = create_lexicon(&[]);
assert!(lexicon.get(*b"abc").is_none());
}
#[test]
fn open_rejects_duplicate_trigrams() {
let tmp = TempDir::new().expect("create temp dir");
let path = tmp.path().join("lexicon.bin");
let mut data = LEXICON_MAGIC.to_vec();
data.extend_from_slice(&2u32.to_le_bytes()); data.extend_from_slice(b"aaa");
data.extend_from_slice(&0u64.to_le_bytes());
data.extend_from_slice(&1u32.to_le_bytes());
data.extend_from_slice(b"aaa"); data.extend_from_slice(&1u64.to_le_bytes());
data.extend_from_slice(&1u32.to_le_bytes());
std::fs::write(&path, &data).expect("write");
let result = Lexicon::open(&path);
assert!(result.is_err());
}
#[test]
fn open_rejects_out_of_order_trigrams() {
let tmp = TempDir::new().expect("create temp dir");
let path = tmp.path().join("lexicon.bin");
let mut data = LEXICON_MAGIC.to_vec();
data.extend_from_slice(&2u32.to_le_bytes()); data.extend_from_slice(b"bbb");
data.extend_from_slice(&0u64.to_le_bytes());
data.extend_from_slice(&1u32.to_le_bytes());
data.extend_from_slice(b"aaa"); data.extend_from_slice(&1u64.to_le_bytes());
data.extend_from_slice(&1u32.to_le_bytes());
std::fs::write(&path, &data).expect("write");
let result = Lexicon::open(&path);
assert!(result.is_err());
}
#[test]
fn open_rejects_non_monotonic_posting_offsets() {
let tmp = TempDir::new().expect("create temp dir");
let path = tmp.path().join("lexicon.bin");
let mut data = LEXICON_MAGIC.to_vec();
data.extend_from_slice(&2u32.to_le_bytes()); data.extend_from_slice(b"aaa");
data.extend_from_slice(&10u64.to_le_bytes());
data.extend_from_slice(&1u32.to_le_bytes());
data.extend_from_slice(b"bbb");
data.extend_from_slice(&5u64.to_le_bytes());
data.extend_from_slice(&1u32.to_le_bytes());
std::fs::write(&path, &data).expect("write");
let result = Lexicon::open(&path);
assert!(result.is_err());
}
#[test]
fn iter_yields_entries_in_stored_order() {
let entries = vec![
make_entry(*b"aaa", 0, 4),
make_entry(*b"bbb", 4, 4),
make_entry(*b"ccc", 8, 4),
];
let lexicon = create_lexicon(&entries);
let collected: Vec<_> = lexicon.iter().collect();
assert_eq!(collected.len(), 3);
assert_eq!(collected[0].trigram, *b"aaa");
assert_eq!(collected[1].trigram, *b"bbb");
assert_eq!(collected[2].trigram, *b"ccc");
}
#[test]
fn into_iterator_for_ref_works() {
let entries = vec![make_entry(*b"abc", 0, 4)];
let lexicon = create_lexicon(&entries);
let collected: Vec<_> = (&lexicon).into_iter().collect();
assert_eq!(collected.len(), 1);
assert_eq!(collected[0].trigram, *b"abc");
}
#[test]
fn open_rejects_bad_magic() {
let tmp = TempDir::new().expect("create temp dir");
let path = tmp.path().join("lexicon.bin");
let mut file = std::fs::File::create(&path).expect("create file");
file.write_all(b"BADMAGIC").expect("write bad magic");
file.write_all(&0u32.to_le_bytes()).expect("write count");
let result = Lexicon::open(&path);
assert!(result.is_err());
}
#[test]
fn open_rejects_truncated_entry_data() {
let tmp = TempDir::new().expect("create temp dir");
let path = tmp.path().join("lexicon.bin");
let mut file = std::fs::File::create(&path).expect("create file");
file.write_all(&LEXICON_MAGIC).expect("write magic");
file.write_all(&1u32.to_le_bytes()).expect("write count 1");
file.write_all(&[0u8; 8]).expect("write only 8 of 15 bytes");
let result = Lexicon::open(&path);
assert!(result.is_err());
}
#[test]
fn open_rejects_truncated_magic() {
let tmp = TempDir::new().expect("create temp dir");
let path = tmp.path().join("lexicon.bin");
std::fs::write(&path, b"SHORT").expect("write short file");
let result = Lexicon::open(&path);
assert!(result.is_err());
}
#[test]
fn get_returns_correct_offset_and_len() {
let entries = vec![make_entry(*b"aaa", 100, 12), make_entry(*b"bbb", 200, 8)];
let lexicon = create_lexicon(&entries);
let entry = lexicon.get(*b"aaa").expect("find aaa");
assert_eq!(entry.offset, 100);
assert_eq!(entry.len, 12);
let entry = lexicon.get(*b"bbb").expect("find bbb");
assert_eq!(entry.offset, 200);
assert_eq!(entry.len, 8);
}
}