#![allow(dead_code)]
use crate::error::{LaurusError, Result};
use crate::lexical::index::structures::dictionary::TermInfo;
use super::block_max_data::BlockMaxData;
use super::front_coding::FrontCodingDecoder;
use super::term_info_block::FixedTermInfoBlock;
pub(super) struct BlockReader<'a> {
pub(super) block_term_count: u16,
pub(super) term_bytes: &'a [u8],
pub(super) term_info_block: FixedTermInfoBlock,
pub(super) block_max_data: BlockMaxData,
}
impl<'a> BlockReader<'a> {
pub(super) fn parse(bytes: &'a [u8], cursor: usize) -> Result<(Self, usize)> {
let mut c = cursor;
let term_count_u64 = read_varint(bytes, &mut c)?;
let term_count = u16::try_from(term_count_u64).map_err(|_| {
LaurusError::index(format!(
"BlockReader: term_count {term_count_u64} exceeds u16::MAX"
))
})?;
if term_count == 0 {
return Err(LaurusError::index("BlockReader: empty block"));
}
let fc_len = read_varint(bytes, &mut c)? as usize;
if c + fc_len > bytes.len() {
return Err(LaurusError::index(
"BlockReader: front-coded section overruns BlockSection",
));
}
let term_bytes = &bytes[c..c + fc_len];
c += fc_len;
const FTIB_FIXED_HEADER_BYTES: usize = 8 + 8 + 8 + 8 + 4 + 5 + 4;
if c + FTIB_FIXED_HEADER_BYTES > bytes.len() {
return Err(LaurusError::index(
"BlockReader: FixedTermInfoBlock header overruns BlockSection",
));
}
let block_min_posting_offset = read_u64_le(bytes, &mut c);
let block_min_posting_length = read_u64_le(bytes, &mut c);
let block_min_doc_frequency = read_u64_le(bytes, &mut c);
let block_min_total_frequency = read_u64_le(bytes, &mut c);
let ref_max_score_factor = read_f32_le(bytes, &mut c);
let posting_offset_nbits = read_u8(bytes, &mut c);
let posting_length_nbits = read_u8(bytes, &mut c);
let doc_frequency_nbits = read_u8(bytes, &mut c);
let total_frequency_nbits = read_u8(bytes, &mut c);
let max_score_factor_size = read_u8(bytes, &mut c);
let payload_len = read_u32_le(bytes, &mut c) as usize;
if c + payload_len > bytes.len() {
return Err(LaurusError::index(
"BlockReader: FixedTermInfoBlock payload overruns BlockSection",
));
}
let payload = bytes[c..c + payload_len].to_vec();
c += payload_len;
let term_info_block = FixedTermInfoBlock {
term_count,
block_min_posting_offset,
block_min_posting_length,
block_min_doc_frequency,
block_min_total_frequency,
ref_max_score_factor,
posting_offset_nbits,
posting_length_nbits,
doc_frequency_nbits,
total_frequency_nbits,
max_score_factor_size,
payload,
};
let offsets_count = term_count as usize + 1;
let offsets_bytes = offsets_count * 4;
if c + offsets_bytes + 4 > bytes.len() {
return Err(LaurusError::index(
"BlockReader: BlockMaxData header overruns BlockSection",
));
}
let mut offsets = Vec::with_capacity(offsets_count);
for _ in 0..offsets_count {
offsets.push(read_u32_le(bytes, &mut c));
}
let bm_data_len = read_u32_le(bytes, &mut c) as usize;
if c + bm_data_len > bytes.len() {
return Err(LaurusError::index(
"BlockReader: BlockMaxData data overruns BlockSection",
));
}
let bm_data = bytes[c..c + bm_data_len].to_vec();
c += bm_data_len;
let block_max_data = BlockMaxData {
offsets,
data: bm_data,
};
Ok((
BlockReader {
block_term_count: term_count,
term_bytes,
term_info_block,
block_max_data,
},
c,
))
}
pub(super) fn lookup(&self, target: &[u8]) -> Option<TermInfo> {
let mut decoder = FrontCodingDecoder::new(self.term_bytes, self.block_term_count as u32);
let mut inner: usize = 0;
while let Some(term) = decoder.next() {
match term.cmp(target) {
std::cmp::Ordering::Equal => {
return Some(self.materialise_term_info(inner));
}
std::cmp::Ordering::Greater => return None,
std::cmp::Ordering::Less => {
inner += 1;
}
}
}
None
}
pub(super) fn iter(&self) -> BlockReaderIter<'_, 'a> {
BlockReaderIter {
decoder: FrontCodingDecoder::new(self.term_bytes, self.block_term_count as u32),
term_info_block: &self.term_info_block,
block_max_data: &self.block_max_data,
inner: 0,
}
}
fn materialise_term_info(&self, inner_offset: usize) -> TermInfo {
let fti = self.term_info_block.decode_at(inner_offset);
let block_max = self.block_max_data.get(inner_offset);
TermInfo {
posting_offset: fti.posting_offset,
posting_length: fti.posting_length,
doc_frequency: fti.doc_frequency,
total_frequency: fti.total_frequency,
max_score_factor: fti.max_score_factor,
block_max,
}
}
}
pub(super) struct BlockReaderIter<'r, 'a: 'r> {
decoder: FrontCodingDecoder<'a>,
term_info_block: &'r FixedTermInfoBlock,
block_max_data: &'r BlockMaxData,
inner: usize,
}
impl<'r, 'a: 'r> Iterator for BlockReaderIter<'r, 'a> {
type Item = (Vec<u8>, TermInfo);
fn next(&mut self) -> Option<Self::Item> {
let term = self.decoder.next()?;
let term_owned = term.to_vec();
let fti = self.term_info_block.decode_at(self.inner);
let block_max = self.block_max_data.get(self.inner);
let info = TermInfo {
posting_offset: fti.posting_offset,
posting_length: fti.posting_length,
doc_frequency: fti.doc_frequency,
total_frequency: fti.total_frequency,
max_score_factor: fti.max_score_factor,
block_max,
};
self.inner += 1;
Some((term_owned, info))
}
}
pub(super) struct BlockSectionIter<'a> {
bytes: &'a [u8],
cursor: usize,
blocks_remaining: u32,
current: Option<CurrentBlock<'a>>,
}
struct CurrentBlock<'a> {
block_term_count: u16,
decoder: FrontCodingDecoder<'a>,
term_info_block: FixedTermInfoBlock,
block_max_data: BlockMaxData,
inner: usize,
}
impl<'a> BlockSectionIter<'a> {
pub(super) fn new(bytes: &'a [u8], block_count: u32) -> Self {
BlockSectionIter {
bytes,
cursor: 0,
blocks_remaining: block_count,
current: None,
}
}
fn advance_block(&mut self) -> Result<bool> {
if self.blocks_remaining == 0 {
self.current = None;
return Ok(false);
}
let (reader, next_cursor) = BlockReader::parse(self.bytes, self.cursor)?;
self.cursor = next_cursor;
self.blocks_remaining -= 1;
let decoder = FrontCodingDecoder::new(reader.term_bytes, reader.block_term_count as u32);
self.current = Some(CurrentBlock {
block_term_count: reader.block_term_count,
decoder,
term_info_block: reader.term_info_block,
block_max_data: reader.block_max_data,
inner: 0,
});
Ok(true)
}
}
impl<'a> Iterator for BlockSectionIter<'a> {
type Item = (String, TermInfo);
fn next(&mut self) -> Option<Self::Item> {
loop {
if self.current.is_none() {
match self.advance_block() {
Ok(true) => {}
Ok(false) => return None,
Err(_) => return None, }
}
let cb = self.current.as_mut().unwrap();
if cb.inner >= cb.block_term_count as usize {
self.current = None;
continue;
}
let term_bytes = match cb.decoder.next() {
Some(b) => b,
None => {
self.current = None;
return None;
}
};
let term_string = String::from_utf8_lossy(term_bytes).into_owned();
let fti = cb.term_info_block.decode_at(cb.inner);
let block_max = cb.block_max_data.get(cb.inner);
let info = TermInfo {
posting_offset: fti.posting_offset,
posting_length: fti.posting_length,
doc_frequency: fti.doc_frequency,
total_frequency: fti.total_frequency,
max_score_factor: fti.max_score_factor,
block_max,
};
cb.inner += 1;
return Some((term_string, info));
}
}
}
fn read_varint(bytes: &[u8], cursor: &mut usize) -> Result<u64> {
let mut result: u64 = 0;
let mut shift: u32 = 0;
loop {
if *cursor >= bytes.len() {
return Err(LaurusError::index("BlockReader: truncated varint"));
}
let byte = bytes[*cursor];
*cursor += 1;
result |= u64::from(byte & 0x7F) << shift;
if byte & 0x80 == 0 {
return Ok(result);
}
shift += 7;
if shift >= 64 {
return Err(LaurusError::index("BlockReader: varint overflow"));
}
}
}
fn read_u64_le(bytes: &[u8], cursor: &mut usize) -> u64 {
let v = u64::from_le_bytes(
bytes[*cursor..*cursor + 8]
.try_into()
.expect("8-byte slice"),
);
*cursor += 8;
v
}
fn read_u32_le(bytes: &[u8], cursor: &mut usize) -> u32 {
let v = u32::from_le_bytes(
bytes[*cursor..*cursor + 4]
.try_into()
.expect("4-byte slice"),
);
*cursor += 4;
v
}
fn read_f32_le(bytes: &[u8], cursor: &mut usize) -> f32 {
let v = f32::from_le_bytes(
bytes[*cursor..*cursor + 4]
.try_into()
.expect("4-byte slice"),
);
*cursor += 4;
v
}
fn read_u8(bytes: &[u8], cursor: &mut usize) -> u8 {
let v = bytes[*cursor];
*cursor += 1;
v
}
#[cfg(test)]
mod tests {
use super::*;
use crate::lexical::index::structures::dictionary::BlockMax;
use super::super::front_coding::encode_block_terms;
use super::super::term_info_block::FixedTermInfo;
fn fti(po: u64, pl: u64, df: u64, tf: u64, msf: f32) -> FixedTermInfo {
FixedTermInfo {
posting_offset: po,
posting_length: pl,
doc_frequency: df,
total_frequency: tf,
max_score_factor: msf,
}
}
fn sample_block_max(doc: u64, factor: f32) -> BlockMax {
BlockMax {
last_doc_id: doc,
max_factor: factor,
}
}
fn encode_block_into(
out: &mut Vec<u8>,
terms: &[&[u8]],
infos: &[(FixedTermInfo, Vec<BlockMax>)],
) {
assert_eq!(terms.len(), infos.len());
let term_count = terms.len();
write_varint_into(out, term_count as u64);
let fc = encode_block_terms(terms);
write_varint_into(out, fc.len() as u64);
out.extend_from_slice(&fc);
let fixed: Vec<FixedTermInfo> = infos.iter().map(|(f, _)| *f).collect();
let ftib = FixedTermInfoBlock::encode(&fixed);
out.extend_from_slice(&ftib.block_min_posting_offset.to_le_bytes());
out.extend_from_slice(&ftib.block_min_posting_length.to_le_bytes());
out.extend_from_slice(&ftib.block_min_doc_frequency.to_le_bytes());
out.extend_from_slice(&ftib.block_min_total_frequency.to_le_bytes());
out.extend_from_slice(&ftib.ref_max_score_factor.to_le_bytes());
out.push(ftib.posting_offset_nbits);
out.push(ftib.posting_length_nbits);
out.push(ftib.doc_frequency_nbits);
out.push(ftib.total_frequency_nbits);
out.push(ftib.max_score_factor_size);
out.extend_from_slice(&(ftib.payload.len() as u32).to_le_bytes());
out.extend_from_slice(&ftib.payload);
let per_term: Vec<Vec<BlockMax>> = infos.iter().map(|(_, bm)| bm.clone()).collect();
let bmd = BlockMaxData::encode(&per_term);
for off in &bmd.offsets {
out.extend_from_slice(&off.to_le_bytes());
}
out.extend_from_slice(&(bmd.data.len() as u32).to_le_bytes());
out.extend_from_slice(&bmd.data);
}
fn write_varint_into(buf: &mut Vec<u8>, mut value: u64) {
while value >= 0x80 {
buf.push(((value as u8) & 0x7F) | 0x80);
value >>= 7;
}
buf.push(value as u8);
}
fn build_two_block_section() -> (Vec<u8>, Vec<(Vec<u8>, TermInfo)>) {
let block0_terms: Vec<&[u8]> = vec![b"alpha", b"beta", b"gamma"];
let block0_infos = vec![
(
fti(0, 100, 5, 50, 1.5),
vec![sample_block_max(10, 1.5), sample_block_max(20, 2.5)],
),
(fti(100, 80, 3, 40, 1.5), Vec::new()),
(fti(180, 200, 8, 100, 2.0), vec![sample_block_max(30, 3.5)]),
];
let block1_terms: Vec<&[u8]> = vec![b"kappa", b"lambda", b"mu"];
let block1_infos = vec![
(fti(380, 50, 1, 1, 0.5), Vec::new()),
(fti(430, 70, 2, 5, 0.5), Vec::new()),
(fti(500, 90, 4, 8, 0.5), vec![sample_block_max(99, 9.9)]),
];
let mut section = Vec::new();
encode_block_into(&mut section, &block0_terms, &block0_infos);
encode_block_into(&mut section, &block1_terms, &block1_infos);
let mut expected: Vec<(Vec<u8>, TermInfo)> = Vec::new();
for (term, (f, bm)) in block0_terms.iter().zip(block0_infos.iter()) {
expected.push((
term.to_vec(),
TermInfo {
posting_offset: f.posting_offset,
posting_length: f.posting_length,
doc_frequency: f.doc_frequency,
total_frequency: f.total_frequency,
max_score_factor: f.max_score_factor,
block_max: bm.clone(),
},
));
}
for (term, (f, bm)) in block1_terms.iter().zip(block1_infos.iter()) {
expected.push((
term.to_vec(),
TermInfo {
posting_offset: f.posting_offset,
posting_length: f.posting_length,
doc_frequency: f.doc_frequency,
total_frequency: f.total_frequency,
max_score_factor: f.max_score_factor,
block_max: bm.clone(),
},
));
}
(section, expected)
}
#[test]
fn parse_single_block_round_trip() {
let terms: Vec<&[u8]> = vec![b"apple", b"banana", b"cherry"];
let infos = vec![
(fti(0, 50, 1, 10, 1.0), Vec::new()),
(fti(50, 60, 2, 20, 1.0), vec![sample_block_max(5, 0.5)]),
(fti(110, 70, 3, 30, 1.0), Vec::new()),
];
let mut section = Vec::new();
encode_block_into(&mut section, &terms, &infos);
let (reader, next_cursor) = BlockReader::parse(§ion, 0).unwrap();
assert_eq!(reader.block_term_count, 3);
assert_eq!(next_cursor, section.len());
let info = reader.lookup(b"banana").unwrap();
assert_eq!(info.posting_offset, 50);
assert_eq!(info.block_max.len(), 1);
assert_eq!(info.block_max[0].last_doc_id, 5);
assert!(reader.lookup(b"aardvark").is_none());
assert!(reader.lookup(b"blueberry").is_none());
assert!(reader.lookup(b"date").is_none());
}
#[test]
fn block_reader_iter_yields_in_order() {
let terms: Vec<&[u8]> = vec![b"alpha", b"beta", b"gamma"];
let infos = vec![
(fti(0, 100, 5, 50, 1.5), Vec::new()),
(fti(100, 80, 3, 40, 1.5), Vec::new()),
(fti(180, 200, 8, 100, 2.0), Vec::new()),
];
let mut section = Vec::new();
encode_block_into(&mut section, &terms, &infos);
let (reader, _) = BlockReader::parse(§ion, 0).unwrap();
let collected: Vec<(Vec<u8>, TermInfo)> = reader.iter().collect();
assert_eq!(collected.len(), 3);
assert_eq!(collected[0].0, b"alpha");
assert_eq!(collected[1].0, b"beta");
assert_eq!(collected[2].0, b"gamma");
}
#[test]
fn block_section_iter_walks_two_blocks() {
let (section, expected) = build_two_block_section();
let iter = BlockSectionIter::new(§ion, 2);
let collected: Vec<(String, TermInfo)> = iter.collect();
assert_eq!(collected.len(), expected.len());
for ((got_term, got_info), (exp_term, exp_info)) in collected.iter().zip(expected.iter()) {
assert_eq!(got_term.as_bytes(), exp_term.as_slice());
assert_eq!(got_info.posting_offset, exp_info.posting_offset);
assert_eq!(got_info.posting_length, exp_info.posting_length);
assert_eq!(got_info.doc_frequency, exp_info.doc_frequency);
assert_eq!(got_info.total_frequency, exp_info.total_frequency);
assert_eq!(
got_info.max_score_factor.to_bits(),
exp_info.max_score_factor.to_bits()
);
assert_eq!(got_info.block_max.len(), exp_info.block_max.len());
}
}
#[test]
fn parse_rejects_truncated_section() {
let terms: Vec<&[u8]> = vec![b"alpha", b"beta"];
let infos = vec![
(fti(0, 50, 1, 10, 1.0), Vec::new()),
(fti(50, 60, 2, 20, 1.0), Vec::new()),
];
let mut section = Vec::new();
encode_block_into(&mut section, &terms, &infos);
let truncated = §ion[..section.len() - 5];
assert!(BlockReader::parse(truncated, 0).is_err());
}
#[test]
fn empty_section_iter() {
let iter = BlockSectionIter::new(&[], 0);
let collected: Vec<(String, TermInfo)> = iter.collect();
assert!(collected.is_empty());
}
#[test]
fn varied_term_lengths_in_block() {
let mut long_term = b"long_".to_vec();
long_term.extend(std::iter::repeat_n(b'a', 100));
let terms: Vec<&[u8]> = vec![long_term.as_slice(), b"short"];
let infos = vec![
(fti(0, 50, 1, 1, 0.0), Vec::new()),
(fti(50, 60, 2, 2, 0.0), Vec::new()),
];
let mut section = Vec::new();
encode_block_into(&mut section, &terms, &infos);
let (reader, _) = BlockReader::parse(§ion, 0).unwrap();
assert!(reader.lookup(&long_term).is_some());
assert!(reader.lookup(b"short").is_some());
assert!(reader.lookup(b"medium").is_none());
}
#[test]
fn lookup_miss_after_last_block_term() {
let terms: Vec<&[u8]> = vec![b"alpha", b"beta"];
let infos = vec![
(fti(0, 50, 1, 1, 0.0), Vec::new()),
(fti(50, 60, 2, 2, 0.0), Vec::new()),
];
let mut section = Vec::new();
encode_block_into(&mut section, &terms, &infos);
let (reader, _) = BlockReader::parse(§ion, 0).unwrap();
assert!(reader.lookup(b"zulu").is_none());
}
}