use std::borrow::Borrow;
use std::hash::Hash;
use std::num::NonZeroUsize;
use lru::LruCache;
pub struct FifoCache<K: Hash + Eq, V> {
entries: LruCache<K, V>,
}
impl<K: Hash + Eq, V: Clone> FifoCache<K, V> {
pub fn new(capacity: usize) -> Self {
Self {
entries: LruCache::new(
NonZeroUsize::new(capacity).expect("cache capacity is a nonzero constant"),
),
}
}
pub fn get(&self, key: &K) -> Option<V> {
self.entries.peek(key).cloned()
}
pub fn get_ref(&self, key: &K) -> Option<&V> {
self.entries.peek(key)
}
pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
self.entries.peek_mut(key)
}
pub fn get_ref_by<Q>(&self, key: &Q) -> Option<&V>
where
K: Borrow<Q>,
Q: Hash + Eq + ?Sized,
{
self.entries.peek(key)
}
pub fn contains(&self, key: &K) -> bool {
self.entries.contains(key)
}
pub fn insert(&mut self, key: K, value: V) {
if let Some(slot) = self.entries.peek_mut(&key) {
*slot = value;
} else {
self.entries.push(key, value);
}
}
pub fn insert_with_previous(&mut self, key: K, value: V) -> (&mut V, Option<V>) {
if self.entries.contains(&key) {
let slot = self.entries.peek_mut(&key).expect("已确认缓存槽存在");
let previous = std::mem::replace(slot, value);
return (slot, Some(previous));
}
let previous = self.entries.push(key, value).map(|(_, value)| value);
let (_, slot) = self.entries.iter_mut().next().expect("新插入的槽必须存在");
(slot, previous)
}
pub fn remove(&mut self, key: &K) -> Option<V> {
self.entries.pop(key)
}
pub fn clear(&mut self) {
self.entries.clear();
}
pub fn retain_mut<F>(&mut self, mut keep: F)
where
K: Clone,
F: FnMut(&K, &mut V) -> bool,
{
let mut removed = Vec::new();
for (key, value) in self.entries.iter_mut() {
if !keep(key, value) {
removed.push(key.clone());
}
}
for key in removed {
self.entries.pop(&key);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn evicts_by_insertion_order_not_use() {
let mut cache = FifoCache::new(2);
cache.insert("a", 1);
cache.insert("b", 2);
assert_eq!(cache.get(&"a"), Some(1));
cache.insert("a", 3);
cache.insert("c", 4);
assert!(!cache.contains(&"a"));
assert_eq!(cache.get(&"b"), Some(2));
assert_eq!(cache.get(&"c"), Some(4));
}
#[test]
fn looks_up_owned_keys_by_borrowed_key() {
let mut cache = FifoCache::new(2);
cache.insert("a".to_owned(), 1);
assert_eq!(cache.get_ref_by("a"), Some(&1));
}
#[test]
fn recovered_values_keep_fifo_age_and_reuse_the_full_cache_node() {
use crate::ime::personal_rerank::allocations::count;
let mut cache = FifoCache::new(2);
let (slot, previous) = cache.insert_with_previous("a", 1);
assert_eq!(previous, None);
*slot = 2;
cache.insert("b", 3);
let ((slot, previous), allocations) = count(|| cache.insert_with_previous("a", 4));
assert_eq!(previous, Some(2));
assert_eq!(allocations, 0);
*slot = 5;
assert_eq!(cache.get(&"a"), Some(5));
let ((slot, previous), allocations) = count(|| cache.insert_with_previous("c", 6));
assert_eq!(previous, Some(5));
assert_eq!(allocations, 0);
*slot = 7;
assert!(!cache.contains(&"a"));
assert_eq!(cache.get(&"b"), Some(3));
assert_eq!(cache.get(&"c"), Some(7));
cache.insert("d", 8);
assert!(!cache.contains(&"b"));
}
}