use expandable_cuckoo_filter::ExpandableCuckooFilter;
use std::sync::atomic::{AtomicBool, Ordering};
use tracing::{debug, info, warn};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FilterTrust {
Untrusted,
Trusted,
}
pub struct FilterManager {
filter: ExpandableCuckooFilter,
trust_state: AtomicBool,
}
impl FilterManager {
#[must_use]
pub fn new(node_id: impl Into<String>, initial_capacity: usize) -> Self {
let node_id = node_id.into();
info!(
node_id = %node_id,
initial_capacity,
"Creating expandable cuckoo filter manager"
);
let filter = ExpandableCuckooFilter::new(node_id, initial_capacity);
Self {
filter,
trust_state: AtomicBool::new(false),
}
}
#[must_use]
#[inline]
pub fn trust_state(&self) -> FilterTrust {
if self.trust_state.load(Ordering::Acquire) {
FilterTrust::Trusted
} else {
FilterTrust::Untrusted
}
}
#[must_use]
#[inline]
pub fn is_trusted(&self) -> bool {
self.trust_state.load(Ordering::Acquire)
}
pub fn mark_trusted(&self) {
info!(
entries = self.filter.len(),
"Marking cuckoo filter as trusted"
);
self.trust_state.store(true, Ordering::Release);
}
pub fn mark_untrusted(&self) {
warn!("Marking cuckoo filter as untrusted");
self.trust_state.store(false, Ordering::Release);
}
#[must_use]
pub fn might_exist(&self, key: &str) -> Option<bool> {
if !self.trust_state.load(Ordering::Acquire) {
return None;
}
Some(self.filter.contains(key))
}
#[must_use]
#[inline]
pub fn should_check_l3(&self, key: &str) -> bool {
match self.might_exist(key) {
None => true, Some(true) => true, Some(false) => false, }
}
pub fn insert(&self, key: &str) {
if !self.filter.contains(key) {
self.filter.insert(key);
}
}
pub fn remove(&self, key: &str) -> bool {
self.filter.remove(key)
}
pub fn bulk_insert(&self, keys: &[String]) {
for key in keys {
if !self.filter.contains(key) {
self.filter.insert(key.as_str());
}
}
debug!(
added = keys.len(),
total = self.filter.len(),
"Bulk inserted keys into cuckoo filter"
);
}
#[must_use]
pub fn export(&self) -> Option<Vec<u8>> {
self.filter.export().ok()
}
pub fn import(&self, data: &[u8]) -> Result<(), String> {
self.filter.import(data).map_err(|e| e.to_string())
}
#[must_use]
pub fn len(&self) -> usize {
self.filter.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.filter.is_empty()
}
#[must_use]
pub fn stats(&self) -> (usize, usize, FilterTrust) {
(
self.filter.len(),
self.filter.capacity(),
self.trust_state(),
)
}
#[must_use]
pub fn memory_usage_bytes(&self) -> usize {
self.filter.capacity().saturating_mul(8).max(1024)
}
#[must_use]
pub fn node_id(&self) -> &str {
self.filter.node_id()
}
#[must_use]
pub fn seed(&self) -> u64 {
self.filter.seed()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_new_filter_is_untrusted() {
let fm = FilterManager::new("test-node", 1000);
assert_eq!(fm.trust_state(), FilterTrust::Untrusted);
assert!(fm.is_empty());
assert_eq!(fm.len(), 0);
}
#[test]
fn test_mark_trusted_and_untrusted() {
let fm = FilterManager::new("test-node", 1000);
assert_eq!(fm.trust_state(), FilterTrust::Untrusted);
fm.mark_trusted();
assert_eq!(fm.trust_state(), FilterTrust::Trusted);
fm.mark_untrusted();
assert_eq!(fm.trust_state(), FilterTrust::Untrusted);
}
#[test]
fn test_insert_and_contains() {
let fm = FilterManager::new("test-node", 1000);
fm.mark_trusted();
fm.insert("key-1");
assert_eq!(fm.len(), 1);
assert_eq!(fm.might_exist("key-1"), Some(true));
assert_eq!(fm.might_exist("key-2"), Some(false));
}
#[test]
fn test_might_exist_returns_none_when_untrusted() {
let fm = FilterManager::new("test-node", 1000);
fm.insert("key-1");
assert_eq!(fm.might_exist("key-1"), None);
assert_eq!(fm.might_exist("key-2"), None);
}
#[test]
fn test_should_check_l3_when_untrusted() {
let fm = FilterManager::new("test-node", 1000);
assert!(fm.should_check_l3("any-key"));
assert!(fm.should_check_l3("another-key"));
}
#[test]
fn test_should_check_l3_when_trusted() {
let fm = FilterManager::new("test-node", 1000);
fm.mark_trusted();
fm.insert("exists");
assert!(fm.should_check_l3("exists"));
assert!(!fm.should_check_l3("nonexistent"));
}
#[test]
fn test_remove() {
let fm = FilterManager::new("test-node", 1000);
fm.mark_trusted();
fm.insert("to-remove");
assert_eq!(fm.might_exist("to-remove"), Some(true));
let removed = fm.remove("to-remove");
assert!(removed || fm.might_exist("to-remove") == Some(false));
}
#[test]
fn test_bulk_insert() {
let fm = FilterManager::new("test-node", 1000);
fm.mark_trusted();
let keys: Vec<String> = (0..100).map(|i| format!("bulk-key-{}", i)).collect();
fm.bulk_insert(&keys);
assert_eq!(fm.len(), 100);
assert_eq!(fm.might_exist("bulk-key-0"), Some(true));
assert_eq!(fm.might_exist("bulk-key-50"), Some(true));
assert_eq!(fm.might_exist("bulk-key-99"), Some(true));
assert_eq!(fm.might_exist("not-inserted"), Some(false));
}
#[test]
fn test_export_import() {
let fm1 = FilterManager::new("test-node", 1000);
fm1.insert("key-a");
fm1.insert("key-b");
fm1.insert("key-c");
let exported = fm1.export();
assert!(exported.is_some());
let fm2 = FilterManager::new("test-node", 1000);
let import_result = fm2.import(&exported.unwrap());
assert!(import_result.is_ok());
assert_eq!(fm2.len(), 3);
fm2.mark_trusted();
assert_eq!(fm2.might_exist("key-a"), Some(true));
assert_eq!(fm2.might_exist("key-b"), Some(true));
assert_eq!(fm2.might_exist("key-c"), Some(true));
}
#[test]
fn test_stats() {
let fm = FilterManager::new("test-node", 1000);
let (count, capacity, trust) = fm.stats();
assert_eq!(count, 0);
assert!(capacity > 0);
assert_eq!(trust, FilterTrust::Untrusted);
fm.insert("key");
fm.mark_trusted();
let (count, _, trust) = fm.stats();
assert_eq!(count, 1);
assert_eq!(trust, FilterTrust::Trusted);
}
#[test]
fn test_memory_usage_bytes() {
let fm = FilterManager::new("test-node", 10000);
let mem = fm.memory_usage_bytes();
assert!(mem >= 1024, "Should report at least 1KB");
assert!(mem < 1024 * 1024, "Should be less than 1MB for 10k capacity");
}
#[test]
fn test_node_id_and_seed() {
let fm = FilterManager::new("my-unique-node", 1000);
assert_eq!(fm.node_id(), "my-unique-node");
assert!(fm.seed() > 0);
}
#[test]
fn test_deterministic_seed_same_node() {
let fm1 = FilterManager::new("same-node", 1000);
let fm2 = FilterManager::new("same-node", 1000);
assert_eq!(fm1.seed(), fm2.seed());
}
#[test]
fn test_different_seeds_different_nodes() {
let fm1 = FilterManager::new("node-a", 1000);
let fm2 = FilterManager::new("node-b", 1000);
assert_ne!(fm1.seed(), fm2.seed());
}
#[test]
fn test_filter_grows_beyond_initial_capacity() {
let fm = FilterManager::new("test-node", 10);
for i in 0..100 {
fm.insert(&format!("key-{}", i));
}
assert!(fm.len() >= 95, "Expected at least 95 entries, got {}", fm.len());
let (count, capacity, _) = fm.stats();
assert!(capacity >= count);
}
#[test]
fn test_is_empty() {
let fm = FilterManager::new("test-node", 1000);
assert!(fm.is_empty());
fm.insert("key");
assert!(!fm.is_empty());
}
#[test]
fn test_filter_trust_display() {
assert_eq!(format!("{:?}", FilterTrust::Untrusted), "Untrusted");
assert_eq!(format!("{:?}", FilterTrust::Trusted), "Trusted");
}
#[test]
fn test_multiple_inserts_same_key() {
let fm = FilterManager::new("test-node", 1000);
fm.insert("same-key");
fm.insert("same-key");
fm.insert("same-key");
fm.mark_trusted();
assert_eq!(fm.might_exist("same-key"), Some(true));
}
#[test]
fn test_false_positive_rate_reasonable() {
let fm = FilterManager::new("test-node", 10000);
fm.mark_trusted();
for i in 0..1000 {
fm.insert(&format!("inserted-{}", i));
}
let mut false_positives = 0;
for i in 0..1000 {
if fm.might_exist(&format!("not-inserted-{}", i)) == Some(true) {
false_positives += 1;
}
}
let fp_rate = false_positives as f64 / 1000.0;
assert!(fp_rate < 0.05, "False positive rate {} is too high", fp_rate);
}
}