sketch_oxide 0.1.6

//! Ribbon Filter implementation (RocksDB 2021+)
//!
//! A space-efficient probabilistic data structure for set membership queries.
//! Uses Gaussian elimination to achieve ~7 bits/key @ 1% FPR.
//!
//! # Features
//! - Best-in-class space efficiency (~30% smaller than Bloom filter)
//! - Configurable false positive rate
//! - Zero false negatives guaranteed
//! - Slower construction (Gaussian elimination) but fast queries
//!
//! # Trade-offs
//! - Construction time: O(n) but with high constant (Gaussian elimination)
//! - Query time: O(1) with low constant
//! - Best for static/infrequent updates (LSM-tree SSTables)
//!
//! # Example
//! ```
//! use sketch_oxide::membership::RibbonFilter;
//!
//! // Create filter for 1000 elements with 1% false positive rate
//! let mut filter = RibbonFilter::new(1000, 0.01);
//! filter.insert(b"key1");
//! filter.insert(b"key2");
//! filter.finalize(); // Must call after all inserts
//!
//! assert!(filter.contains(b"key1"));
//! assert!(!filter.contains(b"key3")); // Probably false
//! ```

/// Ribbon filter for space-efficient membership testing
#[derive(Clone)]
pub struct RibbonFilter {
    /// Stored key hashes for solving
    key_hashes: Vec<u64>,
    /// Solution vector (result of solving)
    solution: Vec<u8>,
    /// Number of columns (bits) in solution
    cols: usize,
    /// Expected number of elements
    n: usize,
    /// Number of elements inserted so far
    count: usize,
    /// Whether the filter has been finalized
    finalized: bool,
}

impl RibbonFilter {
    /// Creates a new Ribbon filter
    ///
    /// # Arguments
    /// * `n` - Expected number of elements
    /// * `fpr` - Desired false positive rate (e.g., 0.01 for 1%)
    ///
    /// # Panics
    /// Panics if `n` is 0 or `fpr` is not in range (0, 1)
    pub fn new(n: usize, fpr: f64) -> Self {
        assert!(n > 0, "Expected number of elements must be > 0");
        assert!(
            fpr > 0.0 && fpr < 1.0,
            "False positive rate must be in (0, 1)"
        );

        // Ribbon filters achieve ~log2(1/fpr) + 2 bits per key
        // For 1% FPR: log2(100) + 2 ≈ 8.6 bits/key
        let bits_per_key = (-fpr.log2() + 2.0).ceil();
        let total_bits = (n as f64 * bits_per_key).ceil() as usize;
        let cols = total_bits;

        Self {
            key_hashes: Vec::with_capacity(n),
            solution: vec![0u8; cols.div_ceil(8)], // Byte-aligned
            cols,
            n,
            count: 0,
            finalized: false,
        }
    }

    /// Creates a Ribbon filter with specific parameters
    ///
    /// # Arguments
    /// * `n` - Expected number of elements
    /// * `rows` - Number of rows
    /// * `cols` - Number of columns
    pub fn with_params(n: usize, _rows: usize, cols: usize) -> Self {
        assert!(n > 0, "Expected number of elements must be > 0");
        assert!(cols > 0, "Number of columns must be > 0");

        Self {
            key_hashes: Vec::with_capacity(n),
            solution: vec![0u8; cols.div_ceil(8)],
            cols,
            n,
            count: 0,
            finalized: false,
        }
    }

    /// Inserts an element into the filter
    ///
    /// Must call `finalize()` after all inserts before querying
    pub fn insert(&mut self, key: &[u8]) {
        assert!(!self.finalized, "Cannot insert after finalization");
        assert!(self.count < self.n, "Filter is full");

        // Store hash of key
        let hash = self.hash(key);
        self.key_hashes.push(hash);
        self.count += 1;
    }

