1extern crate alloc;
12
13use crate::config::LfuCacheConfig;
14use crate::list::{Entry, List};
15use crate::metrics::{CacheMetrics, LfuCacheMetrics};
16use alloc::boxed::Box;
17use alloc::collections::BTreeMap;
18use alloc::string::String;
19use core::borrow::Borrow;
20use core::hash::{BuildHasher, Hash};
21use core::mem;
22use core::num::NonZeroUsize;
23
24type FrequencyMetadata<K, V> = (usize, *mut Entry<(K, V)>);
26
27#[cfg(feature = "hashbrown")]
28use hashbrown::hash_map::DefaultHashBuilder;
29#[cfg(feature = "hashbrown")]
30use hashbrown::HashMap;
31
32#[cfg(not(feature = "hashbrown"))]
33use std::collections::hash_map::RandomState as DefaultHashBuilder;
34#[cfg(not(feature = "hashbrown"))]
35use std::collections::HashMap;
36
37pub(crate) struct LfuSegment<K, V, S = DefaultHashBuilder> {
51 config: LfuCacheConfig,
53
54 min_frequency: usize,
56
57 map: HashMap<K, FrequencyMetadata<K, V>, S>,
59
60 frequency_lists: BTreeMap<usize, List<(K, V)>>,
63
64 metrics: LfuCacheMetrics,
66}
67
68unsafe impl<K: Send, V: Send, S: Send> Send for LfuSegment<K, V, S> {}
71
72unsafe impl<K: Send, V: Send, S: Sync> Sync for LfuSegment<K, V, S> {}
74
75impl<K: Hash + Eq, V: Clone, S: BuildHasher> LfuSegment<K, V, S> {
76 pub(crate) fn with_hasher(cap: NonZeroUsize, hash_builder: S) -> Self {
78 let config = LfuCacheConfig::new(cap);
79 let map_capacity = config.capacity().get().next_power_of_two();
80 let max_cache_size_bytes = config.capacity().get() as u64 * 128;
81 LfuSegment {
82 config,
83 min_frequency: 1,
84 map: HashMap::with_capacity_and_hasher(map_capacity, hash_builder),
85 frequency_lists: BTreeMap::new(),
86 metrics: LfuCacheMetrics::new(max_cache_size_bytes),
87 }
88 }
89
90 #[inline]
92 pub(crate) fn cap(&self) -> NonZeroUsize {
93 self.config.capacity()
94 }
95
96 #[inline]
98 pub(crate) fn len(&self) -> usize {
99 self.map.len()
100 }
101
102 #[inline]
104 pub(crate) fn is_empty(&self) -> bool {
105 self.map.is_empty()
106 }
107
108 fn estimate_object_size(&self, _key: &K, _value: &V) -> u64 {
110 mem::size_of::<K>() as u64 + mem::size_of::<V>() as u64 + 64
111 }
112
113 #[inline]
115 pub(crate) fn metrics(&self) -> &LfuCacheMetrics {
116 &self.metrics
117 }
118
119 unsafe fn update_frequency_by_node(
127 &mut self,
128 node: *mut Entry<(K, V)>,
129 old_frequency: usize,
130 ) -> *mut Entry<(K, V)>
131 where
132 K: Clone + Hash + Eq,
133 {
134 let new_frequency = old_frequency + 1;
135
136 self.metrics
138 .record_frequency_increment(old_frequency, new_frequency);
139
140 let (key_ref, _) = (*node).get_value();
142 let key_cloned = key_ref.clone();
143
144 let (_, node) = self.map.get(&key_cloned).unwrap();
146
147 let boxed_entry = self
149 .frequency_lists
150 .get_mut(&old_frequency)
151 .unwrap()
152 .remove(*node)
153 .unwrap();
154
155 if self.frequency_lists.get(&old_frequency).unwrap().is_empty()
158 && old_frequency == self.min_frequency
159 {
160 self.min_frequency = new_frequency;
161 }
162
163 let entry_ptr = Box::into_raw(boxed_entry);
165
166 let capacity = self.config.capacity();
168 self.frequency_lists
169 .entry(new_frequency)
170 .or_insert_with(|| List::new(capacity));
171
172 self.frequency_lists
174 .get_mut(&new_frequency)
175 .unwrap()
176 .attach_from_other_list(entry_ptr);
177
178 self.map.get_mut(&key_cloned).unwrap().0 = new_frequency;
180 self.map.get_mut(&key_cloned).unwrap().1 = entry_ptr;
181
182 self.metrics.update_frequency_levels(&self.frequency_lists);
184
185 entry_ptr
186 }
187
188 pub(crate) fn get<Q>(&mut self, key: &Q) -> Option<&V>
190 where
191 K: Borrow<Q> + Clone,
192 Q: ?Sized + Hash + Eq,
193 {
194 if let Some(&(frequency, node)) = self.map.get(key) {
195 unsafe {
196 let (key_ref, value) = (*node).get_value();
198 let object_size = self.estimate_object_size(key_ref, value);
199 self.metrics.record_frequency_hit(object_size, frequency);
200
201 let new_node = self.update_frequency_by_node(node, frequency);
202 let (_, value) = (*new_node).get_value();
203 Some(value)
204 }
205 } else {
206 None
207 }
208 }
209
210 pub(crate) fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
212 where
213 K: Borrow<Q> + Clone,
214 Q: ?Sized + Hash + Eq,
215 {
216 if let Some(&(frequency, node)) = self.map.get(key) {
217 unsafe {
218 let (key_ref, value) = (*node).get_value();
220 let object_size = self.estimate_object_size(key_ref, value);
221 self.metrics.record_frequency_hit(object_size, frequency);
222
223 let new_node = self.update_frequency_by_node(node, frequency);
224 let (_, value) = (*new_node).get_value_mut();
225 Some(value)
226 }
227 } else {
228 None
229 }
230 }
231
232 pub(crate) fn put(&mut self, key: K, value: V) -> Option<(K, V)>
234 where
235 K: Clone,
236 {
237 let object_size = self.estimate_object_size(&key, &value);
238
239 if let Some(&(frequency, node)) = self.map.get(&key) {
241 unsafe {
242 let old_entry = self.frequency_lists.get_mut(&frequency).unwrap().update(
244 node,
245 (key.clone(), value),
246 true,
247 );
248
249 self.metrics.core.record_insertion(object_size);
250 return old_entry.0;
251 }
252 }
253
254 let mut evicted = None;
255
256 if self.len() >= self.config.capacity().get() {
258 if let Some(min_freq_list) = self.frequency_lists.get_mut(&self.min_frequency) {
259 if let Some(old_entry) = min_freq_list.remove_last() {
260 unsafe {
261 let entry_ptr = Box::into_raw(old_entry);
262 let (old_key, old_value) = (*entry_ptr).get_value().clone();
263 let evicted_size = self.estimate_object_size(&old_key, &old_value);
264 self.metrics.core.record_eviction(evicted_size);
265 self.map.remove(&old_key);
266 evicted = Some((old_key, old_value));
267 let _ = Box::from_raw(entry_ptr);
268 }
269 }
270 }
271 }
272
273 let frequency = 1;
275 self.min_frequency = 1;
276
277 let capacity = self.config.capacity();
279 self.frequency_lists
280 .entry(frequency)
281 .or_insert_with(|| List::new(capacity));
282
283 if let Some(node) = self
284 .frequency_lists
285 .get_mut(&frequency)
286 .unwrap()
287 .add((key.clone(), value))
288 {
289 self.map.insert(key, (frequency, node));
290 }
291
292 self.metrics.core.record_insertion(object_size);
293 self.metrics.update_frequency_levels(&self.frequency_lists);
294
295 evicted
296 }
297
298 pub(crate) fn remove<Q>(&mut self, key: &Q) -> Option<V>
300 where
301 K: Borrow<Q>,
302 Q: ?Sized + Hash + Eq,
303 V: Clone,
304 {
305 let (frequency, node) = self.map.remove(key)?;
306
307 unsafe {
308 let boxed_entry = self.frequency_lists.get_mut(&frequency)?.remove(node)?;
310 let entry_ptr = Box::into_raw(boxed_entry);
311 let value = (*entry_ptr).get_value().1.clone();
312 let _ = Box::from_raw(entry_ptr);
313
314 if self.frequency_lists.get(&frequency).unwrap().is_empty()
316 && frequency == self.min_frequency
317 {
318 self.min_frequency = self
319 .frequency_lists
320 .keys()
321 .find(|&&f| f > frequency && !self.frequency_lists.get(&f).unwrap().is_empty())
322 .copied()
323 .unwrap_or(1);
324 }
325
326 Some(value)
327 }
328 }
329
330 pub(crate) fn clear(&mut self) {
332 self.map.clear();
333 self.frequency_lists.clear();
334 self.min_frequency = 1;
335 }
336
337 #[inline]
339 pub(crate) fn record_miss(&mut self, object_size: u64) {
340 self.metrics.record_miss(object_size);
341 }
342}
343
344impl<K, V, S> core::fmt::Debug for LfuSegment<K, V, S> {
346 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
347 f.debug_struct("LfuSegment")
348 .field("capacity", &self.config.capacity())
349 .field("len", &self.map.len())
350 .field("min_frequency", &self.min_frequency)
351 .finish()
352 }
353}
354
355#[derive(Debug)]
385pub struct LfuCache<K, V, S = DefaultHashBuilder> {
386 segment: LfuSegment<K, V, S>,
387}
388
389impl<K: Hash + Eq, V: Clone, S: BuildHasher> LfuCache<K, V, S> {
390 pub fn with_hasher(cap: NonZeroUsize, hash_builder: S) -> Self {
405 Self {
406 segment: LfuSegment::with_hasher(cap, hash_builder),
407 }
408 }
409
410 #[inline]
412 pub fn cap(&self) -> NonZeroUsize {
413 self.segment.cap()
414 }
415
416 #[inline]
418 pub fn len(&self) -> usize {
419 self.segment.len()
420 }
421
422 #[inline]
424 pub fn is_empty(&self) -> bool {
425 self.segment.is_empty()
426 }
427
428 #[inline]
436 pub fn get<Q>(&mut self, key: &Q) -> Option<&V>
437 where
438 K: Borrow<Q> + Clone,
439 Q: ?Sized + Hash + Eq,
440 {
441 self.segment.get(key)
442 }
443
444 #[inline]
452 pub fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
453 where
454 K: Borrow<Q> + Clone,
455 Q: ?Sized + Hash + Eq,
456 {
457 self.segment.get_mut(key)
458 }
459
460 #[inline]
469 pub fn put(&mut self, key: K, value: V) -> Option<(K, V)>
470 where
471 K: Clone,
472 {
473 self.segment.put(key, value)
474 }
475
476 #[inline]
482 pub fn remove<Q>(&mut self, key: &Q) -> Option<V>
483 where
484 K: Borrow<Q>,
485 Q: ?Sized + Hash + Eq,
486 V: Clone,
487 {
488 self.segment.remove(key)
489 }
490
491 #[inline]
493 pub fn clear(&mut self) {
494 self.segment.clear()
495 }
496
497 #[inline]
499 pub fn record_miss(&mut self, object_size: u64) {
500 self.segment.record_miss(object_size);
501 }
502}
503
504impl<K: Hash + Eq, V> LfuCache<K, V>
505where
506 V: Clone,
507{
508 pub fn new(cap: NonZeroUsize) -> LfuCache<K, V, DefaultHashBuilder> {
519 let config = LfuCacheConfig::new(cap);
520 LfuCache::with_hasher(config.capacity(), DefaultHashBuilder::default())
521 }
522}
523
524impl<K: Hash + Eq, V: Clone, S: BuildHasher> CacheMetrics for LfuCache<K, V, S> {
525 fn metrics(&self) -> BTreeMap<String, f64> {
526 self.segment.metrics().metrics()
527 }
528
529 fn algorithm_name(&self) -> &'static str {
530 self.segment.metrics().algorithm_name()
531 }
532}
533
534#[cfg(test)]
535mod tests {
536 extern crate std;
537 use std::string::ToString;
538
539 use super::*;
540 use alloc::string::String;
541
542 #[test]
543 fn test_lfu_basic() {
544 let mut cache = LfuCache::new(NonZeroUsize::new(3).unwrap());
545
546 assert_eq!(cache.put("a", 1), None);
548 assert_eq!(cache.put("b", 2), None);
549 assert_eq!(cache.put("c", 3), None);
550
551 assert_eq!(cache.get(&"a"), Some(&1));
553 assert_eq!(cache.get(&"a"), Some(&1));
554
555 assert_eq!(cache.get(&"b"), Some(&2));
557
558 let evicted = cache.put("d", 4);
560 assert!(evicted.is_some());
561 let (evicted_key, evicted_val) = evicted.unwrap();
562 assert_eq!(evicted_key, "c");
563 assert_eq!(evicted_val, 3);
564
565 assert_eq!(cache.get(&"a"), Some(&1)); assert_eq!(cache.get(&"b"), Some(&2)); assert_eq!(cache.get(&"d"), Some(&4)); assert_eq!(cache.get(&"c"), None); }
571
572 #[test]
573 fn test_lfu_frequency_ordering() {
574 let mut cache = LfuCache::new(NonZeroUsize::new(2).unwrap());
575
576 cache.put("a", 1);
578 cache.put("b", 2);
579
580 cache.get(&"a");
582 cache.get(&"a");
583 cache.get(&"a");
584
585 cache.get(&"b");
587
588 let evicted = cache.put("c", 3);
590 assert_eq!(evicted.unwrap().0, "b");
591
592 assert_eq!(cache.get(&"a"), Some(&1));
593 assert_eq!(cache.get(&"c"), Some(&3));
594 assert_eq!(cache.get(&"b"), None);
595 }
596
597 #[test]
598 fn test_lfu_update_existing() {
599 let mut cache = LfuCache::new(NonZeroUsize::new(2).unwrap());
600
601 cache.put("a", 1);
602 cache.get(&"a"); let old_value = cache.put("a", 10);
606 assert_eq!(old_value.unwrap().1, 1);
607
608 cache.put("b", 2);
610 cache.put("c", 3); assert_eq!(cache.get(&"a"), Some(&10));
613 assert_eq!(cache.get(&"c"), Some(&3));
614 assert_eq!(cache.get(&"b"), None);
615 }
616
617 #[test]
618 fn test_lfu_remove() {
619 let mut cache = LfuCache::new(NonZeroUsize::new(3).unwrap());
620
621 cache.put("a", 1);
622 cache.put("b", 2);
623 cache.put("c", 3);
624
625 assert_eq!(cache.remove(&"b"), Some(2));
627 assert_eq!(cache.remove(&"b"), None);
628
629 assert_eq!(cache.get(&"a"), Some(&1));
631 assert_eq!(cache.get(&"c"), Some(&3));
632 assert_eq!(cache.len(), 2);
633 }
634
635 #[test]
636 fn test_lfu_clear() {
637 let mut cache = LfuCache::new(NonZeroUsize::new(3).unwrap());
638
639 cache.put("a", 1);
640 cache.put("b", 2);
641 cache.put("c", 3);
642
643 assert_eq!(cache.len(), 3);
644 cache.clear();
645 assert_eq!(cache.len(), 0);
646 assert!(cache.is_empty());
647
648 cache.put("d", 4);
650 assert_eq!(cache.get(&"d"), Some(&4));
651 }
652
653 #[test]
654 fn test_lfu_get_mut() {
655 let mut cache = LfuCache::new(NonZeroUsize::new(2).unwrap());
656
657 cache.put("a", 1);
658
659 if let Some(value) = cache.get_mut(&"a") {
661 *value = 10;
662 }
663
664 assert_eq!(cache.get(&"a"), Some(&10));
665 }
666
667 #[test]
668 fn test_lfu_complex_values() {
669 let mut cache = LfuCache::new(NonZeroUsize::new(2).unwrap());
670
671 #[derive(Debug, Clone, PartialEq)]
672 struct ComplexValue {
673 id: usize,
674 data: String,
675 }
676
677 cache.put(
679 "a",
680 ComplexValue {
681 id: 1,
682 data: "a-data".to_string(),
683 },
684 );
685
686 cache.put(
687 "b",
688 ComplexValue {
689 id: 2,
690 data: "b-data".to_string(),
691 },
692 );
693
694 if let Some(value) = cache.get_mut(&"a") {
696 value.id = 100;
697 value.data = "a-modified".to_string();
698 }
699
700 let a = cache.get(&"a").unwrap();
702 assert_eq!(a.id, 100);
703 assert_eq!(a.data, "a-modified");
704 }
705
706 #[test]
708 fn test_miri_stacked_borrows_fix() {
709 let mut cache = LfuCache::new(NonZeroUsize::new(10).unwrap());
710
711 cache.put("a", 1);
713 cache.put("b", 2);
714 cache.put("c", 3);
715
716 for _ in 0..3 {
718 assert_eq!(cache.get(&"a"), Some(&1));
719 assert_eq!(cache.get(&"b"), Some(&2));
720 assert_eq!(cache.get(&"c"), Some(&3));
721 }
722
723 assert_eq!(cache.len(), 3);
724
725 if let Some(val) = cache.get_mut(&"a") {
727 *val += 10;
728 }
729 assert_eq!(cache.get(&"a"), Some(&11));
730 }
731
732 #[test]
734 fn test_lfu_segment_directly() {
735 let mut segment: LfuSegment<&str, i32, DefaultHashBuilder> =
736 LfuSegment::with_hasher(NonZeroUsize::new(3).unwrap(), DefaultHashBuilder::default());
737
738 assert_eq!(segment.len(), 0);
739 assert!(segment.is_empty());
740 assert_eq!(segment.cap().get(), 3);
741
742 segment.put("a", 1);
743 segment.put("b", 2);
744 assert_eq!(segment.len(), 2);
745
746 assert_eq!(segment.get(&"a"), Some(&1));
748 assert_eq!(segment.get(&"a"), Some(&1));
749 assert_eq!(segment.get(&"b"), Some(&2));
750 }
751
752 #[test]
753 fn test_lfu_concurrent_access() {
754 extern crate std;
755 use std::sync::{Arc, Mutex};
756 use std::thread;
757 use std::vec::Vec;
758
759 let cache = Arc::new(Mutex::new(LfuCache::new(NonZeroUsize::new(100).unwrap())));
760 let num_threads = 4;
761 let ops_per_thread = 100;
762
763 let mut handles: Vec<std::thread::JoinHandle<()>> = Vec::new();
764
765 for t in 0..num_threads {
766 let cache = Arc::clone(&cache);
767 handles.push(thread::spawn(move || {
768 for i in 0..ops_per_thread {
769 let key = std::format!("key_{}_{}", t, i);
770 let mut guard = cache.lock().unwrap();
771 guard.put(key.clone(), i);
772 if i % 3 == 0 {
774 let _ = guard.get(&key);
775 let _ = guard.get(&key);
776 }
777 }
778 }));
779 }
780
781 for handle in handles {
782 handle.join().unwrap();
783 }
784
785 let mut guard = cache.lock().unwrap();
786 assert!(guard.len() <= 100);
787 guard.clear(); }
789}