Skip to main content

cuckoo_clock/
filter.rs

1use std::{
2    hash::{BuildHasher, Hash, RandomState},
3    io::Read,
4    iter::repeat_with,
5    sync::{
6        Arc, Mutex, MutexGuard,
7        atomic::{AtomicUsize, Ordering},
8    },
9};
10
11use crate::{
12    associated_data::AssociatedData,
13    bucket::{Bucket, InsertValues, LookupValues},
14    config::CuckooConfiguration,
15    data_block::{DataBlock, Fingerprint},
16    exporter::{
17        CuckooFilterExporter, ExportableBuildHasher, ExportableRandomState, import_config,
18        read_hasher_from,
19    },
20};
21
22/// Thread-safe cuckoo filter, with support for TTL, LRU and custom counters associated with the
23/// stored data.
24///
25/// Instances of [`CuckooFilter`] can be cloned and used across different threads. To ensure thread
26/// safety, locks are used, but locking is done per bucket, meaning that 2 separate threads can
27/// freely access different buckets without conflicts. In most cases locks shouldn't block, because
28/// optimal cuckoo filter configuration will have a large number of buckets, reducing the change of
29/// concurrent access to the same bucket.
30///
31/// Instances of [`CuckooFilter`] build using [`CuckooFilter::new_random_exportable`] or with a
32/// [`BuildHasher`] that also implements [`ExportableBuildHasher`] can be exported and imported,
33/// using [`CuckooFilter::exporter`] and [`CuckooFilter::import`]. Note that configuration is
34/// stored in the exported data too and can't be changed, because any changes to the configuration
35/// data would invalidate all of the stored data.
36///
37/// # Examples
38///
39/// Basic cuckoo filter with default configuration
40/// ```
41/// use cuckoo_clock::{CuckooFilter, config::CuckooConfiguration};
42///
43/// let filter = CuckooFilter::new_random(CuckooConfiguration::builder(100_000).build()?);
44///
45/// // None returned from insertion means no entry was evicted
46/// assert!(filter.insert("example_data").is_none());
47///
48/// // Insertion must have been successful
49/// assert!(filter.contains("example_data"));
50///
51/// // Deletion must have been successful
52/// assert!(filter.remove("example_data"));
53/// assert!(!filter.contains("example_data"));
54///
55/// # Ok::<(), Box<dyn std::error::Error>>(())
56/// ```
57///
58/// More complex use-case, with additional options
59/// ```
60/// use cuckoo_clock::{CuckooFilter, config::{CuckooConfiguration, CounterConfig, TtlConfig}};
61///
62/// let filter = CuckooFilter::new_random(
63///     CuckooConfiguration::builder(10_000_000)
64///         .fingerprint_bits(18.try_into()?)
65///         .bucket_size(8.try_into()?)
66///         .with_counter(CounterConfig {
67///             counter_bits: 4.try_into()?,
68///             ..Default::default()
69///         })
70///         .with_ttl(TtlConfig {
71///             ttl: 600.try_into()?,
72///             ttl_bits: 10.try_into()?
73///         })
74///         .build()?
75/// );
76///
77/// // In this case, we use `insert_if_not_present` to ensure no duplicates, because we care about
78/// // the counter
79/// // None returned from insertion means no entry was evicted
80/// assert!(filter.insert_if_not_present("example_data").is_none());
81/// assert!(filter.insert_if_not_present("example_data").is_none());
82///
83/// // Insertion must have been successful
84/// assert!(filter.contains("example_data"));
85///
86/// // Counter should be 4 now.
87/// // We have accessed this item 3 times, but `get_associated_data` also counts as an access.
88/// assert_eq!(filter.get_associated_data("example_data").unwrap().get_counter()?, 4);
89///
90/// # Ok::<(), Box<dyn std::error::Error>>(())
91/// ```
92///
93/// Export/import
94/// ```
95/// use std::{collections::VecDeque, io::Read};
96/// use cuckoo_clock::{CuckooFilter, config::CuckooConfiguration};
97///
98/// let filter = CuckooFilter::new_random_exportable(CuckooConfiguration::builder(100_000).build()?);
99///
100/// // None returned from insertion means no entry was evicted
101/// assert!(filter.insert("example_data").is_none());
102///
103/// let mut buf = Vec::new();
104/// filter.exporter().write_to(&mut buf)?;
105///
106/// let mut buf = VecDeque::from(buf);
107/// let imported_filter = CuckooFilter::import_random_exportable(&mut buf)?;
108///
109/// // The inserted data is available in the imported filter
110/// assert!(filter.contains("example_data"));
111///
112/// # Ok::<(), Box<dyn std::error::Error>>(())
113/// ```
114///
115#[derive(Clone)]
116pub struct CuckooFilter<H: BuildHasher> {
117    configuration: CuckooConfiguration,
118    buckets: Arc<Vec<Mutex<Bucket>>>,
119    build_hasher: H,
120    items: Arc<AtomicUsize>,
121}
122
123impl CuckooFilter<RandomState> {
124    /// Creates a new instance of [`CuckooFilter`], using [`RandomState`] as its [`BuildHasher`].
125    ///
126    /// # Panics
127    ///
128    /// Panics if allocation of buckets fails (if too much memory was requested).
129    #[must_use]
130    pub fn new_random(configuration: CuckooConfiguration) -> Self {
131        Self::new(configuration, RandomState::new())
132    }
133}
134
135impl CuckooFilter<ExportableRandomState> {
136    /// Creates a new instance of [`CuckooFilter`], using [`ExportableBuildHasher`] based on
137    /// [`RandomState`] as its [`BuildHasher`].
138    /// This instance supports export using [`CuckooFilter::export`].
139    ///
140    /// # Panics
141    ///
142    /// Panics if allocation of buckets fails (if too much memory was requested).
143    #[must_use]
144    pub fn new_random_exportable(configuration: CuckooConfiguration) -> Self {
145        Self::new(configuration, ExportableRandomState::new_random())
146    }
147
148    /// Creates a new instance of [`CuckooFilter`], using [`ExportableBuildHasher`] based on
149    /// exported data.
150    ///
151    /// # Panics
152    ///
153    /// Panics if allocation of buckets fails (if too much memory was requested).
154    pub fn import_random_exportable(reader: impl Read) -> Result<Self, crate::ImportError> {
155        Self::import(reader)
156    }
157}
158
159impl<H: ExportableBuildHasher + BuildHasher> CuckooFilter<H> {
160    /// Creates a cuckoo filter from its exported state.
161    ///
162    /// # Panics
163    ///
164    /// Panics if allocation of buckets fails (if too much memory was requested).
165    pub fn import(mut reader: impl Read) -> Result<Self, crate::ImportError> {
166        let hasher = read_hasher_from(&mut reader)?;
167        let configuration = import_config(&mut reader)?;
168        let mut buckets = Vec::with_capacity(configuration.bucket_count);
169
170        let mut item_count = 0;
171        for _ in 0..configuration.bucket_count {
172            let bucket = Bucket::take_from(&mut reader, &configuration)?;
173            item_count += bucket.occupied_count(&configuration);
174            buckets.push(Mutex::new(bucket));
175        }
176
177        Ok(Self {
178            configuration,
179            buckets: Arc::new(buckets),
180            build_hasher: hasher,
181            items: Arc::new(AtomicUsize::new(item_count)),
182        })
183    }
184
185    /// Prepares an exporter for this [`CuckooFilter`], enabling to persist it and import it later
186    /// using [`CuckooFilter::import`].
187    pub fn exporter<'a>(&'a self) -> CuckooFilterExporter<'a, H> {
188        CuckooFilterExporter::new(&self.build_hasher, &self.buckets, &self.configuration)
189    }
190}
191
192impl<H: BuildHasher> CuckooFilter<H> {
193    /// Creates a new instance of [`CuckooFilter`], using provided [`BuildHasher`].
194    ///
195    /// # Panics
196    ///
197    /// Panics if allocation of buckets fails (if too much memory was requested).
198    pub fn new(configuration: CuckooConfiguration, build_hasher: H) -> Self {
199        Self {
200            configuration: configuration.clone(),
201            buckets: repeat_with(|| Bucket::new(&configuration).into())
202                .take(configuration.bucket_count)
203                .collect::<Vec<_>>()
204                .into(),
205            build_hasher,
206            items: Arc::new(AtomicUsize::new(0)),
207        }
208    }
209
210    /// Returns the actual bucket count for this [`CuckooFilter`].
211    ///
212    /// Bucket count is calculated as first next power of two of capacity / bucket_size.
213    /// This means that the actual capacity of the filter is usually bigger than the requested
214    /// capacity.
215    pub const fn get_bucket_count(&self) -> usize {
216        self.configuration.bucket_count
217    }
218
219    /// Returns the actual number of items currently stored in this [`CuckooFilter`].
220    pub fn get_item_count(&self) -> usize {
221        self.items.load(Ordering::Relaxed)
222    }
223
224    /// Returns the configuration for this [`CuckooFilter`].
225    pub fn get_configuration(&self) -> CuckooConfiguration {
226        self.configuration.clone()
227    }
228
229    /// Returns the memory usage of this filter in bytes.
230    pub fn get_memory_usage(&self) -> usize {
231        size_of::<Self>()
232            + size_of::<AtomicUsize>()
233            + size_of::<Vec<Mutex<Bucket>>>()
234            + size_of::<Mutex<Bucket>>() * self.buckets.len()
235            + self.configuration.bucket_byte_size * self.buckets.len()
236    }
237
238    /// Inserts a new item into the filter, only if the filter doesn't contain it already.
239    ///
240    /// This is slower than [`CuckooFilter::insert`], but it ensures that no duplicates are present
241    /// in the filter. That can be useful when [`AssociatedData`] is used, to ensure consistent
242    /// results.
243    ///
244    /// Returns fingerprint of the item that was evicted from the filter, if eviction had to take
245    /// place to finalize the insertion. It is possible that the item that was just inserted gets
246    /// evicted in random kicking process. That can be confirmed using
247    /// [`Fingerprint::matches_key`].
248    pub fn insert_if_not_present<K: Hash + ?Sized>(&self, key: &K) -> Option<Fingerprint> {
249        self.insert_if_not_present_with_update(
250            key,
251            InsertValues::default(),
252            LookupValues::default(),
253        )
254    }
255
256    /// Inserts a new item into the filter, only if the filter doesn't contain it already. Also
257    /// applies provided updates.
258    ///
259    /// This is similar to [`CuckooFilter::insert_if_not_present`], but it also updates found
260    /// values, or starts off values with different values.
261    pub fn insert_if_not_present_with_update<K: Hash + ?Sized>(
262        &self,
263        key: &K,
264        insert_values: InsertValues,
265        lookup_update: LookupValues,
266    ) -> Option<Fingerprint> {
267        let (fp, i1) = self.get_fingerprint_and_index(key);
268
269        let mut contains =
270            self.lock_bucket(i1 as usize)
271                .contains(&fp, &self.configuration, &lookup_update);
272
273        if contains {
274            return None;
275        }
276
277        let i2 = self.alt_index(&fp, i1);
278        contains = self
279            .lock_bucket(i2 as usize)
280            .contains(&fp, &self.configuration, &lookup_update);
281
282        if contains {
283            return None;
284        }
285
286        let mut cur_data_block = self.new_data_block(&fp, insert_values);
287
288        let inserted = self
289            .lock_bucket(i1 as usize)
290            .insert(&cur_data_block, &self.configuration);
291
292        if inserted {
293            self.items.fetch_add(1, Ordering::Relaxed);
294            return None;
295        }
296
297        let inserted = self
298            .lock_bucket(i2 as usize)
299            .insert(&cur_data_block, &self.configuration);
300
301        if inserted {
302            self.items.fetch_add(1, Ordering::Relaxed);
303            return None;
304        }
305
306        let mut cur_index = if rand::random::<bool>() { i1 } else { i2 };
307        for _ in 0..self.configuration.max_kicks {
308            {
309                let mut bucket = self.lock_bucket(cur_index as usize);
310                // Replace a random item first
311                if let Some(lru_config) = self.configuration.lru_field_config.as_ref() {
312                    if !bucket.kick_lru(&mut cur_data_block, &self.configuration, lru_config) {
313                        return Some(cur_data_block.get_fingerprint(&self.configuration));
314                    }
315                } else {
316                    bucket.kick_random(&mut cur_data_block, &self.configuration);
317                }
318                cur_index = self.alt_index(
319                    &cur_data_block.get_fingerprint(&self.configuration),
320                    cur_index,
321                );
322            }
323
324            if self
325                .lock_bucket(cur_index as usize)
326                .insert(&cur_data_block, &self.configuration)
327            {
328                self.items.fetch_add(1, Ordering::Relaxed);
329                // Found an alternative spot for evicted item, done with kicks
330                return None;
331            }
332        }
333
334        // Filter is full
335        Some(cur_data_block.get_fingerprint(&self.configuration))
336    }
337
338    /// Inserts a new item into the filter.
339    ///
340    /// If both target buckets for this item are full, random item is kicked out of one of these 2
341    /// buckets and moved into its alternate bucket, starting a recursive kicking process, which
342    /// stops once an empty slot is found in alternate bucket of a kicked item, or
343    /// [`crate::config::CuckooConfigurationBuilder::max_kicks`] is reached.
344    ///
345    /// Returns fingerprint of the item that was evicted from the filter, if eviction had to take
346    /// place to finalize the insertion. It is possible that the item that was just inserted gets
347    /// evicted in random kicking process. That can be confirmed using
348    /// [`Fingerprint::matches_key`].
349    pub fn insert<K: Hash + ?Sized>(&self, key: &K) -> Option<Fingerprint> {
350        self.insert_with_defaults(key, InsertValues::default())
351    }
352
353    /// Inserts a new item into the filter, with defined defaults for associated data.
354    ///
355    /// This is similar to [`CuckooFilter::insert`], but allows overrides for associated data
356    /// defaults.
357    pub fn insert_with_defaults<K: Hash + ?Sized>(
358        &self,
359        key: &K,
360        default: InsertValues,
361    ) -> Option<Fingerprint> {
362        let (fp, i1) = self.get_fingerprint_and_index(key);
363        let mut cur_data_block = self.new_data_block(&fp, default);
364
365        let inserted = self
366            .lock_bucket(i1 as usize)
367            .insert(&cur_data_block, &self.configuration);
368
369        if inserted {
370            self.items.fetch_add(1, Ordering::Relaxed);
371            return None;
372        }
373
374        let i2 = self.alt_index(&fp, i1);
375
376        let inserted = self
377            .lock_bucket(i2 as usize)
378            .insert(&cur_data_block, &self.configuration);
379
380        if inserted {
381            self.items.fetch_add(1, Ordering::Relaxed);
382            return None;
383        }
384
385        let mut cur_index = i1;
386        for _ in 0..self.configuration.max_kicks {
387            {
388                let mut bucket = self.lock_bucket(cur_index as usize);
389                // Replace a random item first
390                if let Some(lru_config) = self.configuration.lru_field_config.as_ref() {
391                    if !bucket.kick_lru(&mut cur_data_block, &self.configuration, lru_config) {
392                        return Some(cur_data_block.get_fingerprint(&self.configuration));
393                    }
394                } else {
395                    // TODO: this can even kick the newest item, which is not ideal
396                    bucket.kick_random(&mut cur_data_block, &self.configuration);
397                }
398                cur_index = self.alt_index(
399                    &cur_data_block.get_fingerprint(&self.configuration),
400                    cur_index,
401                );
402            }
403
404            if self
405                .lock_bucket(cur_index as usize)
406                .insert(&cur_data_block, &self.configuration)
407            {
408                self.items.fetch_add(1, Ordering::Relaxed);
409                // Found an alternative spot for evicted item, done with kicks
410                return None;
411            }
412        }
413
414        // Filter is full
415        Some(cur_data_block.get_fingerprint(&self.configuration))
416    }
417
418    /// Check if this key is stored in the filter and applies the provided [`LookupValues`].
419    ///
420    /// This is similar to [`CuckooFilter::contains`], but allows overrides for updates on
421    /// successful lookup.
422    pub fn contains_with_update<K: Hash + ?Sized>(&self, key: &K, update: LookupValues) -> bool {
423        let (fp, i1) = self.get_fingerprint_and_index(key);
424
425        let mut contains =
426            self.lock_bucket(i1 as usize)
427                .contains(&fp, &self.configuration, &update);
428
429        if !contains {
430            let i2 = self.alt_index(&fp, i1);
431            contains = self
432                .lock_bucket(i2 as usize)
433                .contains(&fp, &self.configuration, &update);
434        }
435
436        contains
437    }
438
439    /// Check if this key is stored in the filter.
440    ///
441    /// Returns true if this key might be present in the filter. If false is returned, then the key
442    /// is definitely not present.
443    pub fn contains<K: Hash + ?Sized>(&self, key: &K) -> bool {
444        self.contains_with_update(key, LookupValues::default())
445    }
446
447    /// Loads associated data of a key stored in the filter.
448    ///
449    /// Returns None if this filter is not present in the filter. Returns associated data for the
450    /// first item with the fingerprint matching this key's fingerprint. Note that it is
451    /// recommended to use [`CuckooFilter::insert_if_not_present`] if consistent [`AssociatedData`]
452    /// is required.
453    pub fn get_associated_data<K: Hash + ?Sized>(&self, key: &K) -> Option<AssociatedData> {
454        self.get_associated_data_with_update(key, LookupValues::default())
455    }
456
457    /// Loads associated data of a key stored in the filter and applies the provided [`LookupValues`].
458    ///
459    /// This is similar to [`CuckooFilter::get_associated_data`], but allows overrides for updates on
460    /// successful lookup.
461    pub fn get_associated_data_with_update<K: Hash + ?Sized>(
462        &self,
463        key: &K,
464        update: LookupValues,
465    ) -> Option<AssociatedData> {
466        let (fp, i1) = self.get_fingerprint_and_index(key);
467
468        let mut contains =
469            self.lock_bucket(i1 as usize)
470                .get_associated_data(&fp, &self.configuration, &update);
471
472        if contains.is_none() {
473            let i2 = self.alt_index(&fp, i1);
474            contains = self.lock_bucket(i2 as usize).get_associated_data(
475                &fp,
476                &self.configuration,
477                &update,
478            );
479        }
480
481        contains
482    }
483
484    /// Removes this key from the filter, if present.
485    ///
486    /// Returns true if the key was present in the filter.
487    pub fn remove<K: Hash + ?Sized>(&self, key: &K) -> bool {
488        let (fp, i1) = self.get_fingerprint_and_index(key);
489
490        let mut removed = self
491            .lock_bucket(i1 as usize)
492            .remove(&fp, &self.configuration);
493
494        if !removed {
495            let i2 = self.alt_index(&fp, i1);
496            removed = self
497                .lock_bucket(i2 as usize)
498                .remove(&fp, &self.configuration);
499        }
500
501        if removed {
502            self.items.fetch_sub(1, Ordering::Relaxed);
503        }
504
505        removed
506    }
507
508    /// Scans all buckets of this filter and reduces TTL and LRU counters.
509    ///
510    /// If LRU and/or TTL features are used, this must be called periodically.
511    /// Each call to this function will age all the LRU and TTL counters. The frequency of calls
512    /// will affect both LRU and TTL in different ways:
513    /// - TTL will get reduced by 1 on each call, meaning that scanning each second indirectly sets
514    ///   the unit of TTL field to be seconds.
515    /// - LRU will get halved on each call. By scanning more frequently, items will require more
516    ///   frequent usage to stay in the filter.
517    ///
518    /// This is a no-op if both LRU and TTL are disabled.
519    ///
520    /// # Examples
521    ///
522    /// ```
523    /// use cuckoo_clock::{CuckooFilter, config::{CuckooConfiguration, CounterConfig, TtlConfig}};
524    ///
525    /// let filter = CuckooFilter::new_random(
526    ///     CuckooConfiguration::builder(10_000)
527    ///         .with_ttl(TtlConfig {
528    ///             ttl: 3.try_into()?,
529    ///             ttl_bits: 2.try_into()?
530    ///         })
531    ///         .build()?
532    /// );
533    ///
534    /// filter.insert("example_data");
535    ///
536    /// assert!(filter.contains("example_data"));
537    ///
538    /// filter.scan_and_update_full();
539    /// assert!(filter.contains("example_data"));
540    ///
541    /// filter.scan_and_update_full();
542    /// assert!(filter.contains("example_data"));
543    ///
544    /// // The item will get removed now, due to expired TTL
545    /// filter.scan_and_update_full();
546    /// assert!(!filter.contains("example_data"));
547    ///
548    /// # Ok::<(), Box<dyn std::error::Error>>(())
549    /// ```
550    pub fn scan_and_update_full(&self) -> usize {
551        if self.configuration.lru_field_config.is_none()
552            && self.configuration.ttl_field_config.is_none()
553        {
554            return 0;
555        }
556
557        let mut removed = 0;
558        for b in self.buckets.iter() {
559            #[expect(clippy::unwrap_used)]
560            let mut bucket = b.lock().unwrap();
561            if let Some(lru_config) = &self.configuration.lru_field_config {
562                bucket.age_lru_counters(&self.configuration, lru_config);
563            }
564            if let Some(ttl_config) = &self.configuration.ttl_field_config {
565                removed += bucket.age_ttl_counters(&self.configuration, ttl_config);
566            }
567        }
568
569        if removed > 0 {
570            self.items.fetch_sub(removed, Ordering::Relaxed);
571        }
572        removed
573    }
574
575    /// Scans all buckets of this filter and reduces TTL counters.
576    ///
577    /// Similar to [`CuckooFilter::scan_and_update_full`], but updates only TTL counters. This
578    /// allows more control, enabling different update frequency for TTL and LRU.
579    ///
580    /// This is a no-op if TTL is disabled.
581    pub fn scan_and_update_ttl(&self) -> usize {
582        if self.configuration.ttl_field_config.is_none() {
583            return 0;
584        }
585
586        let mut removed = 0;
587        for b in self.buckets.iter() {
588            #[expect(clippy::unwrap_used)]
589            let mut bucket = b.lock().unwrap();
590            if let Some(ttl_config) = &self.configuration.ttl_field_config {
591                removed += bucket.age_ttl_counters(&self.configuration, ttl_config);
592            }
593        }
594
595        if removed > 0 {
596            self.items.fetch_sub(removed, Ordering::Relaxed);
597        }
598        removed
599    }
600
601    /// Scans all buckets of this filter and reduces LRU counters.
602    ///
603    /// Similar to [`CuckooFilter::scan_and_update_full`], but updates only LRU counters. This
604    /// allows more control, enabling different update frequency for TTL and LRU.
605    ///
606    /// This is a no-op if LRU is disabled.
607    pub fn scan_and_update_lru(&self) {
608        if self.configuration.lru_field_config.is_none() {
609            return;
610        }
611
612        for b in self.buckets.iter() {
613            #[expect(clippy::unwrap_used)]
614            let mut bucket = b.lock().unwrap();
615            if let Some(lru_config) = &self.configuration.lru_field_config {
616                bucket.age_lru_counters(&self.configuration, lru_config);
617            }
618        }
619    }
620
621    /// Generates the fingerprint and first index for the provided key.
622    pub(crate) fn get_fingerprint<K: Hash + ?Sized>(&self, key: &K) -> Fingerprint {
623        self.get_fingerprint_and_index(key).0
624    }
625
626    fn new_data_block(&self, fp: &Fingerprint, defaults: InsertValues) -> DataBlock<Vec<u8>> {
627        let data = vec![0u8; self.configuration.data_block_size];
628        let mut cur_data_block = DataBlock::from(data);
629        cur_data_block.store_fingerprint(fp, &self.configuration);
630
631        if let Some(ttl_config) = self.configuration.ttl_field_config.as_ref() {
632            cur_data_block.set_ttl(ttl_config, defaults.ttl.unwrap_or(ttl_config.0.ttl.into()));
633        }
634        if let Some(counter_config) = self.configuration.counter_field_config.as_ref() {
635            cur_data_block.update_counter(
636                counter_config,
637                defaults
638                    .counter
639                    .unwrap_or(counter_config.0.change_on_insert),
640            );
641        }
642        if let Some(lru_config) = self.configuration.lru_field_config.as_ref() {
643            cur_data_block.inc_lru_counter(lru_config);
644        }
645        cur_data_block
646    }
647
648    fn get_fingerprint_and_index<K: Hash + ?Sized>(&self, key: &K) -> (Fingerprint, u32) {
649        let result = self.build_hasher.hash_one(key);
650
651        // Fingeprint bits over 32 are definitely an overkill
652        // We can reduce number of hashes by using one hash as fingerprint and first index
653        let fingerprint = (result >> 32) as u32;
654        // Intentional truncation here
655        #[expect(clippy::cast_possible_truncation)]
656        let index = result as u32 & self.configuration.buckets_mask;
657
658        (
659            Fingerprint::new(
660                fingerprint,
661                self.configuration.fingerprint_field_config.value_mask(),
662            ),
663            index,
664        )
665    }
666
667    // Intentional truncation here
668    #[expect(clippy::cast_possible_truncation)]
669    fn alt_index(&self, fingerprint: &Fingerprint, index: u32) -> u32 {
670        let result = self.build_hasher.hash_one(fingerprint);
671
672        (index ^ ((result as u32) & self.configuration.buckets_mask))
673            & self.configuration.buckets_mask
674    }
675
676    #[expect(clippy::unwrap_used)]
677    fn lock_bucket(&self, index: usize) -> MutexGuard<'_, Bucket> {
678        // Any panic while lock is held should come from this library
679        // Any panic produced while the lock is held is a bug in the library!
680        self.buckets[index].lock().unwrap()
681    }
682}
683
684#[cfg(test)]
685#[expect(clippy::unwrap_used)]
686mod tests {
687    use std::{
688        collections::{HashSet, VecDeque},
689        hash::Hasher,
690        ops::Range,
691    };
692
693    use crate::config::{CounterConfig, LruConfig, TtlConfig};
694
695    use super::*;
696
697    fn get_words(range: Range<usize>) -> Vec<String> {
698        std::fs::read_to_string("/usr/share/dict/words")
699            .unwrap()
700            .split("\n")
701            .skip(range.start)
702            .take(range.len())
703            .map(ToString::to_string)
704            .collect()
705    }
706
707    #[test]
708    fn basic_insertion() {
709        let filter = CuckooFilter::new_random(CuckooConfiguration::builder(1000).build().unwrap());
710
711        filter.insert("basic");
712
713        assert!(filter.contains("basic"));
714    }
715
716    #[test]
717    fn basic_removal() {
718        let filter = CuckooFilter::new_random(CuckooConfiguration::builder(1000).build().unwrap());
719
720        filter.insert("basic");
721
722        assert!(filter.contains("basic"));
723
724        filter.remove("basic");
725
726        assert!(!filter.contains("basic"));
727    }
728
729    struct PredefinedBucketItem(u64);
730    struct TestHasher(u64);
731    impl BuildHasher for TestHasher {
732        type Hasher = TestHasher;
733
734        fn build_hasher(&self) -> Self::Hasher {
735            TestHasher(0)
736        }
737    }
738    impl Hasher for TestHasher {
739        fn finish(&self) -> u64 {
740            self.0
741        }
742
743        fn write(&mut self, bytes: &[u8]) {
744            if bytes.len() == 8 {
745                self.0 = u64::from_ne_bytes(bytes.try_into().unwrap());
746            } else {
747                // Shift fingeprint hashes a bit, to allow control
748                self.0 = 1 - (u32::from_ne_bytes(bytes.try_into().unwrap()) as u64 % 2);
749            }
750        }
751    }
752    impl Hash for PredefinedBucketItem {
753        fn hash<H: Hasher>(&self, state: &mut H) {
754            state.write_u64(self.0);
755        }
756    }
757
758    #[test]
759    fn lru_insertion() {
760        let filter = CuckooFilter::new(
761            CuckooConfiguration::builder(1000)
762                .bucket_size(2.try_into().unwrap())
763                .with_lru(LruConfig {
764                    counter_bits: 8.try_into().unwrap(),
765                })
766                .build()
767                .unwrap(),
768            TestHasher(0),
769        );
770
771        let test_item = PredefinedBucketItem(2 << 32);
772        filter.insert(&test_item);
773        filter.contains(&test_item); // Make it more used than others
774
775        let test_item_2 = PredefinedBucketItem(4 << 32);
776        filter.insert(&test_item_2); // Sharing the same bucket as "test", but less used
777
778        let test_item_3 = PredefinedBucketItem((3 << 32) + 2);
779        filter.insert(&test_item_3); // Another bucket, but also valid for "test" bucket
780        filter.contains(&test_item_3); // Make it more used
781
782        let test_item_4 = PredefinedBucketItem((5 << 32) + 2);
783        filter.insert(&test_item_4); // Takes bucket of "test_item_3", but less used
784
785        // Everything fits now
786        assert!(filter.contains(&test_item));
787        assert!(filter.contains(&test_item_2));
788        assert!(filter.contains(&test_item_3));
789        assert!(filter.contains(&test_item_4));
790
791        let test_item_5 = PredefinedBucketItem((1 << 32) + 2);
792        // Insert a new item which has to take one of the 2 fully occupied buckets
793        filter.insert(&test_item_5);
794
795        assert!(filter.contains(&test_item_2));
796        assert!(filter.contains(&test_item));
797        assert!(filter.contains(&test_item_3));
798
799        assert!(
800            !filter.contains(&test_item_5) || !filter.contains(&test_item_4),
801            "No inserted items are missing, but filter can't hold them all"
802        );
803
804        // Insert both of these items again and confirm the more used ones are still there
805        filter.insert(&test_item_5);
806        filter.insert(&test_item_4);
807        assert!(filter.contains(&test_item));
808        assert!(filter.contains(&test_item_3));
809    }
810
811    #[test]
812    fn alt_index() {
813        let words = get_words(0..200_000);
814        let filter = CuckooFilter::new_random(
815            CuckooConfiguration::builder(200_000)
816                .fingerprint_bits(32.try_into().unwrap())
817                .build()
818                .unwrap(),
819        );
820
821        for word in words {
822            let (fp, index) = filter.get_fingerprint_and_index(&word);
823            let alt_index = filter.alt_index(&fp, index);
824            assert_eq!(index, filter.alt_index(&fp, alt_index));
825        }
826    }
827
828    #[test]
829    fn random_kicks() {
830        let filter = CuckooFilter::new(
831            CuckooConfiguration::builder(1000)
832                .bucket_size(2.try_into().unwrap())
833                .build()
834                .unwrap(),
835            TestHasher(0),
836        );
837
838        let test_item = PredefinedBucketItem(2 << 32);
839        filter.insert(&test_item);
840
841        let test_item_2 = PredefinedBucketItem(4 << 32);
842        filter.insert(&test_item_2); // Sharing the same bucket as "test"
843
844        let test_item_3 = PredefinedBucketItem((3 << 32) + 2);
845        filter.insert(&test_item_3); // Another bucket, but also valid for "test" bucket
846
847        let test_item_4 = PredefinedBucketItem((5 << 32) + 2);
848        filter.insert(&test_item_4); // Takes bucket of "test_item_3"
849
850        // This one should not be kicked, because it takes an unrelated bucket
851        let test_item_unrelated = PredefinedBucketItem((10 << 32) + 10);
852        filter.insert(&test_item_unrelated);
853
854        // Everything fits now
855        assert!(filter.contains(&test_item));
856        assert!(filter.contains(&test_item_2));
857        assert!(filter.contains(&test_item_3));
858        assert!(filter.contains(&test_item_4));
859
860        let test_item_5 = PredefinedBucketItem((1 << 32) + 2);
861        // Insert a new item which has to take one of the 2 fully occupied buckets
862        let kicked = filter.insert(&test_item_5);
863        assert!(kicked.is_some(), "An item had to be kicked");
864        assert!(filter.contains(&test_item_5));
865        assert!(filter.contains(&test_item_unrelated));
866
867        for item in [&test_item, &test_item_2, &test_item_3, &test_item_4]
868            .iter()
869            .filter(|i| !kicked.as_ref().unwrap().matches_key(i, &filter))
870        {
871            assert!(filter.contains(item), "Only one item should be kicked");
872        }
873    }
874
875    #[test]
876    fn overriding_defaults() {
877        let filter = CuckooFilter::new_random(
878            CuckooConfiguration::builder(1000)
879                .with_ttl(TtlConfig {
880                    ttl: 30.try_into().unwrap(),
881                    ttl_bits: 8.try_into().unwrap(),
882                })
883                .with_counter(CounterConfig::default())
884                .build()
885                .unwrap(),
886        );
887
888        filter.insert_with_defaults(
889            "basic",
890            InsertValues {
891                ttl: Some(50),
892                counter: Some(10),
893            },
894        );
895
896        assert!(filter.contains("basic"));
897        assert_eq!(
898            filter
899                .get_associated_data("basic")
900                .unwrap()
901                .get_stored_ttl_value()
902                .unwrap(),
903            50
904        );
905        assert_eq!(
906            filter
907                .get_associated_data("basic")
908                .unwrap()
909                .get_counter()
910                .unwrap(),
911            13 // initial 10 + 1 on contains + 2x1 on get_associated_data
912        );
913    }
914
915    #[test]
916    fn overriding_updates() {
917        let filter = CuckooFilter::new_random(
918            CuckooConfiguration::builder(1000)
919                .with_ttl(TtlConfig {
920                    ttl: 30.try_into().unwrap(),
921                    ttl_bits: 8.try_into().unwrap(),
922                })
923                .with_counter(CounterConfig::default())
924                .build()
925                .unwrap(),
926        );
927
928        filter.insert_with_defaults(
929            "basic",
930            InsertValues {
931                ttl: Some(5),
932                counter: Some(1),
933            },
934        );
935
936        assert!(filter.contains_with_update(
937            "basic",
938            LookupValues {
939                ttl: Some(50),
940                counter_diff: Some(10),
941            },
942        ));
943        assert_eq!(
944            filter
945                .get_associated_data("basic")
946                .unwrap()
947                .get_stored_ttl_value()
948                .unwrap(),
949            50
950        );
951        assert_eq!(
952            filter
953                .get_associated_data("basic")
954                .unwrap()
955                .get_counter()
956                .unwrap(),
957            13 // initial 1 + 10 on contains + 2x1 on get_associated_data
958        );
959    }
960
961    #[test]
962    fn scan_and_update_full() {
963        let words = get_words(0..100_000);
964        let filter = CuckooFilter::new_random(
965            CuckooConfiguration::builder(100_000)
966                .fingerprint_bits(32.try_into().unwrap())
967                .with_lru(LruConfig::default())
968                .with_ttl(TtlConfig {
969                    ttl: 3.try_into().unwrap(),
970                    ttl_bits: 2.try_into().unwrap(),
971                })
972                .build()
973                .unwrap(),
974        );
975
976        assert_eq!(filter.get_item_count(), 0);
977
978        let mut stored_words = HashSet::new();
979
980        for (index, word) in words.iter().enumerate() {
981            stored_words.insert(word);
982            if let Some(evicted_fp) = filter.insert(word) {
983                words[0..=index]
984                    .iter()
985                    .filter(|w| evicted_fp.matches_key(w, &filter))
986                    .for_each(|evicted_word| {
987                        stored_words.remove(evicted_word);
988                    });
989            }
990        }
991
992        assert_eq!(filter.get_item_count(), stored_words.len());
993
994        for _ in 0..2 {
995            assert_eq!(filter.scan_and_update_full(), 0);
996        }
997
998        assert_eq!(filter.get_item_count(), stored_words.len());
999        for word in stored_words.iter() {
1000            assert!(
1001                filter.contains(word),
1002                "Word: {word} expected in the filter, but not found"
1003            );
1004        }
1005
1006        // TTL should remove all entries now
1007        assert_eq!(filter.scan_and_update_full(), stored_words.len());
1008        for word in &words {
1009            assert!(
1010                !filter.contains(word),
1011                "Filter contained {word}, but shouldn't have"
1012            );
1013        }
1014        assert_eq!(filter.get_item_count(), 0);
1015    }
1016
1017    #[test]
1018    fn export_import() {
1019        let words = get_words(0..100_000);
1020        let filter = CuckooFilter::new_random_exportable(
1021            CuckooConfiguration::builder(100_000)
1022                .fingerprint_bits(32.try_into().unwrap())
1023                .with_lru(LruConfig::default())
1024                .with_ttl(TtlConfig {
1025                    ttl: 3.try_into().unwrap(),
1026                    ttl_bits: 2.try_into().unwrap(),
1027                })
1028                .build()
1029                .unwrap(),
1030        );
1031
1032        assert_eq!(filter.get_item_count(), 0);
1033
1034        let mut stored_words = HashSet::new();
1035
1036        for (index, word) in words.iter().enumerate() {
1037            stored_words.insert(word);
1038            if let Some(evicted_fp) = filter.insert(word) {
1039                words[0..=index]
1040                    .iter()
1041                    .filter(|w| evicted_fp.matches_key(w, &filter))
1042                    .for_each(|evicted_word| {
1043                        stored_words.remove(evicted_word);
1044                    });
1045            }
1046        }
1047
1048        assert_eq!(filter.get_item_count(), stored_words.len());
1049
1050        let exported_buf = filter.exporter().snapshot().unwrap();
1051        let mut readable_buf = VecDeque::from(exported_buf);
1052
1053        let imported_filter = CuckooFilter::import_random_exportable(&mut readable_buf).unwrap();
1054
1055        assert_eq!(
1056            imported_filter.get_configuration(),
1057            filter.get_configuration()
1058        );
1059
1060        assert_eq!(imported_filter.get_item_count(), stored_words.len());
1061        for word in stored_words.iter() {
1062            assert!(
1063                imported_filter.contains(word),
1064                "Word: {word} expected in the filter, but not found"
1065            );
1066        }
1067    }
1068}