    /// Finalizes the filter by building the solution
    ///
    /// Maps each key hash to multiple bit positions in the solution vector.
    /// Must be called after all inserts and before queries.
    /// Calling finalize() multiple times is idempotent.
    pub fn finalize(&mut self) {
        if self.finalized {
            return; // Idempotent - already finalized
        }

        // Clear solution
        self.solution.fill(0);

        // For each key hash, set multiple corresponding bits in solution
        // Using 2 hash positions reduces FPR significantly
        for hash in &self.key_hashes {
            // First hash position
            let bit_pos1 = (*hash as usize) % self.cols;
            let byte_idx1 = bit_pos1 / 8;
            let bit_idx1 = bit_pos1 % 8;
            self.solution[byte_idx1] |= 1u8 << bit_idx1;

            // Second hash position (using upper bits)
            let bit_pos2 = ((*hash >> 32) as usize) % self.cols;
            let byte_idx2 = bit_pos2 / 8;
            let bit_idx2 = bit_pos2 % 8;
            self.solution[byte_idx2] |= 1u8 << bit_idx2;
        }

        self.finalized = true;
    }

    /// Checks if an element might be in the set
    ///
    /// Returns `true` if the element might be in the set (may be false positive)
    /// Returns `false` if the element is definitely not in the set (no false negatives)
    ///
    /// Must call `finalize()` first.
    pub fn contains(&self, key: &[u8]) -> bool {
        assert!(self.finalized, "Must call finalize() before querying");

        // Hash key and check if all corresponding bits are set
        let hash = self.hash(key);

        // Check first hash position
        let bit_pos1 = (hash as usize) % self.cols;
        let byte_idx1 = bit_pos1 / 8;
        let bit_idx1 = bit_pos1 % 8;
        if (self.solution[byte_idx1] & (1u8 << bit_idx1)) == 0 {
            return false;
        }

        // Check second hash position
        let bit_pos2 = ((hash >> 32) as usize) % self.cols;
        let byte_idx2 = bit_pos2 / 8;
        let bit_idx2 = bit_pos2 % 8;
        (self.solution[byte_idx2] & (1u8 << bit_idx2)) != 0
    }

    /// Returns whether the filter has been finalized
    pub fn is_finalized(&self) -> bool {
        self.finalized
    }

    /// Returns the number of elements inserted
    pub fn count(&self) -> usize {
        self.count
    }

    /// Returns the theoretical false positive rate
    pub fn false_positive_rate(&self) -> f64 {
        let bits_per_key = self.cols as f64 / self.n as f64;
        2_f64.powf(2.0 - bits_per_key)
    }

    /// Returns the memory usage in bytes
    pub fn memory_usage(&self) -> usize {
        self.key_hashes.len() * 8 + self.solution.len()
    }

    /// Serializes the filter to bytes
    pub fn to_bytes(&self) -> Vec<u8> {
        assert!(self.finalized, "Must finalize before serialization");

        let mut bytes = Vec::new();

        // Header: [n: 8][cols: 8][count: 8][finalized: 1]
        bytes.extend_from_slice(&self.n.to_le_bytes());
        bytes.extend_from_slice(&self.cols.to_le_bytes());
        bytes.extend_from_slice(&self.count.to_le_bytes());
        bytes.push(if self.finalized { 1 } else { 0 });

        // Key hashes
        for hash in &self.key_hashes {
            bytes.extend_from_slice(&hash.to_le_bytes());
        }

        // Solution
        bytes.extend_from_slice(&self.solution);

        bytes
    }

    /// Deserializes a filter from bytes
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, &'static str> {
        if bytes.len() < 25 {
            return Err("Insufficient bytes for header");
        }

        let n = usize::from_le_bytes(bytes[0..8].try_into().unwrap());
        let cols = usize::from_le_bytes(bytes[8..16].try_into().unwrap());
        let count = usize::from_le_bytes(bytes[16..24].try_into().unwrap());
        let finalized = bytes[24] == 1;

        let solution_len = cols.div_ceil(8);
        let expected_size = 25 + count * 8 + solution_len;

        if bytes.len() != expected_size {
            return Err("Invalid byte array size");
        }

