expandable_cuckoo_filter/

lib.rs

//! # Expandable Cuckoo Filter
//!
//! An expandable, serializable, and orthogonal Cuckoo Filter implementation.
//!
//! ## Features
//! - **Expandable**: Automatically grows capacity when load exceeds 80%.
//! - **Orthogonal**: Supports deterministic seeding for distributed consensus.
//! - **Persistent**: Export/Import with Zstd compression and integrity checks.
//!
//! ## Example
//! ```
//! use expandable_cuckoo_filter::ExpandableCuckooFilter;
//!
//! let filter = ExpandableCuckooFilter::new("my-filter", 100);
//! filter.insert("hello");
//! assert!(filter.contains("hello"));
//! ```

use byteorder::{ByteOrder, LittleEndian};
use cuckoofilter::{CuckooFilter, ExportedCuckooFilter};
use parking_lot::RwLock;
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
use std::sync::Arc;
use thiserror::Error;

#[derive(Error, Debug)]
pub enum CuckooError {
    #[error("Invalid data format")]
    InvalidData,
}

struct FilterNode {
    filter: CuckooFilter<DefaultHasher>,
    capacity: usize,
}

/// An expandable Cuckoo Filter.
///
/// It maintains a chain of standard Cuckoo Filters. When the active filter reaches
/// a high load factor (80%), a new, larger filter is appended.
#[derive(Clone)]
pub struct ExpandableCuckooFilter {
    node_id: String,
    seed: u64,
    filters: Arc<RwLock<Vec<FilterNode>>>,
    initial_capacity: usize,
}

impl ExpandableCuckooFilter {
    /// Creates a new `ExpandableCuckooFilter`.
    ///
    /// - `node_id`: Unique identifier used to derive the seeding salt.
    /// - `initial_capacity`: Starting capacity. Should be at least 64 for stability.
    ///
    /// # Example
    /// ```
    /// use expandable_cuckoo_filter::ExpandableCuckooFilter;
    /// let filter = ExpandableCuckooFilter::new("test-node-1", 1000);
    /// ```
    pub fn new(node_id: impl Into<String>, initial_capacity: usize) -> Self {
        let node_id = node_id.into();

        let mut hasher = DefaultHasher::new();
        node_id.hash(&mut hasher);
        let seed = hasher.finish();

        let filter = CuckooFilter::with_capacity(initial_capacity);

        Self {
            node_id,
            seed,
            filters: Arc::new(RwLock::new(vec![FilterNode {
                filter,
                capacity: initial_capacity,
            }])),
            initial_capacity,
        }
    }

    /// Returns the node ID of this filter.
    pub fn node_id(&self) -> &str {
        &self.node_id
    }

    /// Returns the derived seed used for orthogonality.
    pub fn seed(&self) -> u64 {
        self.seed
    }

    /// Generates a salted key tuple for orthogonality.
    ///
    /// ### Why manual salting?
    /// We explored using the internal hasher of the `cuckoofilter` crate, but discovered
    /// it uses `PhantomData<H>`. This means the filter does not store a hasher instance,
    /// but instantiates a new one via `H::default()` for every operation.
    ///
    /// To ensure deterministic hashing across imports/exports and to guarantee
    /// orthogonality between nodes, we must manually salt the key here with our
    /// stable `seed` before passing it to the underlying filter.
    ///
    /// Uses `(seed, item_hash)` for maximum performance.
    fn hash_key<T: Hash + ?Sized>(&self, item: &T) -> (u64, u64) {
        let mut hasher = DefaultHasher::new();
        item.hash(&mut hasher);
        (self.seed, hasher.finish())
    }

