seagrep-index 0.8.0

//! Delta-bitpacked posting lists for the sparse strategy (Lucene/tantivy's
//! layout): blocks of 128 doc-id deltas, each block packed at the width of
//! its own largest delta behind a one-byte header. Dense lists collapse to
//! one or two bits per id where the fixed-width codec paid the full
//! `ceil(log2(doc_count))`. The encoded length is not derivable from the
//! count alone, so the dictionary carries it (measured net win on real
//! prose: −17.9% of postings).

use anyhow::{Context, Result};
use seagrep_core::DocId;

pub(crate) const BLOCK_IDS: usize = 128;

/// Encode strictly-ascending unique ids. The first delta of the first block
/// is the raw first id; every later delta is `id - previous - 1`, so
/// consecutive runs pack at width 1 (width 0 is invalid, keeping headers
/// honest and corruption detectable).
pub(crate) fn encode_delta_blocks(ids: &[DocId], out: &mut Vec<u8>) {
    let mut previous: Option<DocId> = None;
    for chunk in ids.chunks(BLOCK_IDS) {
        let mut deltas = [0u32; BLOCK_IDS];
        let mut max_delta = 0u32;
        for (slot, &id) in chunk.iter().enumerate() {
            let delta = match previous {
                None => id,
                Some(previous) => id - previous - 1,
            };
            deltas[slot] = delta;
            max_delta = max_delta.max(delta);
            previous = Some(id);
        }
        let width = bits_for(max_delta);
        out.push(width as u8);
        pack_fixed(&deltas[..chunk.len()], width, out);
    }
}

/// Decode `count` ids, validating widths, bounds, and strict ascent — a
/// block that fails any check is a corrupt index, reported loudly.
pub(crate) fn decode_delta_blocks(bytes: &[u8], count: u32, doc_count: u32) -> Result<Vec<DocId>> {
    anyhow::ensure!(
        count <= doc_count,
        "posting count {count} exceeds doc count {doc_count}"
    );
    let mut ids = Vec::with_capacity(count as usize);
    let mut cursor = 0usize;
    let mut previous: Option<DocId> = None;
    let mut remaining = count as usize;
    while remaining > 0 {
        let width = usize::from(
            *bytes
                .get(cursor)
                .context("posting block ends before its width header")?,
        );
        cursor += 1;
        anyhow::ensure!(
            (1..=32).contains(&width),
            "posting block width {width} is invalid"
        );
        let block_ids = remaining.min(BLOCK_IDS);
        let block_bytes = (block_ids * width).div_ceil(8);
        let packed = bytes
            .get(cursor..cursor + block_bytes)
            .context("posting block ends inside its packed ids")?;
        cursor += block_bytes;
        unpack_fixed(packed, block_ids, width, |delta| {
            let id = match previous {
                None => delta,
                Some(previous) => previous
                    .checked_add(delta)
                    .and_then(|sum| sum.checked_add(1))
                    .context("posting delta overflows the id space")?,
            };
            anyhow::ensure!(id < doc_count, "posting id {id} is out of bounds");
            previous = Some(id);
            ids.push(id);
            Ok(())
        })?;
        remaining -= block_ids;
    }
    anyhow::ensure!(
        cursor == bytes.len(),
        "posting list has {} trailing bytes",
        bytes.len() - cursor
    );
    Ok(ids)
}

/// Exact encoded length without materializing the encoding; tests use it to
/// pin the writer's length accounting (the build path measures the actual
/// buffer).
#[cfg(test)]
pub(crate) fn encoded_len(ids: &[DocId]) -> u64 {
    let mut total = 0u64;
    let mut previous: Option<DocId> = None;
    for chunk in ids.chunks(BLOCK_IDS) {
        let mut max_delta = 0u32;
        for &id in chunk {
            let delta = match previous {
                None => id,
                Some(previous) => id - previous - 1,
            };
            max_delta = max_delta.max(delta);
            previous = Some(id);
        }
        let width = bits_for(max_delta);
        total += 1 + ((chunk.len() * width).div_ceil(8)) as u64;
    }
    total
}

