steeldb/

bitmap.rs

//! Postings — an exact set of unsigned situation ids with roaring-compatible set algebra.
//!
//! The whole engine is bit-parallel logic over these: AND (∩), OR (∪), AND-NOT (−), popcount.
//! `Postings` is the trait the query layer is generic over, so we can benchmark the historical
//! `Set<number>` baseline (`SetPostings`, a `HashSet<u32>`) against real roaring containers
//! (`RoarPostings`) without touching a single caller — exactly the swap the TS `Bitmap` was
//! designed for ("swap in a native roaring library later without touching callers").
//!
//! `serialize_deltagap` uses the same varint delta-gap (LEB128 over sorted gaps) as the TS/Python
//! parts store, so posting blobs stay byte-compatible and portable across engines.

use roaring::RoaringBitmap;
use std::collections::HashSet;

/// Set-algebra surface the tokenql evaluator and index build against.
pub trait Postings: Clone {
    fn empty() -> Self;
    /// Build from strictly-ascending, de-duplicated ids (the natural order of an inverted index).
    fn from_sorted(ids: &[u32]) -> Self;
    fn insert(&mut self, id: u32);
    fn contains(&self, id: u32) -> bool;
    fn len(&self) -> usize;
    fn is_empty(&self) -> bool {
        self.len() == 0
    }
    /// ∩
    fn and(&self, other: &Self) -> Self;
    /// ∪
    fn or(&self, other: &Self) -> Self;
    /// − (this and-not other)
    fn and_not(&self, other: &Self) -> Self;
    /// in-place ∪ (the `|=` of the build/OR loop)
    fn or_inplace(&mut self, other: &Self);
    /// ids in ascending order
    fn to_sorted(&self) -> Vec<u32>;
    /// byte size of this set's native serialized form (memory proxy)
    fn native_bytes(&self) -> usize;

    /// varint delta-gap of the sorted ids — byte-identical to the TS/Python parts store.
    fn serialize_deltagap(&self) -> Vec<u8> {
        let ids = self.to_sorted();
        let mut out = Vec::new();
        let mut prev: u32 = 0;
        for id in ids {
            let mut n = id.wrapping_sub(prev);
            prev = id;
            loop {
                let b = (n & 0x7f) as u8;
                n >>= 7;
                if n != 0 {
                    out.push(b | 0x80);
                } else {
                    out.push(b);
                    break;
                }
            }
        }
        out
    }

    /// inverse of `serialize_deltagap`
    fn deserialize_deltagap(bytes: &[u8]) -> Self
    where
        Self: Sized,
    {
        let mut ids = Vec::new();
        let mut prev: u32 = 0;
        let mut i = 0;
        while i < bytes.len() {
            let mut shift = 0u32;
            let mut val: u32 = 0;
            loop {
                let b = bytes[i];
                i += 1;
                val |= ((b & 0x7f) as u32) << shift;
                if b & 0x80 == 0 {
                    break;
                }
                shift += 7;
            }
            prev = prev.wrapping_add(val);
            ids.push(prev);
        }
        Self::from_sorted(&ids)
    }
}

// ── Baseline: HashSet<u32> (the TS `Set<number>` engine, ported faithfully) ──────────────────

#[derive(Clone, Default)]
pub struct SetPostings(pub HashSet<u32>);