    /// Inserts an item into the filter.
    ///
    /// If the filter is full (load > 80% or collision), it expands automatically.
    ///
    /// # Returns
    /// `true` if inserted successfully (which should be always due to expansion).
    pub fn insert<T: Hash + ?Sized>(&self, item: &T) -> bool {
        let hashed_keys = self.hash_key(item);

        let mut filters = self.filters.write();
        let last_idx = filters.len() - 1;

        // Check load factor of active filter
        let current_node = &filters[last_idx];
        let load_factor = current_node.filter.len() as f64 / current_node.capacity as f64;

        // Conservative threshold to prevent "NotEnoughSpace" data loss.
        // We expand early because if cuckoofilter::add fails, it drops the victim data.
        if load_factor > 0.80 {
            // Expand immediately
            self.expand(&mut filters);
            // Insert into new last
            let last_idx = filters.len() - 1;
            // We can ignore result here mostly since it's empty, but strict safety:
            let _ = filters[last_idx].filter.add(&hashed_keys);
            return true;
        }

        let result = filters[last_idx].filter.add(&hashed_keys);

        match result {
            Ok(_) => true,
            Err(_) => {
                // Collision/Full -> Expand
                //
                // We expand by doubling the capacity of the new stage.
                // This geometric growth keeps the number of chained filters (and thus
                // the number of hash checks in `contains`) logarithmic relative to
                // the total items.
                let last_capacity = filters[last_idx].capacity;
                let new_capacity = last_capacity * 2;
                let mut new_filter = CuckooFilter::with_capacity(new_capacity);

                if let Err(e) = new_filter.add(&hashed_keys) {
                    eprintln!(
                        "CRITICAL: Failed to insert into NEW filter (cap: {}). Error: {:?}",
                        new_capacity, e
                    );
                    return false;
                }

                filters.push(FilterNode {
                    filter: new_filter,
                    capacity: new_capacity,
                });
                true
            }
        }
    }

    fn expand(&self, filters: &mut Vec<FilterNode>) {
        let last_capacity = filters
            .last()
            .map(|n| n.capacity)
            .unwrap_or(self.initial_capacity);
        let new_capacity = last_capacity * 2;
        let new_filter = CuckooFilter::with_capacity(new_capacity);
        filters.push(FilterNode {
            filter: new_filter,
            capacity: new_capacity,
        });
    }

    /// Checks if an item is possibly in the filter.
    ///
    /// # Returns
    /// `true` if the item is in the filter (with small false positive probability).
    /// `false` if the item is definitely NOT in the filter.
    pub fn contains<T: Hash + ?Sized>(&self, item: &T) -> bool {
        let hashed_keys = self.hash_key(item);
        let filters = self.filters.read();

        // We search in reverse order (newest filters first).
        // Statistically, if an item was recently added, it will be in a
        // later stage, potentially failing faster.
        for node in filters.iter().rev() {
            if node.filter.contains(&hashed_keys) {
                return true;
            }
        }
        false
    }

    /// Removes an item from the filter.
    ///
    /// **Note**: Only safe if you KNOW the item was inserted.
    /// Removing non-existent items may delete a colliding fingerprint (false deletion).
    pub fn remove<T: Hash + ?Sized>(&self, item: &T) -> bool {
        let hashed_keys = self.hash_key(item);
        let mut filters = self.filters.write();

        // Note: We only remove from one stage. In the unlikely event of
        // duplicate insertions across stages (e.g. if a duplicate was
        // inserted before/after an expansion), only the latest one is removed.
        for node in filters.iter_mut().rev() {
            if node.filter.contains(&hashed_keys) && node.filter.delete(&hashed_keys) {
                return true;
            }
        }
        false
    }

    /// Returns the total number of items in the filter.
    pub fn len(&self) -> usize {
        let filters = self.filters.read();
        filters.iter().map(|f| f.filter.len()).sum()
    }

    /// Returns true if the filter contains no items.
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Returns the total capacity (sum of all chained filters).
    pub fn capacity(&self) -> usize {
        let filters = self.filters.read();
        filters.iter().map(|f| f.capacity).sum()
    }

    const MAGIC_BYTES: &[u8; 4] = b"ECF1";

    /// Exports the filter state as a byte vector.
    ///
    /// Format: `ECF1` magic bytes + Count + [Filter Data...]
    pub fn export(&self) -> Result<Vec<u8>, CuckooError> {
        let filters = self.filters.read();
        let mut buffer = Vec::new();

        // Write Magic Bytes
        buffer.extend_from_slice(Self::MAGIC_BYTES);

        let mut u64_buf = [0u8; 8];
        LittleEndian::write_u64(&mut u64_buf, filters.len() as u64);
        buffer.extend_from_slice(&u64_buf);

        for node in filters.iter() {
            let exported = node.filter.export();

            LittleEndian::write_u64(&mut u64_buf, node.capacity as u64);
            buffer.extend_from_slice(&u64_buf);

            LittleEndian::write_u64(&mut u64_buf, exported.length as u64);
            buffer.extend_from_slice(&u64_buf);

            LittleEndian::write_u64(&mut u64_buf, exported.values.len() as u64);
            buffer.extend_from_slice(&u64_buf);

            buffer.extend_from_slice(&exported.values);
        }

        Ok(buffer)
    }

