1extern crate alloc;
54
55use crate::config::LfudaCacheConfig;
56use crate::list::{Entry, List};
57use crate::metrics::{CacheMetrics, LfudaCacheMetrics};
58use alloc::boxed::Box;
59use alloc::collections::BTreeMap;
60use alloc::string::String;
61use core::borrow::Borrow;
62use core::hash::{BuildHasher, Hash};
63use core::mem;
64use core::num::NonZeroUsize;
65
66#[cfg(feature = "hashbrown")]
67use hashbrown::hash_map::DefaultHashBuilder;
68#[cfg(feature = "hashbrown")]
69use hashbrown::HashMap;
70
71#[cfg(not(feature = "hashbrown"))]
72use std::collections::hash_map::RandomState as DefaultHashBuilder;
73#[cfg(not(feature = "hashbrown"))]
74use std::collections::HashMap;
75
76#[derive(Debug, Clone, Copy)]
78struct EntryMetadata<K, V> {
79 frequency: usize,
81 age_at_insertion: usize,
83 node: *mut Entry<(K, V)>,
85}
86
87pub(crate) struct LfudaSegment<K, V, S = DefaultHashBuilder> {
101 config: LfudaCacheConfig,
103
104 global_age: usize,
106
107 min_priority: usize,
109
110 map: HashMap<K, EntryMetadata<K, V>, S>,
112
113 priority_lists: BTreeMap<usize, List<(K, V)>>,
116
117 metrics: LfudaCacheMetrics,
119}
120
121unsafe impl<K: Send, V: Send, S: Send> Send for LfudaSegment<K, V, S> {}
124
125unsafe impl<K: Send, V: Send, S: Sync> Sync for LfudaSegment<K, V, S> {}
127
128impl<K: Hash + Eq, V: Clone, S: BuildHasher> LfudaSegment<K, V, S> {
129 pub(crate) fn with_hasher(cap: NonZeroUsize, hash_builder: S) -> Self {
131 let config = LfudaCacheConfig::new(cap);
132 let map_capacity = config.capacity().get().next_power_of_two();
133 let max_cache_size_bytes = config.capacity().get() as u64 * 128;
134
135 LfudaSegment {
136 config,
137 global_age: config.initial_age(),
138 min_priority: 0,
139 map: HashMap::with_capacity_and_hasher(map_capacity, hash_builder),
140 priority_lists: BTreeMap::new(),
141 metrics: LfudaCacheMetrics::new(max_cache_size_bytes),
142 }
143 }
144
145 #[inline]
147 pub(crate) fn cap(&self) -> NonZeroUsize {
148 self.config.capacity()
149 }
150
151 #[inline]
153 pub(crate) fn len(&self) -> usize {
154 self.map.len()
155 }
156
157 #[inline]
159 pub(crate) fn is_empty(&self) -> bool {
160 self.map.is_empty()
161 }
162
163 #[inline]
165 pub(crate) fn global_age(&self) -> usize {
166 self.global_age
167 }
168
169 #[inline]
171 pub(crate) fn metrics(&self) -> &LfudaCacheMetrics {
172 &self.metrics
173 }
174
175 fn estimate_object_size(&self, _key: &K, _value: &V) -> u64 {
177 mem::size_of::<K>() as u64 + mem::size_of::<V>() as u64 + 64
178 }
179
180 #[inline]
182 pub(crate) fn record_miss(&mut self, object_size: u64) {
183 self.metrics.core.record_miss(object_size);
184 }
185
186 unsafe fn update_priority_by_node(&mut self, node: *mut Entry<(K, V)>) -> *mut Entry<(K, V)>
193 where
194 K: Clone + Hash + Eq,
195 {
196 let (key_ref, _) = (*node).get_value();
198 let key_cloned = key_ref.clone();
199
200 let metadata = self.map.get_mut(&key_cloned).unwrap();
201 let old_priority = metadata.frequency + metadata.age_at_insertion;
202
203 metadata.frequency += 1;
205 let new_priority = metadata.frequency + metadata.age_at_insertion;
206
207 if old_priority == new_priority {
209 let node = metadata.node;
210 self.priority_lists
211 .get_mut(&old_priority)
212 .unwrap()
213 .move_to_front(node);
214 return node;
215 }
216
217 let node = metadata.node;
219 let boxed_entry = self
220 .priority_lists
221 .get_mut(&old_priority)
222 .unwrap()
223 .remove(node)
224 .unwrap();
225
226 if self.priority_lists.get(&old_priority).unwrap().is_empty()
229 && old_priority == self.min_priority
230 {
231 self.min_priority = new_priority;
232 }
233
234 let entry_ptr = Box::into_raw(boxed_entry);
236
237 let capacity = self.config.capacity();
239 self.priority_lists
240 .entry(new_priority)
241 .or_insert_with(|| List::new(capacity));
242
243 self.priority_lists
245 .get_mut(&new_priority)
246 .unwrap()
247 .attach_from_other_list(entry_ptr);
248
249 metadata.node = entry_ptr;
251
252 entry_ptr
253 }
254
255 pub(crate) fn get<Q>(&mut self, key: &Q) -> Option<&V>
257 where
258 K: Borrow<Q> + Clone,
259 Q: ?Sized + Hash + Eq,
260 {
261 if let Some(metadata) = self.map.get(key) {
262 let node = metadata.node;
263 unsafe {
264 let (key_ref, value) = (*node).get_value();
266 let object_size = self.estimate_object_size(key_ref, value);
267 self.metrics.core.record_hit(object_size);
268
269 let new_node = self.update_priority_by_node(node);
270 let (_, value) = (*new_node).get_value();
271 Some(value)
272 }
273 } else {
274 None
275 }
276 }
277
278 pub(crate) fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
280 where
281 K: Borrow<Q> + Clone,
282 Q: ?Sized + Hash + Eq,
283 {
284 if let Some(metadata) = self.map.get(key) {
285 let node = metadata.node;
286 unsafe {
287 let (key_ref, value) = (*node).get_value();
289 let object_size = self.estimate_object_size(key_ref, value);
290 self.metrics.core.record_hit(object_size);
291
292 let old_frequency = metadata.frequency;
293 let new_priority = (old_frequency + 1) + metadata.age_at_insertion;
294 self.metrics.record_frequency_increment(new_priority as u64);
295
296 let new_node = self.update_priority_by_node(node);
297 let (_, value) = (*new_node).get_value_mut();
298 Some(value)
299 }
300 } else {
301 None
302 }
303 }
304
305 pub(crate) fn put(&mut self, key: K, value: V) -> Option<(K, V)>
307 where
308 K: Clone,
309 {
310 let object_size = self.estimate_object_size(&key, &value);
311
312 if let Some(metadata) = self.map.get(&key) {
314 let node = metadata.node;
315 let priority = metadata.frequency + metadata.age_at_insertion;
316
317 unsafe {
318 let old_entry = self.priority_lists.get_mut(&priority).unwrap().update(
320 node,
321 (key.clone(), value),
322 true,
323 );
324
325 self.metrics.core.record_insertion(object_size);
326 return old_entry.0;
327 }
328 }
329
330 let mut evicted = None;
331
332 let frequency = 1;
334 let age_at_insertion = self.global_age;
335 let priority = frequency + age_at_insertion;
336
337 if self.len() >= self.config.capacity().get() {
339 if let Some(min_priority_list) = self.priority_lists.get_mut(&self.min_priority) {
340 if let Some(old_entry) = min_priority_list.remove_last() {
341 unsafe {
342 let entry_ptr = Box::into_raw(old_entry);
343 let (old_key, old_value) = (*entry_ptr).get_value().clone();
344
345 let evicted_size = self.estimate_object_size(&old_key, &old_value);
346 self.metrics.core.record_eviction(evicted_size);
347
348 if let Some(evicted_metadata) = self.map.get(&old_key) {
350 self.global_age =
351 evicted_metadata.frequency + evicted_metadata.age_at_insertion;
352 self.metrics.record_aging_event(self.global_age as u64);
353 }
354
355 self.map.remove(&old_key);
356 evicted = Some((old_key, old_value));
357 let _ = Box::from_raw(entry_ptr);
358
359 if self
361 .priority_lists
362 .get(&self.min_priority)
363 .unwrap()
364 .is_empty()
365 {
366 self.min_priority = self
367 .priority_lists
368 .keys()
369 .find(|&&p| {
370 p > self.min_priority
371 && !self.priority_lists.get(&p).unwrap().is_empty()
372 })
373 .copied()
374 .unwrap_or(priority);
375 }
376 }
377 }
378 }
379 }
380
381 self.min_priority = if self.is_empty() {
382 priority
383 } else {
384 core::cmp::min(self.min_priority, priority)
385 };
386
387 let capacity = self.config.capacity();
389 self.priority_lists
390 .entry(priority)
391 .or_insert_with(|| List::new(capacity));
392
393 if let Some(node) = self
394 .priority_lists
395 .get_mut(&priority)
396 .unwrap()
397 .add((key.clone(), value))
398 {
399 let metadata = EntryMetadata {
400 frequency,
401 age_at_insertion,
402 node,
403 };
404
405 self.map.insert(key, metadata);
406
407 self.metrics.core.record_insertion(object_size);
408 self.metrics.record_frequency_increment(priority as u64);
409 if age_at_insertion > 0 {
410 self.metrics.record_aging_benefit(age_at_insertion as u64);
411 }
412 }
413
414 evicted
415 }
416
417 pub(crate) fn remove<Q>(&mut self, key: &Q) -> Option<V>
419 where
420 K: Borrow<Q>,
421 Q: ?Sized + Hash + Eq,
422 V: Clone,
423 {
424 let metadata = self.map.remove(key)?;
425 let priority = metadata.frequency + metadata.age_at_insertion;
426
427 unsafe {
428 let boxed_entry = self
430 .priority_lists
431 .get_mut(&priority)?
432 .remove(metadata.node)?;
433 let entry_ptr = Box::into_raw(boxed_entry);
434 let value = (*entry_ptr).get_value().1.clone();
435 let _ = Box::from_raw(entry_ptr);
436
437 if self.priority_lists.get(&priority).unwrap().is_empty()
439 && priority == self.min_priority
440 {
441 self.min_priority = self
442 .priority_lists
443 .keys()
444 .find(|&&p| p > priority && !self.priority_lists.get(&p).unwrap().is_empty())
445 .copied()
446 .unwrap_or(self.global_age);
447 }
448
449 Some(value)
450 }
451 }
452
453 pub(crate) fn clear(&mut self) {
455 self.map.clear();
456 self.priority_lists.clear();
457 self.global_age = 0;
458 self.min_priority = 0;
459 }
460}
461
462impl<K, V, S> core::fmt::Debug for LfudaSegment<K, V, S> {
464 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
465 f.debug_struct("LfudaSegment")
466 .field("capacity", &self.config.capacity())
467 .field("len", &self.map.len())
468 .field("global_age", &self.global_age)
469 .field("min_priority", &self.min_priority)
470 .finish()
471 }
472}
473
474#[derive(Debug)]
507pub struct LfudaCache<K, V, S = DefaultHashBuilder> {
508 segment: LfudaSegment<K, V, S>,
509}
510
511impl<K: Hash + Eq, V: Clone, S: BuildHasher> LfudaCache<K, V, S> {
512 pub fn with_hasher(cap: NonZeroUsize, hash_builder: S) -> Self {
527 Self {
528 segment: LfudaSegment::with_hasher(cap, hash_builder),
529 }
530 }
531
532 #[inline]
534 pub fn cap(&self) -> NonZeroUsize {
535 self.segment.cap()
536 }
537
538 #[inline]
540 pub fn len(&self) -> usize {
541 self.segment.len()
542 }
543
544 #[inline]
546 pub fn is_empty(&self) -> bool {
547 self.segment.is_empty()
548 }
549
550 #[inline]
552 pub fn global_age(&self) -> usize {
553 self.segment.global_age()
554 }
555
556 #[inline]
558 pub fn record_miss(&mut self, object_size: u64) {
559 self.segment.record_miss(object_size);
560 }
561
562 #[inline]
570 pub fn get<Q>(&mut self, key: &Q) -> Option<&V>
571 where
572 K: Borrow<Q> + Clone,
573 Q: ?Sized + Hash + Eq,
574 {
575 self.segment.get(key)
576 }
577
578 #[inline]
586 pub fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
587 where
588 K: Borrow<Q> + Clone,
589 Q: ?Sized + Hash + Eq,
590 {
591 self.segment.get_mut(key)
592 }
593
594 #[inline]
603 pub fn put(&mut self, key: K, value: V) -> Option<(K, V)>
604 where
605 K: Clone,
606 {
607 self.segment.put(key, value)
608 }
609
610 #[inline]
616 pub fn remove<Q>(&mut self, key: &Q) -> Option<V>
617 where
618 K: Borrow<Q>,
619 Q: ?Sized + Hash + Eq,
620 V: Clone,
621 {
622 self.segment.remove(key)
623 }
624
625 #[inline]
628 pub fn clear(&mut self) {
629 self.segment.clear()
630 }
631}
632
633impl<K: Hash + Eq, V: Clone, S: BuildHasher> CacheMetrics for LfudaCache<K, V, S> {
634 fn metrics(&self) -> BTreeMap<String, f64> {
635 self.segment.metrics().metrics()
636 }
637
638 fn algorithm_name(&self) -> &'static str {
639 self.segment.metrics().algorithm_name()
640 }
641}
642
643impl<K: Hash + Eq, V> LfudaCache<K, V>
644where
645 V: Clone,
646{
647 pub fn new(cap: NonZeroUsize) -> LfudaCache<K, V, DefaultHashBuilder> {
658 let config = LfudaCacheConfig::new(cap);
659 LfudaCache::with_hasher(config.capacity(), DefaultHashBuilder::default())
660 }
661}
662
663#[cfg(test)]
664mod tests {
665 extern crate std;
666 use alloc::string::ToString;
667
668 use super::*;
669 use alloc::string::String;
670
671 #[test]
672 fn test_lfuda_basic() {
673 let mut cache = LfudaCache::new(NonZeroUsize::new(3).unwrap());
674
675 assert_eq!(cache.put("a", 1), None);
677 assert_eq!(cache.put("b", 2), None);
678 assert_eq!(cache.put("c", 3), None);
679
680 assert_eq!(cache.get(&"a"), Some(&1));
682 assert_eq!(cache.get(&"a"), Some(&1));
683
684 assert_eq!(cache.get(&"b"), Some(&2));
686
687 let evicted = cache.put("d", 4);
689 assert!(evicted.is_some());
690 let (evicted_key, evicted_val) = evicted.unwrap();
691 assert_eq!(evicted_key, "c");
692 assert_eq!(evicted_val, 3);
693
694 assert_eq!(cache.get(&"a"), Some(&1));
696 assert_eq!(cache.get(&"b"), Some(&2));
697 assert_eq!(cache.get(&"d"), Some(&4));
698 assert_eq!(cache.get(&"c"), None); }
700
701 #[test]
702 fn test_lfuda_aging() {
703 let mut cache = LfudaCache::new(NonZeroUsize::new(2).unwrap());
704
705 cache.put("a", 1);
707 cache.put("b", 2);
708
709 for _ in 0..10 {
711 cache.get(&"a");
712 }
713
714 assert_eq!(cache.global_age(), 0);
716
717 let evicted = cache.put("c", 3);
719 assert!(evicted.is_some());
720
721 assert!(cache.global_age() > 0);
723
724 cache.put("d", 4);
726
727 assert!(cache.len() <= cache.cap().get());
729 }
730
731 #[test]
732 fn test_lfuda_priority_calculation() {
733 let mut cache = LfudaCache::new(NonZeroUsize::new(3).unwrap());
734
735 cache.put("a", 1);
736 assert_eq!(cache.global_age(), 0);
737
738 cache.get(&"a");
740
741 cache.put("b", 2);
743 cache.put("c", 3);
744
745 let evicted = cache.put("d", 4);
747 assert!(evicted.is_some());
748
749 assert!(cache.global_age() > 0);
751 }
752
753 #[test]
754 fn test_lfuda_update_existing() {
755 let mut cache = LfudaCache::new(NonZeroUsize::new(2).unwrap());
756
757 cache.put("a", 1);
758 cache.get(&"a"); let old_value = cache.put("a", 10);
762 assert_eq!(old_value.unwrap().1, 1);
763
764 cache.put("b", 2);
766 cache.put("c", 3); assert_eq!(cache.get(&"a"), Some(&10));
769 assert_eq!(cache.get(&"c"), Some(&3));
770 assert_eq!(cache.get(&"b"), None);
771 }
772
773 #[test]
774 fn test_lfuda_remove() {
775 let mut cache = LfudaCache::new(NonZeroUsize::new(3).unwrap());
776
777 cache.put("a", 1);
778 cache.put("b", 2);
779 cache.put("c", 3);
780
781 assert_eq!(cache.remove(&"b"), Some(2));
783 assert_eq!(cache.remove(&"b"), None);
784
785 assert_eq!(cache.get(&"a"), Some(&1));
787 assert_eq!(cache.get(&"c"), Some(&3));
788 assert_eq!(cache.len(), 2);
789 }
790
791 #[test]
792 fn test_lfuda_clear() {
793 let mut cache = LfudaCache::new(NonZeroUsize::new(3).unwrap());
794
795 cache.put("a", 1);
796 cache.put("b", 2);
797 cache.put("c", 3);
798
799 cache.get(&"a");
801 cache.put("d", 4); let age_before_clear = cache.global_age();
804 assert!(age_before_clear > 0);
805
806 cache.clear();
807 assert_eq!(cache.len(), 0);
808 assert!(cache.is_empty());
809 assert_eq!(cache.global_age(), 0); cache.put("e", 5);
813 assert_eq!(cache.get(&"e"), Some(&5));
814 }
815
816 #[test]
817 fn test_lfuda_get_mut() {
818 let mut cache = LfudaCache::new(NonZeroUsize::new(2).unwrap());
819
820 cache.put("a", 1);
821
822 if let Some(value) = cache.get_mut(&"a") {
824 *value = 10;
825 }
826
827 assert_eq!(cache.get(&"a"), Some(&10));
828 }
829
830 #[test]
831 fn test_lfuda_complex_values() {
832 let mut cache = LfudaCache::new(NonZeroUsize::new(2).unwrap());
833
834 #[derive(Debug, Clone, PartialEq)]
835 struct ComplexValue {
836 id: usize,
837 data: String,
838 }
839
840 cache.put(
842 "a",
843 ComplexValue {
844 id: 1,
845 data: "a-data".to_string(),
846 },
847 );
848
849 cache.put(
850 "b",
851 ComplexValue {
852 id: 2,
853 data: "b-data".to_string(),
854 },
855 );
856
857 if let Some(value) = cache.get_mut(&"a") {
859 value.id = 100;
860 value.data = "a-modified".to_string();
861 }
862
863 let a = cache.get(&"a").unwrap();
865 assert_eq!(a.id, 100);
866 assert_eq!(a.data, "a-modified");
867 }
868
869 #[test]
870 fn test_lfuda_aging_advantage() {
871 let mut cache = LfudaCache::new(NonZeroUsize::new(2).unwrap());
872
873 cache.put("old", 1);
875 for _ in 0..100 {
876 cache.get(&"old");
877 }
878
879 cache.put("temp", 2);
881
882 let _evicted = cache.put("new1", 3);
884 let age_after_first_eviction = cache.global_age();
885
886 let _evicted = cache.put("new2", 4);
888 let age_after_second_eviction = cache.global_age();
889
890 assert!(age_after_second_eviction >= age_after_first_eviction);
892
893 cache.put("newer", 5);
895
896 assert!(cache.len() <= cache.cap().get());
898 }
899
900 #[test]
902 fn test_miri_stacked_borrows_fix() {
903 let mut cache = LfudaCache::new(NonZeroUsize::new(10).unwrap());
904
905 cache.put("a", 1);
907 cache.put("b", 2);
908 cache.put("c", 3);
909
910 for _ in 0..3 {
912 assert_eq!(cache.get(&"a"), Some(&1));
913 assert_eq!(cache.get(&"b"), Some(&2));
914 assert_eq!(cache.get(&"c"), Some(&3));
915 }
916
917 assert_eq!(cache.len(), 3);
918
919 if let Some(val) = cache.get_mut(&"a") {
921 *val += 10;
922 }
923 assert_eq!(cache.get(&"a"), Some(&11));
924 }
925
926 #[test]
928 fn test_lfuda_segment_directly() {
929 let mut segment: LfudaSegment<&str, i32, DefaultHashBuilder> =
930 LfudaSegment::with_hasher(NonZeroUsize::new(3).unwrap(), DefaultHashBuilder::default());
931
932 assert_eq!(segment.len(), 0);
933 assert!(segment.is_empty());
934 assert_eq!(segment.cap().get(), 3);
935 assert_eq!(segment.global_age(), 0);
936
937 segment.put("a", 1);
938 segment.put("b", 2);
939 assert_eq!(segment.len(), 2);
940
941 assert_eq!(segment.get(&"a"), Some(&1));
943 assert_eq!(segment.get(&"b"), Some(&2));
944 }
945
946 #[test]
947 fn test_lfuda_concurrent_access() {
948 extern crate std;
949 use std::sync::{Arc, Mutex};
950 use std::thread;
951 use std::vec::Vec;
952
953 let cache = Arc::new(Mutex::new(LfudaCache::new(NonZeroUsize::new(100).unwrap())));
954 let num_threads = 4;
955 let ops_per_thread = 100;
956
957 let mut handles: Vec<std::thread::JoinHandle<()>> = Vec::new();
958
959 for t in 0..num_threads {
960 let cache = Arc::clone(&cache);
961 handles.push(thread::spawn(move || {
962 for i in 0..ops_per_thread {
963 let key = std::format!("key_{}_{}", t, i);
964 let mut guard = cache.lock().unwrap();
965 guard.put(key.clone(), i);
966 let _ = guard.get(&key);
967 }
968 }));
969 }
970
971 for handle in handles {
972 handle.join().unwrap();
973 }
974
975 let mut guard = cache.lock().unwrap();
976 assert!(guard.len() <= 100);
977 guard.clear(); }
979}