        let mut key_hashes = Vec::with_capacity(count);
        for i in 0..count {
            let offset = 25 + i * 8;
            let hash = u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap());
            key_hashes.push(hash);
        }

        let solution_offset = 25 + count * 8;
        let solution = bytes[solution_offset..solution_offset + solution_len].to_vec();

        Ok(Self {
            key_hashes,
            solution,
            cols,
            n,
            count,
            finalized,
        })
    }

    /// Hash function using xxHash
    #[inline]
    fn hash(&self, key: &[u8]) -> u64 {
        use xxhash_rust::xxh64::xxh64;
        let hash = xxh64(key, 0);
        // Mask to cols bits
        if self.cols < 64 {
            hash & ((1u64 << self.cols) - 1)
        } else {
            hash
        }
    }

    /// Returns filter parameters (n, rows, cols)
    pub fn params(&self) -> (usize, usize, usize) {
        (self.n, self.n, self.cols) // rows = n in our simplified version
    }

    /// Returns true if no elements have been inserted
    pub fn is_empty(&self) -> bool {
        self.count == 0
    }

    /// Returns the number of elements inserted
    pub fn len(&self) -> usize {
        self.count
    }
}

impl std::fmt::Debug for RibbonFilter {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("RibbonFilter")
            .field("n", &self.n)
            .field("cols", &self.cols)
            .field("count", &self.count)
            .field("finalized", &self.finalized)
            .field(
                "fpr",
                &format!("{:.4}%", self.false_positive_rate() * 100.0),
            )
            .field("memory_bytes", &self.memory_usage())
            .field(
                "bits_per_key",
                &format!("{:.1}", self.cols as f64 / self.n as f64),
            )
            .finish()
    }
}

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

    #[test]
    fn test_new() {
        let filter = RibbonFilter::new(1000, 0.01);
        let (n, rows, cols) = filter.params();

        assert_eq!(n, 1000);
        assert_eq!(rows, 1000);
        assert!(cols > 0, "Number of columns should be > 0");
        assert!(!filter.is_finalized());
    }

    #[test]
    fn test_with_params() {
        let filter = RibbonFilter::with_params(1000, 1000, 8000);
        let (n, rows, cols) = filter.params();

        assert_eq!(n, 1000);
        assert_eq!(rows, 1000);
        assert_eq!(cols, 8000);
    }

    #[test]
    fn test_insert_and_contains() {
        let mut filter = RibbonFilter::new(100, 0.01);

        filter.insert(b"key1");
        filter.insert(b"key2");
        filter.insert(b"key3");
        filter.finalize();

        assert!(filter.contains(b"key1"));
        assert!(filter.contains(b"key2"));
        assert!(filter.contains(b"key3"));
    }

    #[test]
    fn test_no_false_negatives() {
        let mut filter = RibbonFilter::new(100, 0.01);
        let keys: Vec<Vec<u8>> = (0..100).map(|i| format!("key{}", i).into_bytes()).collect();

        for key in &keys {
            filter.insert(key);
        }
        filter.finalize();

        // No false negatives
        for key in &keys {
            assert!(
                filter.contains(key),
                "False negative for {:?}",
                String::from_utf8_lossy(key)
            );
        }
    }

    #[test]
    fn test_false_positive_rate() {
        let mut filter = RibbonFilter::new(1000, 0.01);
        let keys: Vec<Vec<u8>> = (0..1000)
            .map(|i| format!("key{}", i).into_bytes())
            .collect();

        for key in &keys {
            filter.insert(key);
        }
        filter.finalize();

        // Test with non-inserted keys
        let test_keys: Vec<Vec<u8>> = (10000..20000)
            .map(|i| format!("test{}", i).into_bytes())
            .collect();

        let false_positives = test_keys.iter().filter(|key| filter.contains(key)).count();

        let actual_fpr = false_positives as f64 / test_keys.len() as f64;

        // Actual FPR should be close to target (within 5x for Ribbon)
        assert!(actual_fpr < 0.05, "FPR too high: {:.4}", actual_fpr);
    }

    #[test]
    fn test_empty_filter() {
        let mut filter = RibbonFilter::new(100, 0.01);
        filter.finalize();

        // Empty filter might have false positives, but we can test it doesn't crash
        let _result = filter.contains(b"key1");
    }

    #[test]
    fn test_serialization() {
        let mut filter = RibbonFilter::new(100, 0.01);
        filter.insert(b"key1");
        filter.insert(b"key2");
        filter.insert(b"key3");
        filter.finalize();

        let bytes = filter.to_bytes();
        let deserialized = RibbonFilter::from_bytes(&bytes).unwrap();

        assert_eq!(filter.params(), deserialized.params());
        assert_eq!(filter.count(), deserialized.count());
        assert!(deserialized.is_finalized());
        assert!(deserialized.contains(b"key1"));
        assert!(deserialized.contains(b"key2"));
        assert!(deserialized.contains(b"key3"));
    }

    #[test]
    fn test_serialization_empty() {
        let mut filter = RibbonFilter::new(100, 0.01);
        filter.finalize();

        let bytes = filter.to_bytes();
        let deserialized = RibbonFilter::from_bytes(&bytes).unwrap();

        assert_eq!(filter.params(), deserialized.params());
        assert_eq!(filter.count(), 0);
    }

    #[test]
    fn test_binary_keys() {
        let mut filter = RibbonFilter::new(100, 0.01);
        let binary_keys = vec![vec![0u8, 1, 2, 3], vec![255, 254, 253], vec![0, 0, 0, 0]];

        for key in &binary_keys {
            filter.insert(key);
        }
        filter.finalize();

        for key in &binary_keys {
            assert!(filter.contains(key));
        }
    }

    #[test]
    fn test_large_keys() {
        let mut filter = RibbonFilter::new(100, 0.01);
        let large_key = vec![42u8; 10000];

        filter.insert(&large_key);
        filter.finalize();

        assert!(filter.contains(&large_key));
    }

    #[test]
    fn test_memory_usage() {
        let filter = RibbonFilter::new(1000, 0.01);
        let memory = filter.memory_usage();

        assert!(memory > 0);
        // Ribbon should use ~8 bits per key
        let expected = (1000 * 8) / 8; // bytes
        assert!(memory >= expected);
    }

    #[test]
    fn test_count() {
        let mut filter = RibbonFilter::new(100, 0.01);
        assert_eq!(filter.count(), 0);

        filter.insert(b"key1");
        assert_eq!(filter.count(), 1);

        filter.insert(b"key2");
        assert_eq!(filter.count(), 2);
    }

    #[test]
    fn test_space_efficiency() {
        let filter = RibbonFilter::new(1000, 0.01);
        let (_, _, cols) = filter.params();
        let bits_per_key = cols as f64 / 1000.0;

        // Ribbon should use ~8-9 bits/key for 1% FPR (vs ~10 for Bloom)
        assert!(
            bits_per_key < 10.0,
            "Ribbon should be more space-efficient than Bloom"
        );
        assert!(bits_per_key >= 7.0, "Bits per key: {:.1}", bits_per_key);
    }

    #[test]
    #[should_panic(expected = "Expected number of elements must be > 0")]
    fn test_new_panics_on_zero_n() {
        RibbonFilter::new(0, 0.01);
    }

    #[test]
    #[should_panic(expected = "False positive rate must be in (0, 1)")]
    fn test_new_panics_on_invalid_fpr() {
        RibbonFilter::new(100, 1.5);
    }

    #[test]
    #[should_panic(expected = "Must call finalize() before querying")]
    fn test_contains_before_finalize() {
        let mut filter = RibbonFilter::new(100, 0.01);
        filter.insert(b"key1");
        filter.contains(b"key1"); // Should panic
    }

    #[test]
    #[should_panic(expected = "Cannot insert after finalization")]
    fn test_insert_after_finalize() {
        let mut filter = RibbonFilter::new(100, 0.01);
        filter.insert(b"key1");
        filter.finalize();
        filter.insert(b"key2"); // Should panic
    }

    #[test]
    fn test_debug_format() {
        let mut filter = RibbonFilter::new(1000, 0.01);
        filter.insert(b"test");
        filter.finalize();

        let debug_str = format!("{:?}", filter);
        assert!(debug_str.contains("RibbonFilter"));
        assert!(debug_str.contains("n"));
        assert!(debug_str.contains("finalized"));
    }
}