use std::collections::HashMap;
use std::hash::Hash;
use std::ptr::NonNull;
use std::time::{Duration, Instant};
struct Node<K, V> {
key: K,
value: V,
expires_at: Option<Instant>,
next: Option<NonNull<Node<K, V>>>,
prev: Option<NonNull<Node<K, V>>>,
}
pub enum CleanupMode {
OnAccess,
OnDemand,
}
pub struct LruCache<K, V> {
map: HashMap<K, NonNull<Node<K, V>>>,
head: Option<NonNull<Node<K, V>>>,
tail: Option<NonNull<Node<K, V>>>,
capacity: usize,
cleanup_mode: CleanupMode,
}
impl<K: Eq + Hash + Clone, V> LruCache<K, V> {
pub fn new(capacity: usize, cleanup_mode: CleanupMode) -> Self {
assert!(capacity > 0);
LruCache {
map: HashMap::with_capacity(capacity),
head: None,
tail: None,
capacity,
cleanup_mode,
}
}
pub fn put(&mut self, key: K, value: V, ttl: Option<Duration>) {
if matches!(self.cleanup_mode, CleanupMode::OnAccess) {
self.evict_expired();
}
let expires_at = ttl.map(|d| Instant::now() + d);
if let Some(&node_ptr) = self.map.get(&key) {
unsafe {
let node = node_ptr.as_ptr().as_mut().unwrap();
node.value = value;
node.expires_at = expires_at;
self.remove_node(node_ptr);
self.push_front(node_ptr);
}
return;
}
if self.map.len() >= self.capacity {
self.remove_last();
}
let node = Box::new(Node {
key: key.clone(),
value,
expires_at,
next: self.head,
prev: None,
});
let node_ptr = unsafe { NonNull::new_unchecked(Box::into_raw(node)) };
if let Some(mut head) = self.head {
unsafe { head.as_mut().prev = Some(node_ptr) };
} else {
self.tail = Some(node_ptr);
}
self.head = Some(node_ptr);
self.map.insert(key, node_ptr);
}
pub fn get(&mut self, key: &K) -> Option<&V> {
if matches!(self.cleanup_mode, CleanupMode::OnAccess) {
self.evict_expired();
}
let node_ptr = *self.map.get(key)?;
unsafe {
let node = node_ptr.as_ptr().as_ref().unwrap();
if node.expired() {
self.map.remove(key);
self.remove_node(node_ptr);
let _ = Box::from_raw(node_ptr.as_ptr());
return None;
}
self.remove_node(node_ptr);
self.push_front(node_ptr);
Some(&(*node_ptr.as_ptr()).value)
}
}
pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
if matches!(self.cleanup_mode, CleanupMode::OnAccess) {
self.evict_expired();
}
let node_ptr = *self.map.get(key)?;
unsafe {
let node = node_ptr.as_ptr().as_mut().unwrap();
if node.expired() {
self.map.remove(key);
self.remove_node(node_ptr);
let _ = Box::from_raw(node_ptr.as_ptr());
return None;
}
self.remove_node(node_ptr);
self.push_front(node_ptr);
Some(&mut (*node_ptr.as_ptr()).value)
}
}
fn remove_node(&mut self, node_ptr: NonNull<Node<K, V>>) {
unsafe {
let node = node_ptr.as_ptr();
if let Some(prev) = (*node).prev {
(*prev.as_ptr()).next = (*node).next;
} else {
self.head = (*node).next;
}
if let Some(next) = (*node).next {
(*next.as_ptr()).prev = (*node).prev;
} else {
self.tail = (*node).prev;
}
}
}
fn push_front(&mut self, node_ptr: NonNull<Node<K, V>>) {
unsafe {
(*node_ptr.as_ptr()).next = self.head;
(*node_ptr.as_ptr()).prev = None;
if let Some(head) = self.head {
let head_mut = head.as_ptr() as *mut Node<K, V>;
(*head_mut).prev = Some(node_ptr);
} else {
self.tail = Some(node_ptr);
}
self.head = Some(node_ptr);
}
}
fn remove_last(&mut self) {
if let Some(tail_ptr) = self.tail {
unsafe {
let key = &(*tail_ptr.as_ptr()).key;
let prev = (*tail_ptr.as_ptr()).prev;
self.map.remove(key);
match prev {
Some(prev) => {
let prev_mut = prev.as_ptr() as *mut Node<K, V>;
(*prev_mut).next = None;
self.tail = Some(prev);
}
None => {
self.head = None;
self.tail = None;
}
}
let _ = Box::from_raw(tail_ptr.as_ptr());
}
}
}
pub fn evict_expired(&mut self) {
let now = Instant::now();
let mut current = self.head;
while let Some(node_ptr) = current {
unsafe {
let node = node_ptr.as_ptr();
current = (*node).next;
if (*node).expired_at(now) {
self.map.remove(&(*node).key);
self.remove_node(node_ptr);
let _ = Box::from_raw(node);
}
}
}
}
pub fn len(&self) -> usize {
self.map.len()
}
pub fn is_empty(&self) -> bool {
self.map.is_empty()
}
pub fn capacity(&self) -> usize {
self.capacity
}
}
impl<K, V> Drop for LruCache<K, V> {
fn drop(&mut self) {
let mut current = self.head;
while let Some(node_ptr) = current {
unsafe {
current = (*node_ptr.as_ptr()).next;
let _ = Box::from_raw(node_ptr.as_ptr());
}
}
}
}
impl<K, V> Node<K, V> {
fn expired(&self) -> bool {
self.expires_at.map_or(false, |e| e <= Instant::now())
}
fn expired_at(&self, now: Instant) -> bool {
self.expires_at.map_or(false, |e| e <= now)
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::thread;
#[test]
fn test_basic_operations() {
let mut cache = LruCache::new(2, CleanupMode::OnAccess);
cache.put("a", 1, None);
cache.put("b", 2, None);
assert_eq!(cache.get(&"a"), Some(&1));
assert_eq!(cache.get(&"b"), Some(&2));
assert_eq!(cache.get(&"c"), None);
cache.put("c", 3, None);
assert_eq!(cache.get(&"a"), None);
assert_eq!(cache.get(&"b"), Some(&2));
assert_eq!(cache.get(&"c"), Some(&3));
}
#[test]
fn test_ttl_expiration_auto() {
let mut cache = LruCache::new(2, CleanupMode::OnDemand);
cache.put("a", 1, Some(Duration::from_millis(150)));
cache.put("b", 2, None);
assert_eq!(cache.get(&"a"), Some(&1));
assert_eq!(cache.get(&"b"), Some(&2));
thread::sleep(Duration::from_millis(200));
assert_eq!(cache.get(&"a"), None);
assert_eq!(cache.get(&"b"), Some(&2));
}
#[test]
fn test_ttl_expiration() {
let mut cache = LruCache::new(2, CleanupMode::OnAccess);
cache.put("a", 1, Some(Duration::from_millis(150)));
cache.put("b", 2, None);
assert_eq!(cache.get(&"a"), Some(&1));
assert_eq!(cache.get(&"b"), Some(&2));
thread::sleep(Duration::from_millis(200));
cache.evict_expired();
assert_eq!(cache.get(&"a"), None);
assert_eq!(cache.get(&"b"), Some(&2));
}
#[test]
fn test_lru_eviction() {
let mut cache = LruCache::new(3, CleanupMode::OnAccess);
cache.put("a", 1, None);
cache.put("b", 2, None);
cache.put("c", 3, None);
cache.get(&"a");
cache.put("d", 4, None);
assert_eq!(cache.get(&"b"), None);
assert_eq!(cache.get(&"a"), Some(&1));
assert_eq!(cache.get(&"c"), Some(&3));
assert_eq!(cache.get(&"d"), Some(&4));
}
#[test]
fn test_no_memory_leaks() {
let mut cache = LruCache::new(2, CleanupMode::OnAccess);
for i in 0..1000 {
cache.put(i, Box::new([0u8; 1024]), None);
}
assert_eq!(cache.len(), 2);
}
}