    /// Imports filter state from a byte slice.
    ///
    /// # Errors
    /// Returns `CuckooError::InvalidData` if magic bytes don't match or data is corrupt.
    pub fn import(&self, data: &[u8]) -> Result<(), CuckooError> {
        let mut cursor = 0;

        // Check Magic Bytes
        if data.len() < 4 || &data[0..4] != Self::MAGIC_BYTES {
            return Err(CuckooError::InvalidData);
        }
        cursor += 4;

        if cursor + 8 > data.len() {
            return Err(CuckooError::InvalidData);
        }
        let count = LittleEndian::read_u64(&data[cursor..cursor + 8]) as usize;
        cursor += 8;

        let mut new_filters = Vec::with_capacity(count);

        for _i in 0..count {
            if cursor + 8 > data.len() {
                return Err(CuckooError::InvalidData);
            }
            let capacity = LittleEndian::read_u64(&data[cursor..cursor + 8]) as usize;
            cursor += 8;

            if cursor + 8 > data.len() {
                return Err(CuckooError::InvalidData);
            }
            let item_count = LittleEndian::read_u64(&data[cursor..cursor + 8]) as usize;
            cursor += 8;

            if cursor + 8 > data.len() {
                return Err(CuckooError::InvalidData);
            }
            let values_len = LittleEndian::read_u64(&data[cursor..cursor + 8]) as usize;
            cursor += 8;

            if cursor + values_len > data.len() {
                eprintln!(
                    "IMPORT: Invalid values_len {} vs rem {}",
                    values_len,
                    data.len() - cursor
                );
                return Err(CuckooError::InvalidData);
            }
            let values = data[cursor..cursor + values_len].to_vec();
            cursor += values_len;

            let exported = ExportedCuckooFilter {
                length: item_count,
                values,
            };

            new_filters.push(FilterNode {
                filter: CuckooFilter::from(exported),
                capacity,
            });
        }

        if new_filters.is_empty() {
            new_filters.push(FilterNode {
                filter: CuckooFilter::with_capacity(self.initial_capacity),
                capacity: self.initial_capacity,
            });
        }

        *self.filters.write() = new_filters;

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

    #[test]
    fn test_basic_lifecycle() {
        let filter = ExpandableCuckooFilter::new("test-basic", 100);
        assert!(filter.is_empty());

        filter.insert("item1");
        filter.insert("item2");

        assert!(filter.contains("item1"));
        assert!(filter.contains("item2"));
        assert!(!filter.contains("item3"));
        assert_eq!(filter.len(), 2);

        filter.remove("item1");
        assert!(!filter.contains("item1"));
        assert!(filter.contains("item2"));
        assert_eq!(filter.len(), 1);
    }

    #[test]
    fn test_orthogonality() {
        let f1 = ExpandableCuckooFilter::new("node-A", 1000);
        let f2 = ExpandableCuckooFilter::new("node-B", 1000);

        assert_ne!(f1.seed(), f2.seed());
    }

    #[test]
    fn test_expansion() {
        // Start with capacity 4
        let initial_cap = 4;
        let filter = ExpandableCuckooFilter::new("test-expand", initial_cap);

        // Insert 100 items (force expansions)
        for i in 0..100 {
            filter.insert(&format!("item-{}", i));
        }

        // Verify presence
        for i in 0..100 {
            let key = format!("item-{}", i);
            assert!(filter.contains(&key), "Missing key {}", key);
        }

        assert!(filter.capacity() > initial_cap);
        assert_eq!(filter.len(), 100);
    }

    #[test]
    fn test_expansion_remove() {
        // Note: With small capacity, FP rate increases due to small buckets?
        // Actually using a slightly larger base capacity helps avoid trivial collisions loops
        // which might affect test stability regarding FPs.
        let filter = ExpandableCuckooFilter::new("test-expand-remove", 16);

        for i in 0..100 {
            filter.insert(&format!("val-{}", i));
        }

        let key0 = "val-0";
        assert!(filter.contains(key0));
        assert!(filter.remove(key0));
        assert!(!filter.contains(key0));

        let key99 = "val-99";
        assert!(filter.contains(key99));
        assert!(filter.remove(key99));
        assert!(!filter.contains(key99));
    }

    #[test]
    fn test_export_import() {
        let filter = ExpandableCuckooFilter::new("test-persist", 100);
        for i in 0..20 {
            filter.insert(&i.to_string());
        }

        let restored = ExpandableCuckooFilter::new("test-persist", 100);
        let bytes = filter.export().expect("Export failed");
        restored.import(&bytes).expect("Import failed");

        assert_eq!(restored.len(), 20);
        for i in 0..20 {
            assert!(restored.contains(&i.to_string()));
        }
    }

    #[test]
    fn test_import_invalid_data() {
        let filter = ExpandableCuckooFilter::new("test-bad", 100);

        // Garbage data (Wrong Magic)
        let result = filter.import(&[1, 2, 3, 4]);
        assert!(result.is_err(), "Should fail on garbage logic (Bad Magic)");

        // Valid buffer but missing magic bytes (All zeros)
        let _payload = vec![0u8; 100];
        assert!(filter.import(&_payload).is_err());
    }

    #[test]
    fn test_coverage_gap_fillers() {
        let filter = ExpandableCuckooFilter::new("gap-fill", 100);
        assert_eq!(filter.node_id(), "gap-fill");
        assert!(filter.seed() > 0);
        assert_eq!(filter.capacity(), 100);
        assert!(filter.is_empty());

        // Test truncated import data
        let bytes = filter.export().unwrap();
        for i in 0..bytes.len() - 1 {
            assert!(
                filter.import(&bytes[..i]).is_err(),
                "Truncated at {} should fail",
                i
            );
        }

        // Test empty filters array import (malformed but valid structure)
        let mut empty_data = Vec::new();
        empty_data.extend_from_slice(b"ECF1");
        empty_data.extend_from_slice(&0u64.to_le_bytes()); // Count = 0
        filter
            .import(&empty_data)
            .expect("Should handle empty filter count by re-initializing");
        assert_eq!(filter.len(), 0);
        assert_eq!(filter.capacity(), 100);
    }

    mod proptests {
        use super::*;
        use proptest::prelude::*;
        use std::collections::HashSet;

        proptest! {
            #[test]
            fn test_fuzz_insert_contains(keys in proptest::collection::vec(".*", 1..100)) {
                let filter = ExpandableCuckooFilter::new("fuzz-test", 64);
                for key in &keys {
                    prop_assert!(filter.insert(key), "Insert failed check for key: {:?}", key);
                }
                for key in &keys {
                    prop_assert!(filter.contains(key));
                }
            }

            #[test]
            fn test_fuzz_persistence_roundtrip(keys in proptest::collection::vec(".*", 1..50)) {
                let filter = ExpandableCuckooFilter::new("persist-fuzz", 20);
                for key in &keys {
                    filter.insert(key);
                }
                let bytes = filter.export().unwrap();
                let restored = ExpandableCuckooFilter::new("persist-fuzz", 20);
                restored.import(&bytes).unwrap();
                for key in &keys {
                    prop_assert!(restored.contains(key));
                }
                prop_assert_eq!(filter.len(), restored.len());
            }

            #[test]
            fn test_fuzz_set_semantics(ops in proptest::collection::vec((0..2u8, ".*"), 1..100)) {
                let filter = ExpandableCuckooFilter::new("set-fuzz", 50);
                let mut shadow_set = HashSet::new();

                for (op_code, key) in ops {
                    match op_code {
                        0 => { // Insert
                            filter.insert(&key);
                            shadow_set.insert(key);
                        },
                        1 => { // Remove
                            // Cuckoo filters support deletion, BUT if we try to remove a key that wasn't inserted,
                            // we risk a "false deletion" (colliding fingerprint removed).
                            // So we only remove if we think it's there (shadow set agrees).
                            if shadow_set.contains(&key) {
                                filter.remove(&key);
                                shadow_set.remove(&key);
                            }
                        },
                        _ => unreachable!(),
                    }
                }

                for key in &shadow_set {
                    prop_assert!(filter.contains(key), "Filter missing key: {}", key);
                }
            }
        }
    }
}