impl Postings for SetPostings {
    fn empty() -> Self {
        SetPostings(HashSet::new())
    }
    fn from_sorted(ids: &[u32]) -> Self {
        SetPostings(ids.iter().copied().collect())
    }
    fn insert(&mut self, id: u32) {
        self.0.insert(id);
    }
    fn contains(&self, id: u32) -> bool {
        self.0.contains(&id)
    }
    fn len(&self) -> usize {
        self.0.len()
    }
    fn and(&self, other: &Self) -> Self {
        // iterate the smaller, probe the larger (matches the TS `and`)
        let (small, big) = if self.0.len() <= other.0.len() {
            (&self.0, &other.0)
        } else {
            (&other.0, &self.0)
        };
        SetPostings(small.iter().copied().filter(|id| big.contains(id)).collect())
    }
    fn or(&self, other: &Self) -> Self {
        let mut out = self.0.clone();
        out.extend(other.0.iter().copied());
        SetPostings(out)
    }
    fn and_not(&self, other: &Self) -> Self {
        SetPostings(self.0.iter().copied().filter(|id| !other.0.contains(id)).collect())
    }
    fn or_inplace(&mut self, other: &Self) {
        self.0.extend(other.0.iter().copied());
    }
    fn to_sorted(&self) -> Vec<u32> {
        let mut v: Vec<u32> = self.0.iter().copied().collect();
        v.sort_unstable();
        v
    }
    fn native_bytes(&self) -> usize {
        // hashbrown table: ~ capacity * (4 byte key + 1 control byte), load-factor ~7/8
        let cap = (self.0.len() as f64 / 0.875).ceil() as usize;
        cap * 5
    }
}

// ── Roaring containers (the recommendation) ──────────────────────────────────────────────────

#[derive(Clone, Default, Debug)]
pub struct RoarPostings(pub RoaringBitmap);

impl Postings for RoarPostings {
    fn empty() -> Self {
        RoarPostings(RoaringBitmap::new())
    }
    fn from_sorted(ids: &[u32]) -> Self {
        // ids are ascending & unique → the fast bulk path that builds optimal containers
        match RoaringBitmap::from_sorted_iter(ids.iter().copied()) {
            Ok(b) => RoarPostings(b),
            Err(_) => {
                let mut b = RoaringBitmap::new();
                b.extend(ids.iter().copied());
                RoarPostings(b)
            }
        }
    }
    fn insert(&mut self, id: u32) {
        self.0.insert(id);
    }
    fn contains(&self, id: u32) -> bool {
        self.0.contains(id)
    }
    fn len(&self) -> usize {
        self.0.len() as usize
    }
    fn and(&self, other: &Self) -> Self {
        RoarPostings(&self.0 & &other.0)
    }
    fn or(&self, other: &Self) -> Self {
        RoarPostings(&self.0 | &other.0)
    }
    fn and_not(&self, other: &Self) -> Self {
        RoarPostings(&self.0 - &other.0)
    }
    fn or_inplace(&mut self, other: &Self) {
        self.0 |= &other.0;
    }
    fn to_sorted(&self) -> Vec<u32> {
        self.0.iter().collect()
    }
    fn native_bytes(&self) -> usize {
        self.0.serialized_size()
    }
}

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

    fn roundtrip<B: Postings>() {
        let ids = [0u32, 1, 2, 5, 300, 70000, 70001, 1_000_000];
        let b = B::from_sorted(&ids);
        let bytes = b.serialize_deltagap();
        let back = B::deserialize_deltagap(&bytes);
        assert_eq!(back.to_sorted(), ids.to_vec());
    }

    #[test]
    fn deltagap_roundtrip_both_backends() {
        roundtrip::<SetPostings>();
        roundtrip::<RoarPostings>();
    }

    #[test]
    fn set_algebra_parity() {
        let a_ids = [1u32, 2, 3, 4, 5, 100, 200];
        let b_ids = [3u32, 4, 5, 6, 200, 300];
        let (sa, sb) = (SetPostings::from_sorted(&a_ids), SetPostings::from_sorted(&b_ids));
        let (ra, rb) = (RoarPostings::from_sorted(&a_ids), RoarPostings::from_sorted(&b_ids));
        assert_eq!(sa.and(&sb).to_sorted(), ra.and(&rb).to_sorted());
        assert_eq!(sa.or(&sb).to_sorted(), ra.or(&rb).to_sorted());
        assert_eq!(sa.and_not(&sb).to_sorted(), ra.and_not(&rb).to_sorted());
    }
}