Skip to main content

cuckoo_clock/
filter.rs

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