use std::{
hash::{BuildHasher, Hash, RandomState},
io::Read,
iter::repeat_with,
sync::{
Arc, Mutex, MutexGuard,
atomic::{AtomicUsize, Ordering},
},
};
use crate::{
associated_data::AssociatedData,
bucket::{Bucket, InsertValues, LookupValues},
config::CuckooConfiguration,
data_block::{DataBlock, Fingerprint},
exporter::{
CuckooFilterExporter, ExportableBuildHasher, ExportableRandomState, import_config,
read_hasher_from,
},
};
#[derive(Clone)]
pub struct CuckooFilter<H: BuildHasher> {
configuration: CuckooConfiguration,
buckets: Arc<Vec<Mutex<Bucket>>>,
build_hasher: H,
items: Arc<AtomicUsize>,
}
impl CuckooFilter<RandomState> {
#[must_use]
pub fn new_random(configuration: CuckooConfiguration) -> Self {
Self::new(configuration, RandomState::new())
}
}
impl CuckooFilter<ExportableRandomState> {
#[must_use]
pub fn new_random_exportable(configuration: CuckooConfiguration) -> Self {
Self::new(configuration, ExportableRandomState::new_random())
}
pub fn import_random_exportable(reader: impl Read) -> Result<Self, crate::ImportError> {
Self::import(reader)
}
}
impl<H: ExportableBuildHasher + BuildHasher> CuckooFilter<H> {
pub fn import(mut reader: impl Read) -> Result<Self, crate::ImportError> {
let hasher = read_hasher_from(&mut reader)?;
let configuration = import_config(&mut reader)?;
let mut buckets = Vec::with_capacity(configuration.bucket_count);
let mut item_count = 0;
for _ in 0..configuration.bucket_count {
let bucket = Bucket::take_from(&mut reader, &configuration)?;
item_count += bucket.occupied_count(&configuration);
buckets.push(Mutex::new(bucket));
}
Ok(Self {
configuration,
buckets: Arc::new(buckets),
build_hasher: hasher,
items: Arc::new(AtomicUsize::new(item_count)),
})
}
pub fn exporter<'a>(&'a self) -> CuckooFilterExporter<'a, H> {
CuckooFilterExporter::new(&self.build_hasher, &self.buckets, &self.configuration)
}
}
impl<H: BuildHasher> CuckooFilter<H> {
pub fn new(configuration: CuckooConfiguration, build_hasher: H) -> Self {
Self {
configuration: configuration.clone(),
buckets: repeat_with(|| Bucket::new(&configuration).into())
.take(configuration.bucket_count)
.collect::<Vec<_>>()
.into(),
build_hasher,
items: Arc::new(AtomicUsize::new(0)),
}
}
pub const fn get_bucket_count(&self) -> usize {
self.configuration.bucket_count
}
pub fn get_item_count(&self) -> usize {
self.items.load(Ordering::Relaxed)
}
pub fn get_configuration(&self) -> CuckooConfiguration {
self.configuration.clone()
}
pub fn get_memory_usage(&self) -> usize {
size_of::<Self>()
+ size_of::<AtomicUsize>()
+ size_of::<Vec<Mutex<Bucket>>>()
+ size_of::<Mutex<Bucket>>() * self.buckets.len()
+ self.configuration.bucket_byte_size * self.buckets.len()
}
pub fn insert_if_not_present<K: Hash + ?Sized>(&self, key: &K) -> Option<Fingerprint> {
self.insert_if_not_present_with_update(
key,
InsertValues::default(),
LookupValues::default(),
)
}
pub fn insert_if_not_present_with_update<K: Hash + ?Sized>(
&self,
key: &K,
insert_values: InsertValues,
lookup_update: LookupValues,
) -> Option<Fingerprint> {
let (fp, i1) = self.get_fingerprint_and_index(key);
let mut contains =
self.lock_bucket(i1 as usize)
.contains(&fp, &self.configuration, &lookup_update);
if contains {
return None;
}
let i2 = self.alt_index(&fp, i1);
contains = self
.lock_bucket(i2 as usize)
.contains(&fp, &self.configuration, &lookup_update);
if contains {
return None;
}
let mut cur_data_block = self.new_data_block(&fp, insert_values);
let inserted = self
.lock_bucket(i1 as usize)
.insert(&cur_data_block, &self.configuration);
if inserted {
self.items.fetch_add(1, Ordering::Relaxed);
return None;
}
let inserted = self
.lock_bucket(i2 as usize)
.insert(&cur_data_block, &self.configuration);
if inserted {
self.items.fetch_add(1, Ordering::Relaxed);
return None;
}
let mut cur_index = if rand::random::<bool>() { i1 } else { i2 };
for _ in 0..self.configuration.max_kicks {
{
let mut bucket = self.lock_bucket(cur_index as usize);
if let Some(lru_config) = self.configuration.lru_field_config.as_ref() {
if !bucket.kick_lru(&mut cur_data_block, &self.configuration, lru_config) {
return Some(cur_data_block.get_fingerprint(&self.configuration));
}
} else {
bucket.kick_random(&mut cur_data_block, &self.configuration);
}
cur_index = self.alt_index(
&cur_data_block.get_fingerprint(&self.configuration),
cur_index,
);
}
if self
.lock_bucket(cur_index as usize)
.insert(&cur_data_block, &self.configuration)
{
self.items.fetch_add(1, Ordering::Relaxed);
return None;
}
}
Some(cur_data_block.get_fingerprint(&self.configuration))
}
pub fn insert<K: Hash + ?Sized>(&self, key: &K) -> Option<Fingerprint> {
self.insert_with_defaults(key, InsertValues::default())
}
pub fn insert_with_defaults<K: Hash + ?Sized>(
&self,
key: &K,
default: InsertValues,
) -> Option<Fingerprint> {
let (fp, i1) = self.get_fingerprint_and_index(key);
let mut cur_data_block = self.new_data_block(&fp, default);
let inserted = self
.lock_bucket(i1 as usize)
.insert(&cur_data_block, &self.configuration);
if inserted {
self.items.fetch_add(1, Ordering::Relaxed);
return None;
}
let i2 = self.alt_index(&fp, i1);
let inserted = self
.lock_bucket(i2 as usize)
.insert(&cur_data_block, &self.configuration);
if inserted {
self.items.fetch_add(1, Ordering::Relaxed);
return None;
}
let mut cur_index = i1;
for _ in 0..self.configuration.max_kicks {
{
let mut bucket = self.lock_bucket(cur_index as usize);
if let Some(lru_config) = self.configuration.lru_field_config.as_ref() {
if !bucket.kick_lru(&mut cur_data_block, &self.configuration, lru_config) {
return Some(cur_data_block.get_fingerprint(&self.configuration));
}
} else {
bucket.kick_random(&mut cur_data_block, &self.configuration);
}
cur_index = self.alt_index(
&cur_data_block.get_fingerprint(&self.configuration),
cur_index,
);
}
if self
.lock_bucket(cur_index as usize)
.insert(&cur_data_block, &self.configuration)
{
self.items.fetch_add(1, Ordering::Relaxed);
return None;
}
}
Some(cur_data_block.get_fingerprint(&self.configuration))
}
pub fn contains_with_update<K: Hash + ?Sized>(&self, key: &K, update: LookupValues) -> bool {
let (fp, i1) = self.get_fingerprint_and_index(key);
let mut contains =
self.lock_bucket(i1 as usize)
.contains(&fp, &self.configuration, &update);
if !contains {
let i2 = self.alt_index(&fp, i1);
contains = self
.lock_bucket(i2 as usize)
.contains(&fp, &self.configuration, &update);
}
contains
}
pub fn contains<K: Hash + ?Sized>(&self, key: &K) -> bool {
self.contains_with_update(key, LookupValues::default())
}
pub fn get_associated_data<K: Hash + ?Sized>(&self, key: &K) -> Option<AssociatedData> {
self.get_associated_data_with_update(key, LookupValues::default())
}
pub fn get_associated_data_with_update<K: Hash + ?Sized>(
&self,
key: &K,
update: LookupValues,
) -> Option<AssociatedData> {
let (fp, i1) = self.get_fingerprint_and_index(key);
let mut contains =
self.lock_bucket(i1 as usize)
.get_associated_data(&fp, &self.configuration, &update);
if contains.is_none() {
let i2 = self.alt_index(&fp, i1);
contains = self.lock_bucket(i2 as usize).get_associated_data(
&fp,
&self.configuration,
&update,
);
}
contains
}
pub fn remove<K: Hash + ?Sized>(&self, key: &K) -> bool {
let (fp, i1) = self.get_fingerprint_and_index(key);
let mut removed = self
.lock_bucket(i1 as usize)
.remove(&fp, &self.configuration);
if !removed {
let i2 = self.alt_index(&fp, i1);
removed = self
.lock_bucket(i2 as usize)
.remove(&fp, &self.configuration);
}
if removed {
self.items.fetch_sub(1, Ordering::Relaxed);
}
removed
}
pub fn scan_and_update_full(&self) -> usize {
if self.configuration.lru_field_config.is_none()
&& self.configuration.ttl_field_config.is_none()
{
return 0;
}
let mut removed = 0;
for b in self.buckets.iter() {
#[expect(clippy::unwrap_used)]
let mut bucket = b.lock().unwrap();
if let Some(lru_config) = &self.configuration.lru_field_config {
bucket.age_lru_counters(&self.configuration, lru_config);
}
if let Some(ttl_config) = &self.configuration.ttl_field_config {
removed += bucket.age_ttl_counters(&self.configuration, ttl_config);
}
}
if removed > 0 {
self.items.fetch_sub(removed, Ordering::Relaxed);
}
removed
}
pub fn scan_and_update_ttl(&self) -> usize {
if self.configuration.ttl_field_config.is_none() {
return 0;
}
let mut removed = 0;
for b in self.buckets.iter() {
#[expect(clippy::unwrap_used)]
let mut bucket = b.lock().unwrap();
if let Some(ttl_config) = &self.configuration.ttl_field_config {
removed += bucket.age_ttl_counters(&self.configuration, ttl_config);
}
}
if removed > 0 {
self.items.fetch_sub(removed, Ordering::Relaxed);
}
removed
}
pub fn scan_and_update_lru(&self) {
if self.configuration.lru_field_config.is_none() {
return;
}
for b in self.buckets.iter() {
#[expect(clippy::unwrap_used)]
let mut bucket = b.lock().unwrap();
if let Some(lru_config) = &self.configuration.lru_field_config {
bucket.age_lru_counters(&self.configuration, lru_config);
}
}
}
pub(crate) fn get_fingerprint<K: Hash + ?Sized>(&self, key: &K) -> Fingerprint {
self.get_fingerprint_and_index(key).0
}
fn new_data_block(&self, fp: &Fingerprint, defaults: InsertValues) -> DataBlock<Vec<u8>> {
let data = vec![0u8; self.configuration.data_block_size];
let mut cur_data_block = DataBlock::from(data);
cur_data_block.store_fingerprint(fp, &self.configuration);
if let Some(ttl_config) = self.configuration.ttl_field_config.as_ref() {
cur_data_block.set_ttl(ttl_config, defaults.ttl.unwrap_or(ttl_config.0.ttl.into()));
}
if let Some(counter_config) = self.configuration.counter_field_config.as_ref() {
cur_data_block.update_counter(
counter_config,
defaults
.counter
.unwrap_or(counter_config.0.change_on_insert),
);
}
if let Some(lru_config) = self.configuration.lru_field_config.as_ref() {
cur_data_block.inc_lru_counter(lru_config);
}
cur_data_block
}
fn get_fingerprint_and_index<K: Hash + ?Sized>(&self, key: &K) -> (Fingerprint, u32) {
let result = self.build_hasher.hash_one(key);
let fingerprint = (result >> 32) as u32;
#[expect(clippy::cast_possible_truncation)]
let index = result as u32 & self.configuration.buckets_mask;
(
Fingerprint::new(
fingerprint,
self.configuration.fingerprint_field_config.value_mask(),
),
index,
)
}
#[expect(clippy::cast_possible_truncation)]
fn alt_index(&self, fingerprint: &Fingerprint, index: u32) -> u32 {
let result = self.build_hasher.hash_one(fingerprint);
(index ^ ((result as u32) & self.configuration.buckets_mask))
& self.configuration.buckets_mask
}
#[expect(clippy::unwrap_used)]
fn lock_bucket(&self, index: usize) -> MutexGuard<'_, Bucket> {
self.buckets[index].lock().unwrap()
}
}
#[cfg(test)]
#[expect(clippy::unwrap_used)]
mod tests {
use std::{
collections::{HashSet, VecDeque},
hash::Hasher,
ops::Range,
};
use crate::config::{CounterConfig, LruConfig, TtlConfig};
use super::*;
fn get_words(range: Range<usize>) -> Vec<String> {
std::fs::read_to_string("/usr/share/dict/words")
.unwrap()
.split("\n")
.skip(range.start)
.take(range.len())
.map(ToString::to_string)
.collect()
}
#[test]
fn basic_insertion() {
let filter = CuckooFilter::new_random(CuckooConfiguration::builder(1000).build().unwrap());
filter.insert("basic");
assert!(filter.contains("basic"));
}
#[test]
fn basic_removal() {
let filter = CuckooFilter::new_random(CuckooConfiguration::builder(1000).build().unwrap());
filter.insert("basic");
assert!(filter.contains("basic"));
filter.remove("basic");
assert!(!filter.contains("basic"));
}
struct PredefinedBucketItem(u64);
struct TestHasher(u64);
impl BuildHasher for TestHasher {
type Hasher = TestHasher;
fn build_hasher(&self) -> Self::Hasher {
TestHasher(0)
}
}
impl Hasher for TestHasher {
fn finish(&self) -> u64 {
self.0
}
fn write(&mut self, bytes: &[u8]) {
if bytes.len() == 8 {
self.0 = u64::from_ne_bytes(bytes.try_into().unwrap());
} else {
self.0 = 1 - (u32::from_ne_bytes(bytes.try_into().unwrap()) as u64 % 2);
}
}
}
impl Hash for PredefinedBucketItem {
fn hash<H: Hasher>(&self, state: &mut H) {
state.write_u64(self.0);
}
}
#[test]
fn lru_insertion() {
let filter = CuckooFilter::new(
CuckooConfiguration::builder(1000)
.bucket_size(2.try_into().unwrap())
.with_lru(LruConfig {
counter_bits: 8.try_into().unwrap(),
})
.build()
.unwrap(),
TestHasher(0),
);
let test_item = PredefinedBucketItem(2 << 32);
filter.insert(&test_item);
filter.contains(&test_item);
let test_item_2 = PredefinedBucketItem(4 << 32);
filter.insert(&test_item_2);
let test_item_3 = PredefinedBucketItem((3 << 32) + 2);
filter.insert(&test_item_3); filter.contains(&test_item_3);
let test_item_4 = PredefinedBucketItem((5 << 32) + 2);
filter.insert(&test_item_4);
assert!(filter.contains(&test_item));
assert!(filter.contains(&test_item_2));
assert!(filter.contains(&test_item_3));
assert!(filter.contains(&test_item_4));
let test_item_5 = PredefinedBucketItem((1 << 32) + 2);
filter.insert(&test_item_5);
assert!(filter.contains(&test_item_2));
assert!(filter.contains(&test_item));
assert!(filter.contains(&test_item_3));
assert!(
!filter.contains(&test_item_5) || !filter.contains(&test_item_4),
"No inserted items are missing, but filter can't hold them all"
);
filter.insert(&test_item_5);
filter.insert(&test_item_4);
assert!(filter.contains(&test_item));
assert!(filter.contains(&test_item_3));
}
#[test]
fn alt_index() {
let words = get_words(0..200_000);
let filter = CuckooFilter::new_random(
CuckooConfiguration::builder(200_000)
.fingerprint_bits(32.try_into().unwrap())
.build()
.unwrap(),
);
for word in words {
let (fp, index) = filter.get_fingerprint_and_index(&word);
let alt_index = filter.alt_index(&fp, index);
assert_eq!(index, filter.alt_index(&fp, alt_index));
}
}
#[test]
fn random_kicks() {
let filter = CuckooFilter::new(
CuckooConfiguration::builder(1000)
.bucket_size(2.try_into().unwrap())
.build()
.unwrap(),
TestHasher(0),
);
let test_item = PredefinedBucketItem(2 << 32);
filter.insert(&test_item);
let test_item_2 = PredefinedBucketItem(4 << 32);
filter.insert(&test_item_2);
let test_item_3 = PredefinedBucketItem((3 << 32) + 2);
filter.insert(&test_item_3);
let test_item_4 = PredefinedBucketItem((5 << 32) + 2);
filter.insert(&test_item_4);
let test_item_unrelated = PredefinedBucketItem((10 << 32) + 10);
filter.insert(&test_item_unrelated);
assert!(filter.contains(&test_item));
assert!(filter.contains(&test_item_2));
assert!(filter.contains(&test_item_3));
assert!(filter.contains(&test_item_4));
let test_item_5 = PredefinedBucketItem((1 << 32) + 2);
let kicked = filter.insert(&test_item_5);
assert!(kicked.is_some(), "An item had to be kicked");
assert!(filter.contains(&test_item_5));
assert!(filter.contains(&test_item_unrelated));
for item in [&test_item, &test_item_2, &test_item_3, &test_item_4]
.iter()
.filter(|i| !kicked.as_ref().unwrap().matches_key(i, &filter))
{
assert!(filter.contains(item), "Only one item should be kicked");
}
}
#[test]
fn overriding_defaults() {
let filter = CuckooFilter::new_random(
CuckooConfiguration::builder(1000)
.with_ttl(TtlConfig {
ttl: 30.try_into().unwrap(),
ttl_bits: 8.try_into().unwrap(),
})
.with_counter(CounterConfig::default())
.build()
.unwrap(),
);
filter.insert_with_defaults(
"basic",
InsertValues {
ttl: Some(50),
counter: Some(10),
},
);
assert!(filter.contains("basic"));
assert_eq!(
filter
.get_associated_data("basic")
.unwrap()
.get_stored_ttl_value()
.unwrap(),
50
);
assert_eq!(
filter
.get_associated_data("basic")
.unwrap()
.get_counter()
.unwrap(),
13 );
}
#[test]
fn overriding_updates() {
let filter = CuckooFilter::new_random(
CuckooConfiguration::builder(1000)
.with_ttl(TtlConfig {
ttl: 30.try_into().unwrap(),
ttl_bits: 8.try_into().unwrap(),
})
.with_counter(CounterConfig::default())
.build()
.unwrap(),
);
filter.insert_with_defaults(
"basic",
InsertValues {
ttl: Some(5),
counter: Some(1),
},
);
assert!(filter.contains_with_update(
"basic",
LookupValues {
ttl: Some(50),
counter_diff: Some(10),
},
));
assert_eq!(
filter
.get_associated_data("basic")
.unwrap()
.get_stored_ttl_value()
.unwrap(),
50
);
assert_eq!(
filter
.get_associated_data("basic")
.unwrap()
.get_counter()
.unwrap(),
13 );
}
#[test]
fn scan_and_update_full() {
let words = get_words(0..100_000);
let filter = CuckooFilter::new_random(
CuckooConfiguration::builder(100_000)
.fingerprint_bits(32.try_into().unwrap())
.with_lru(LruConfig::default())
.with_ttl(TtlConfig {
ttl: 3.try_into().unwrap(),
ttl_bits: 2.try_into().unwrap(),
})
.build()
.unwrap(),
);
assert_eq!(filter.get_item_count(), 0);
let mut stored_words = HashSet::new();
for (index, word) in words.iter().enumerate() {
stored_words.insert(word);
if let Some(evicted_fp) = filter.insert(word) {
words[0..=index]
.iter()
.filter(|w| evicted_fp.matches_key(w, &filter))
.for_each(|evicted_word| {
stored_words.remove(evicted_word);
});
}
}
assert_eq!(filter.get_item_count(), stored_words.len());
for _ in 0..2 {
assert_eq!(filter.scan_and_update_full(), 0);
}
assert_eq!(filter.get_item_count(), stored_words.len());
for word in stored_words.iter() {
assert!(
filter.contains(word),
"Word: {word} expected in the filter, but not found"
);
}
assert_eq!(filter.scan_and_update_full(), stored_words.len());
for word in &words {
assert!(
!filter.contains(word),
"Filter contained {word}, but shouldn't have"
);
}
assert_eq!(filter.get_item_count(), 0);
}
#[test]
fn export_import() {
let words = get_words(0..100_000);
let filter = CuckooFilter::new_random_exportable(
CuckooConfiguration::builder(100_000)
.fingerprint_bits(32.try_into().unwrap())
.with_lru(LruConfig::default())
.with_ttl(TtlConfig {
ttl: 3.try_into().unwrap(),
ttl_bits: 2.try_into().unwrap(),
})
.build()
.unwrap(),
);
assert_eq!(filter.get_item_count(), 0);
let mut stored_words = HashSet::new();
for (index, word) in words.iter().enumerate() {
stored_words.insert(word);
if let Some(evicted_fp) = filter.insert(word) {
words[0..=index]
.iter()
.filter(|w| evicted_fp.matches_key(w, &filter))
.for_each(|evicted_word| {
stored_words.remove(evicted_word);
});
}
}
assert_eq!(filter.get_item_count(), stored_words.len());
let exported_buf = filter.exporter().snapshot().unwrap();
let mut readable_buf = VecDeque::from(exported_buf);
let imported_filter = CuckooFilter::import_random_exportable(&mut readable_buf).unwrap();
assert_eq!(
imported_filter.get_configuration(),
filter.get_configuration()
);
assert_eq!(imported_filter.get_item_count(), stored_words.len());
for word in stored_words.iter() {
assert!(
imported_filter.contains(word),
"Word: {word} expected in the filter, but not found"
);
}
}
}