src/libstd/collections/lru_cache.rs

   1 // Copyright 2013 The Rust Project Developers. See the COPYRIGHT
   2 // file at the top-level directory of this distribution and at
   3 // http://rust-lang.org/COPYRIGHT.
   4 //
   5 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
   6 // http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
   7 // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
   8 // option. This file may not be copied, modified, or distributed
   9 // except according to those terms.
  10
  11
  12 //! A cache that holds a limited number of key-value pairs. When the
  13 //! capacity of the cache is exceeded, the least-recently-used
  14 //! (where "used" means a look-up or putting the pair into the cache)
  15 //! pair is automatically removed.
  16 //!
  17 //! # Example
  18 //!
  19 //! ```rust
  20 //! use std::collections::LruCache;
  21 //!
  22 //! let mut cache: LruCache<int, int> = LruCache::new(2);
  23 //! cache.put(1, 10);
  24 //! cache.put(2, 20);
  25 //! cache.put(3, 30);
  26 //! assert!(cache.get(&1).is_none());
  27 //! assert_eq!(*cache.get(&2).unwrap(), 20);
  28 //! assert_eq!(*cache.get(&3).unwrap(), 30);
  29 //!
  30 //! cache.put(2, 22);
  31 //! assert_eq!(*cache.get(&2).unwrap(), 22);
  32 //!
  33 //! cache.put(6, 60);
  34 //! assert!(cache.get(&3).is_none());
  35 //!
  36 //! cache.change_capacity(1);
  37 //! assert!(cache.get(&2).is_none());
  38 //! ```
  39
  40 use cmp::{PartialEq, Eq};
  41 use collections::{HashMap, Collection, Mutable, MutableMap};
  42 use fmt;
  43 use hash::Hash;
  44 use iter::{range, Iterator};
  45 use mem;
  46 use ops::Drop;
  47 use option::{Some, None, Option};
  48 use owned::Box;
  49 use ptr;
  50 use result::{Ok, Err};
  51
  52 struct KeyRef<K> { k: *K }
  53
  54 struct LruEntry<K, V> {
  55     next: *mut LruEntry<K, V>,
  56     prev: *mut LruEntry<K, V>,
  57     key: K,
  58     value: V,
  59 }
  60
  61 /// An LRU Cache.
  62 pub struct LruCache<K, V> {
  63     map: HashMap<KeyRef<K>, Box<LruEntry<K, V>>>,
  64     max_size: uint,
  65     head: *mut LruEntry<K, V>,
  66 }
  67
  68 impl<S, K: Hash<S>> Hash<S> for KeyRef<K> {
  69     fn hash(&self, state: &mut S) {
  70         unsafe { (*self.k).hash(state) }
  71     }
  72 }
  73
  74 impl<K: PartialEq> PartialEq for KeyRef<K> {
  75     fn eq(&self, other: &KeyRef<K>) -> bool {
  76         unsafe{ (*self.k).eq(&*other.k) }
  77     }
  78 }
  79
  80 impl<K: Eq> Eq for KeyRef<K> {}
  81
  82 impl<K, V> LruEntry<K, V> {
  83     fn new(k: K, v: V) -> LruEntry<K, V> {
  84         LruEntry {
  85             key: k,
  86             value: v,
  87             next: ptr::mut_null(),
  88             prev: ptr::mut_null(),
  89         }
  90     }
  91 }
  92
  93 impl<K: Hash + Eq, V> LruCache<K, V> {
  94     /// Create an LRU Cache that holds at most `capacity` items.
  95     pub fn new(capacity: uint) -> LruCache<K, V> {
  96         let cache = LruCache {
  97             map: HashMap::new(),
  98             max_size: capacity,
  99             head: unsafe{ mem::transmute(box mem::uninitialized::<LruEntry<K, V>>()) },
 100         };
 101         unsafe {
 102             (*cache.head).next = cache.head;
 103             (*cache.head).prev = cache.head;
 104         }
 105         return cache;
 106     }
 107
 108     /// Put a key-value pair into cache.
 109     pub fn put(&mut self, k: K, v: V) {
 110         let (node_ptr, node_opt) = match self.map.find_mut(&KeyRef{k: &k}) {
 111             Some(node) => {
 112                 node.value = v;
 113                 let node_ptr: *mut LruEntry<K, V> = &mut **node;
 114                 (node_ptr, None)
 115             }
 116             None => {
 117                 let mut node = box LruEntry::new(k, v);
 118                 let node_ptr: *mut LruEntry<K, V> = &mut *node;
 119                 (node_ptr, Some(node))
 120             }
 121         };
 122         match node_opt {
 123             None => {
 124                 // Existing node, just update LRU position
 125                 self.detach(node_ptr);
 126                 self.attach(node_ptr);
 127             }
 128             Some(node) => {
 129                 let keyref = unsafe { &(*node_ptr).key };
 130                 self.map.swap(KeyRef{k: keyref}, node);
 131                 self.attach(node_ptr);
 132                 if self.len() > self.capacity() {
 133                     self.remove_lru();
 134                 }
 135             }
 136         }
 137     }
 138
 139     /// Return a value corresponding to the key in the cache.
 140     pub fn get<'a>(&'a mut self, k: &K) -> Option<&'a V> {
 141         let (value, node_ptr_opt) = match self.map.find_mut(&KeyRef{k: k}) {
 142             None => (None, None),
 143             Some(node) => {
 144                 let node_ptr: *mut LruEntry<K, V> = &mut **node;
 145                 (Some(unsafe { &(*node_ptr).value }), Some(node_ptr))
 146             }
 147         };
 148         match node_ptr_opt {
 149             None => (),
 150             Some(node_ptr) => {
 151                 self.detach(node_ptr);
 152                 self.attach(node_ptr);
 153             }
 154         }
 155         return value;
 156     }
 157
 158     /// Remove and return a value corresponding to the key from the cache.
 159     pub fn pop(&mut self, k: &K) -> Option<V> {
 160         match self.map.pop(&KeyRef{k: k}) {
 161             None => None,
 162             Some(lru_entry) => Some(lru_entry.value)
 163         }
 164     }
 165
 166     /// Return the maximum number of key-value pairs the cache can hold.
 167     pub fn capacity(&self) -> uint {
 168         self.max_size
 169     }
 170
 171     /// Change the number of key-value pairs the cache can hold. Remove
 172     /// least-recently-used key-value pairs if necessary.
 173     pub fn change_capacity(&mut self, capacity: uint) {
 174         for _ in range(capacity, self.len()) {
 175             self.remove_lru();
 176         }
 177         self.max_size = capacity;
 178     }
 179
 180     #[inline]
 181     fn remove_lru(&mut self) {
 182         if self.len() > 0 {
 183             let lru = unsafe { (*self.head).prev };
 184             self.detach(lru);
 185             self.map.pop(&KeyRef{k: unsafe { &(*lru).key }});
 186         }
 187     }
 188
 189     #[inline]
 190     fn detach(&mut self, node: *mut LruEntry<K, V>) {
 191         unsafe {
 192             (*(*node).prev).next = (*node).next;
 193             (*(*node).next).prev = (*node).prev;
 194         }
 195     }
 196
 197     #[inline]
 198     fn attach(&mut self, node: *mut LruEntry<K, V>) {
 199         unsafe {
 200             (*node).next = (*self.head).next;
 201             (*node).prev = self.head;
 202             (*self.head).next = node;
 203             (*(*node).next).prev = node;
 204         }
 205     }
 206 }
 207
 208 impl<A: fmt::Show + Hash + Eq, B: fmt::Show> fmt::Show for LruCache<A, B> {
 209     /// Return a string that lists the key-value pairs from most-recently
 210     /// used to least-recently used.
 211     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 212         try!(write!(f, "{{"));
 213         let mut cur = self.head;
 214         for i in range(0, self.len()) {
 215             if i > 0 { try!(write!(f, ", ")) }
 216             unsafe {
 217                 cur = (*cur).next;
 218                 try!(write!(f, "{}", (*cur).key));
 219             }
 220             try!(write!(f, ": "));
 221             unsafe {
 222                 try!(write!(f, "{}", (*cur).value));
 223             }
 224         }
 225         write!(f, r"}}")
 226     }
 227 }
 228
 229 impl<K: Hash + Eq, V> Collection for LruCache<K, V> {
 230     /// Return the number of key-value pairs in the cache.
 231     fn len(&self) -> uint {
 232         self.map.len()
 233     }
 234 }
 235
 236 impl<K: Hash + Eq, V> Mutable for LruCache<K, V> {
 237     /// Clear the cache of all key-value pairs.
 238     fn clear(&mut self) {
 239         self.map.clear();
 240     }
 241 }
 242
 243 #[unsafe_destructor]
 244 impl<K, V> Drop for LruCache<K, V> {
 245     fn drop(&mut self) {
 246         unsafe {
 247             let node: Box<LruEntry<K, V>> = mem::transmute(self.head);
 248             // Prevent compiler from trying to drop the un-initialized field in the sigil node.
 249             let box LruEntry { key: k, value: v, .. } = node;
 250             mem::forget(k);
 251             mem::forget(v);
 252         }
 253     }
 254 }
 255
 256 #[cfg(test)]
 257 mod tests {
 258     use prelude::*;
 259     use super::LruCache;
 260
 261     fn assert_opt_eq<V: PartialEq>(opt: Option<&V>, v: V) {
 262         assert!(opt.is_some());
 263         assert!(opt.unwrap() == &v);
 264     }
 265
 266     #[test]
 267     fn test_put_and_get() {
 268         let mut cache: LruCache<int, int> = LruCache::new(2);
 269         cache.put(1, 10);
 270         cache.put(2, 20);
 271         assert_opt_eq(cache.get(&1), 10);
 272         assert_opt_eq(cache.get(&2), 20);
 273         assert_eq!(cache.len(), 2);
 274     }
 275
 276     #[test]
 277     fn test_put_update() {
 278         let mut cache: LruCache<String, Vec<u8>> = LruCache::new(1);
 279         cache.put("1".to_string(), vec![10, 10]);
 280         cache.put("1".to_string(), vec![10, 19]);
 281         assert_opt_eq(cache.get(&"1".to_string()), vec![10, 19]);
 282         assert_eq!(cache.len(), 1);
 283     }
 284
 285     #[test]
 286     fn test_expire_lru() {
 287         let mut cache: LruCache<String, String> = LruCache::new(2);
 288         cache.put("foo1".to_string(), "bar1".to_string());
 289         cache.put("foo2".to_string(), "bar2".to_string());
 290         cache.put("foo3".to_string(), "bar3".to_string());
 291         assert!(cache.get(&"foo1".to_string()).is_none());
 292         cache.put("foo2".to_string(), "bar2update".to_string());
 293         cache.put("foo4".to_string(), "bar4".to_string());
 294         assert!(cache.get(&"foo3".to_string()).is_none());
 295     }
 296
 297     #[test]
 298     fn test_pop() {
 299         let mut cache: LruCache<int, int> = LruCache::new(2);
 300         cache.put(1, 10);
 301         cache.put(2, 20);
 302         assert_eq!(cache.len(), 2);
 303         let opt1 = cache.pop(&1);
 304         assert!(opt1.is_some());
 305         assert_eq!(opt1.unwrap(), 10);
 306         assert!(cache.get(&1).is_none());
 307         assert_eq!(cache.len(), 1);
 308     }
 309
 310     #[test]
 311     fn test_change_capacity() {
 312         let mut cache: LruCache<int, int> = LruCache::new(2);
 313         assert_eq!(cache.capacity(), 2);
 314         cache.put(1, 10);
 315         cache.put(2, 20);
 316         cache.change_capacity(1);
 317         assert!(cache.get(&1).is_none());
 318         assert_eq!(cache.capacity(), 1);
 319     }
 320
 321     #[test]
 322     fn test_to_str() {
 323         let mut cache: LruCache<int, int> = LruCache::new(3);
 324         cache.put(1, 10);
 325         cache.put(2, 20);
 326         cache.put(3, 30);
 327         assert_eq!(cache.to_str(), "{3: 30, 2: 20, 1: 10}".to_string());
 328         cache.put(2, 22);
 329         assert_eq!(cache.to_str(), "{2: 22, 3: 30, 1: 10}".to_string());
 330         cache.put(6, 60);
 331         assert_eq!(cache.to_str(), "{6: 60, 2: 22, 3: 30}".to_string());
 332         cache.get(&3);
 333         assert_eq!(cache.to_str(), "{3: 30, 6: 60, 2: 22}".to_string());
 334         cache.change_capacity(2);
 335         assert_eq!(cache.to_str(), "{3: 30, 6: 60}".to_string());
 336     }
 337
 338     #[test]
 339     fn test_clear() {
 340         let mut cache: LruCache<int, int> = LruCache::new(2);
 341         cache.put(1, 10);
 342         cache.put(2, 20);
 343         cache.clear();
 344         assert!(cache.get(&1).is_none());
 345         assert!(cache.get(&2).is_none());
 346         assert_eq!(cache.to_str(), "{}".to_string());
 347     }
 348 }