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.insert(1, 10);
  24 //! cache.insert(2, 20);
  25 //! cache.insert(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.insert(2, 22);
  31 //! assert_eq!(*cache.get(&2).unwrap(), 22);
  32 //!
  33 //! cache.insert(6, 60);
  34 //! assert!(cache.get(&3).is_none());
  35 //!
  36 //! cache.set_capacity(1);
  37 //! assert!(cache.get(&2).is_none());
  38 //! ```
  39
  40 use cmp::{PartialEq, Eq};
  41 use collections::HashMap;
  42 use fmt;
  43 use hash::Hash;
  44 use iter::{range, Iterator, Extend};
  45 use mem;
  46 use ops::Drop;
  47 use option::Option;
  48 use option::Option::{Some, None};
  49 use boxed::Box;
  50 use ptr;
  51 use result::Result::{Ok, Err};
  52
  53 // FIXME(conventions): implement iterators?
  54 // FIXME(conventions): implement indexing?
  55
  56 struct KeyRef<K> { k: *const K }
  57
  58 struct LruEntry<K, V> {
  59     next: *mut LruEntry<K, V>,
  60     prev: *mut LruEntry<K, V>,
  61     key: K,
  62     value: V,
  63 }
  64
  65 /// An LRU Cache.
  66 pub struct LruCache<K, V> {
  67     map: HashMap<KeyRef<K>, Box<LruEntry<K, V>>>,
  68     max_size: uint,
  69     head: *mut LruEntry<K, V>,
  70 }
  71
  72 impl<S, K: Hash<S>> Hash<S> for KeyRef<K> {
  73     fn hash(&self, state: &mut S) {
  74         unsafe { (*self.k).hash(state) }
  75     }
  76 }
  77
  78 impl<K: PartialEq> PartialEq for KeyRef<K> {
  79     fn eq(&self, other: &KeyRef<K>) -> bool {
  80         unsafe{ (*self.k).eq(&*other.k) }
  81     }
  82 }
  83
  84 impl<K: Eq> Eq for KeyRef<K> {}
  85
  86 impl<K, V> LruEntry<K, V> {
  87     fn new(k: K, v: V) -> LruEntry<K, V> {
  88         LruEntry {
  89             key: k,
  90             value: v,
  91             next: ptr::null_mut(),
  92             prev: ptr::null_mut(),
  93         }
  94     }
  95 }
  96
  97 impl<K: Hash + Eq, V> LruCache<K, V> {
  98     /// Create an LRU Cache that holds at most `capacity` items.
  99     ///
 100     /// # Example
 101     ///
 102     /// ```
 103     /// use std::collections::LruCache;
 104     /// let mut cache: LruCache<int, &str> = LruCache::new(10);
 105     /// ```
 106     #[unstable = "matches collection reform specification, waiting for dust to settle"]
 107     pub fn new(capacity: uint) -> LruCache<K, V> {
 108         let cache = LruCache {
 109             map: HashMap::new(),
 110             max_size: capacity,
 111             head: unsafe{ mem::transmute(box mem::uninitialized::<LruEntry<K, V>>()) },
 112         };
 113         unsafe {
 114             (*cache.head).next = cache.head;
 115             (*cache.head).prev = cache.head;
 116         }
 117         return cache;
 118     }
 119
 120     /// Deprecated: Replaced with `insert`.
 121     #[deprecated = "Replaced with `insert`"]
 122     pub fn put(&mut self, k: K, v: V) {
 123         self.insert(k, v);
 124     }
 125
 126     /// Inserts a key-value pair into the cache. If the key already existed, the old value is
 127     /// returned.
 128     ///
 129     /// # Example
 130     ///
 131     /// ```
 132     /// use std::collections::LruCache;
 133     /// let mut cache = LruCache::new(2);
 134     ///
 135     /// cache.insert(1i, "a");
 136     /// cache.insert(2, "b");
 137     /// assert_eq!(cache.get(&1), Some(&"a"));
 138     /// assert_eq!(cache.get(&2), Some(&"b"));
 139     /// ```
 140     #[unstable = "matches collection reform specification, waiting for dust to settle"]
 141     pub fn insert(&mut self, k: K, v: V) -> Option<V> {
 142         let (node_ptr, node_opt, old_val) = match self.map.get_mut(&KeyRef{k: &k}) {
 143             Some(node) => {
 144                 let old_val = mem::replace(&mut node.value, v);
 145                 let node_ptr: *mut LruEntry<K, V> = &mut **node;
 146                 (node_ptr, None, Some(old_val))
 147             }
 148             None => {
 149                 let mut node = box LruEntry::new(k, v);
 150                 let node_ptr: *mut LruEntry<K, V> = &mut *node;
 151                 (node_ptr, Some(node), None)
 152             }
 153         };
 154         match node_opt {
 155             None => {
 156                 // Existing node, just update LRU position
 157                 self.detach(node_ptr);
 158                 self.attach(node_ptr);
 159             }
 160             Some(node) => {
 161                 let keyref = unsafe { &(*node_ptr).key };
 162                 self.map.insert(KeyRef{k: keyref}, node);
 163                 self.attach(node_ptr);
 164                 if self.len() > self.capacity() {
 165                     self.remove_lru();
 166                 }
 167             }
 168         }
 169         old_val
 170     }
 171
 172     /// Return a value corresponding to the key in the cache.
 173     ///
 174     /// # Example
 175     ///
 176     /// ```
 177     /// use std::collections::LruCache;
 178     /// let mut cache = LruCache::new(2);
 179     ///
 180     /// cache.insert(1i, "a");
 181     /// cache.insert(2, "b");
 182     /// cache.insert(2, "c");
 183     /// cache.insert(3, "d");
 184     ///
 185     /// assert_eq!(cache.get(&1), None);
 186     /// assert_eq!(cache.get(&2), Some(&"c"));
 187     /// ```
 188     #[unstable = "matches collection reform specification, waiting for dust to settle"]
 189     pub fn get(&mut self, k: &K) -> Option<&V> {
 190         let (value, node_ptr_opt) = match self.map.get_mut(&KeyRef{k: k}) {
 191             None => (None, None),
 192             Some(node) => {
 193                 let node_ptr: *mut LruEntry<K, V> = &mut **node;
 194                 (Some(unsafe { &(*node_ptr).value }), Some(node_ptr))
 195             }
 196         };
 197         match node_ptr_opt {
 198             None => (),
 199             Some(node_ptr) => {
 200                 self.detach(node_ptr);
 201                 self.attach(node_ptr);
 202             }
 203         }
 204         return value;
 205     }
 206
 207     /// Deprecated: Renamed to `remove`.
 208     #[deprecated = "Renamed to `remove`"]
 209     pub fn pop(&mut self, k: &K) -> Option<V> {
 210         self.remove(k)
 211     }
 212
 213     /// Remove and return a value corresponding to the key from the cache.
 214     ///
 215     /// # Example
 216     ///
 217     /// ```
 218     /// use std::collections::LruCache;
 219     /// let mut cache = LruCache::new(2);
 220     ///
 221     /// cache.insert(2i, "a");
 222     ///
 223     /// assert_eq!(cache.remove(&1), None);
 224     /// assert_eq!(cache.remove(&2), Some("a"));
 225     /// assert_eq!(cache.remove(&2), None);
 226     /// assert_eq!(cache.len(), 0);
 227     /// ```
 228     #[unstable = "matches collection reform specification, waiting for dust to settle"]
 229     pub fn remove(&mut self, k: &K) -> Option<V> {
 230         match self.map.remove(&KeyRef{k: k}) {
 231             None => None,
 232             Some(lru_entry) => Some(lru_entry.value)
 233         }
 234     }
 235
 236     /// Return the maximum number of key-value pairs the cache can hold.
 237     ///
 238     /// # Example
 239     ///
 240     /// ```
 241     /// use std::collections::LruCache;
 242     /// let mut cache: LruCache<int, &str> = LruCache::new(2);
 243     /// assert_eq!(cache.capacity(), 2);
 244     /// ```
 245     #[unstable = "matches collection reform specification, waiting for dust to settle"]
 246     pub fn capacity(&self) -> uint {
 247         self.max_size
 248     }
 249
 250     /// Deprecated: Renamed to `set_capacity`.
 251     #[deprecated = "Renamed to `set_capacity`"]
 252     pub fn change_capacity(&mut self, capacity: uint) {
 253         self.set_capacity(capacity)
 254     }
 255
 256     /// Change the number of key-value pairs the cache can hold. Remove
 257     /// least-recently-used key-value pairs if necessary.
 258     ///
 259     /// # Example
 260     ///
 261     /// ```
 262     /// use std::collections::LruCache;
 263     /// let mut cache = LruCache::new(2);
 264     ///
 265     /// cache.insert(1i, "a");
 266     /// cache.insert(2, "b");
 267     /// cache.insert(3, "c");
 268     ///
 269     /// assert_eq!(cache.get(&1), None);
 270     /// assert_eq!(cache.get(&2), Some(&"b"));
 271     /// assert_eq!(cache.get(&3), Some(&"c"));
 272     ///
 273     /// cache.set_capacity(3);
 274     /// cache.insert(1i, "a");
 275     /// cache.insert(2, "b");
 276     ///
 277     /// assert_eq!(cache.get(&1), Some(&"a"));
 278     /// assert_eq!(cache.get(&2), Some(&"b"));
 279     /// assert_eq!(cache.get(&3), Some(&"c"));
 280     ///
 281     /// cache.set_capacity(1);
 282     ///
 283     /// assert_eq!(cache.get(&1), None);
 284     /// assert_eq!(cache.get(&2), None);
 285     /// assert_eq!(cache.get(&3), Some(&"c"));
 286     /// ```
 287     #[unstable = "matches collection reform specification, waiting for dust to settle"]
 288     pub fn set_capacity(&mut self, capacity: uint) {
 289         for _ in range(capacity, self.len()) {
 290             self.remove_lru();
 291         }
 292         self.max_size = capacity;
 293     }
 294
 295     #[inline]
 296     fn remove_lru(&mut self) {
 297         if self.len() > 0 {
 298             let lru = unsafe { (*self.head).prev };
 299             self.detach(lru);
 300             self.map.remove(&KeyRef{k: unsafe { &(*lru).key }});
 301         }
 302     }
 303
 304     #[inline]
 305     fn detach(&mut self, node: *mut LruEntry<K, V>) {
 306         unsafe {
 307             (*(*node).prev).next = (*node).next;
 308             (*(*node).next).prev = (*node).prev;
 309         }
 310     }
 311
 312     #[inline]
 313     fn attach(&mut self, node: *mut LruEntry<K, V>) {
 314         unsafe {
 315             (*node).next = (*self.head).next;
 316             (*node).prev = self.head;
 317             (*self.head).next = node;
 318             (*(*node).next).prev = node;
 319         }
 320     }
 321
 322     /// Return the number of key-value pairs in the cache.
 323     #[unstable = "matches collection reform specification, waiting for dust to settle"]
 324     pub fn len(&self) -> uint { self.map.len() }
 325
 326     /// Returns whether the cache is currently empty.
 327     #[unstable = "matches collection reform specification, waiting for dust to settle"]
 328     pub fn is_empty(&self) -> bool { self.len() == 0 }
 329
 330     /// Clear the cache of all key-value pairs.
 331     #[unstable = "matches collection reform specification, waiting for dust to settle"]
 332     pub fn clear(&mut self) { self.map.clear(); }
 333
 334 }
 335
 336 impl<K: Hash + Eq, V> Extend<(K, V)> for LruCache<K, V> {
 337     fn extend<T: Iterator<(K, V)>>(&mut self, mut iter: T) {
 338         for (k, v) in iter{
 339             self.insert(k, v);
 340         }
 341     }
 342 }
 343
 344 impl<A: fmt::Show + Hash + Eq, B: fmt::Show> fmt::Show for LruCache<A, B> {
 345     /// Return a string that lists the key-value pairs from most-recently
 346     /// used to least-recently used.
 347     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 348         try!(write!(f, "{{"));
 349         let mut cur = self.head;
 350         for i in range(0, self.len()) {
 351             if i > 0 { try!(write!(f, ", ")) }
 352             unsafe {
 353                 cur = (*cur).next;
 354                 try!(write!(f, "{}", (*cur).key));
 355             }
 356             try!(write!(f, ": "));
 357             unsafe {
 358                 try!(write!(f, "{}", (*cur).value));
 359             }
 360         }
 361         write!(f, r"}}")
 362     }
 363 }
 364
 365 #[unsafe_destructor]
 366 impl<K, V> Drop for LruCache<K, V> {
 367     fn drop(&mut self) {
 368         unsafe {
 369             let node: Box<LruEntry<K, V>> = mem::transmute(self.head);
 370             // Prevent compiler from trying to drop the un-initialized field in the sigil node.
 371             let box internal_node = node;
 372             let LruEntry { next: _, prev: _, key: k, value: v } = internal_node;
 373             mem::forget(k);
 374             mem::forget(v);
 375         }
 376     }
 377 }
 378
 379 #[cfg(test)]
 380 mod tests {
 381     use prelude::*;
 382     use super::LruCache;
 383
 384     fn assert_opt_eq<V: PartialEq>(opt: Option<&V>, v: V) {
 385         assert!(opt.is_some());
 386         assert!(opt.unwrap() == &v);
 387     }
 388
 389     #[test]
 390     fn test_put_and_get() {
 391         let mut cache: LruCache<int, int> = LruCache::new(2);
 392         cache.insert(1, 10);
 393         cache.insert(2, 20);
 394         assert_opt_eq(cache.get(&1), 10);
 395         assert_opt_eq(cache.get(&2), 20);
 396         assert_eq!(cache.len(), 2);
 397     }
 398
 399     #[test]
 400     fn test_put_update() {
 401         let mut cache: LruCache<String, Vec<u8>> = LruCache::new(1);
 402         cache.insert("1".to_string(), vec![10, 10]);
 403         cache.insert("1".to_string(), vec![10, 19]);
 404         assert_opt_eq(cache.get(&"1".to_string()), vec![10, 19]);
 405         assert_eq!(cache.len(), 1);
 406     }
 407
 408     #[test]
 409     fn test_expire_lru() {
 410         let mut cache: LruCache<String, String> = LruCache::new(2);
 411         cache.insert("foo1".to_string(), "bar1".to_string());
 412         cache.insert("foo2".to_string(), "bar2".to_string());
 413         cache.insert("foo3".to_string(), "bar3".to_string());
 414         assert!(cache.get(&"foo1".to_string()).is_none());
 415         cache.insert("foo2".to_string(), "bar2update".to_string());
 416         cache.insert("foo4".to_string(), "bar4".to_string());
 417         assert!(cache.get(&"foo3".to_string()).is_none());
 418     }
 419
 420     #[test]
 421     fn test_pop() {
 422         let mut cache: LruCache<int, int> = LruCache::new(2);
 423         cache.insert(1, 10);
 424         cache.insert(2, 20);
 425         assert_eq!(cache.len(), 2);
 426         let opt1 = cache.remove(&1);
 427         assert!(opt1.is_some());
 428         assert_eq!(opt1.unwrap(), 10);
 429         assert!(cache.get(&1).is_none());
 430         assert_eq!(cache.len(), 1);
 431     }
 432
 433     #[test]
 434     fn test_change_capacity() {
 435         let mut cache: LruCache<int, int> = LruCache::new(2);
 436         assert_eq!(cache.capacity(), 2);
 437         cache.insert(1, 10);
 438         cache.insert(2, 20);
 439         cache.set_capacity(1);
 440         assert!(cache.get(&1).is_none());
 441         assert_eq!(cache.capacity(), 1);
 442     }
 443
 444     #[test]
 445     fn test_to_string() {
 446         let mut cache: LruCache<int, int> = LruCache::new(3);
 447         cache.insert(1, 10);
 448         cache.insert(2, 20);
 449         cache.insert(3, 30);
 450         assert_eq!(cache.to_string(), "{3: 30, 2: 20, 1: 10}");
 451         cache.insert(2, 22);
 452         assert_eq!(cache.to_string(), "{2: 22, 3: 30, 1: 10}");
 453         cache.insert(6, 60);
 454         assert_eq!(cache.to_string(), "{6: 60, 2: 22, 3: 30}");
 455         cache.get(&3);
 456         assert_eq!(cache.to_string(), "{3: 30, 6: 60, 2: 22}");
 457         cache.set_capacity(2);
 458         assert_eq!(cache.to_string(), "{3: 30, 6: 60}");
 459     }
 460
 461     #[test]
 462     fn test_clear() {
 463         let mut cache: LruCache<int, int> = LruCache::new(2);
 464         cache.insert(1, 10);
 465         cache.insert(2, 20);
 466         cache.clear();
 467         assert!(cache.get(&1).is_none());
 468         assert!(cache.get(&2).is_none());
 469         assert_eq!(cache.to_string(), "{}");
 470     }
 471 }