use std::collections::HashMap;
use std::hash::Hash;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{PoisonError, RwLock, RwLockWriteGuard};
#[derive(Clone)]
struct LruNode<K, V> {
value: V,
prev: Option<K>,
next: Option<K>,
}
struct LruCache<K, V> {
map: HashMap<K, LruNode<K, V>>,
head: Option<K>, tail: Option<K>, max_capacity: Option<usize>,
}
impl<K, V> LruCache<K, V>
where
K: Eq + Hash + Clone,
V: Clone,
{
fn new(max_capacity: Option<usize>) -> Self {
Self {
map: HashMap::new(),
head: None,
tail: None,
max_capacity,
}
}
fn len(&self) -> usize {
self.map.len()
}
fn is_empty(&self) -> bool {
self.map.is_empty()
}
fn capacity(&self) -> usize {
self.map.capacity()
}
fn contains_key(&self, key: &K) -> bool {
self.map.contains_key(key)
}
fn peek(&self, key: &K) -> Option<&V> {
self.map.get(key).map(|node| &node.value)
}
fn get(&mut self, key: &K) -> Option<V> {
if !self.map.contains_key(key) {
return None;
}
self.move_to_tail(key);
self.map.get(key).map(|node| node.value.clone())
}
fn insert(&mut self, key: K, value: V) -> Option<(K, V)> {
let mut evicted = None;
if self.map.contains_key(&key) {
if let Some(node) = self.map.get_mut(&key) {
node.value = value;
}
self.move_to_tail(&key);
} else {
if self
.max_capacity
.is_some_and(|max| max > 0 && self.map.len() >= max)
{
evicted = self.evict_lru();
}
let node = LruNode {
value,
prev: self.tail.clone(),
next: None,
};
if let Some(tail_node) = self.tail.as_ref().and_then(|k| self.map.get_mut(k)) {
tail_node.next = Some(key.clone());
}
self.map.insert(key.clone(), node);
self.tail = Some(key.clone());
if self.head.is_none() {
self.head = Some(key);
}
}
evicted
}
fn remove(&mut self, key: &K) -> Option<V> {
if let Some(node) = self.map.remove(key) {
self.unlink(&node);
Some(node.value)
} else {
None
}
}
fn clear(&mut self) {
self.map.clear();
self.head = None;
self.tail = None;
}
fn keys(&self) -> Vec<K> {
self.map.keys().cloned().collect()
}
fn values(&self) -> Vec<V> {
self.map.values().map(|n| n.value.clone()).collect()
}
fn reserve(&mut self, additional: usize) {
self.map.reserve(additional);
}
fn shrink_to_fit(&mut self) {
self.map.shrink_to_fit();
}
fn move_to_tail(&mut self, key: &K) {
if self.tail.as_ref() == Some(key) {
return; }
let (prev, next) = {
let node = match self.map.get(key) {
Some(n) => n,
None => return,
};
(node.prev.clone(), node.next.clone())
};
if let Some(ref prev_key) = prev {
if let Some(prev_node) = self.map.get_mut(prev_key) {
prev_node.next = next.clone();
}
} else {
self.head = next.clone();
}
if let Some(next_node) = next.as_ref().and_then(|k| self.map.get_mut(k)) {
next_node.prev = prev.clone();
}
if let Some(tail_node) = self
.tail
.as_ref()
.filter(|k| *k != key)
.and_then(|k| self.map.get_mut(k))
{
tail_node.next = Some(key.clone());
}
if let Some(node) = self.map.get_mut(key) {
node.prev = self.tail.clone();
node.next = None;
}
self.tail = Some(key.clone());
}
fn unlink(&mut self, node: &LruNode<K, V>) {
if let Some(ref prev_key) = node.prev {
if let Some(prev_node) = self.map.get_mut(prev_key) {
prev_node.next = node.next.clone();
}
} else {
self.head = node.next.clone();
}
if let Some(ref next_key) = node.next {
if let Some(next_node) = self.map.get_mut(next_key) {
next_node.prev = node.prev.clone();
}
} else {
self.tail = node.prev.clone();
}
}
fn evict_lru(&mut self) -> Option<(K, V)> {
let head_key = self.head.clone()?;
let node = self.map.remove(&head_key)?;
self.head = node.next.clone();
if let Some(ref new_head_key) = self.head {
if let Some(new_head_node) = self.map.get_mut(new_head_key) {
new_head_node.prev = None;
}
} else {
self.tail = None;
}
Some((head_key, node.value))
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct CacheStats {
pub hits: u64,
pub misses: u64,
pub evictions: u64,
}
impl CacheStats {
#[must_use]
pub fn hit_rate(&self) -> f64 {
let total = self.hits + self.misses;
if total == 0 {
0.0
} else {
self.hits as f64 / total as f64
}
}
#[must_use]
pub fn total_accesses(&self) -> u64 {
self.hits + self.misses
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MemoizerError {
pub message: String,
}
pub type InsertEvictionResult<V, K> = (Option<V>, Option<K>, Option<V>);
impl std::fmt::Display for MemoizerError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "MemoizerError: {}", self.message)
}
}
impl std::error::Error for MemoizerError {}
impl<T> From<PoisonError<T>> for MemoizerError {
fn from(err: PoisonError<T>) -> Self {
MemoizerError {
message: format!("Lock poisoned: {}", err),
}
}
}
pub struct Memoizer<K, V> {
cache: RwLock<LruCache<K, V>>,
hits: AtomicU64,
misses: AtomicU64,
evictions: AtomicU64,
}
impl<K, V> Default for Memoizer<K, V>
where
K: Eq + Hash + Clone,
V: Clone,
{
fn default() -> Self {
Self::new()
}
}
impl<K, V> Memoizer<K, V>
where
K: Eq + Hash + Clone,
V: Clone,
{
pub fn new() -> Self {
Memoizer {
cache: RwLock::new(LruCache::new(None)),
hits: AtomicU64::new(0),
misses: AtomicU64::new(0),
evictions: AtomicU64::new(0),
}
}
pub fn with_capacity(max_capacity: usize) -> Self {
Memoizer {
cache: RwLock::new(LruCache::new(Some(max_capacity))),
hits: AtomicU64::new(0),
misses: AtomicU64::new(0),
evictions: AtomicU64::new(0),
}
}
pub fn stats(&self) -> CacheStats {
CacheStats {
hits: self.hits.load(Ordering::Relaxed),
misses: self.misses.load(Ordering::Relaxed),
evictions: self.evictions.load(Ordering::Relaxed),
}
}
pub fn reset_stats(&self) {
self.hits.store(0, Ordering::Relaxed);
self.misses.store(0, Ordering::Relaxed);
self.evictions.store(0, Ordering::Relaxed);
}
pub fn max_capacity(&self) -> Option<usize> {
self.cache.read().unwrap().max_capacity
}
pub fn get_or_compute<F>(&self, key: K, f: F) -> V
where
F: FnOnce(&K) -> V,
{
let mut cache = self.cache.write().unwrap();
if let Some(v) = cache.get(&key) {
self.hits.fetch_add(1, Ordering::Relaxed);
return v;
}
self.misses.fetch_add(1, Ordering::Relaxed);
let value = f(&key);
if let Some(_evicted) = cache.insert(key, value.clone()) {
self.evictions.fetch_add(1, Ordering::Relaxed);
}
value
}
pub fn get_or_compute_optimistic<F>(&self, key: K, f: F) -> V
where
F: FnOnce(&K) -> V,
{
{
let mut cache = self.cache.write().unwrap();
if let Some(v) = cache.get(&key) {
self.hits.fetch_add(1, Ordering::Relaxed);
return v;
}
}
self.misses.fetch_add(1, Ordering::Relaxed);
let value = f(&key);
let mut cache = self.cache.write().unwrap();
if let Some(existing) = cache.get(&key) {
return existing;
}
if cache.insert(key, value.clone()).is_some() {
self.evictions.fetch_add(1, Ordering::Relaxed);
}
value
}
pub fn try_get_or_compute_optimistic<F>(&self, key: K, f: F) -> Result<V, MemoizerError>
where
F: FnOnce(&K) -> V,
{
{
let mut cache = self.write_cache()?;
if let Some(v) = cache.get(&key) {
self.hits.fetch_add(1, Ordering::Relaxed);
return Ok(v);
}
}
self.misses.fetch_add(1, Ordering::Relaxed);
let value = f(&key);
let mut cache = self.write_cache()?;
if let Some(existing) = cache.get(&key) {
return Ok(existing);
}
if cache.insert(key, value.clone()).is_some() {
self.evictions.fetch_add(1, Ordering::Relaxed);
}
Ok(value)
}
pub fn clear(&self) {
let mut cache = self.cache.write().unwrap();
cache.clear();
}
fn write_cache(&self) -> Result<RwLockWriteGuard<'_, LruCache<K, V>>, MemoizerError> {
self.cache.write().map_err(|e| e.into())
}
pub fn try_get_or_compute<F>(&self, key: K, f: F) -> Result<V, MemoizerError>
where
F: FnOnce(&K) -> V,
{
let mut cache = self.write_cache()?;
if let Some(v) = cache.get(&key) {
self.hits.fetch_add(1, Ordering::Relaxed);
return Ok(v);
}
self.misses.fetch_add(1, Ordering::Relaxed);
let value = f(&key);
if let Some(_evicted) = cache.insert(key, value.clone()) {
self.evictions.fetch_add(1, Ordering::Relaxed);
}
Ok(value)
}
pub fn insert(&self, key: K, value: V) -> Option<V> {
let mut cache = self.cache.write().unwrap();
let old = cache.peek(&key).cloned();
if cache.insert(key, value).is_some() {
self.evictions.fetch_add(1, Ordering::Relaxed);
}
old
}
pub fn insert_with_eviction_info(&self, key: K, value: V) -> InsertEvictionResult<V, K> {
let mut cache = self.cache.write().unwrap();
let old = cache.peek(&key).cloned();
let evicted = cache.insert(key, value);
let (evicted_key, evicted_value) = if let Some((k, v)) = evicted {
self.evictions.fetch_add(1, Ordering::Relaxed);
(Some(k), Some(v))
} else {
(None, None)
};
(old, evicted_key, evicted_value)
}
pub fn try_insert(&self, key: K, value: V) -> Result<Option<V>, MemoizerError> {
let mut cache = self.write_cache()?;
let old = cache.peek(&key).cloned();
if cache.insert(key, value).is_some() {
self.evictions.fetch_add(1, Ordering::Relaxed);
}
Ok(old)
}
pub fn try_insert_with_eviction_info(
&self, key: K, value: V,
) -> Result<InsertEvictionResult<V, K>, MemoizerError> {
let mut cache = self.write_cache()?;
let old = cache.peek(&key).cloned();
let evicted = cache.insert(key, value);
let (evicted_key, evicted_value) = if let Some((k, v)) = evicted {
self.evictions.fetch_add(1, Ordering::Relaxed);
(Some(k), Some(v))
} else {
(None, None)
};
Ok((old, evicted_key, evicted_value))
}
pub fn get_or_try_compute<F, E>(&self, key: K, f: F) -> Result<V, E>
where
F: FnOnce(&K) -> Result<V, E>,
{
let mut cache = self.cache.write().unwrap();
if let Some(v) = cache.get(&key) {
self.hits.fetch_add(1, Ordering::Relaxed);
return Ok(v);
}
self.misses.fetch_add(1, Ordering::Relaxed);
let value = f(&key)?;
if cache.insert(key, value.clone()).is_some() {
self.evictions.fetch_add(1, Ordering::Relaxed);
}
Ok(value)
}
pub fn touch(&self, key: &K) -> bool {
let mut cache = self.cache.write().unwrap();
cache.get(key).is_some()
}
pub fn try_touch(&self, key: &K) -> Result<bool, MemoizerError> {
let mut cache = self.write_cache()?;
Ok(cache.get(key).is_some())
}
pub fn len(&self) -> usize {
self.cache.read().unwrap().len()
}
pub fn try_len(&self) -> Result<usize, MemoizerError> {
Ok(self.cache.read().map_err(MemoizerError::from)?.len())
}
pub fn is_empty(&self) -> bool {
self.cache.read().unwrap().is_empty()
}
pub fn try_is_empty(&self) -> Result<bool, MemoizerError> {
Ok(self.cache.read().map_err(MemoizerError::from)?.is_empty())
}
pub fn contains_key(&self, key: &K) -> bool {
self.cache.read().unwrap().contains_key(key)
}
pub fn try_contains_key(&self, key: &K) -> Result<bool, MemoizerError> {
Ok(self
.cache
.read()
.map_err(MemoizerError::from)?
.contains_key(key))
}
pub fn remove(&self, key: &K) -> Option<V> {
self.cache.write().unwrap().remove(key)
}
pub fn try_remove(&self, key: &K) -> Result<Option<V>, MemoizerError> {
Ok(self.write_cache()?.remove(key))
}
pub fn get(&self, key: &K) -> Option<V> {
self.cache.read().unwrap().peek(key).cloned()
}
pub fn try_get(&self, key: &K) -> Result<Option<V>, MemoizerError> {
Ok(self
.cache
.read()
.map_err(MemoizerError::from)?
.peek(key)
.cloned())
}
pub fn get_with_lru_update(&self, key: &K) -> Option<V> {
let mut cache = self.cache.write().unwrap();
cache.get(key)
}
pub fn try_get_with_lru_update(&self, key: &K) -> Result<Option<V>, MemoizerError> {
let mut cache = self.write_cache()?;
Ok(cache.get(key))
}
pub fn reserve(&self, additional: usize) {
self.cache.write().unwrap().reserve(additional)
}
pub fn try_reserve(&self, additional: usize) -> Result<(), MemoizerError> {
self.write_cache()?.reserve(additional);
Ok(())
}
pub fn shrink_to_fit(&self) {
self.cache.write().unwrap().shrink_to_fit()
}
pub fn try_shrink_to_fit(&self) -> Result<(), MemoizerError> {
self.write_cache()?.shrink_to_fit();
Ok(())
}
pub fn keys(&self) -> Vec<K> {
self.cache.read().unwrap().keys()
}
pub fn try_keys(&self) -> Result<Vec<K>, MemoizerError> {
Ok(self.cache.read().map_err(MemoizerError::from)?.keys())
}
pub fn values(&self) -> Vec<V> {
self.cache.read().unwrap().values()
}
pub fn try_values(&self) -> Result<Vec<V>, MemoizerError> {
Ok(self.cache.read().map_err(MemoizerError::from)?.values())
}
pub fn capacity(&self) -> usize {
self.cache.read().unwrap().capacity()
}
pub fn try_capacity(&self) -> Result<usize, MemoizerError> {
Ok(self.cache.read().map_err(MemoizerError::from)?.capacity())
}
pub fn try_clear(&self) -> Result<(), MemoizerError> {
self.write_cache()?.clear();
Ok(())
}
}