fn bits_for(value: u32) -> usize {
    (32 - value.leading_zeros()).max(1) as usize
}

fn pack_fixed(values: &[u32], width: usize, out: &mut Vec<u8>) {
    let mut acc: u64 = 0;
    let mut filled: usize = 0;
    for &value in values {
        acc |= u64::from(value) << filled;
        filled += width;
        while filled >= 8 {
            out.push(acc as u8);
            acc >>= 8;
            filled -= 8;
        }
    }
    if filled > 0 {
        out.push(acc as u8);
    }
}

fn unpack_fixed(
    bytes: &[u8],
    n: usize,
    width: usize,
    mut visit: impl FnMut(u32) -> Result<()>,
) -> Result<()> {
    let mut acc: u64 = 0;
    let mut filled: usize = 0;
    let mut input = bytes.iter();
    let mask: u64 = (1u64 << width) - 1;
    for _ in 0..n {
        while filled < width {
            acc |= u64::from(*input.next().context("packed ids are truncated")?) << filled;
            filled += 8;
        }
        visit((acc & mask) as u32)?;
        acc >>= width;
        filled -= width;
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    fn round_trip(ids: &[DocId], doc_count: u32) {
        let mut encoded = Vec::new();
        encode_delta_blocks(ids, &mut encoded);
        assert_eq!(
            encoded.len() as u64,
            encoded_len(ids),
            "length prediction must match the encoding"
        );
        let decoded = decode_delta_blocks(&encoded, ids.len() as u32, doc_count).unwrap();
        assert_eq!(decoded, ids);
    }

    #[test]
    fn round_trips_every_shape() {
        round_trip(&[], 10);
        round_trip(&[0], 10);
        round_trip(&[9], 10);
        round_trip(&(0..128).collect::<Vec<_>>(), 200); // exactly one full block, width 1
        round_trip(&(0..129).collect::<Vec<_>>(), 200); // full block + 1
        round_trip(&(0..1000).collect::<Vec<_>>(), 1000); // dense: every doc
        round_trip(&[0, 1_000_000, 4_000_000 - 1], 4_000_000); // huge deltas
        let sparse: Vec<DocId> = (0..3000).map(|i| i * 7 + 3).collect();
        round_trip(&sparse, 30_000);
        // adversarial widths: one huge delta inside an otherwise-dense block
        let mut mixed: Vec<DocId> = (0..127).collect();
        mixed.push(3_999_999);
        round_trip(&mixed, 4_000_000);
    }

    #[test]
    fn dense_lists_shrink_versus_fixed_width() {
        // 1000 consecutive ids in a 4M-doc segment: fixed-width pays 22
        // bits/id; deltas pay 1 bit/id plus headers.
        let ids: Vec<DocId> = (500..1500).collect();
        let mut encoded = Vec::new();
        encode_delta_blocks(&ids, &mut encoded);
        assert!(
            encoded.len() < 1000 * 22 / 8 / 4,
            "expected a big win, got {} bytes",
            encoded.len()
        );
    }

    #[test]
    fn corrupt_blocks_fail_loudly() {
        let ids: Vec<DocId> = (0..300).map(|i| i * 3).collect();
        let mut encoded = Vec::new();
        encode_delta_blocks(&ids, &mut encoded);

        // truncated
        assert!(decode_delta_blocks(&encoded[..encoded.len() - 1], 300, 1000).is_err());
        // trailing garbage
        let mut long = encoded.clone();
        long.push(0);
        assert!(decode_delta_blocks(&long, 300, 1000).is_err());
        // invalid width header
        let mut bad = encoded.clone();
        bad[0] = 0;
        assert!(decode_delta_blocks(&bad, 300, 1000).is_err());
        bad[0] = 33;
        assert!(decode_delta_blocks(&bad, 300, 1000).is_err());
        // out-of-bounds id
        assert!(decode_delta_blocks(&encoded, 300, 500).is_err());
        // count larger than doc_count
        assert!(decode_delta_blocks(&encoded, 1001, 1000).is_err());
    }
}