src/tools/lru_cache.rs

   1 //! Least recently used (LRU) cache
   2 //!
   3 //! Implements a cache with least recently used cache replacement policy.
   4 //! A HashMap is used for fast access by a given key and a doubly linked list
   5 //! is used to keep track of the cache access order.
   6
   7 use std::collections::{HashMap, hash_map::Entry};
   8 use std::marker::PhantomData;
   9
  10 /// Interface for getting values on cache misses.
  11 pub trait Cacher<K, V> {
  12     /// Fetch a value for key on cache miss.
  13     ///
  14     /// Whenever a cache miss occurs, the fetch method provides a corresponding value.
  15     /// If no value can be obtained for the given key, None is returned, the cache is
  16     /// not updated in that case.
  17     fn fetch(&mut self, key: K) -> Result<Option<V>, failure::Error>;
  18 }
  19
  20 /// Node of the doubly linked list storing key and value
  21 struct CacheNode<K, V> {
  22     // We need to additionally store the key to be able to remove it
  23     // from the HashMap when removing the tail.
  24     key: K,
  25     value: V,
  26     prev: *mut CacheNode<K, V>,
  27     next: *mut CacheNode<K, V>,
  28     // Dropcheck marker. See the phantom-data section in the rustonomicon.
  29     _marker: PhantomData<Box<CacheNode<K, V>>>,
  30 }
  31
  32 impl<K, V> CacheNode<K, V> {
  33     fn new(key: K, value: V) -> Self {
  34         Self {
  35             key,
  36             value,
  37             prev: std::ptr::null_mut(),
  38             next: std::ptr::null_mut(),
  39             _marker: PhantomData,
  40         }
  41     }
  42 }
  43
  44 /// LRU cache instance.
  45 ///
  46 /// # Examples:
  47 /// ```
  48 /// # use self::proxmox_backup::tools::lru_cache::{Cacher, LruCache};
  49 /// # fn main() -> Result<(), failure::Error> {
  50 /// struct LruCacher {};
  51 ///
  52 /// impl Cacher<u64, u64> for LruCacher {
  53 ///     fn fetch(&mut self, key: u64) -> Result<Option<u64>, failure::Error> {
  54 ///         Ok(Some(key))
  55 ///     }
  56 /// }
  57 ///
  58 /// let mut cache = LruCache::new(3);
  59 ///
  60 /// assert_eq!(cache.get_mut(1), None);
  61 /// assert_eq!(cache.len(), 0);
  62 ///
  63 /// cache.insert(1, 1);
  64 /// cache.insert(2, 2);
  65 /// cache.insert(3, 3);
  66 /// cache.insert(4, 4);
  67 /// assert_eq!(cache.len(), 3);
  68 ///
  69 /// assert_eq!(cache.get_mut(1), None);
  70 /// assert_eq!(cache.get_mut(2), Some(&mut 2));
  71 /// assert_eq!(cache.get_mut(3), Some(&mut 3));
  72 /// assert_eq!(cache.get_mut(4), Some(&mut 4));
  73 ///
  74 /// cache.remove(4);
  75 /// cache.remove(3);
  76 /// cache.remove(2);
  77 /// assert_eq!(cache.len(), 0);
  78 /// assert_eq!(cache.get_mut(2), None);
  79 /// // access will fill in missing cache entry by fetching from LruCacher
  80 /// assert_eq!(cache.access(2, &mut LruCacher {}).unwrap(), Some(&mut 2));
  81 ///
  82 /// cache.insert(1, 1);
  83 /// assert_eq!(cache.get_mut(1), Some(&mut 1));
  84 ///
  85 /// cache.clear();
  86 /// assert_eq!(cache.len(), 0);
  87 /// assert_eq!(cache.get_mut(1), None);
  88 /// # Ok(())
  89 /// # }
  90 /// ```
  91 pub struct LruCache<K, V> {
  92     /// Quick access to individual nodes via the node pointer.
  93     map: HashMap<K, *mut CacheNode<K, V>>,
  94     /// Actual nodes stored in a linked list.
  95     list: LinkedList<K, V>,
  96     /// Max nodes the cache can hold, temporarily exceeded by 1 due to
  97     /// implementation details.
  98     capacity: usize,
  99     // Dropcheck marker. See the phantom-data section in the rustonomicon.
 100     _marker: PhantomData<Box<CacheNode<K, V>>>,
 101 }
 102
 103 unsafe impl<K, V> Send for LruCache<K, V> {}
 104
 105 impl<K: std::cmp::Eq + std::hash::Hash + Copy, V> LruCache<K, V> {
 106     /// Create LRU cache instance which holds up to `capacity` nodes at once.
 107     pub fn new(capacity: usize) -> Self {
 108         Self {
 109             map: HashMap::with_capacity(capacity),
 110             list: LinkedList::new(),
 111             capacity,
 112             _marker: PhantomData,
 113         }
 114     }
 115
 116     /// Clear all the entries from the cache.
 117     pub fn clear(&mut self) {
 118         // This frees only the HashMap with the node pointers.
 119         self.map.clear();
 120         // This frees the actual nodes and resets the list head and tail.
 121         self.list.clear();
 122     }
 123
 124     /// Insert or update an entry identified by `key` with the given `value`.
 125     /// This entry is placed as the most recently used node at the head.
 126     pub fn insert(&mut self, key: K, value: V) {
 127         match self.map.entry(key) {
 128             Entry::Occupied(mut o) => {
 129                 // Node present, update value
 130                 let node_ptr = *o.get_mut();
 131                 self.list.bring_to_front(node_ptr);
 132                 let mut node = unsafe { Box::from_raw(node_ptr) };
 133                 node.value = value;
 134                 let _node_ptr = Box::into_raw(node);
 135             }
 136             Entry::Vacant(v) => {
 137                 // Node not present, insert a new one
 138                 // Unfortunately we need a copy of the key here, therefore it has
 139                 // to impl the copy trait
 140                 let node = Box::new(CacheNode::new(key, value));
 141                 let node_ptr = Box::into_raw(node);
 142                 self.list.push_front(node_ptr);
 143                 v.insert(node_ptr);
 144                 // If we have more elements than capacity,
 145                 // delete the lists tail node (= oldest node).
 146                 // This needs to be executed after the insert in order to
 147                 // avoid borrow conflict. This means there are temporarily
 148                 // self.capacity + 1 cache nodes.
 149                 if self.map.len() > self.capacity {
 150                     self.pop_tail();
 151                 }
 152             }
 153         }
 154     }
 155
 156     /// Remove the given `key` and its `value` from the cache.
 157     pub fn remove(&mut self, key: K) -> Option<V> {
 158         // Remove node pointer from the HashMap and get ownership of the node
 159         let node_ptr = self.map.remove(&key)?;
 160         let node = self.list.remove(node_ptr);
 161         Some(node.value)
 162     }
 163
 164     /// Remove the least recently used node from the cache.
 165     fn pop_tail(&mut self) {
 166         if let Some(old_tail) = self.list.pop_tail() {
 167             // Remove HashMap entry for old tail
 168             self.map.remove(&old_tail.key);
 169         }
 170     }
 171
 172     /// Get a mutable reference to the value identified by `key`.
 173     /// This will update the cache entry to be the most recently used entry.
 174     /// On cache misses, None is returned.
 175     pub fn get_mut<'a>(&'a mut self, key: K) -> Option<&'a mut V> {
 176         let node_ptr = self.map.get(&key)?;
 177         self.list.bring_to_front(*node_ptr);
 178         Some(unsafe { &mut (*self.list.head).value })
 179     }
 180
 181     /// Number of entries in the cache.
 182     pub fn len(&self) -> usize {
 183         self.map.len()
 184     }
 185
 186     /// Get a mutable reference to the value identified by `key`.
 187     /// This will update the cache entry to be the most recently used entry.
 188     /// On cache misses, the cachers fetch method is called to get a corresponding
 189     /// value.
 190     /// If fetch returns a value, it is inserted as the most recently used entry
 191     /// in the cache.
 192     pub fn access<'a>(&'a mut self, key: K, cacher: &mut dyn Cacher<K, V>) -> Result<Option<&'a mut V>, failure::Error> {
 193         match self.map.entry(key) {
 194             Entry::Occupied(mut o) => {
 195                 // Cache hit, birng node to front of list
 196                 let node_ptr = *o.get_mut();
 197                 self.list.bring_to_front(node_ptr);
 198             }
 199             Entry::Vacant(v) => {
 200                 // Cache miss, try to fetch from cacher and insert at the front
 201                 match cacher.fetch(key)? {
 202                     None => return Ok(None),
 203                     Some(value) => {
 204                         // Unfortunately we need a copy of the key here, therefore it has
 205                         // to impl the copy trait
 206                         let node = Box::new(CacheNode::new(key, value));
 207                         let node_ptr = Box::into_raw(node);
 208                         self.list.push_front(node_ptr);
 209                         v.insert(node_ptr);
 210                         // If we have more elements than capacity,
 211                         // delete the lists tail node (= oldest node).
 212                         // This needs to be executed after the insert in order to
 213                         // avoid borrow conflict. This means there are temporarily
 214                         // self.capacity + 1 cache nodes.
 215                         if self.map.len() > self.capacity {
 216                             self.pop_tail();
 217                         }
 218                     }
 219                 }
 220             }
 221         }
 222
 223         Ok(Some(unsafe { &mut (*self.list.head).value }))
 224     }
 225 }
 226
 227 /// Linked list holding the nodes of the LruCache.
 228 ///
 229 /// This struct actually holds the CacheNodes via the raw linked list pointers
 230 /// and allows to define the access sequence of these via the list sequence.
 231 /// The LinkedList of the standard library unfortunately does not implement
 232 /// an efficient way to bring list entries to the front, therefore we need our own.
 233 struct LinkedList<K, V> {
 234     head: *mut CacheNode<K, V>,
 235     tail: *mut CacheNode<K, V>,
 236 }
 237
 238 impl<K, V> LinkedList<K, V> {
 239     /// Create a new empty linked list.
 240     fn new() -> Self {
 241         Self {
 242             head: std::ptr::null_mut(),
 243             tail: std::ptr::null_mut(),
 244         }
 245     }
 246
 247     /// Bring the CacheNode referenced by `node_ptr` to the front of the linked list.
 248     fn bring_to_front(&mut self, node_ptr: *mut CacheNode<K, V>) {
 249         if node_ptr == self.head {
 250             // node is already head, just return
 251             return;
 252         }
 253
 254         let mut node = unsafe { Box::from_raw(node_ptr) };
 255         // Update the prev node to point to next (or null if current node is tail)
 256         unsafe { (*node.prev).next = node.next };
 257
 258         // Update the next node or otherwise the tail
 259         if !node.next.is_null() {
 260             unsafe { (*node.next).prev = node.prev };
 261         } else {
 262             // No next node means this was the tail
 263             self.tail = node.prev;
 264         }
 265
 266         node.prev = std::ptr::null_mut();
 267         node.next = self.head;
 268         // update the head and release ownership of the node again
 269         let node_ptr = Box::into_raw(node);
 270         // Update current head
 271         unsafe { (*self.head).prev = node_ptr };
 272         // Update to new head
 273         self.head = node_ptr;
 274     }
 275
 276     /// Insert a new node at the front of the linked list.
 277     fn push_front(&mut self, node_ptr: *mut CacheNode<K, V>) {
 278         let mut node = unsafe { Box::from_raw(node_ptr) };
 279
 280         // Old head gets new heads next
 281         node.next = self.head;
 282         // Release ownership of node, rest can be handled with just the pointer.
 283         let node_ptr = Box::into_raw(node);
 284
 285         // Update the prev for the old head
 286         if !self.head.is_null() {
 287             unsafe { (*self.head).prev = node_ptr };
 288         }
 289
 290         // Update the head to the new node pointer
 291         self.head = node_ptr;
 292
 293         // If there was no old tail, this node will be the new tail too
 294         if self.tail.is_null() {
 295             self.tail = node_ptr;
 296         }
 297     }
 298
 299     /// Remove the node referenced by `node_ptr` from the linke list and return it.
 300     fn remove(&mut self, node_ptr: *mut CacheNode<K, V>) -> Box<CacheNode<K, V>> {
 301         let node = unsafe { Box::from_raw(node_ptr) };
 302
 303         // Update the previous node or otherwise the head
 304         if !node.prev.is_null() {
 305             unsafe { (*node.prev).next = node.next };
 306         } else {
 307             // No previous node means this was the head
 308             self.head = node.next;
 309         }
 310
 311         // Update the next node or otherwise the tail
 312         if !node.next.is_null() {
 313             unsafe { (*node.next).prev = node.prev };
 314         } else {
 315             // No next node means this was the tail
 316             self.tail = node.prev;
 317         }
 318         node
 319     }
 320
 321     /// Remove the tail node from the linked list and return it.
 322     fn pop_tail(&mut self) -> Option<Box<CacheNode<K, V>>> {
 323         if self.tail.is_null() {
 324             return None;
 325         }
 326
 327         let old_tail = unsafe { Box::from_raw(self.tail) };
 328         self.tail = old_tail.prev;
 329         // Update next node for new tail
 330         if !self.tail.is_null() {
 331             unsafe { (*self.tail).next = std::ptr::null_mut() };
 332         }
 333         Some(old_tail)
 334     }
 335
 336     /// Clear the linked list and free all the nodes.
 337     fn clear(&mut self) {
 338         let mut next = self.head;
 339         while !next.is_null() {
 340             // Taking ownership of node and drop it at the end of the block.
 341             let current = unsafe { Box::from_raw(next) };
 342             next = current.next;
 343         }
 344         // Reset head and tail pointers
 345         self.head = std::ptr::null_mut();
 346         self.tail = std::ptr::null_mut();
 347     }
 348 }
 349
 350 #[test]
 351 fn test_linked_list() {
 352     let mut list = LinkedList::new();
 353     for idx in 0..3 {
 354         let node = Box::new(CacheNode::new(idx, idx + 1));
 355         // Get pointer, release ownership.
 356         let node_ptr = Box::into_raw(node);
 357         list.push_front(node_ptr);
 358     }
 359     assert_eq!(unsafe { (*list.head).key }, 2);
 360     assert_eq!(unsafe { (*list.head).value }, 3);
 361     assert_eq!(unsafe { (*list.tail).key }, 0);
 362     assert_eq!(unsafe { (*list.tail).value }, 1);
 363
 364     list.bring_to_front(list.tail);
 365     assert_eq!(unsafe { (*list.head).key }, 0);
 366     assert_eq!(unsafe { (*list.head).value }, 1);
 367     assert_eq!(unsafe { (*list.tail).key }, 1);
 368     assert_eq!(unsafe { (*list.tail).value }, 2);
 369
 370     list.bring_to_front(list.tail);
 371     assert_eq!(unsafe { (*list.head).key }, 1);
 372     assert_eq!(unsafe { (*list.head).value }, 2);
 373     assert_eq!(unsafe { (*list.tail).key }, 2);
 374     assert_eq!(unsafe { (*list.tail).value }, 3);
 375
 376     let tail = list.pop_tail().unwrap();
 377     assert_eq!(tail.key, 2);
 378     assert_eq!(tail.value, 3);
 379     assert_eq!(unsafe { (*list.head).key }, 1);
 380     assert_eq!(unsafe { (*list.head).value }, 2);
 381     assert_eq!(unsafe { (*list.tail).key }, 0);
 382     assert_eq!(unsafe { (*list.tail).value }, 1);
 383
 384     list.clear();
 385     assert!(list.head.is_null());
 386     assert!(list.tail.is_null());
 387 }