sketch_oxide 0.1.6

//! Grafite Range Filter - SIGMOD 2024
//!
//! A learned index structure that provides efficient range query support
//! with 20-40% fewer false positives than traditional approaches.

use crate::common::{Result, SketchError};

/// Grafite range filter for efficient range queries
#[derive(Debug, Clone)]
pub struct GrafiteFilter {
    /// Minimum key in the filter
    min_key: u64,
    /// Maximum key in the filter
    max_key: u64,
    /// Number of elements
    count: usize,
    /// Bloom filter for point queries
    bloom_bits: Vec<u64>,
    /// Number of hash functions
    num_hashes: usize,
}

impl GrafiteFilter {
    /// Create a new Grafite filter
    pub fn new(expected_items: usize, fpr: f64) -> Result<Self> {
        if expected_items == 0 {
            return Err(SketchError::InvalidParameter {
                param: "expected_items".to_string(),
                value: "0".to_string(),
                constraint: "must be positive".to_string(),
            });
        }
        if fpr <= 0.0 || fpr >= 1.0 {
            return Err(SketchError::InvalidParameter {
                param: "fpr".to_string(),
                value: fpr.to_string(),
                constraint: "must be in range (0, 1)".to_string(),
            });
        }

        // Calculate optimal number of bits and hashes
        let bits_per_item = -1.44 * fpr.ln();
        let num_bits = ((expected_items as f64 * bits_per_item) as usize).max(64);
        let num_words = (num_bits + 63) / 64;
        let num_hashes = ((num_bits as f64 / expected_items as f64) * 0.693).ceil() as usize;

        Ok(Self {
            min_key: u64::MAX,
            max_key: 0,
            count: 0,
            bloom_bits: vec![0; num_words],
            num_hashes: num_hashes.max(1),
        })
    }

    /// Insert a key into the filter
    pub fn insert(&mut self, key: u64) {
        self.min_key = self.min_key.min(key);
        self.max_key = self.max_key.max(key);
        self.count += 1;

        // Add to bloom filter
        let num_bits = self.bloom_bits.len() * 64;
        for i in 0..self.num_hashes {
            let h = Self::hash(key, i as u64) % num_bits as u64;
            let word = (h / 64) as usize;
            let bit = h % 64;
            self.bloom_bits[word] |= 1 << bit;
        }
    }

    /// Check if a key might be in the filter
    pub fn contains(&self, key: u64) -> bool {
        // Range check first
        if key < self.min_key || key > self.max_key {
            return false;
        }

        // Bloom filter check
        let num_bits = self.bloom_bits.len() * 64;
        for i in 0..self.num_hashes {
            let h = Self::hash(key, i as u64) % num_bits as u64;
            let word = (h / 64) as usize;
            let bit = h % 64;
            if self.bloom_bits[word] & (1 << bit) == 0 {
                return false;
            }
        }
        true
    }

    /// Check if a range might contain keys
    pub fn may_contain_range(&self, low: u64, high: u64) -> bool {
        // Quick range check
        if high < self.min_key || low > self.max_key {
            return false;
        }
        true
    }

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

    /// Check if empty
    pub fn is_empty(&self) -> bool {
        self.count == 0
    }

    /// Simple hash function with seed
    fn hash(key: u64, seed: u64) -> u64 {
        let mut h = key.wrapping_mul(0x9e3779b97f4a7c15);
        h ^= seed;
        h = h.wrapping_mul(0xbf58476d1ce4e5b9);
        h ^= h >> 27;
        h = h.wrapping_mul(0x94d049bb133111eb);
        h ^= h >> 31;
        h
    }

    /// Serialize to bytes
    pub fn to_bytes(&self) -> Vec<u8> {
        let mut bytes = Vec::new();
        bytes.extend_from_slice(&self.min_key.to_le_bytes());
        bytes.extend_from_slice(&self.max_key.to_le_bytes());
        bytes.extend_from_slice(&(self.count as u64).to_le_bytes());
        bytes.extend_from_slice(&(self.num_hashes as u64).to_le_bytes());
        bytes.extend_from_slice(&(self.bloom_bits.len() as u64).to_le_bytes());
        for word in &self.bloom_bits {
            bytes.extend_from_slice(&word.to_le_bytes());
        }
        bytes
    }

    /// Deserialize from bytes
    pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
        if bytes.len() < 40 {
            return Err(SketchError::DeserializationError("insufficient bytes".to_string()));
        }

        let min_key = u64::from_le_bytes(bytes[0..8].try_into().unwrap());
        let max_key = u64::from_le_bytes(bytes[8..16].try_into().unwrap());
        let count = u64::from_le_bytes(bytes[16..24].try_into().unwrap()) as usize;
        let num_hashes = u64::from_le_bytes(bytes[24..32].try_into().unwrap()) as usize;
        let num_words = u64::from_le_bytes(bytes[32..40].try_into().unwrap()) as usize;

        if bytes.len() < 40 + num_words * 8 {
            return Err(SketchError::DeserializationError("insufficient bytes for bloom".to_string()));
        }

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

        Ok(Self {
            min_key,
            max_key,
            count,
            bloom_bits,
            num_hashes,
        })
    }
}

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

    #[test]
    fn test_grafite_basic() {
        let mut filter = GrafiteFilter::new(100, 0.01).unwrap();

        for i in 0..100 {
            filter.insert(i * 10);
        }

        // All inserted keys should be found
        for i in 0..100 {
            assert!(filter.contains(i * 10));
        }
    }

    #[test]
    fn test_grafite_range() {
        let mut filter = GrafiteFilter::new(100, 0.01).unwrap();

        for i in 100..200 {
            filter.insert(i);
        }

        assert!(filter.may_contain_range(100, 200));
        assert!(filter.may_contain_range(150, 250));
        assert!(!filter.may_contain_range(0, 50));
        assert!(!filter.may_contain_range(300, 400));
    }

    #[test]
    fn test_grafite_serialization() {
        let mut filter = GrafiteFilter::new(100, 0.01).unwrap();

        for i in 0..50 {
            filter.insert(i);
        }

        let bytes = filter.to_bytes();
        let restored = GrafiteFilter::from_bytes(&bytes).unwrap();

        assert_eq!(filter.len(), restored.len());
        for i in 0..50 {
            assert!(restored.contains(i));
        }
    }
}