use std::borrow::Borrow;
use std::collections::hash_map;
use std::collections::HashMap;
use std::fmt::{Debug, Formatter};
use std::hash::Hash;
use std::hint::unreachable_unchecked;
use std::iter::{FromIterator, FusedIterator};
use std::num::NonZeroUsize;
use std::ptr::NonNull;
use std::rc::Rc;
use crate::{Entry, LfuCacheIter};
pub(self) use freq_list::FrequencyList;
pub(self) use lfu_entry::LfuEntry;
pub(self) use node::Node;
use self::entry::{OccupiedEntry, VacantEntry};
use self::node::WithFrequency;
pub mod entry;
mod freq_list;
mod lfu_entry;
mod node;
mod util;
#[derive(Eq, PartialEq)]
#[allow(clippy::module_name_repetitions)]
pub struct LfuCache<Key: Hash + Eq, Value> {
lookup: LookupMap<Key, Value>,
freq_list: FrequencyList<Key, Value>,
capacity: Option<NonZeroUsize>,
len: usize,
}
#[derive(Eq, PartialEq)]
struct LookupMap<Key: Hash + Eq, Value>(HashMap<Rc<Key>, NonNull<LfuEntry<Key, Value>>>);
#[cfg(not(tarpaulin_include))]
impl<Key: Hash + Eq + Debug, Value> Debug for LookupMap<Key, Value> {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
let mut dbg = f.debug_struct("LookupMap");
for (key, value) in &self.0 {
dbg.field(&format!("{:?}", key), &unsafe {
value.as_ref().owner.as_ref().frequency
});
}
dbg.finish()
}
}
unsafe impl<Key: Hash + Eq, Value> Send for LfuCache<Key, Value> {}
unsafe impl<Key: Hash + Eq, Value> Sync for LfuCache<Key, Value> {}
impl<Key: Hash + Eq, Value> Drop for LookupMap<Key, Value> {
fn drop(&mut self) {
for (_, v) in self.0.drain() {
unsafe { Box::from_raw(v.as_ptr()) };
}
}
}
impl<Key: Hash + Eq, Value> LfuCache<Key, Value> {
#[inline]
#[must_use]
pub fn with_capacity(capacity: usize) -> Self {
Self {
lookup: LookupMap(HashMap::with_capacity(capacity)),
freq_list: FrequencyList::new(),
capacity: NonZeroUsize::new(capacity),
len: 0,
}
}
#[inline]
#[must_use]
pub fn unbounded() -> Self {
Self::with_capacity(0)
}
pub fn set_capacity(&mut self, new_capacity: usize) {
self.capacity = NonZeroUsize::new(new_capacity);
if let Some(capacity) = self.capacity {
while self.len > capacity.get() {
self.pop_lfu();
}
}
}
#[inline]
pub fn insert(&mut self, key: Key, value: Value) -> Option<Value> {
self.insert_rc(Rc::new(key), value)
}
pub(crate) fn insert_rc(&mut self, key: Rc<Key>, value: Value) -> Option<Value> {
let mut evicted = self.remove(&key);
if let Some(capacity) = self.capacity {
if capacity.get() <= self.len {
evicted = self.pop_lfu();
}
}
self.lookup.0.insert(Rc::clone(&key), NonNull::dangling());
let v = self.lookup.0.get_mut(&key).unwrap();
*v = self.freq_list.insert(key, value);
self.len += 1;
evicted
}
#[inline]
pub fn get(&mut self, key: &Key) -> Option<&Value> {
let entry = self.lookup.0.get_mut(key)?;
self.freq_list.update(*entry);
Some(&unsafe { entry.as_ref() }.value)
}
pub(crate) fn get_rc_key_value(&mut self, key: &Key) -> Option<(Rc<Key>, &Value)> {
let entry = self.lookup.0.get_mut(key)?;
self.freq_list.update(*entry);
let entry = unsafe { entry.as_ref() };
Some((Rc::clone(&entry.key), &entry.value))
}
#[inline]
pub fn get_mut(&mut self, key: &Key) -> Option<&mut Value> {
self.get_rc_key_value_mut(key).map(|(_, v)| v)
}
pub(crate) fn get_rc_key_value_mut(&mut self, key: &Key) -> Option<(Rc<Key>, &mut Value)> {
let entry = self.lookup.0.get_mut(key)?;
self.freq_list.update(*entry);
let entry = unsafe { entry.as_mut() };
Some((Rc::clone(&entry.key), &mut entry.value))
}
#[inline]
pub fn remove(&mut self, key: &Key) -> Option<Value> {
self.lookup.0.remove(key).map(|node| {
let mut freq_node = unsafe { node.as_ref() }.owner;
let detached = unsafe { freq_node.as_mut() }.remove(node);
if unsafe { freq_node.as_ref() }.elements.is_none() {
let freq_head = unsafe { Box::from_raw(freq_node.as_ptr()) };
if self.freq_list.head == Some(NonNull::from(&*freq_head)) {
self.freq_list.head = freq_head.next;
}
freq_head.detach();
}
self.len -= 1;
detached.value
})
}
#[inline]
pub fn entry(&mut self, key: Key) -> Entry<'_, Key, Value> {
let key = Rc::new(key);
match self.lookup.0.entry(Rc::clone(&key)) {
hash_map::Entry::Occupied(mut entry) => {
self.freq_list.update(*entry.get_mut());
Entry::Occupied(OccupiedEntry::new(entry, &mut self.len))
}
hash_map::Entry::Vacant(entry) => Entry::Vacant(VacantEntry::new(
entry,
key,
&mut self.freq_list,
self.capacity,
&mut self.len,
)),
}
}
#[inline]
pub fn pop_lfu(&mut self) -> Option<Value> {
self.pop_lfu_key_value_frequency().map(|(_, v, _)| v)
}
#[inline]
pub fn pop_lfu_key_value(&mut self) -> Option<(Key, Value)> {
self.pop_lfu_key_value_frequency().map(|(k, v, _)| (k, v))
}
#[inline]
pub fn pop_lfu_key_value_frequency(&mut self) -> Option<(Key, Value, usize)> {
self.freq_list
.pop_lfu()
.map(|WithFrequency(freq, detached)| {
let key = detached.key.as_ref();
self.lookup.0.remove(key);
self.len -= 1;
let key = match Rc::try_unwrap(detached.key) {
Ok(k) => k,
Err(_) => unsafe { unreachable_unchecked() },
};
(key, detached.value, freq)
})
}
pub fn clear(&mut self) -> LfuCacheIter<Key, Value> {
let mut to_return = Self::with_capacity(self.capacity.map_or(0, NonZeroUsize::get));
std::mem::swap(&mut to_return, self);
to_return.into_iter()
}
#[inline]
#[must_use]
pub fn peek_lfu(&self) -> Option<&Value> {
self.freq_list.peek_lfu()
}
#[inline]
#[must_use]
pub const fn capacity(&self) -> Option<NonZeroUsize> {
self.capacity
}
#[inline]
#[must_use]
pub const fn len(&self) -> usize {
self.len
}
#[inline]
#[must_use]
pub const fn is_empty(&self) -> bool {
self.len == 0
}
#[inline]
#[must_use]
pub const fn is_unbounded(&self) -> bool {
self.capacity.is_none()
}
#[inline]
#[must_use]
pub fn frequencies(&self) -> Vec<usize> {
self.freq_list.frequencies().collect()
}
#[inline]
pub fn shrink_to_fit(&mut self) {
self.set_capacity(self.len);
}
#[inline]
pub fn keys(&self) -> impl Iterator<Item = &Key> + FusedIterator + '_ {
self.lookup.0.keys().map(Borrow::borrow)
}
#[inline]
pub fn peek_values(&self) -> impl Iterator<Item = &Value> + FusedIterator + '_ {
self.lookup
.0
.values()
.map(|value| &unsafe { value.as_ref() }.value)
}
#[inline]
pub fn peek_iter(&self) -> impl Iterator<Item = (&Key, &Value)> + FusedIterator + '_ {
self.lookup
.0
.iter()
.map(|(key, value)| (key.borrow(), &unsafe { value.as_ref() }.value))
}
}
#[cfg(not(tarpaulin_include))]
impl<Key: Hash + Eq + Debug, Value> Debug for LfuCache<Key, Value> {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
let mut dbg = f.debug_struct("LfuCache");
dbg.field("len", &self.len);
dbg.field("capacity", &self.capacity);
dbg.field("freq_list", &self.freq_list);
dbg.field("lookup_map", &self.lookup);
dbg.finish()
}
}
impl<Key: Hash + Eq, Value> FromIterator<(Key, Value)> for LfuCache<Key, Value> {
fn from_iter<T: IntoIterator<Item = (Key, Value)>>(iter: T) -> Self {
let mut cache = Self::unbounded();
for (k, v) in iter {
cache.insert(k, v);
}
cache.shrink_to_fit();
cache
}
}
impl<Key: Hash + Eq, Value> Extend<(Key, Value)> for LfuCache<Key, Value> {
fn extend<T: IntoIterator<Item = (Key, Value)>>(&mut self, iter: T) {
for (k, v) in iter {
self.insert(k, v);
}
}
}
impl<Key: Hash + Eq, Value> IntoIterator for LfuCache<Key, Value> {
type Item = (Key, Value);
type IntoIter = LfuCacheIter<Key, Value>;
fn into_iter(self) -> Self::IntoIter {
LfuCacheIter(self)
}
}
#[cfg(test)]
mod get {
use super::LfuCache;
#[test]
fn empty() {
let mut cache = LfuCache::<u64, u64>::unbounded();
for i in 0..100 {
assert!(cache.get(&i).is_none())
}
}
#[test]
fn get_mut() {
let mut cache = LfuCache::unbounded();
cache.insert(1, 2);
assert_eq!(cache.frequencies(), vec![0]);
*cache.get_mut(&1).unwrap() = 3;
assert_eq!(cache.frequencies(), vec![1]);
assert_eq!(cache.get(&1), Some(&3));
}
#[test]
fn getting_is_ok_after_adding_other_value() {
let mut cache = LfuCache::unbounded();
cache.insert(1, 2);
assert_eq!(cache.get(&1), Some(&2));
cache.insert(3, 4);
assert_eq!(cache.get(&1), Some(&2));
}
#[test]
fn bounded_alternating_values() {
let mut cache = LfuCache::with_capacity(8);
cache.insert(1, 1);
cache.insert(2, 2);
for _ in 0..100 {
cache.get(&1);
cache.get(&2);
}
assert_eq!(cache.len(), 2);
assert_eq!(cache.frequencies(), vec![100]);
}
}
#[cfg(test)]
mod insert {
use super::LfuCache;
#[test]
fn insert_unbounded() {
let mut cache = LfuCache::unbounded();
for i in 0..100 {
cache.insert(i, i + 100);
}
for i in 0..100 {
assert_eq!(cache.get(&i), Some(&(i + 100)));
assert!(cache.get(&(i + 100)).is_none());
}
}
#[test]
fn reinsertion_of_same_key_resets_freq() {
let mut cache = LfuCache::unbounded();
cache.insert(1, 1);
cache.get(&1);
cache.insert(1, 1);
assert_eq!(cache.frequencies(), vec![0]);
}
#[test]
fn insert_bounded() {
let mut cache = LfuCache::with_capacity(20);
for i in 0..100 {
cache.insert(i, i + 100);
}
}
#[test]
fn insert_returns_evicted() {
let mut cache = LfuCache::with_capacity(1);
assert_eq!(cache.insert(1, 2), None);
for _ in 0..10 {
assert_eq!(cache.insert(3, 4), Some(2));
assert_eq!(cache.insert(1, 2), Some(4));
}
}
}
#[cfg(test)]
mod pop {
use super::LfuCache;
#[test]
fn pop() {
let mut cache = LfuCache::unbounded();
for i in 0..100 {
cache.insert(i, i + 100);
}
for i in 0..100 {
assert_eq!(cache.lookup.0.len(), 100 - i);
assert_eq!(cache.pop_lfu(), Some(200 - i - 1));
}
}
#[test]
fn pop_empty() {
let mut cache = LfuCache::<i32, i32>::unbounded();
assert_eq!(None, cache.pop_lfu());
assert_eq!(None, cache.pop_lfu_key_value());
}
#[test]
fn set_capacity_evicts_multiple() {
let mut cache = LfuCache::unbounded();
for i in 0..100 {
cache.insert(i, i + 100);
}
cache.set_capacity(10);
assert_eq!(cache.len(), 10);
}
#[test]
fn pop_multiple_varying_frequencies() {
let mut cache = LfuCache::unbounded();
for i in 0..10 {
cache.insert(i, i + 10);
}
for i in 0..10 {
for _ in 0..i * i {
cache.get(&i).unwrap();
}
}
for i in 0..10 {
assert_eq!(10 - i, cache.len());
assert_eq!(Some(i + 10), cache.pop_lfu());
}
}
}
#[cfg(test)]
mod remove {
use super::LfuCache;
#[test]
fn remove_to_empty() {
let mut cache = LfuCache::unbounded();
cache.insert(1, 2);
assert_eq!(cache.remove(&1), Some(2));
assert!(cache.is_empty());
assert_eq!(cache.freq_list.len(), 0);
}
#[test]
fn remove_empty() {
let mut cache = LfuCache::<usize, usize>::unbounded();
assert!(cache.remove(&1).is_none());
}
#[test]
fn remove_to_nonempty() {
let mut cache = LfuCache::unbounded();
cache.insert(1, 2);
cache.insert(3, 4);
assert_eq!(cache.remove(&1), Some(2));
assert!(!cache.is_empty());
assert_eq!(cache.remove(&3), Some(4));
assert!(cache.is_empty());
assert_eq!(cache.freq_list.len(), 0);
}
#[test]
fn remove_middle() {
let mut cache = LfuCache::unbounded();
cache.insert(1, 2);
cache.insert(3, 4);
cache.insert(5, 6);
cache.insert(7, 8);
cache.insert(9, 10);
cache.insert(11, 12);
cache.get(&7);
cache.get(&9);
cache.get(&11);
assert_eq!(cache.frequencies(), vec![0, 1]);
assert_eq!(cache.len(), 6);
cache.remove(&9);
assert!(cache.get(&7).is_some());
assert!(cache.get(&11).is_some());
cache.remove(&3);
assert!(cache.get(&1).is_some());
assert!(cache.get(&5).is_some());
}
#[test]
fn remove_end() {
let mut cache = LfuCache::unbounded();
cache.insert(1, 2);
cache.insert(3, 4);
cache.insert(5, 6);
cache.insert(7, 8);
cache.insert(9, 10);
cache.insert(11, 12);
cache.get(&7);
cache.get(&9);
cache.get(&11);
assert_eq!(cache.frequencies(), vec![0, 1]);
assert_eq!(cache.len(), 6);
cache.remove(&7);
assert!(cache.get(&9).is_some());
assert!(cache.get(&11).is_some());
cache.remove(&1);
assert!(cache.get(&3).is_some());
assert!(cache.get(&5).is_some());
}
#[test]
fn remove_start() {
let mut cache = LfuCache::unbounded();
cache.insert(1, 2);
cache.insert(3, 4);
cache.insert(5, 6);
cache.insert(7, 8);
cache.insert(9, 10);
cache.insert(11, 12);
cache.get(&7);
cache.get(&9);
cache.get(&11);
assert_eq!(cache.frequencies(), vec![0, 1]);
assert_eq!(cache.len(), 6);
cache.remove(&11);
assert!(cache.get(&9).is_some());
assert!(cache.get(&7).is_some());
cache.remove(&5);
assert!(cache.get(&3).is_some());
assert!(cache.get(&1).is_some());
}
#[test]
fn remove_connects_next_owner() {
let mut cache = LfuCache::unbounded();
cache.insert(1, 1);
cache.insert(2, 2);
assert_eq!(cache.get(&1), Some(&1));
assert_eq!(cache.remove(&2), Some(2));
assert_eq!(cache.get(&1), Some(&1));
}
}
#[cfg(test)]
mod bookkeeping {
use std::num::NonZeroUsize;
use super::LfuCache;
#[test]
fn getting_one_element_has_constant_freq_list_size() {
let mut cache = LfuCache::unbounded();
cache.insert(1, 2);
assert_eq!(cache.freq_list.len(), 1);
for _ in 0..100 {
cache.get(&1);
assert_eq!(cache.freq_list.len(), 1);
}
}
#[test]
fn freq_list_node_merges() {
let mut cache = LfuCache::unbounded();
cache.insert(1, 2);
cache.insert(3, 4);
assert_eq!(cache.freq_list.len(), 1);
assert!(cache.get(&1).is_some());
assert_eq!(cache.freq_list.len(), 2);
assert!(cache.get(&3).is_some());
assert_eq!(cache.freq_list.len(), 1);
}
#[test]
fn freq_list_multi_items() {
let mut cache = LfuCache::unbounded();
cache.insert(1, 2);
cache.get(&1);
cache.get(&1);
cache.insert(3, 4);
assert_eq!(cache.freq_list.len(), 2);
cache.get(&3);
assert_eq!(cache.freq_list.len(), 2);
cache.get(&3);
assert_eq!(cache.freq_list.len(), 1);
}
#[test]
fn unbounded_is_unbounded() {
assert!(LfuCache::<i32, i32>::unbounded().is_unbounded());
assert!(!LfuCache::<i32, i32>::with_capacity(3).is_unbounded());
}
#[test]
fn capacity_reports_internal_capacity() {
assert_eq!(LfuCache::<i32, i32>::unbounded().capacity(), None);
assert_eq!(
LfuCache::<i32, i32>::with_capacity(3).capacity(),
Some(NonZeroUsize::new(3).unwrap())
);
}
#[test]
fn clear_is_ok() {
let mut cache = LfuCache::unbounded();
for i in 0..10 {
cache.insert(i, i);
}
assert!(!cache.is_empty());
cache.clear();
assert!(cache.is_empty());
for i in 0..10 {
assert!(cache.get(&i).is_none());
}
}
}