use std::collections::HashMap;
use std::time::Instant;
use ulid::Ulid;
use crate::error::{Error, Result};
use crate::format::{PAGE_HEADER_LEN, PageHeader, PageType, stamp_page_checksum};
use crate::index::dict;
use crate::storage::btree::PageSource;
use crate::storage::pager::Txn;
const BM25_K1: f32 = 1.2;
const BM25_B: f32 = 0.75;
const MAX_TERM_LEN: usize = 128;
const FTS_DICT: dict::DictSpec = dict::DictSpec {
dict: PageType::FtsDict,
overflow: PageType::FtsPostings,
max_key_len: MAX_TERM_LEN,
};
const POSTING_LEN: usize = 20;
const DELTA_VARINT_MIN_FORMAT_VERSION: u32 = 4;
const SKIP_MIN_FORMAT_VERSION: u32 = 5;
const SKIP_BOUND_MIN_FORMAT_VERSION: u32 = 6;
const FRAME_OF_REFERENCE_MIN_FORMAT_VERSION: u32 = 8;
const SKIP_BLOCK_SIZE: usize = 128;
const SKIP_MIN_DOC_FREQ: usize = 4 * SKIP_BLOCK_SIZE;
const SKIP_ENTRY_LEN_V5: usize = 24;
const SKIP_ENTRY_LEN_V6: usize = 40;
const MAX_VARINT_LEN: usize = 19;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum PostingsLayout {
FixedWidth,
DeltaVarint,
DeltaVarintSkip(SkipEntry),
FrameOfReference,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum SkipEntry {
V5,
V6,
}
impl SkipEntry {
fn len(self) -> usize {
match self {
SkipEntry::V5 => SKIP_ENTRY_LEN_V5,
SkipEntry::V6 => SKIP_ENTRY_LEN_V6,
}
}
fn tail_off(self) -> usize {
match self {
SkipEntry::V5 => 16,
SkipEntry::V6 => 32,
}
}
}
impl PostingsLayout {
fn for_format_version(version: u32) -> Self {
if version >= FRAME_OF_REFERENCE_MIN_FORMAT_VERSION {
PostingsLayout::FrameOfReference
} else if version >= SKIP_BOUND_MIN_FORMAT_VERSION {
PostingsLayout::DeltaVarintSkip(SkipEntry::V6)
} else if version >= SKIP_MIN_FORMAT_VERSION {
PostingsLayout::DeltaVarintSkip(SkipEntry::V5)
} else if version >= DELTA_VARINT_MIN_FORMAT_VERSION {
PostingsLayout::DeltaVarint
} else {
PostingsLayout::FixedWidth
}
}
}
fn put_varint(out: &mut Vec<u8>, mut v: u128) {
loop {
let byte = (v & 0x7F) as u8;
v >>= 7;
if v == 0 {
out.push(byte);
return;
}
out.push(byte | 0x80);
}
}
fn read_varint(body: &[u8], off: &mut usize, page_no: u64) -> Result<u128> {
let mut value: u128 = 0;
for i in 0..MAX_VARINT_LEN {
let byte = *body
.get(*off + i)
.ok_or_else(|| malformed(page_no, "fts varint truncated"))?;
let bits = u128::from(byte & 0x7F);
let shift = 7 * i as u32;
if bits != 0 && (bits << shift) >> shift != bits {
return Err(malformed(page_no, "fts varint overflow"));
}
value |= bits << shift;
if byte & 0x80 == 0 {
*off += i + 1;
return Ok(value);
}
}
Err(malformed(page_no, "fts varint too long"))
}
fn malformed(page_no: u64, what: &'static str) -> Error {
Error::MalformedPage { page_no, what }
}
fn read_u128_le(body: &[u8], off: usize, page_no: u64) -> Result<u128> {
let bytes: [u8; 16] = body
.get(off..off + 16)
.and_then(|b| b.try_into().ok())
.ok_or_else(|| malformed(page_no, "fts u128 truncated"))?;
Ok(u128::from_le_bytes(bytes))
}
fn encode_delta_run(out: &mut Vec<u8>, entries: &[Posting]) {
let mut prev: u128 = 0;
for p in entries {
let id = u128::from(p.record_id);
put_varint(out, id.wrapping_sub(prev));
put_varint(out, u128::from(p.term_freq));
prev = id;
}
}
fn reject_hostile_delta_count(body: &[u8], count: usize, page_no: u64) -> Result<()> {
let min_need = 4usize
.checked_add(
count
.checked_mul(2)
.ok_or_else(|| malformed(page_no, "fts postings count overflow"))?,
)
.ok_or_else(|| malformed(page_no, "fts postings length overflow"))?;
if body.len() < min_need {
return Err(malformed(page_no, "fts postings truncated"));
}
Ok(())
}
fn decode_delta_run(
body: &[u8],
off: &mut usize,
count: usize,
page_no: u64,
entries: &mut Vec<Posting>,
) -> Result<()> {
let mut prev: u128 = 0;
for i in 0..count {
let delta = read_varint(body, off, page_no)?;
if i > 0 && delta == 0 {
return Err(malformed(page_no, "unsorted fts postings"));
}
let id = if i == 0 {
delta
} else {
prev.checked_add(delta)
.ok_or_else(|| malformed(page_no, "fts posting id overflow"))?
};
prev = id;
let term_freq = u32::try_from(read_varint(body, off, page_no)?)
.map_err(|_| malformed(page_no, "fts posting term_freq overflow"))?;
if term_freq == 0 {
return Err(malformed(page_no, "fts posting zero term_freq"));
}
entries.push(Posting {
record_id: Ulid::from(id),
term_freq,
});
}
Ok(())
}
fn decode_block_run(
layout: PostingsLayout,
body: &[u8],
off: &mut usize,
count: usize,
page_no: u64,
entries: &mut Vec<Posting>,
) -> Result<()> {
match layout {
PostingsLayout::FrameOfReference => decode_for_block(body, off, count, page_no, entries),
_ => decode_delta_run(body, off, count, page_no, entries),
}
}
fn byte_width(v: u128) -> usize {
let bits = 128 - v.leading_zeros() as usize;
bits.div_ceil(8).max(1)
}
fn put_fixed(out: &mut Vec<u8>, v: u128, width: usize) {
out.extend_from_slice(&v.to_le_bytes()[..width]);
}
fn read_fixed(body: &[u8], off: &mut usize, width: usize, page_no: u64) -> Result<u128> {
let end = off
.checked_add(width)
.ok_or_else(|| malformed(page_no, "fts fixed-width overflow"))?;
let slice = body
.get(*off..end)
.ok_or_else(|| malformed(page_no, "fts fixed-width truncated"))?;
let mut buf = [0u8; 16];
buf[..width].copy_from_slice(slice);
*off = end;
Ok(u128::from_le_bytes(buf))
}
fn encode_for_block(out: &mut Vec<u8>, entries: &[Posting]) {
let mut prev: u128 = 0;
let mut max_delta: u128 = 0;
let mut max_tf: u32 = 0;
for p in entries {
let id = u128::from(p.record_id);
max_delta = max_delta.max(id.wrapping_sub(prev));
max_tf = max_tf.max(p.term_freq);
prev = id;
}
let delta_width = byte_width(max_delta);
let tf_width = byte_width(u128::from(max_tf));
out.push(delta_width as u8);
out.push(tf_width as u8);
let mut prev: u128 = 0;
for p in entries {
let id = u128::from(p.record_id);
put_fixed(out, id.wrapping_sub(prev), delta_width);
prev = id;
}
for p in entries {
put_fixed(out, u128::from(p.term_freq), tf_width);
}
}
fn decode_for_block(
body: &[u8],
off: &mut usize,
count: usize,
page_no: u64,
entries: &mut Vec<Posting>,
) -> Result<()> {
let delta_width = *body
.get(*off)
.ok_or_else(|| malformed(page_no, "fts for widths truncated"))?
as usize;
let tf_width = *body
.get(*off + 1)
.ok_or_else(|| malformed(page_no, "fts for widths truncated"))? as usize;
if !(1..=16).contains(&delta_width) || !(1..=4).contains(&tf_width) {
return Err(malformed(page_no, "fts for width out of range"));
}
let stream_bytes = count
.checked_mul(delta_width + tf_width)
.ok_or_else(|| malformed(page_no, "fts for stream overflow"))?;
let need = off
.checked_add(2)
.and_then(|h| h.checked_add(stream_bytes))
.ok_or_else(|| malformed(page_no, "fts for stream overflow"))?;
if body.len() < need {
return Err(malformed(page_no, "fts for stream truncated"));
}
let mut delta_off = *off + 2;
let mut tf_off = delta_off + count * delta_width;
let mut prev: u128 = 0;
for i in 0..count {
let delta = read_fixed(body, &mut delta_off, delta_width, page_no)?;
if i > 0 && delta == 0 {
return Err(malformed(page_no, "unsorted fts postings"));
}
let id = if i == 0 {
delta
} else {
prev.checked_add(delta)
.ok_or_else(|| malformed(page_no, "fts posting id overflow"))?
};
prev = id;
let term_freq = u32::try_from(read_fixed(body, &mut tf_off, tf_width, page_no)?)
.map_err(|_| malformed(page_no, "fts posting term_freq overflow"))?;
if term_freq == 0 {
return Err(malformed(page_no, "fts posting zero term_freq"));
}
entries.push(Posting {
record_id: Ulid::from(id),
term_freq,
});
}
*off = tf_off;
Ok(())
}
fn lookup_via_skip(
body: &[u8],
page_no: u64,
target: Ulid,
layout: PostingsLayout,
) -> Result<Option<u32>> {
let entry = match layout {
PostingsLayout::DeltaVarintSkip(e) => e,
PostingsLayout::FrameOfReference => SkipEntry::V6,
_ => return Err(malformed(page_no, "fts skip lookup on non-skip layout")),
};
let count = dict::read_u32(body, 0, page_no)? as usize;
let block_count = dict::read_u32(body, 4, page_no)? as usize;
if block_count == 0 {
let decoded = Postings::decode(body, page_no, layout)?;
return Ok(decoded
.entries
.binary_search_by(|p| p.record_id.cmp(&target))
.ok()
.map(|i| decoded.entries[i].term_freq));
}
let expected_blocks = count.div_ceil(SKIP_BLOCK_SIZE);
if count < SKIP_MIN_DOC_FREQ || block_count != expected_blocks {
return Err(malformed(page_no, "fts skip block_count mismatch"));
}
let entry_len = entry.len();
let index_len = block_count
.checked_mul(entry_len)
.ok_or_else(|| malformed(page_no, "fts skip index overflow"))?;
let blocks_start = 8usize
.checked_add(index_len)
.ok_or_else(|| malformed(page_no, "fts skip index overflow"))?;
if body.len() < blocks_start {
return Err(malformed(page_no, "fts skip index truncated"));
}
let target_id = u128::from(target);
let first_id_of = |b: usize| -> Result<u128> { read_u128_le(body, 8 + b * entry_len, page_no) };
let mut lo = 0usize;
let mut hi = block_count;
while lo < hi {
let mid = lo + (hi - lo) / 2;
if first_id_of(mid)? <= target_id {
lo = mid + 1;
} else {
hi = mid;
}
}
if lo == 0 {
return Ok(None); }
let b = lo - 1;
let stored_offset =
dict::read_u32(body, 8 + b * entry_len + entry.tail_off(), page_no)? as usize;
let block_len = if b + 1 == block_count {
count - b * SKIP_BLOCK_SIZE
} else {
SKIP_BLOCK_SIZE
};
let mut block = Vec::with_capacity(block_len);
let mut off = blocks_start
.checked_add(stored_offset)
.ok_or_else(|| malformed(page_no, "fts skip byte_offset overflow"))?;
decode_block_run(layout, body, &mut off, block_len, page_no, &mut block)?;
Ok(block
.binary_search_by(|p| p.record_id.cmp(&target))
.ok()
.map(|i| block[i].term_freq))
}
pub fn tokenize(text: &str) -> Vec<String> {
let mut out = Vec::new();
let mut cur = String::new();
for ch in text.chars() {
if ch.is_alphanumeric() {
cur.extend(ch.to_lowercase());
} else if !cur.is_empty() {
out.push(std::mem::take(&mut cur));
}
}
if !cur.is_empty() {
out.push(cur);
}
out
}
pub fn doc_len(text: &str) -> u32 {
let mut count = 0u32;
let mut in_token = false;
for ch in text.chars() {
if ch.is_alphanumeric() {
in_token = true;
} else if in_token {
count = count.saturating_add(1);
in_token = false;
}
}
if in_token {
count = count.saturating_add(1);
}
count
}
fn clip_term(term: &str) -> &str {
if term.len() <= MAX_TERM_LEN {
return term;
}
let mut end = MAX_TERM_LEN;
while end > 0 && !term.is_char_boundary(end) {
end -= 1;
}
&term[..end]
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct FtsMeta {
doc_count: u64,
total_tokens: u64,
dict_root: u64,
}
impl FtsMeta {
fn empty() -> Self {
FtsMeta {
doc_count: 0,
total_tokens: 0,
dict_root: 0,
}
}
fn encode(&self, page_size: u32) -> Result<Vec<u8>> {
let mut page = vec![0u8; page_size as usize];
PageHeader {
page_type: PageType::FtsDict,
entry_count: 0,
next_page: 0,
}
.encode_into(&mut page);
let mut off = PAGE_HEADER_LEN;
page[off] = dict::NODE_META;
off += 1;
dict::put_u64(&mut page, &mut off, self.doc_count);
dict::put_u64(&mut page, &mut off, self.total_tokens);
dict::put_u64(&mut page, &mut off, self.dict_root);
stamp_page_checksum(&mut page);
Ok(page)
}
fn decode(page: &[u8], page_no: u64) -> Result<Self> {
let header =
PageHeader::decode(page).ok_or_else(|| malformed(page_no, "fts page header"))?;
if header.page_type != PageType::FtsDict {
return Err(malformed(page_no, "not an FTS page"));
}
let mut off = PAGE_HEADER_LEN;
if page.get(off).copied() != Some(dict::NODE_META) {
return Err(malformed(page_no, "not an FTS meta page"));
}
off += 1;
let doc_count = dict::get_u64(page, &mut off, page_no)?;
let total_tokens = dict::get_u64(page, &mut off, page_no)?;
let dict_root = dict::get_u64(page, &mut off, page_no)?;
Ok(FtsMeta {
doc_count,
total_tokens,
dict_root,
})
}
}
fn load_meta(src: &dyn PageSource, root: u64) -> Result<Option<FtsMeta>> {
if root == 0 {
return Ok(None);
}
let page = src.page(root)?;
Ok(Some(FtsMeta::decode(&page, root)?))
}
fn save_meta(txn: &mut Txn<'_>, meta: &FtsMeta) -> Result<()> {
let page_no = match txn.fts_root_page() {
0 => txn.allocate_page()?,
p => p,
};
let page = meta.encode(txn.page_size())?;
txn.write_page(page_no, &page)?;
txn.set_fts_root_page(page_no);
Ok(())
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct Posting {
record_id: Ulid,
term_freq: u32,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
struct Postings {
entries: Vec<Posting>,
}
impl Postings {
fn upsert(&mut self, record_id: Ulid, term_freq: u32) {
match self
.entries
.binary_search_by(|p| p.record_id.cmp(&record_id))
{
Ok(i) => self.entries[i].term_freq = term_freq,
Err(i) => self.entries.insert(
i,
Posting {
record_id,
term_freq,
},
),
}
}
fn encode(&self, layout: PostingsLayout) -> Vec<u8> {
let mut out = Vec::with_capacity(4 + self.entries.len() * POSTING_LEN);
out.extend_from_slice(&(self.entries.len() as u32).to_le_bytes());
match layout {
PostingsLayout::FixedWidth => {
for p in &self.entries {
out.extend_from_slice(&p.record_id.to_bytes());
out.extend_from_slice(&p.term_freq.to_le_bytes());
}
}
PostingsLayout::DeltaVarint => {
encode_delta_run(&mut out, &self.entries);
}
PostingsLayout::DeltaVarintSkip(entry) => {
self.encode_skip(&mut out, entry, false);
}
PostingsLayout::FrameOfReference => {
self.encode_skip(&mut out, SkipEntry::V6, true);
}
}
out
}
fn encode_skip(&self, out: &mut Vec<u8>, entry: SkipEntry, frame_of_reference: bool) {
if self.entries.len() < SKIP_MIN_DOC_FREQ {
out.extend_from_slice(&0u32.to_le_bytes());
encode_delta_run(out, &self.entries);
return;
}
let blocks: Vec<&[Posting]> = self.entries.chunks(SKIP_BLOCK_SIZE).collect();
out.extend_from_slice(&(blocks.len() as u32).to_le_bytes());
let mut blocks_body: Vec<u8> = Vec::new();
let mut skip_index: Vec<u8> = Vec::with_capacity(blocks.len() * entry.len());
for block in &blocks {
let first_id = u128::from(block[0].record_id);
let byte_offset = blocks_body.len() as u32;
let max_tf = block.iter().map(|p| p.term_freq).max().unwrap_or(0);
skip_index.extend_from_slice(&first_id.to_le_bytes());
if entry == SkipEntry::V6 {
let last_id = u128::from(block[block.len() - 1].record_id);
skip_index.extend_from_slice(&last_id.to_le_bytes());
}
skip_index.extend_from_slice(&byte_offset.to_le_bytes());
skip_index.extend_from_slice(&max_tf.to_le_bytes());
if frame_of_reference {
encode_for_block(&mut blocks_body, block);
} else {
encode_delta_run(&mut blocks_body, block);
}
}
out.extend_from_slice(&skip_index);
out.extend_from_slice(&blocks_body);
}
fn decode(body: &[u8], page_no: u64, layout: PostingsLayout) -> Result<Self> {
match layout {
PostingsLayout::FixedWidth => Self::decode_fixed_width(body, page_no),
PostingsLayout::DeltaVarint => Self::decode_delta_varint(body, page_no),
PostingsLayout::DeltaVarintSkip(entry) => {
Self::decode_skip(body, page_no, entry, layout)
}
PostingsLayout::FrameOfReference => {
Self::decode_skip(body, page_no, SkipEntry::V6, layout)
}
}
}
fn decode_fixed_width(body: &[u8], page_no: u64) -> Result<Self> {
let count = dict::read_u32(body, 0, page_no)? as usize;
let need = 4usize
.checked_add(
count
.checked_mul(POSTING_LEN)
.ok_or_else(|| malformed(page_no, "fts postings count overflow"))?,
)
.ok_or_else(|| malformed(page_no, "fts postings length overflow"))?;
if body.len() < need {
return Err(malformed(page_no, "fts postings truncated"));
}
let mut entries = Vec::with_capacity(count);
let mut prev: Option<Ulid> = None;
let mut off = 4;
for _ in 0..count {
let id_bytes: [u8; 16] = body
.get(off..off + 16)
.and_then(|b| b.try_into().ok())
.ok_or_else(|| malformed(page_no, "fts posting id"))?;
let record_id = Ulid::from_bytes(id_bytes);
if prev.is_some_and(|p| p >= record_id) {
return Err(malformed(page_no, "unsorted fts postings"));
}
prev = Some(record_id);
let term_freq = dict::read_u32(body, off + 16, page_no)?;
if term_freq == 0 {
return Err(malformed(page_no, "fts posting zero term_freq"));
}
entries.push(Posting {
record_id,
term_freq,
});
off += POSTING_LEN;
}
Ok(Postings { entries })
}
fn decode_delta_varint(body: &[u8], page_no: u64) -> Result<Self> {
let count = dict::read_u32(body, 0, page_no)? as usize;
reject_hostile_delta_count(body, count, page_no)?;
let mut entries = Vec::with_capacity(count);
let mut off = 4;
decode_delta_run(body, &mut off, count, page_no, &mut entries)?;
Ok(Postings { entries })
}
fn decode_skip(
body: &[u8],
page_no: u64,
entry: SkipEntry,
layout: PostingsLayout,
) -> Result<Self> {
let count = dict::read_u32(body, 0, page_no)? as usize;
let block_count = dict::read_u32(body, 4, page_no)? as usize;
if block_count == 0 {
reject_hostile_delta_count(body, count, page_no)?;
let mut entries = Vec::with_capacity(count);
let mut off = 8;
decode_delta_run(body, &mut off, count, page_no, &mut entries)?;
return Ok(Postings { entries });
}
let expected_blocks = count.div_ceil(SKIP_BLOCK_SIZE);
if count < SKIP_MIN_DOC_FREQ || block_count != expected_blocks {
return Err(malformed(page_no, "fts skip block_count mismatch"));
}
let entry_len = entry.len();
let index_len = block_count
.checked_mul(entry_len)
.ok_or_else(|| malformed(page_no, "fts skip index overflow"))?;
let blocks_start = 8usize
.checked_add(index_len)
.ok_or_else(|| malformed(page_no, "fts skip index overflow"))?;
if body.len() < blocks_start {
return Err(malformed(page_no, "fts skip index truncated"));
}
reject_hostile_delta_count(&body[blocks_start.saturating_sub(4)..], count, page_no)?;
let mut entries = Vec::with_capacity(count);
let mut prev_block_last: Option<u128> = None;
let mut off = blocks_start;
for b in 0..block_count {
let block_len = if b + 1 == block_count {
count - b * SKIP_BLOCK_SIZE
} else {
SKIP_BLOCK_SIZE
};
let idx_off = 8 + b * entry_len;
let stored_first = read_u128_le(body, idx_off, page_no)?;
let stored_last = match entry {
SkipEntry::V6 => Some(read_u128_le(body, idx_off + 16, page_no)?),
SkipEntry::V5 => None,
};
let tail = idx_off + entry.tail_off();
let stored_offset = dict::read_u32(body, tail, page_no)? as usize;
let stored_max_tf = dict::read_u32(body, tail + 4, page_no)?;
if blocks_start + stored_offset != off {
return Err(malformed(page_no, "fts skip byte_offset mismatch"));
}
let block_start_entry = entries.len();
decode_block_run(layout, body, &mut off, block_len, page_no, &mut entries)?;
let block = &entries[block_start_entry..];
let first_id = u128::from(block[0].record_id);
if stored_first != first_id {
return Err(malformed(page_no, "fts skip first_id mismatch"));
}
let last_id = u128::from(block[block.len() - 1].record_id);
if stored_last.is_some_and(|s| s != last_id) {
return Err(malformed(page_no, "fts skip last_id mismatch"));
}
let max_tf = block.iter().map(|p| p.term_freq).max().unwrap_or(0);
if stored_max_tf != max_tf {
return Err(malformed(page_no, "fts skip max_term_freq mismatch"));
}
if prev_block_last.is_some_and(|prev_last| first_id <= prev_last) {
return Err(malformed(page_no, "unsorted fts postings"));
}
prev_block_last = Some(last_id);
}
Ok(Postings { entries })
}
}
fn postings_for(src: &dyn PageSource, root: u64, term: &[u8]) -> Result<Option<Postings>> {
let layout = PostingsLayout::for_format_version(src.format_version());
match dict::get(src, FTS_DICT, root, term)? {
Some((body, page_no)) => Ok(Some(Postings::decode(&body, page_no, layout)?)),
None => Ok(None),
}
}
pub fn index_document(txn: &mut Txn<'_>, record_id: Ulid, content: &str) -> Result<()> {
let tokens = tokenize(content);
let mut freqs: HashMap<String, u32> = HashMap::new();
for token in tokens {
let term = clip_term(&token).to_owned();
if term.is_empty() {
continue;
}
*freqs.entry(term).or_insert(0) += 1;
}
let mut meta = load_meta(txn, txn.fts_root_page())?.unwrap_or_else(FtsMeta::empty);
let doc_tokens: u64 = freqs.values().map(|&f| u64::from(f)).sum();
let mut terms: Vec<(String, u32)> = freqs.into_iter().collect();
terms.sort_by(|a, b| a.0.cmp(&b.0));
let layout = PostingsLayout::for_format_version(txn.format_version());
let mut root = meta.dict_root;
for (term, tf) in terms {
let term_bytes = term.as_bytes();
let mut postings = postings_for(txn, root, term_bytes)?.unwrap_or_default();
postings.upsert(record_id, tf);
root = dict::upsert(txn, FTS_DICT, root, term_bytes, &postings.encode(layout))?;
}
meta.dict_root = root;
meta.doc_count = meta.doc_count.saturating_add(1);
meta.total_tokens = meta.total_tokens.saturating_add(doc_tokens);
save_meta(txn, &meta)?;
Ok(())
}
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Hit {
pub record_id: Ulid,
pub score: f32,
}
pub fn search(
src: &dyn PageSource,
fts_root_page: u64,
query: &str,
k: usize,
keep: impl FnMut(Ulid) -> bool,
doc_len: impl FnMut(Ulid) -> Result<Option<u32>>,
) -> Result<Vec<Hit>> {
if PostingsLayout::for_format_version(src.format_version())
== PostingsLayout::DeltaVarintSkip(SkipEntry::V6)
{
Ok(search_bmw_counted(src, fts_root_page, query, k, keep, doc_len)?.0)
} else {
search_linear(src, fts_root_page, query, k, keep, doc_len)
}
}
fn search_linear(
src: &dyn PageSource,
fts_root_page: u64,
query: &str,
k: usize,
mut keep: impl FnMut(Ulid) -> bool,
mut doc_len: impl FnMut(Ulid) -> Result<Option<u32>>,
) -> Result<Vec<Hit>> {
if k == 0 {
return Ok(Vec::new());
}
let Some(meta) = load_meta(src, fts_root_page)? else {
return Ok(Vec::new());
};
if meta.doc_count == 0 || meta.dict_root == 0 {
return Ok(Vec::new());
}
let mut query_terms: Vec<String> = tokenize(query)
.into_iter()
.map(|t| clip_term(&t).to_owned())
.filter(|t| !t.is_empty())
.collect();
query_terms.sort();
query_terms.dedup();
if query_terms.is_empty() {
return Ok(Vec::new());
}
let n = meta.doc_count as f32;
let avgdl = if meta.doc_count == 0 {
0.0
} else {
meta.total_tokens as f32 / meta.doc_count as f32
};
let mut matched: Vec<(f32, Postings)> = Vec::with_capacity(query_terms.len());
let mut bounds: HashMap<Ulid, f32> = HashMap::new();
for term in &query_terms {
let Some(postings) = postings_for(src, meta.dict_root, term.as_bytes())? else {
continue;
};
let df = postings.entries.len() as f32;
if df == 0.0 {
continue;
}
let idf = (1.0 + (n - df + 0.5) / (df + 0.5)).ln();
for p in &postings.entries {
*bounds.entry(p.record_id).or_insert(0.0) += bound_contribution(idf, p.term_freq);
}
matched.push((idf, postings));
}
let mut candidates: Vec<(Ulid, f32)> = bounds.into_iter().collect();
candidates.sort_by(|a, b| {
b.1.partial_cmp(&a.1)
.unwrap_or(std::cmp::Ordering::Equal)
.then_with(|| a.0.cmp(&b.0))
});
let mut hits: Vec<Hit> = Vec::with_capacity(k.saturating_add(1).min(candidates.len()));
for (id, bound) in candidates {
if hits.len() == k && bound < hits[k - 1].score {
break;
}
if !keep(id) {
continue;
}
let Some(dl) = doc_len(id)? else {
continue; };
let mut score = 0.0f32;
for (idf, postings) in &matched {
let Ok(i) = postings.entries.binary_search_by(|p| p.record_id.cmp(&id)) else {
continue;
};
let tf = postings.entries[i].term_freq as f32;
let norm = tf + BM25_K1 * (1.0 - BM25_B + BM25_B * dl as f32 / avgdl.max(1.0));
score += idf * (tf * (BM25_K1 + 1.0)) / norm.max(f32::MIN_POSITIVE);
}
if score <= 0.0 {
continue;
}
let pos =
hits.partition_point(|h| h.score > score || (h.score == score && h.record_id < id));
if pos < k {
hits.insert(
pos,
Hit {
record_id: id,
score,
},
);
hits.truncate(k);
}
}
Ok(hits)
}
fn bound_contribution(idf: f32, tf: u32) -> f32 {
let tf = tf as f32;
let norm = tf + BM25_K1 * (1.0 - BM25_B);
idf * (tf * (BM25_K1 + 1.0)) / norm.max(f32::MIN_POSITIVE)
}
fn bound_slack(terms: usize) -> f64 {
1.0 + terms as f64 * 1.2e-7
}
struct BmwBlock {
first_id: u128,
last_id: u128,
offset: usize,
len: usize,
max_tf: u32,
ub: f32,
}
#[doc(hidden)]
#[derive(Debug, Clone, Copy, Default)]
pub struct BmwCounters {
pub blocks_total: u64,
pub blocks_decoded: u64,
pub docs_evaluated: u64,
pub pivot_skips: u64,
pub decode_varint_ns: u64,
pub decode_revalidate_ns: u64,
}
impl BmwCounters {
pub fn blocks_skipped(&self) -> u64 {
self.blocks_total.saturating_sub(self.blocks_decoded)
}
}
struct BmwCursor {
idf: f32,
body: Vec<u8>,
page_no: u64,
layout: PostingsLayout,
blocks_start: usize,
blocks: Vec<BmwBlock>,
cur_block: usize,
entries: Vec<Posting>,
pos: usize,
cur_id: u128,
term_ub: f32,
exhausted: bool,
}
impl BmwCursor {
fn open(
idf: f32,
body: Vec<u8>,
page_no: u64,
layout: PostingsLayout,
counters: &mut BmwCounters,
) -> Result<Option<BmwCursor>> {
let count = dict::read_u32(&body, 0, page_no)? as usize;
let block_count = dict::read_u32(&body, 4, page_no)? as usize;
if block_count == 0 {
let decoded = Postings::decode(&body, page_no, layout)?;
let Some((first, last)) = decoded.entries.first().zip(decoded.entries.last()) else {
return Ok(None);
};
let max_tf = decoded
.entries
.iter()
.map(|p| p.term_freq)
.max()
.unwrap_or(0);
let ub = bound_contribution(idf, max_tf);
let block = BmwBlock {
first_id: u128::from(first.record_id),
last_id: u128::from(last.record_id),
offset: 0,
len: decoded.entries.len(),
max_tf,
ub,
};
counters.blocks_total += 1;
counters.blocks_decoded += 1;
let cur_id = block.first_id;
return Ok(Some(BmwCursor {
idf,
body,
page_no,
layout,
blocks_start: 0,
blocks: vec![block],
cur_block: 0,
entries: decoded.entries,
pos: 0,
cur_id,
term_ub: ub,
exhausted: false,
}));
}
let expected_blocks = count.div_ceil(SKIP_BLOCK_SIZE);
if count < SKIP_MIN_DOC_FREQ || block_count != expected_blocks {
return Err(malformed(page_no, "fts skip block_count mismatch"));
}
let entry_len = SkipEntry::V6.len();
let index_len = block_count
.checked_mul(entry_len)
.ok_or_else(|| malformed(page_no, "fts skip index overflow"))?;
let blocks_start = 8usize
.checked_add(index_len)
.ok_or_else(|| malformed(page_no, "fts skip index overflow"))?;
if body.len() < blocks_start {
return Err(malformed(page_no, "fts skip index truncated"));
}
reject_hostile_delta_count(&body[blocks_start.saturating_sub(4)..], count, page_no)?;
let mut blocks: Vec<BmwBlock> = Vec::with_capacity(block_count);
let mut term_ub = 0.0f32;
for b in 0..block_count {
let idx = 8 + b * entry_len;
let first_id = read_u128_le(&body, idx, page_no)?;
let last_id = read_u128_le(&body, idx + 16, page_no)?;
let tail = idx + SkipEntry::V6.tail_off();
let offset = dict::read_u32(&body, tail, page_no)? as usize;
let max_tf = dict::read_u32(&body, tail + 4, page_no)?;
let len = if b + 1 == block_count {
count - b * SKIP_BLOCK_SIZE
} else {
SKIP_BLOCK_SIZE
};
if first_id > last_id
|| max_tf == 0
|| (b == 0 && offset != 0)
|| blocks
.last()
.is_some_and(|p| p.last_id >= first_id || p.offset >= offset)
{
return Err(malformed(page_no, "fts skip index inconsistent"));
}
let ub = bound_contribution(idf, max_tf);
term_ub = term_ub.max(ub);
blocks.push(BmwBlock {
first_id,
last_id,
offset,
len,
max_tf,
ub,
});
}
counters.blocks_total += block_count as u64;
let cur_id = blocks[0].first_id;
Ok(Some(BmwCursor {
idf,
body,
page_no,
layout,
blocks_start,
blocks,
cur_block: 0,
entries: Vec::new(),
pos: 0,
cur_id,
term_ub,
exhausted: false,
}))
}
fn decode_block(&mut self, b: usize, counters: &mut BmwCounters) -> Result<()> {
let (offset, len, first_id, last_id, max_tf) = {
let m = &self.blocks[b];
(m.offset, m.len, m.first_id, m.last_id, m.max_tf)
};
let mut off = self
.blocks_start
.checked_add(offset)
.ok_or_else(|| malformed(self.page_no, "fts skip byte_offset overflow"))?;
self.entries.clear();
self.entries.reserve(len);
let varint_started = Instant::now();
decode_block_run(
self.layout,
&self.body,
&mut off,
len,
self.page_no,
&mut self.entries,
)?;
counters.decode_varint_ns += varint_started.elapsed().as_nanos() as u64;
let revalidate_started = Instant::now();
if self.entries.first().map(|p| u128::from(p.record_id)) != Some(first_id) {
return Err(malformed(self.page_no, "fts skip first_id mismatch"));
}
if self.entries.last().map(|p| u128::from(p.record_id)) != Some(last_id) {
return Err(malformed(self.page_no, "fts skip last_id mismatch"));
}
if self.entries.iter().map(|p| p.term_freq).max() != Some(max_tf) {
return Err(malformed(self.page_no, "fts skip max_term_freq mismatch"));
}
counters.decode_revalidate_ns += revalidate_started.elapsed().as_nanos() as u64;
self.cur_block = b;
counters.blocks_decoded += 1;
Ok(())
}
fn advance_to(&mut self, target: u128, counters: &mut BmwCounters) -> Result<()> {
if self.exhausted || target <= self.cur_id {
return Ok(());
}
let rel = self.blocks[self.cur_block..].partition_point(|m| m.last_id < target);
let b = self.cur_block + rel;
if b >= self.blocks.len() {
self.exhausted = true;
return Ok(());
}
if self.blocks[b].first_id >= target {
if b != self.cur_block {
self.entries.clear();
self.cur_block = b;
}
self.pos = 0;
self.cur_id = self.blocks[b].first_id;
return Ok(());
}
if b != self.cur_block || self.entries.is_empty() {
self.decode_block(b, counters)?;
}
let pos = self
.entries
.partition_point(|p| u128::from(p.record_id) < target);
let posting = self
.entries
.get(pos)
.ok_or_else(|| malformed(self.page_no, "fts skip cursor past block"))?;
self.pos = pos;
self.cur_id = u128::from(posting.record_id);
Ok(())
}
fn advance_past(&mut self, id: u128, counters: &mut BmwCounters) -> Result<()> {
match id.checked_add(1) {
Some(next) => self.advance_to(next, counters),
None => {
self.exhausted = true;
Ok(())
}
}
}
fn covering_block(&self, target: u128) -> Option<&BmwBlock> {
let rel = self.blocks[self.cur_block..].partition_point(|m| m.last_id < target);
self.blocks.get(self.cur_block + rel)
}
fn current_tf(&mut self, counters: &mut BmwCounters) -> Result<u32> {
if self.entries.is_empty() {
self.decode_block(self.cur_block, counters)?;
self.pos = 0; }
let posting = self
.entries
.get(self.pos)
.ok_or_else(|| malformed(self.page_no, "fts skip cursor past block"))?;
if u128::from(posting.record_id) != self.cur_id {
return Err(malformed(self.page_no, "fts skip cursor desync"));
}
Ok(posting.term_freq)
}
}
#[doc(hidden)]
pub fn search_bmw_counted(
src: &dyn PageSource,
fts_root_page: u64,
query: &str,
k: usize,
mut keep: impl FnMut(Ulid) -> bool,
mut doc_len: impl FnMut(Ulid) -> Result<Option<u32>>,
) -> Result<(Vec<Hit>, BmwCounters)> {
let mut counters = BmwCounters::default();
if k == 0 {
return Ok((Vec::new(), counters));
}
let Some(meta) = load_meta(src, fts_root_page)? else {
return Ok((Vec::new(), counters));
};
if meta.doc_count == 0 || meta.dict_root == 0 {
return Ok((Vec::new(), counters));
}
let mut query_terms: Vec<String> = tokenize(query)
.into_iter()
.map(|t| clip_term(&t).to_owned())
.filter(|t| !t.is_empty())
.collect();
query_terms.sort();
query_terms.dedup();
if query_terms.is_empty() {
return Ok((Vec::new(), counters));
}
let n = meta.doc_count as f32;
let avgdl = if meta.doc_count == 0 {
0.0
} else {
meta.total_tokens as f32 / meta.doc_count as f32
};
let mut cursors: Vec<BmwCursor> = Vec::with_capacity(query_terms.len());
for term in &query_terms {
let Some((body, page_no)) = dict::get(src, FTS_DICT, meta.dict_root, term.as_bytes())?
else {
continue;
};
let df = dict::read_u32(&body, 0, page_no)? as usize;
if df == 0 {
continue;
}
let idf = (1.0 + (n - df as f32 + 0.5) / (df as f32 + 0.5)).ln();
let layout = PostingsLayout::for_format_version(src.format_version());
if let Some(cursor) = BmwCursor::open(idf, body, page_no, layout, &mut counters)? {
cursors.push(cursor);
}
}
if cursors.is_empty() {
return Ok((Vec::new(), counters));
}
let slack = bound_slack(cursors.len());
let mut hits: Vec<Hit> = Vec::new();
let mut order: Vec<usize> = (0..cursors.len()).collect();
loop {
order.retain(|&i| !cursors[i].exhausted);
if order.is_empty() {
break;
}
order.sort_by(|&a, &b| {
cursors[a]
.cur_id
.cmp(&cursors[b].cur_id)
.then_with(|| a.cmp(&b))
});
let theta: Option<f32> = (hits.len() == k).then(|| hits[k - 1].score);
let mut acc = 0.0f64;
let mut pivot: Option<usize> = None;
for (oi, &ci) in order.iter().enumerate() {
acc += f64::from(cursors[ci].term_ub);
if theta.is_none_or(|t| acc * slack > f64::from(t)) {
pivot = Some(oi);
break;
}
}
let Some(p) = pivot else {
break;
};
let pivot_id = cursors[order[p]].cur_id;
let prefix_end = p + order[p + 1..]
.iter()
.take_while(|&&ci| cursors[ci].cur_id == pivot_id)
.count();
let mut block_acc = 0.0f64;
let mut min_block_last = u128::MAX;
for &ci in &order[..=prefix_end] {
if let Some(block) = cursors[ci].covering_block(pivot_id) {
block_acc += f64::from(block.ub);
min_block_last = min_block_last.min(block.last_id);
}
}
if theta.is_some_and(|t| block_acc * slack <= f64::from(t)) {
counters.pivot_skips += 1;
let next = match (min_block_last.checked_add(1), order.get(prefix_end + 1)) {
(Some(n), Some(&after)) => Some(n.min(cursors[after].cur_id)),
(None, Some(&after)) => Some(cursors[after].cur_id),
(n, None) => n,
};
match next {
Some(n) => {
for &ci in &order[..=prefix_end] {
cursors[ci].advance_to(n, &mut counters)?;
}
}
None => {
for &ci in &order[..=prefix_end] {
cursors[ci].exhausted = true;
}
}
}
continue;
}
if cursors[order[0]].cur_id != pivot_id {
for &ci in &order[..p] {
cursors[ci].advance_to(pivot_id, &mut counters)?;
}
continue;
}
counters.docs_evaluated += 1;
let id = Ulid::from(pivot_id);
if keep(id)
&& let Some(dl) = doc_len(id)?
{
let mut score = 0.0f32;
for c in cursors.iter_mut() {
if c.exhausted || c.cur_id != pivot_id {
continue;
}
let tf = c.current_tf(&mut counters)? as f32;
let norm = tf + BM25_K1 * (1.0 - BM25_B + BM25_B * dl as f32 / avgdl.max(1.0));
score += c.idf * (tf * (BM25_K1 + 1.0)) / norm.max(f32::MIN_POSITIVE);
}
if score > 0.0 {
let pos = hits
.partition_point(|h| h.score > score || (h.score == score && h.record_id < id));
if pos < k {
hits.insert(
pos,
Hit {
record_id: id,
score,
},
);
hits.truncate(k);
}
}
}
for &ci in &order[..=prefix_end] {
cursors[ci].advance_past(pivot_id, &mut counters)?;
}
}
Ok((hits, counters))
}
#[doc(hidden)]
#[derive(Debug, Clone, Copy, Default)]
pub struct BmwPhaseTimings {
pub cursor_open_ns: u64,
pub bound_ns: u64,
pub decode_ns: u64,
pub keep_ns: u64,
pub doc_len_ns: u64,
pub scoring_ns: u64,
pub docs_evaluated: u64,
pub pivot_skips: u64,
pub blocks_decoded: u64,
pub blocks_skipped: u64,
pub decode_varint_ns: u64,
pub decode_revalidate_ns: u64,
}
#[doc(hidden)]
pub fn search_bmw_profiled(
src: &dyn PageSource,
fts_root_page: u64,
query: &str,
k: usize,
mut keep: impl FnMut(Ulid) -> bool,
mut doc_len: impl FnMut(Ulid) -> Result<Option<u32>>,
) -> Result<(Vec<Hit>, BmwPhaseTimings)> {
let mut timings = BmwPhaseTimings::default();
let mut counters = BmwCounters::default();
if k == 0 {
return Ok((Vec::new(), timings));
}
let Some(meta) = load_meta(src, fts_root_page)? else {
return Ok((Vec::new(), timings));
};
if meta.doc_count == 0 || meta.dict_root == 0 {
return Ok((Vec::new(), timings));
}
let mut query_terms: Vec<String> = tokenize(query)
.into_iter()
.map(|t| clip_term(&t).to_owned())
.filter(|t| !t.is_empty())
.collect();
query_terms.sort();
query_terms.dedup();
if query_terms.is_empty() {
return Ok((Vec::new(), timings));
}
let n = meta.doc_count as f32;
let avgdl = if meta.doc_count == 0 {
0.0
} else {
meta.total_tokens as f32 / meta.doc_count as f32
};
let cursor_open_started = Instant::now();
let mut cursors: Vec<BmwCursor> = Vec::with_capacity(query_terms.len());
for term in &query_terms {
let Some((body, page_no)) = dict::get(src, FTS_DICT, meta.dict_root, term.as_bytes())?
else {
continue;
};
let df = dict::read_u32(&body, 0, page_no)? as usize;
if df == 0 {
continue;
}
let idf = (1.0 + (n - df as f32 + 0.5) / (df as f32 + 0.5)).ln();
let layout = PostingsLayout::for_format_version(src.format_version());
if let Some(cursor) = BmwCursor::open(idf, body, page_no, layout, &mut counters)? {
cursors.push(cursor);
}
}
timings.cursor_open_ns += cursor_open_started.elapsed().as_nanos() as u64;
if cursors.is_empty() {
return Ok((Vec::new(), timings));
}
let slack = bound_slack(cursors.len());
let mut hits: Vec<Hit> = Vec::new();
let mut order: Vec<usize> = (0..cursors.len()).collect();
loop {
let bound_started = Instant::now();
order.retain(|&i| !cursors[i].exhausted);
if order.is_empty() {
timings.bound_ns += bound_started.elapsed().as_nanos() as u64;
break;
}
order.sort_by(|&a, &b| {
cursors[a]
.cur_id
.cmp(&cursors[b].cur_id)
.then_with(|| a.cmp(&b))
});
let theta: Option<f32> = (hits.len() == k).then(|| hits[k - 1].score);
let mut acc = 0.0f64;
let mut pivot: Option<usize> = None;
for (oi, &ci) in order.iter().enumerate() {
acc += f64::from(cursors[ci].term_ub);
if theta.is_none_or(|t| acc * slack > f64::from(t)) {
pivot = Some(oi);
break;
}
}
let Some(p) = pivot else {
timings.bound_ns += bound_started.elapsed().as_nanos() as u64;
break;
};
let pivot_id = cursors[order[p]].cur_id;
let prefix_end = p + order[p + 1..]
.iter()
.take_while(|&&ci| cursors[ci].cur_id == pivot_id)
.count();
let mut block_acc = 0.0f64;
let mut min_block_last = u128::MAX;
for &ci in &order[..=prefix_end] {
if let Some(block) = cursors[ci].covering_block(pivot_id) {
block_acc += f64::from(block.ub);
min_block_last = min_block_last.min(block.last_id);
}
}
let skip = theta.is_some_and(|t| block_acc * slack <= f64::from(t));
timings.bound_ns += bound_started.elapsed().as_nanos() as u64;
if skip {
counters.pivot_skips += 1;
let next = match (min_block_last.checked_add(1), order.get(prefix_end + 1)) {
(Some(n), Some(&after)) => Some(n.min(cursors[after].cur_id)),
(None, Some(&after)) => Some(cursors[after].cur_id),
(n, None) => n,
};
let decode_started = Instant::now();
match next {
Some(n) => {
for &ci in &order[..=prefix_end] {
cursors[ci].advance_to(n, &mut counters)?;
}
}
None => {
for &ci in &order[..=prefix_end] {
cursors[ci].exhausted = true;
}
}
}
timings.decode_ns += decode_started.elapsed().as_nanos() as u64;
continue;
}
if cursors[order[0]].cur_id != pivot_id {
let decode_started = Instant::now();
for &ci in &order[..p] {
cursors[ci].advance_to(pivot_id, &mut counters)?;
}
timings.decode_ns += decode_started.elapsed().as_nanos() as u64;
continue;
}
counters.docs_evaluated += 1;
let id = Ulid::from(pivot_id);
let keep_started = Instant::now();
let kept = keep(id);
timings.keep_ns += keep_started.elapsed().as_nanos() as u64;
if kept {
let doc_len_started = Instant::now();
let dl = doc_len(id)?;
timings.doc_len_ns += doc_len_started.elapsed().as_nanos() as u64;
if let Some(dl) = dl {
let scoring_started = Instant::now();
let mut score = 0.0f32;
for c in cursors.iter_mut() {
if c.exhausted || c.cur_id != pivot_id {
continue;
}
let decode_started = Instant::now();
let tf = c.current_tf(&mut counters)? as f32;
timings.decode_ns += decode_started.elapsed().as_nanos() as u64;
let norm = tf + BM25_K1 * (1.0 - BM25_B + BM25_B * dl as f32 / avgdl.max(1.0));
score += c.idf * (tf * (BM25_K1 + 1.0)) / norm.max(f32::MIN_POSITIVE);
}
if score > 0.0 {
let pos = hits.partition_point(|h| {
h.score > score || (h.score == score && h.record_id < id)
});
if pos < k {
hits.insert(
pos,
Hit {
record_id: id,
score,
},
);
hits.truncate(k);
}
}
timings.scoring_ns += scoring_started.elapsed().as_nanos() as u64;
}
}
let decode_started = Instant::now();
for &ci in &order[..=prefix_end] {
cursors[ci].advance_past(pivot_id, &mut counters)?;
}
timings.decode_ns += decode_started.elapsed().as_nanos() as u64;
}
timings.docs_evaluated = counters.docs_evaluated;
timings.pivot_skips = counters.pivot_skips;
timings.blocks_decoded = counters.blocks_decoded;
timings.blocks_skipped = counters.blocks_skipped();
timings.decode_varint_ns = counters.decode_varint_ns;
timings.decode_revalidate_ns = counters.decode_revalidate_ns;
Ok((hits, timings))
}
pub fn indexed_documents(src: &dyn PageSource, fts_root_page: u64) -> Result<u64> {
Ok(load_meta(src, fts_root_page)?.map_or(0, |m| m.doc_count))
}
#[doc(hidden)]
#[derive(Debug, Clone, Copy, Default)]
pub struct SearchPhaseTimings {
pub postings_lookup_ns: u64,
pub keep_ns: u64,
pub doc_len_ns: u64,
pub scoring_ns: u64,
pub terms_matched: u32,
pub postings_visited: u64,
pub keep_outcomes: KeepOutcomeCounts,
}
#[doc(hidden)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum KeepOutcome {
Accepted,
Tombstoned,
OutOfScope,
FilteredOut,
}
#[doc(hidden)]
#[derive(Debug, Clone, Copy, Default)]
pub struct KeepOutcomeCounts {
pub accepted: u64,
pub tombstoned: u64,
pub out_of_scope: u64,
pub filtered_out: u64,
}
impl KeepOutcomeCounts {
fn record(&mut self, outcome: KeepOutcome) {
match outcome {
KeepOutcome::Accepted => self.accepted += 1,
KeepOutcome::Tombstoned => self.tombstoned += 1,
KeepOutcome::OutOfScope => self.out_of_scope += 1,
KeepOutcome::FilteredOut => self.filtered_out += 1,
}
}
pub fn total(&self) -> u64 {
self.accepted + self.tombstoned + self.out_of_scope + self.filtered_out
}
pub fn rejected(&self) -> u64 {
self.tombstoned + self.out_of_scope + self.filtered_out
}
}
#[doc(hidden)]
pub fn search_profiled(
src: &dyn PageSource,
fts_root_page: u64,
query: &str,
k: usize,
mut keep: impl FnMut(Ulid) -> KeepOutcome,
mut doc_len: impl FnMut(Ulid) -> Result<Option<u32>>,
) -> Result<(Vec<Hit>, SearchPhaseTimings)> {
let mut timings = SearchPhaseTimings::default();
if k == 0 {
return Ok((Vec::new(), timings));
}
let Some(meta) = load_meta(src, fts_root_page)? else {
return Ok((Vec::new(), timings));
};
if meta.doc_count == 0 || meta.dict_root == 0 {
return Ok((Vec::new(), timings));
}
let mut query_terms: Vec<String> = tokenize(query)
.into_iter()
.map(|t| clip_term(&t).to_owned())
.filter(|t| !t.is_empty())
.collect();
query_terms.sort();
query_terms.dedup();
if query_terms.is_empty() {
return Ok((Vec::new(), timings));
}
let n = meta.doc_count as f32;
let avgdl = if meta.doc_count == 0 {
0.0
} else {
meta.total_tokens as f32 / meta.doc_count as f32
};
let mut scores: HashMap<Ulid, f32> = HashMap::new();
let mut lengths: HashMap<Ulid, Option<u32>> = HashMap::new();
let mut kept: HashMap<Ulid, bool> = HashMap::new();
for term in &query_terms {
let lookup_started = std::time::Instant::now();
let postings = postings_for(src, meta.dict_root, term.as_bytes())?;
timings.postings_lookup_ns += lookup_started.elapsed().as_nanos() as u64;
let Some(postings) = postings else {
continue;
};
let df = postings.entries.len() as f32;
if df == 0.0 {
continue;
}
timings.terms_matched += 1;
let idf = (1.0 + (n - df + 0.5) / (df + 0.5)).ln();
for p in &postings.entries {
timings.postings_visited += 1;
let id = p.record_id;
let is_kept = match kept.get(&id) {
Some(&v) => v,
None => {
let keep_started = std::time::Instant::now();
let outcome = keep(id);
timings.keep_ns += keep_started.elapsed().as_nanos() as u64;
timings.keep_outcomes.record(outcome);
let v = outcome == KeepOutcome::Accepted;
kept.insert(id, v);
v
}
};
if !is_kept {
continue;
}
let dl = match lengths.get(&id) {
Some(&v) => v,
None => {
let doc_len_started = std::time::Instant::now();
let v = doc_len(id)?;
timings.doc_len_ns += doc_len_started.elapsed().as_nanos() as u64;
lengths.insert(id, v);
v
}
};
let Some(dl) = dl else {
continue;
};
let scoring_started = std::time::Instant::now();
let tf = p.term_freq as f32;
let norm = tf + BM25_K1 * (1.0 - BM25_B + BM25_B * dl as f32 / avgdl.max(1.0));
let contribution = idf * (tf * (BM25_K1 + 1.0)) / norm.max(f32::MIN_POSITIVE);
*scores.entry(id).or_insert(0.0) += contribution;
timings.scoring_ns += scoring_started.elapsed().as_nanos() as u64;
}
}
let scoring_started = std::time::Instant::now();
let mut hits: Vec<Hit> = scores
.into_iter()
.filter(|&(_, s)| s > 0.0)
.map(|(record_id, score)| Hit { record_id, score })
.collect();
hits.sort_by(|a, b| {
b.score
.partial_cmp(&a.score)
.unwrap_or(std::cmp::Ordering::Equal)
.then_with(|| a.record_id.cmp(&b.record_id))
});
hits.truncate(k);
timings.scoring_ns += scoring_started.elapsed().as_nanos() as u64;
Ok((hits, timings))
}
#[doc(hidden)]
pub fn fuzz_decode_page(page: &[u8]) {
dict::fuzz_decode_node(page, FTS_DICT);
let _ = FtsMeta::decode(page, 1);
for layout in [
PostingsLayout::FixedWidth,
PostingsLayout::DeltaVarint,
PostingsLayout::DeltaVarintSkip(SkipEntry::V5),
PostingsLayout::DeltaVarintSkip(SkipEntry::V6),
PostingsLayout::FrameOfReference,
] {
if page.len() > PAGE_HEADER_LEN {
let _ = Postings::decode(&page[PAGE_HEADER_LEN..], 1, layout);
}
let _ = Postings::decode(page, 1, layout);
}
let target = Ulid::from_parts(0x1234_5678, 0x9abc_def0);
for layout in [
PostingsLayout::DeltaVarintSkip(SkipEntry::V5),
PostingsLayout::DeltaVarintSkip(SkipEntry::V6),
PostingsLayout::FrameOfReference,
] {
let _ = lookup_via_skip(page, 1, target, layout);
if page.len() > PAGE_HEADER_LEN {
let _ = lookup_via_skip(&page[PAGE_HEADER_LEN..], 1, target, layout);
}
}
for layout in [
PostingsLayout::DeltaVarintSkip(SkipEntry::V6),
PostingsLayout::FrameOfReference,
] {
fuzz_bmw_cursor(page, layout);
if page.len() > PAGE_HEADER_LEN {
fuzz_bmw_cursor(&page[PAGE_HEADER_LEN..], layout);
}
}
}
fn fuzz_bmw_cursor(body: &[u8], layout: PostingsLayout) {
let mut counters = BmwCounters::default();
let Ok(Some(mut cursor)) = BmwCursor::open(1.0, body.to_vec(), 1, layout, &mut counters) else {
return;
};
let _ = cursor.covering_block(cursor.cur_id);
for _ in 0..64 {
if cursor.exhausted {
return;
}
let _ = cursor.current_tf(&mut counters);
if cursor.advance_past(cursor.cur_id, &mut counters).is_err() {
return;
}
}
let _ = cursor.advance_to(u128::MAX, &mut counters);
}
#[cfg(test)]
mod tests {
#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
use std::path::Path;
use std::sync::Arc;
use super::*;
use crate::storage::pager::{Pager, PagerOptions};
use crate::storage::sim::{SimVfs, SplitMix64};
use crate::storage::vfs::Vfs;
fn pager(page_size: u32) -> Pager {
let vfs: Arc<dyn Vfs> = Arc::new(SimVfs::new());
Pager::create(
vfs,
Path::new("memory.mind"),
PagerOptions {
page_size,
..Default::default()
},
)
.unwrap()
}
fn index_all(pager: &mut Pager, docs: &[&str]) -> Vec<Ulid> {
let mut ids = Vec::new();
let mut txn = pager.begin().unwrap();
for doc in docs {
let id = Ulid::new();
index_document(&mut txn, id, doc).unwrap();
ids.push(id);
}
txn.commit().unwrap();
ids
}
fn len_of<'a>(
map: &'a std::collections::HashMap<Ulid, String>,
) -> impl FnMut(Ulid) -> Result<Option<u32>> + 'a {
move |id| Ok(map.get(&id).map(|c| doc_len(c)))
}
#[test]
fn tokenize_is_lowercase_unicode_and_splits_on_punctuation() {
assert_eq!(tokenize("Hello, WORLD!"), vec!["hello", "world"]);
assert_eq!(
tokenize("Memória número 1 — teste"),
vec!["memória", "número", "1", "teste"]
);
assert!(tokenize(" ... ").is_empty());
assert_eq!(doc_len("Hello, WORLD! foo"), 3);
assert_eq!(doc_len(""), 0);
}
#[test]
fn postings_roundtrip_and_reject_unsorted() {
let mut p = Postings::default();
p.upsert(Ulid::from_parts(2, 0), 3);
p.upsert(Ulid::from_parts(1, 0), 1);
p.upsert(Ulid::from_parts(1, 0), 5); assert_eq!(p.entries.len(), 2);
for layout in [PostingsLayout::FixedWidth, PostingsLayout::DeltaVarint] {
let body = p.encode(layout);
assert_eq!(Postings::decode(&body, 1, layout).unwrap(), p);
let mut bad = 1_000_000u32.to_le_bytes().to_vec();
bad.extend_from_slice(&[0u8; 4]);
assert!(matches!(
Postings::decode(&bad, 1, layout),
Err(Error::MalformedPage { .. })
));
}
}
#[test]
fn delta_varint_roundtrips_and_shrinks_realistic_postings() {
let mut rng = SplitMix64(0x5EED_5EED);
let mut p = Postings::default();
for i in 0..500u64 {
let id = Ulid::from_parts(1_700_000_000_000 + i * 37, rng.next_u64().into());
p.upsert(id, 1 + (rng.next_u64() % 7) as u32);
}
let compact = p.encode(PostingsLayout::DeltaVarint);
let fixed = p.encode(PostingsLayout::FixedWidth);
assert_eq!(
Postings::decode(&compact, 1, PostingsLayout::DeltaVarint).unwrap(),
p
);
assert_eq!(
Postings::decode(&fixed, 1, PostingsLayout::FixedWidth).unwrap(),
p
);
assert!(
compact.len() < fixed.len(),
"delta+varint ({}) must beat fixed-width ({})",
compact.len(),
fixed.len()
);
let mut edge = Postings::default();
edge.upsert(Ulid::from(0u128), 1);
edge.upsert(Ulid::from(1u128), 2);
edge.upsert(Ulid::from(u128::MAX), u32::MAX);
let body = edge.encode(PostingsLayout::DeltaVarint);
assert_eq!(
Postings::decode(&body, 1, PostingsLayout::DeltaVarint).unwrap(),
edge
);
}
#[test]
fn delta_varint_rejects_hostile_bodies() {
let reject = |body: &[u8], what: &str| {
assert!(
matches!(
Postings::decode(body, 1, PostingsLayout::DeltaVarint),
Err(Error::MalformedPage { .. })
),
"must reject: {what}"
);
};
let mut body = 2u32.to_le_bytes().to_vec();
body.push(0x05); body.push(0x01); body.push(0x03); reject(&body, "truncated after last delta");
let mut body = 2u32.to_le_bytes().to_vec();
body.extend_from_slice(&[0x05, 0x01, 0x00, 0x01]);
reject(&body, "zero delta (duplicate id)");
let mut body = 1u32.to_le_bytes().to_vec();
body.extend_from_slice(&[0x05, 0x00]);
reject(&body, "zero term_freq");
let mut body = 1u32.to_le_bytes().to_vec();
body.push(0x05);
put_varint(&mut body, u128::from(u32::MAX) + 1);
reject(&body, "term_freq overflow");
let mut body = 2u32.to_le_bytes().to_vec();
put_varint(&mut body, u128::MAX);
body.push(0x01);
body.extend_from_slice(&[0x01, 0x01]); reject(&body, "id overflow");
let mut body = 1u32.to_le_bytes().to_vec();
body.extend_from_slice(&[0x80; 20]);
reject(&body, "varint too long");
let mut body = 1u32.to_le_bytes().to_vec();
body.extend_from_slice(&[0x80; 18]);
body.push(0x7F); reject(&body, "varint overflow bits");
}
fn realistic_postings(n: u64, seed: u64) -> Postings {
let mut rng = SplitMix64(seed);
let mut p = Postings::default();
for i in 0..n {
let id = Ulid::from_parts(1_700_000_000_000 + i * 37, rng.next_u64().into());
p.upsert(id, 1 + (rng.next_u64() % 9) as u32);
}
p
}
const SKIP_ENTRIES: [SkipEntry; 2] = [SkipEntry::V5, SkipEntry::V6];
const SKIP_LAYOUTS: [PostingsLayout; 3] = [
PostingsLayout::DeltaVarintSkip(SkipEntry::V5),
PostingsLayout::DeltaVarintSkip(SkipEntry::V6),
PostingsLayout::FrameOfReference,
];
#[test]
fn frame_of_reference_roundtrips_and_equals_delta_varint() {
let for_layout = PostingsLayout::FrameOfReference;
let v6 = PostingsLayout::DeltaVarintSkip(SkipEntry::V6);
for n in [
5u64,
20,
SKIP_MIN_DOC_FREQ as u64,
SKIP_MIN_DOC_FREQ as u64 + 1,
SKIP_MIN_DOC_FREQ as u64 + 200,
3 * SKIP_BLOCK_SIZE as u64,
] {
let p = realistic_postings(n, 0xF0 ^ n);
let for_body = p.encode(for_layout);
let decoded = Postings::decode(&for_body, 1, for_layout).unwrap();
assert_eq!(decoded.entries, p.entries, "FOR round-trip mismatch n={n}");
let v6_decoded = Postings::decode(&p.encode(v6), 1, v6).unwrap();
assert_eq!(decoded.entries, v6_decoded.entries, "FOR≠v6 for n={n}");
if (n as usize) < SKIP_MIN_DOC_FREQ {
assert_eq!(le_block_count(&for_body), 0);
assert_eq!(for_body, p.encode(v6), "small FOR body must equal v6");
} else {
assert!(le_block_count(&for_body) > 0);
}
}
}
#[test]
fn skip_layout_roundtrips_small_and_large() {
for entry in SKIP_ENTRIES {
let layout = PostingsLayout::DeltaVarintSkip(entry);
let small = realistic_postings(10, 0xA1);
let skip_body = small.encode(layout);
assert_eq!(Postings::decode(&skip_body, 1, layout).unwrap(), small);
let plain = small.encode(PostingsLayout::DeltaVarint);
assert_eq!(
&skip_body[8..],
&plain[4..],
"small body is plain past count ({entry:?})"
);
assert_eq!(le_block_count(&skip_body), 0);
assert_eq!(
skip_body,
small.encode(PostingsLayout::DeltaVarintSkip(SkipEntry::V5)),
"small body must not depend on entry width"
);
let large = realistic_postings(SKIP_MIN_DOC_FREQ as u64 + 55, 0xB2);
let body = large.encode(layout);
assert_eq!(Postings::decode(&body, 1, layout).unwrap(), large);
let expected_blocks = large.entries.len().div_ceil(SKIP_BLOCK_SIZE);
assert_eq!(le_block_count(&body) as usize, expected_blocks);
assert!(
expected_blocks >= 4,
"test corpus should span several blocks"
);
}
}
#[test]
fn v6_skip_entry_is_16_bytes_wider_per_block() {
let large = realistic_postings(SKIP_MIN_DOC_FREQ as u64 + 200, 0xF6);
let v5 = large.encode(PostingsLayout::DeltaVarintSkip(SkipEntry::V5));
let v6 = large.encode(PostingsLayout::DeltaVarintSkip(SkipEntry::V6));
let blocks = large.entries.len().div_ceil(SKIP_BLOCK_SIZE);
assert_eq!(v6.len(), v5.len() + 16 * blocks);
assert_eq!(SKIP_ENTRY_LEN_V6 - SKIP_ENTRY_LEN_V5, 16);
}
fn le_block_count(body: &[u8]) -> u32 {
u32::from_le_bytes(body[4..8].try_into().unwrap())
}
#[test]
fn lookup_via_skip_matches_linear_scan_for_every_id() {
for layout in SKIP_LAYOUTS {
let large = realistic_postings(SKIP_MIN_DOC_FREQ as u64 + 200, 0xC3);
let body = large.encode(layout);
for p in &large.entries {
assert_eq!(
lookup_via_skip(&body, 1, p.record_id, layout).unwrap(),
Some(p.term_freq),
"present id must be found via skip ({layout:?})"
);
}
let below = Ulid::from(0u128);
assert_eq!(lookup_via_skip(&body, 1, below, layout).unwrap(), None);
let above = Ulid::from(u128::MAX);
assert_eq!(lookup_via_skip(&body, 1, above, layout).unwrap(), None);
for w in large.entries.windows(2) {
let a = u128::from(w[0].record_id);
let b = u128::from(w[1].record_id);
if b - a > 1 {
let mid = Ulid::from(a + 1);
assert_eq!(lookup_via_skip(&body, 1, mid, layout).unwrap(), None);
break;
}
}
let small = realistic_postings(20, 0xD4);
let sbody = small.encode(layout);
for p in &small.entries {
assert_eq!(
lookup_via_skip(&sbody, 1, p.record_id, layout).unwrap(),
Some(p.term_freq)
);
}
assert_eq!(
lookup_via_skip(&sbody, 1, Ulid::from(u128::MAX), layout).unwrap(),
None
);
}
}
#[test]
fn skip_layout_rejects_hostile_bodies() {
for entry in SKIP_ENTRIES {
let layout = PostingsLayout::DeltaVarintSkip(entry);
let large = realistic_postings(SKIP_MIN_DOC_FREQ as u64 + 10, 0xE5);
let good = large.encode(layout);
let reject = |body: &[u8], what: &str| {
assert!(
matches!(
Postings::decode(body, 1, layout),
Err(Error::MalformedPage { .. })
),
"must reject: {what} ({entry:?})"
);
};
let mut bad = good.clone();
bad[4..8].copy_from_slice(&1u32.to_le_bytes());
reject(&bad, "block_count mismatch");
let mut bad = good.clone();
bad[8] ^= 0xFF;
reject(&bad, "first_id mismatch");
let tail = 8 + entry.tail_off();
let mut bad = good.clone();
bad[tail] = bad[tail].wrapping_add(1);
reject(&bad, "byte_offset mismatch");
let mut bad = good.clone();
bad[tail + 4] = bad[tail + 4].wrapping_add(1);
reject(&bad, "max_term_freq mismatch");
if entry == SkipEntry::V6 {
let mut bad = good.clone();
bad[8 + 16] ^= 0xFF; reject(&bad, "last_id mismatch");
}
let mut bad = 10_000_000u32.to_le_bytes().to_vec();
bad.extend_from_slice(&78_125u32.to_le_bytes()); reject(&bad, "skip index truncated");
}
}
#[test]
fn lookup_via_skip_rejects_corpus_crash_input() {
let path = Path::new(env!("CARGO_MANIFEST_DIR"))
.join("../../fuzz/corpus/fuzz_fts_page/crash-9116d630a5fae3ac97551c97104213cd2f5f4e9a");
let data = std::fs::read(&path).unwrap_or_else(|e| panic!("{}: {e}", path.display()));
crate::fuzz::fuzz_fts_page(&data);
let target = Ulid::from_parts(0x1234_5678, 0x9abc_def0);
for layout in SKIP_LAYOUTS {
let _ = lookup_via_skip(&data, 1, target, layout);
if data.len() > PAGE_HEADER_LEN {
let _ = lookup_via_skip(&data[PAGE_HEADER_LEN..], 1, target, layout);
}
}
}
#[test]
fn fuzz_fts_page_survives_every_corpus_seed() {
let dir = Path::new(env!("CARGO_MANIFEST_DIR")).join("../../fuzz/corpus/fuzz_fts_page");
let mut seen_v6 = false;
for ent in std::fs::read_dir(&dir).unwrap_or_else(|e| panic!("{}: {e}", dir.display())) {
let path = ent.unwrap().path();
if !path.is_file() {
continue;
}
let data = std::fs::read(&path).unwrap();
crate::fuzz::fuzz_fts_page(&data); if path
.file_name()
.and_then(|n| n.to_str())
.is_some_and(|n| n.contains("v6"))
{
seen_v6 = true;
}
}
assert!(seen_v6, "corpus must ship at least one v6 skip seed");
}
#[test]
fn index_and_search_ranks_by_relevance() {
let mut pager = pager(4096);
let ids = index_all(
&mut pager,
&[
"the rust compiler enforces memory safety",
"python is a dynamic language",
"rust rust rust is about memory and safety in rust",
],
);
let mut contents = std::collections::HashMap::new();
contents.insert(
ids[0],
"the rust compiler enforces memory safety".to_owned(),
);
contents.insert(ids[1], "python is a dynamic language".to_owned());
contents.insert(
ids[2],
"rust rust rust is about memory and safety in rust".to_owned(),
);
let root = pager.header().fts_root_page;
let hits = search(&pager, root, "rust memory", 10, |_| true, len_of(&contents)).unwrap();
assert_eq!(hits.len(), 2);
assert_eq!(hits[0].record_id, ids[2]);
assert_eq!(hits[1].record_id, ids[0]);
assert!(hits.iter().all(|h| h.score > 0.0));
assert!(hits[0].score > hits[1].score);
}
#[test]
fn search_profiled_matches_search_exactly() {
let mut pager = pager(4096);
let ids = index_all(
&mut pager,
&[
"the rust compiler enforces memory safety",
"python is a dynamic language",
"rust rust rust is about memory and safety in rust",
],
);
let mut contents = std::collections::HashMap::new();
contents.insert(
ids[0],
"the rust compiler enforces memory safety".to_owned(),
);
contents.insert(ids[1], "python is a dynamic language".to_owned());
contents.insert(
ids[2],
"rust rust rust is about memory and safety in rust".to_owned(),
);
let root = pager.header().fts_root_page;
let plain = search(&pager, root, "rust memory", 10, |_| true, len_of(&contents)).unwrap();
let (profiled, timings) = search_profiled(
&pager,
root,
"rust memory",
10,
|_| KeepOutcome::Accepted,
len_of(&contents),
)
.unwrap();
assert_eq!(plain, profiled);
assert_eq!(timings.terms_matched, 2); assert!(timings.postings_visited >= 3); }
#[test]
fn keep_outcomes_are_counted_once_per_distinct_candidate_by_reason() {
let mut pager = pager(4096);
let ids = index_all(
&mut pager,
&[
"the rust compiler enforces memory safety", "python is a dynamic language", "rust rust rust is about memory and safety", "rust engines are fast", ],
);
let mut contents = std::collections::HashMap::new();
for (i, text) in [
"the rust compiler enforces memory safety",
"python is a dynamic language",
"rust rust rust is about memory and safety",
"rust engines are fast",
]
.into_iter()
.enumerate()
{
contents.insert(ids[i], text.to_owned());
}
let root = pager.header().fts_root_page;
let (accepted_id, filtered_id, tombstoned_id, out_of_scope_id) =
(ids[0], ids[1], ids[2], ids[3]);
let (_, timings) = search_profiled(
&pager,
root,
"rust memory python engines",
10,
|id| {
if id == accepted_id {
KeepOutcome::Accepted
} else if id == filtered_id {
KeepOutcome::FilteredOut
} else if id == tombstoned_id {
KeepOutcome::Tombstoned
} else if id == out_of_scope_id {
KeepOutcome::OutOfScope
} else {
KeepOutcome::Tombstoned
}
},
len_of(&contents),
)
.unwrap();
assert_eq!(timings.keep_outcomes.accepted, 1);
assert_eq!(timings.keep_outcomes.filtered_out, 1);
assert_eq!(timings.keep_outcomes.tombstoned, 1);
assert_eq!(timings.keep_outcomes.out_of_scope, 1);
assert_eq!(timings.keep_outcomes.total(), 4);
assert_eq!(timings.keep_outcomes.rejected(), 3);
}
#[test]
fn early_termination_matches_exhaustive_scan_on_larger_corpus() {
let mut pager = pager(4096);
let mut docs: Vec<String> = Vec::new();
let corpus = SKIP_MIN_DOC_FREQ + 133;
for i in 0..corpus {
let mut d = "common ".repeat(1 + i % 7);
d.push_str(&"filler ".repeat(i % 23));
if i % 11 == 0 {
d.push_str("rare ");
}
if i % 37 == 0 {
d.push_str("rarest");
}
docs.push(d);
}
let doc_refs: Vec<&str> = docs.iter().map(String::as_str).collect();
let ids = index_all(&mut pager, &doc_refs);
let mut contents = std::collections::HashMap::new();
for (id, d) in ids.iter().zip(&docs) {
contents.insert(*id, d.clone());
}
let mut partial = contents.clone();
for id in ids.iter().step_by(5) {
partial.remove(id);
}
let dropped: std::collections::HashSet<Ulid> = ids
.iter()
.copied()
.enumerate()
.filter_map(|(i, id)| (i % 3 == 0).then_some(id))
.collect();
let root = pager.header().fts_root_page;
for query in [
"common",
"common rare",
"rare rarest common",
"filler common",
] {
for k in [1, 3, 10, 500] {
let plain = search(&pager, root, query, k, |_| true, len_of(&contents)).unwrap();
let (full, _) = search_profiled(
&pager,
root,
query,
k,
|_| KeepOutcome::Accepted,
len_of(&contents),
)
.unwrap();
assert_eq!(plain, full, "query={query:?} k={k}");
let lin =
search_linear(&pager, root, query, k, |_| true, len_of(&contents)).unwrap();
assert_eq!(plain, lin, "query={query:?} k={k} (linear)");
let plain = search(
&pager,
root,
query,
k,
|id| !dropped.contains(&id),
len_of(&contents),
)
.unwrap();
let (full, _) = search_profiled(
&pager,
root,
query,
k,
|id| {
if dropped.contains(&id) {
KeepOutcome::Tombstoned
} else {
KeepOutcome::Accepted
}
},
len_of(&contents),
)
.unwrap();
assert_eq!(plain, full, "query={query:?} k={k} (keep filter)");
let lin = search_linear(
&pager,
root,
query,
k,
|id| !dropped.contains(&id),
len_of(&contents),
)
.unwrap();
assert_eq!(plain, lin, "query={query:?} k={k} (keep filter, linear)");
let plain = search(&pager, root, query, k, |_| true, len_of(&partial)).unwrap();
let (full, _) = search_profiled(
&pager,
root,
query,
k,
|_| KeepOutcome::Accepted,
len_of(&partial),
)
.unwrap();
assert_eq!(plain, full, "query={query:?} k={k} (vanished records)");
let lin =
search_linear(&pager, root, query, k, |_| true, len_of(&partial)).unwrap();
assert_eq!(plain, lin, "query={query:?} k={k} (vanished, linear)");
}
}
}
#[test]
fn bmw_breaks_boundary_ties_exactly_like_the_oracle() {
let mut pager = pager(4096);
let corpus = SKIP_MIN_DOC_FREQ + 64;
let docs: Vec<String> = (0..corpus)
.map(|_| "twin memory entry".to_owned())
.collect();
let doc_refs: Vec<&str> = docs.iter().map(String::as_str).collect();
let ids = index_all(&mut pager, &doc_refs);
let mut contents = std::collections::HashMap::new();
for (id, d) in ids.iter().zip(&docs) {
contents.insert(*id, d.clone());
}
let root = pager.header().fts_root_page;
for k in [1, 2, 5, 100, corpus] {
let bmw = search(&pager, root, "twin entry", k, |_| true, len_of(&contents)).unwrap();
let (oracle, _) = search_profiled(
&pager,
root,
"twin entry",
k,
|_| KeepOutcome::Accepted,
len_of(&contents),
)
.unwrap();
assert_eq!(bmw, oracle, "k={k}");
let lin =
search_linear(&pager, root, "twin entry", k, |_| true, len_of(&contents)).unwrap();
assert_eq!(bmw, lin, "k={k} (linear)");
let mut expected = ids.clone();
expected.sort();
expected.truncate(k);
let got: Vec<Ulid> = bmw.iter().map(|h| h.record_id).collect();
assert_eq!(got, expected, "k={k} boundary ids");
}
}
#[test]
fn bmw_skips_blocks_and_still_matches_the_oracle() {
let mut pager = pager(4096);
let corpus = SKIP_MIN_DOC_FREQ * 3;
let mut docs: Vec<String> = Vec::new();
for i in 0..corpus {
if i == 3 {
docs.push("common ".repeat(30));
} else {
docs.push(format!("common {}", "filler ".repeat(50)));
}
}
let doc_refs: Vec<&str> = docs.iter().map(String::as_str).collect();
let ids = index_all(&mut pager, &doc_refs);
let mut contents = std::collections::HashMap::new();
for (id, d) in ids.iter().zip(&docs) {
contents.insert(*id, d.clone());
}
let root = pager.header().fts_root_page;
let (hits, counters) =
search_bmw_counted(&pager, root, "common", 1, |_| true, len_of(&contents)).unwrap();
let (oracle, _) = search_profiled(
&pager,
root,
"common",
1,
|_| KeepOutcome::Accepted,
len_of(&contents),
)
.unwrap();
assert_eq!(hits, oracle);
assert_eq!(hits[0].record_id, ids[3]);
let total_blocks = (corpus.div_ceil(SKIP_BLOCK_SIZE)) as u64;
assert_eq!(counters.blocks_total, total_blocks);
assert!(
counters.pivot_skips > 0,
"block-max check never fired: {counters:?}"
);
assert!(
counters.blocks_skipped() > counters.blocks_total / 2,
"BMW must skip most blocks: {counters:?}"
);
assert!(
counters.docs_evaluated <= (SKIP_BLOCK_SIZE + 8) as u64,
"documents outside the dominant block must never be evaluated: {counters:?}"
);
}
proptest::proptest! {
#![proptest_config(proptest::prelude::ProptestConfig::with_cases(40))]
#[test]
fn bmw_equals_oracle_on_random_corpora(
docs in proptest::collection::vec(
proptest::collection::vec(0usize..8, 1..24),
1..48,
),
query in proptest::collection::vec(0usize..8, 1..5),
k in 1usize..12,
keep_mask in proptest::prelude::any::<u64>(),
vanish_mask in proptest::prelude::any::<u64>(),
) {
const VOCAB: [&str; 8] = [
"alpha", "beta", "gamma", "delta", "memo", "rust", "engine", "wal",
];
let mut pager = pager(4096);
let rendered: Vec<String> = docs
.iter()
.map(|d| d.iter().map(|&w| VOCAB[w]).collect::<Vec<_>>().join(" "))
.collect();
let doc_refs: Vec<&str> = rendered.iter().map(String::as_str).collect();
let ids = index_all(&mut pager, &doc_refs);
let mut contents = std::collections::HashMap::new();
for (i, (id, d)) in ids.iter().zip(&rendered).enumerate() {
if vanish_mask & (1 << (i % 64)) == 0 {
contents.insert(*id, d.clone());
}
}
let kept: std::collections::HashSet<Ulid> = ids
.iter()
.enumerate()
.filter_map(|(i, id)| (keep_mask & (1 << (i % 64)) != 0).then_some(*id))
.collect();
let q = query.iter().map(|&w| VOCAB[w]).collect::<Vec<_>>().join(" ");
let root = pager.header().fts_root_page;
let keep = |id: Ulid| kept.contains(&id);
let bmw = search(&pager, root, &q, k, keep, len_of(&contents)).unwrap();
let (oracle, _) = search_profiled(
&pager,
root,
&q,
k,
|id| {
if keep(id) {
KeepOutcome::Accepted
} else {
KeepOutcome::Tombstoned
}
},
len_of(&contents),
)
.unwrap();
proptest::prop_assert_eq!(&bmw, &oracle);
let lin = search_linear(&pager, root, &q, k, keep, len_of(&contents)).unwrap();
proptest::prop_assert_eq!(&bmw, &lin);
}
}
#[test]
fn search_on_empty_or_missing_index_is_empty_not_error() {
let pager = pager(4096);
let none = std::collections::HashMap::new();
let hits = search(&pager, 0, "anything", 5, |_| true, len_of(&none)).unwrap();
assert!(hits.is_empty());
}
#[test]
fn keep_filter_excludes_tombstoned_documents() {
let mut pager = pager(4096);
let ids = index_all(&mut pager, &["shared term here", "shared term also"]);
let mut contents = std::collections::HashMap::new();
contents.insert(ids[0], "shared term here".to_owned());
contents.insert(ids[1], "shared term also".to_owned());
let excluded = ids[0];
let root = pager.header().fts_root_page;
let hits = search(
&pager,
root,
"shared term",
10,
|id| id != excluded,
len_of(&contents),
)
.unwrap();
assert_eq!(hits.len(), 1);
assert_eq!(hits[0].record_id, ids[1]);
}
#[test]
fn large_vocabulary_forces_splits_and_stays_correct() {
let mut pager = pager(512);
let mut docs = Vec::new();
for i in 0..200 {
docs.push(format!("term{i:04} common"));
}
let doc_refs: Vec<&str> = docs.iter().map(String::as_str).collect();
let ids = index_all(&mut pager, &doc_refs);
let mut contents = std::collections::HashMap::new();
for (id, doc) in ids.iter().zip(&docs) {
contents.insert(*id, doc.clone());
}
let root = pager.header().fts_root_page;
for (i, id) in ids.iter().enumerate() {
let q = format!("term{i:04}");
let hits = search(&pager, root, &q, 5, |_| true, len_of(&contents)).unwrap();
assert!(!hits.is_empty(), "term{i:04} not found");
assert_eq!(hits[0].record_id, *id, "term{i:04} ranked wrong");
}
let hits = search(&pager, root, "common", 500, |_| true, len_of(&contents)).unwrap();
assert_eq!(hits.len(), 200);
assert_eq!(indexed_documents(&pager, root).unwrap(), 200);
}
#[test]
fn survives_reopen() {
let vfs: Arc<dyn Vfs> = Arc::new(SimVfs::new());
let mut pager = Pager::create(
Arc::clone(&vfs),
Path::new("memory.mind"),
PagerOptions::default(),
)
.unwrap();
let ids = index_all(&mut pager, &["alpha beta", "beta gamma", "gamma delta"]);
let mut contents = std::collections::HashMap::new();
contents.insert(ids[0], "alpha beta".to_owned());
contents.insert(ids[1], "beta gamma".to_owned());
contents.insert(ids[2], "gamma delta".to_owned());
pager.close().unwrap();
let pager = Pager::open(vfs, Path::new("memory.mind"), PagerOptions::default()).unwrap();
let root = pager.header().fts_root_page;
let hits = search(&pager, root, "beta", 10, |_| true, len_of(&contents)).unwrap();
let found: std::collections::HashSet<Ulid> = hits.iter().map(|h| h.record_id).collect();
assert!(found.contains(&ids[0]));
assert!(found.contains(&ids[1]));
assert!(!found.contains(&ids[2]));
}
#[test]
fn format_version_3_file_reads_and_writes_fixed_width_postings() {
let vfs: Arc<dyn Vfs> = Arc::new(SimVfs::new());
let opts = PagerOptions {
page_size: 4096,
format_version: 3,
..Default::default()
};
let mut pager = Pager::create(Arc::clone(&vfs), Path::new("memory.mind"), opts).unwrap();
let ids = index_all(&mut pager, &["rust memory engine", "python memory model"]);
let mut contents = std::collections::HashMap::new();
contents.insert(ids[0], "rust memory engine".to_owned());
contents.insert(ids[1], "python memory model".to_owned());
let meta = load_meta(&pager, pager.header().fts_root_page)
.unwrap()
.unwrap();
let (body, page_no) = dict::get(&pager, FTS_DICT, meta.dict_root, b"memory")
.unwrap()
.unwrap();
assert_eq!(body.len(), 4 + 2 * POSTING_LEN);
let decoded = Postings::decode(&body, page_no, PostingsLayout::FixedWidth).unwrap();
assert_eq!(decoded.entries.len(), 2);
let root = pager.header().fts_root_page;
let hits = search(&pager, root, "memory", 10, |_| true, len_of(&contents)).unwrap();
assert_eq!(hits.len(), 2);
pager.close().unwrap();
let mut pager = Pager::open(
Arc::clone(&vfs),
Path::new("memory.mind"),
PagerOptions::default(),
)
.unwrap();
assert_eq!(pager.header().format_version, 3);
let new_ids = index_all(&mut pager, &["fresh memory entry"]);
contents.insert(new_ids[0], "fresh memory entry".to_owned());
let meta = load_meta(&pager, pager.header().fts_root_page)
.unwrap()
.unwrap();
let (body, page_no) = dict::get(&pager, FTS_DICT, meta.dict_root, b"memory")
.unwrap()
.unwrap();
assert_eq!(body.len(), 4 + 3 * POSTING_LEN);
assert!(Postings::decode(&body, page_no, PostingsLayout::FixedWidth).is_ok());
let root = pager.header().fts_root_page;
let hits = search(&pager, root, "memory", 10, |_| true, len_of(&contents)).unwrap();
assert_eq!(hits.len(), 3);
}
#[test]
fn new_files_store_postings_as_delta_varint_skip() {
let mut pager = pager(4096);
assert_eq!(pager.header().format_version, crate::format::FORMAT_VERSION);
assert_eq!(pager.header().format_version, 8);
let ids = index_all(&mut pager, &["memory one", "memory two", "memory three"]);
let meta = load_meta(&pager, pager.header().fts_root_page)
.unwrap()
.unwrap();
let (body, page_no) = dict::get(&pager, FTS_DICT, meta.dict_root, b"memory")
.unwrap()
.unwrap();
assert_eq!(le_block_count(&body), 0);
assert!(body.len() < 4 + 3 * POSTING_LEN);
let decoded = Postings::decode(&body, page_no, PostingsLayout::FrameOfReference).unwrap();
let got: Vec<Ulid> = decoded.entries.iter().map(|p| p.record_id).collect();
let mut expected = ids.clone();
expected.sort();
assert_eq!(got, expected);
}
#[test]
fn format_version_5_file_reads_and_writes_v5_skip_entries() {
let vfs: Arc<dyn Vfs> = Arc::new(SimVfs::new());
let opts = PagerOptions {
page_size: 4096,
format_version: 5,
..Default::default()
};
let mut pager = Pager::create(Arc::clone(&vfs), Path::new("memory.mind"), opts).unwrap();
let mut docs = Vec::new();
for i in 0..(SKIP_MIN_DOC_FREQ + 20) {
docs.push(format!("shared unique{i:04}"));
}
let doc_refs: Vec<&str> = docs.iter().map(String::as_str).collect();
index_all(&mut pager, &doc_refs);
let meta = load_meta(&pager, pager.header().fts_root_page)
.unwrap()
.unwrap();
let (body, page_no) = dict::get(&pager, FTS_DICT, meta.dict_root, b"shared")
.unwrap()
.unwrap();
assert!(
le_block_count(&body) > 0,
"v5 file should carry a skip index"
);
let decoded = Postings::decode(
&body,
page_no,
PostingsLayout::DeltaVarintSkip(SkipEntry::V5),
)
.unwrap();
assert_eq!(decoded.entries.len(), SKIP_MIN_DOC_FREQ + 20);
for p in &decoded.entries {
assert_eq!(
lookup_via_skip(
&body,
page_no,
p.record_id,
PostingsLayout::DeltaVarintSkip(SkipEntry::V5)
)
.unwrap(),
Some(p.term_freq)
);
}
}
#[test]
fn format_version_4_file_reads_and_writes_skipless_postings() {
let vfs: Arc<dyn Vfs> = Arc::new(SimVfs::new());
let opts = PagerOptions {
page_size: 4096,
format_version: 4,
..Default::default()
};
let mut pager = Pager::create(Arc::clone(&vfs), Path::new("memory.mind"), opts).unwrap();
let mut docs = Vec::new();
for i in 0..(SKIP_MIN_DOC_FREQ + 20) {
docs.push(format!("shared unique{i:04}"));
}
let doc_refs: Vec<&str> = docs.iter().map(String::as_str).collect();
let ids = index_all(&mut pager, &doc_refs);
let mut contents = std::collections::HashMap::new();
for (id, d) in ids.iter().zip(&docs) {
contents.insert(*id, d.clone());
}
let meta = load_meta(&pager, pager.header().fts_root_page)
.unwrap()
.unwrap();
let (body, page_no) = dict::get(&pager, FTS_DICT, meta.dict_root, b"shared")
.unwrap()
.unwrap();
let decoded = Postings::decode(&body, page_no, PostingsLayout::DeltaVarint).unwrap();
assert_eq!(decoded.entries.len(), SKIP_MIN_DOC_FREQ + 20);
let root = pager.header().fts_root_page;
let hits = search(&pager, root, "shared", 5000, |_| true, len_of(&contents)).unwrap();
assert_eq!(hits.len(), SKIP_MIN_DOC_FREQ + 20);
pager.close().unwrap();
let mut pager = Pager::open(
Arc::clone(&vfs),
Path::new("memory.mind"),
PagerOptions::default(),
)
.unwrap();
assert_eq!(pager.header().format_version, 4);
let new_ids = index_all(&mut pager, &["shared fresh entry"]);
contents.insert(new_ids[0], "shared fresh entry".to_owned());
let meta = load_meta(&pager, pager.header().fts_root_page)
.unwrap()
.unwrap();
let (body, page_no) = dict::get(&pager, FTS_DICT, meta.dict_root, b"shared")
.unwrap()
.unwrap();
assert!(Postings::decode(&body, page_no, PostingsLayout::DeltaVarint).is_ok());
let root = pager.header().fts_root_page;
let hits = search(&pager, root, "shared", 5000, |_| true, len_of(&contents)).unwrap();
assert_eq!(hits.len(), SKIP_MIN_DOC_FREQ + 21);
}
#[test]
fn new_files_store_large_postings_with_a_skip_index() {
let mut pager = pager(4096);
let mut docs = Vec::new();
for i in 0..(SKIP_MIN_DOC_FREQ + 40) {
docs.push(format!("shared unique{i:04}"));
}
let doc_refs: Vec<&str> = docs.iter().map(String::as_str).collect();
let ids = index_all(&mut pager, &doc_refs);
let mut contents = std::collections::HashMap::new();
for (id, d) in ids.iter().zip(&docs) {
contents.insert(*id, d.clone());
}
let meta = load_meta(&pager, pager.header().fts_root_page)
.unwrap()
.unwrap();
let (body, _) = dict::get(&pager, FTS_DICT, meta.dict_root, b"shared")
.unwrap()
.unwrap();
assert!(
le_block_count(&body) > 0,
"large term must carry a skip index"
);
let root = pager.header().fts_root_page;
let hits = search(&pager, root, "shared", 5000, |_| true, len_of(&contents)).unwrap();
assert_eq!(hits.len(), SKIP_MIN_DOC_FREQ + 40);
}
#[test]
fn delta_varint_postings_survive_overflow_chains() {
let mut pager = pager(512);
let mut docs = Vec::new();
for i in 0..120 {
docs.push(format!("shared corpus word plus unique{i:03}"));
}
let doc_refs: Vec<&str> = docs.iter().map(String::as_str).collect();
let ids = index_all(&mut pager, &doc_refs);
let mut contents = std::collections::HashMap::new();
for (id, doc) in ids.iter().zip(&docs) {
contents.insert(*id, doc.clone());
}
let root = pager.header().fts_root_page;
let hits = search(&pager, root, "shared", 500, |_| true, len_of(&contents)).unwrap();
assert_eq!(hits.len(), 120);
}
#[test]
fn index_within_txn_is_searchable_before_commit() {
let mut pager = pager(4096);
let id = Ulid::new();
let mut txn = pager.begin().unwrap();
index_document(&mut txn, id, "in flight transaction text").unwrap();
let root = txn.fts_root_page();
let mut contents = std::collections::HashMap::new();
contents.insert(id, "in flight transaction text".to_owned());
let hits = search(&txn, root, "flight", 5, |_| true, len_of(&contents)).unwrap();
assert_eq!(hits.len(), 1);
assert_eq!(hits[0].record_id, id);
drop(txn); assert_eq!(pager.header().fts_root_page, 0);
}
#[test]
fn decode_never_panics_on_arbitrary_bytes() {
let mut rng = SplitMix64(0xF75_u64);
for _ in 0..3000 {
let len = [64usize, 200, 512, 4096][(rng.next_u64() % 4) as usize];
let mut page = vec![0u8; len];
for b in &mut page {
*b = rng.next_u64() as u8;
}
fuzz_decode_page(&page); }
let entry = dict::LeafEntry {
key: b"rust".to_vec(),
value: dict::Value::Overflow {
total_len: 40,
first_page: 5,
},
};
let valid = dict::encode_leaf(std::slice::from_ref(&entry), 512, FTS_DICT).unwrap();
let meta = FtsMeta {
doc_count: 3,
total_tokens: 30,
dict_root: 7,
}
.encode(512)
.unwrap();
for base in [valid, meta] {
for _ in 0..2000 {
let mut page = base.clone();
let i = (rng.next_u64() as usize) % page.len();
page[i] ^= (rng.next_u64() as u8) | 1;
fuzz_decode_page(&page);
}
}
}
}