use std::collections::HashMap; use std::hash::Hash; /// A least-recently-used cache with a fixed capacity. /// /// Both `get` and `get_mut` mark an entry as recently used. pub struct LruCache { capacity: usize, entries: HashMap, nodes: Vec>>, free: Vec, head: Option, // Most recently used. tail: Option, // Least recently used. } struct Node { key: K, value: V, previous: Option, next: Option, } impl LruCache where K: Eq + Hash + Clone, { pub fn new(capacity: usize) -> Self { Self { capacity, entries: HashMap::with_capacity(capacity), nodes: Vec::with_capacity(capacity), free: Vec::new(), head: None, tail: None, } } pub fn len(&self) -> usize { self.entries.len() } pub fn is_empty(&self) -> bool { self.entries.is_empty() } pub fn capacity(&self) -> usize { self.capacity } pub fn contains_key(&self, key: &K) -> bool { self.entries.contains_key(key) } pub fn get(&mut self, key: &K) -> Option<&V> { let index = *self.entries.get(key)?; self.mark_recent(index); Some(&self.nodes[index].as_ref().unwrap().value) } pub fn get_mut(&mut self, key: &K) -> Option<&mut V> { let index = *self.entries.get(key)?; self.mark_recent(index); Some(&mut self.nodes[index].as_mut().unwrap().value) } /// Inserts a value, returning the previous value for the key if present. pub fn put(&mut self, key: K, value: V) -> Option { if let Some(&index) = self.entries.get(&key) { let old_value = { let node = self.nodes[index].as_mut().unwrap(); std::mem::replace(&mut node.value, value) }; self.mark_recent(index); return Some(old_value); } if self.capacity == 0 { return None; } if self.len() == self.capacity { let least_recent = self.tail.unwrap(); self.remove_index(least_recent); } let index = if let Some(index) = self.free.pop() { index } else { self.nodes.push(None); self.nodes.len() - 1 }; self.nodes[index] = Some(Node { key: key.clone(), value, previous: None, next: self.head, }); if let Some(old_head) = self.head { self.nodes[old_head].as_mut().unwrap().previous = Some(index); } else { self.tail = Some(index); } self.head = Some(index); self.entries.insert(key, index); None } pub fn remove(&mut self, key: &K) -> Option { let index = *self.entries.get(key)?; Some(self.remove_index(index).value) } pub fn clear(&mut self) { self.entries.clear(); self.nodes.clear(); self.free.clear(); self.head = None; self.tail = None; } fn mark_recent(&mut self, index: usize) { if self.head == Some(index) { return; } self.detach(index); let old_head = self.head; { let node = self.nodes[index].as_mut().unwrap(); node.previous = None; node.next = old_head; } if let Some(old_head) = old_head { self.nodes[old_head].as_mut().unwrap().previous = Some(index); } else { self.tail = Some(index); } self.head = Some(index); } fn detach(&mut self, index: usize) { let (previous, next) = { let node = self.nodes[index].as_ref().unwrap(); (node.previous, node.next) }; if let Some(previous) = previous { self.nodes[previous].as_mut().unwrap().next = next; } else { self.head = next; } if let Some(next) = next { self.nodes[next].as_mut().unwrap().previous = previous; } else { self.tail = previous; } } fn remove_index(&mut self, index: usize) -> Node { self.detach(index); let node = self.nodes[index].take().unwrap(); self.entries.remove(&node.key); self.free.push(index); node } } #[cfg(test)] mod tests { use super::LruCache; #[test] fn evicts_the_least_recently_used_entry() { let mut cache = LruCache::new(2); cache.put("a", 1); cache.put("b", 2); assert_eq!(cache.get(&"a"), Some(&1)); cache.put("c", 3); assert_eq!(cache.get(&"b"), None); assert_eq!(cache.get(&"a"), Some(&1)); assert_eq!(cache.get(&"c"), Some(&3)); } #[test] fn updating_an_entry_preserves_capacity_and_marks_it_recent() { let mut cache = LruCache::new(2); cache.put("a", 1); cache.put("b", 2); assert_eq!(cache.put("a", 10), Some(1)); cache.put("c", 3); assert_eq!(cache.len(), 2); assert_eq!(cache.get(&"a"), Some(&10)); assert_eq!(cache.get(&"b"), None); } #[test] fn zero_capacity_never_stores_entries() { let mut cache = LruCache::new(0); cache.put("a", 1); assert!(cache.is_empty()); assert_eq!(cache.get(&"a"), None); } }