src/libcollections/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 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 std::cast;
  41 use std::container::Container;
  42 use std::hash::{Hash, sip};
  43 use std::fmt;
  44 use std::ptr;
  45
  46 use HashMap;
  47
  48 struct KeyRef<K> { k: *K }
  49
  50 struct LruEntry<K, V> {
  51     key: Option<K>,
  52     value: Option<V>,
  53     next: *mut LruEntry<K, V>,
  54     prev: *mut LruEntry<K, V>,
  55 }
  56
  57 /// An LRU Cache.
  58 pub struct LruCache<K, V> {
  59     priv map: HashMap<KeyRef<K>, ~LruEntry<K, V>>,
  60     priv max_size: uint,
  61     priv head: *mut LruEntry<K, V>,
  62     priv tail: *mut LruEntry<K, V>,
  63 }
  64
  65 impl<K: Hash> Hash for KeyRef<K> {
  66     fn hash(&self, s: &mut sip::SipState) {
  67         unsafe {(*self.k).hash(s)}
  68     }
  69 }
  70
  71 impl<K: Eq> Eq for KeyRef<K> {
  72     fn eq(&self, other: &KeyRef<K>) -> bool {
  73         unsafe{ (*self.k).eq(&*other.k) }
  74     }
  75 }
  76
  77 impl<K, V> LruEntry<K, V> {
  78     fn new() -> LruEntry<K, V> {
  79         LruEntry {
  80             key: None,
  81             value: None,
  82             next: ptr::mut_null(),
  83             prev: ptr::mut_null(),
  84         }
  85     }
  86
  87     fn with_key_value(k: K, v: V) -> LruEntry<K, V> {
  88         LruEntry {
  89             key: Some(k),
  90             value: Some(v),
  91             next: ptr::mut_null(),
  92             prev: ptr::mut_null(),
  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     pub fn new(capacity: uint) -> LruCache<K, V> {
 100         let cache = LruCache {
 101             map: HashMap::new(),
 102             max_size: capacity,
 103             head: unsafe{ cast::transmute(~LruEntry::<K, V>::new()) },
 104             tail: unsafe{ cast::transmute(~LruEntry::<K, V>::new()) },
 105         };
 106         unsafe {
 107             (*cache.head).next = cache.tail;
 108             (*cache.tail).prev = cache.head;
 109         }
 110         return cache;
 111     }
 112
 113     /// Put a key-value pair into cache.
 114     pub fn put(&mut self, k: K, v: V) {
 115         let mut key_existed = false;
 116         let (node_ptr, node_opt) = match self.map.find_mut(&KeyRef{k: &k}) {
 117             Some(node) => {
 118                 key_existed = true;
 119                 node.value = Some(v);
 120                 let node_ptr: *mut LruEntry<K, V> = &mut **node;
 121                 (node_ptr, None)
 122             }
 123             None => {
 124                 let mut node = ~LruEntry::with_key_value(k, v);
 125                 let node_ptr: *mut LruEntry<K, V> = &mut *node;
 126                 (node_ptr, Some(node))
 127             }
 128         };
 129         if key_existed {
 130             self.detach(node_ptr);
 131             self.attach(node_ptr);
 132         } else {
 133             let keyref = unsafe { (*node_ptr).key.as_ref().unwrap() };
 134             self.map.swap(KeyRef{k: keyref}, node_opt.unwrap());
 135             self.attach(node_ptr);
 136             if self.len() > self.capacity() {
 137                 self.remove_lru();
 138             }
 139         }
 140     }
 141
 142     /// Return a value corresponding to the key in the cache.
 143     pub fn get<'a>(&'a mut self, k: &K) -> Option<&'a V> {
 144         let (value, node_ptr_opt) = match self.map.find_mut(&KeyRef{k: k}) {
 145             None => (None, None),
 146             Some(node) => {
 147                 let node_ptr: *mut LruEntry<K, V> = &mut **node;
 148                 unsafe {
 149                     match (*node_ptr).value {
 150                         None => (None, None),
 151                         Some(ref value) => (Some(value), Some(node_ptr))
 152                     }
 153                 }
 154             }
 155         };
 156         match node_ptr_opt {
 157             None => (),
 158             Some(node_ptr) => {
 159                 self.detach(node_ptr);
 160                 self.attach(node_ptr);
 161             }
 162         }
 163         return value;
 164     }
 165
 166     /// Remove and return a value corresponding to the key from the cache.
 167     pub fn pop(&mut self, k: &K) -> Option<V> {
 168         match self.map.pop(&KeyRef{k: k}) {
 169             None => None,
 170             Some(lru_entry) => lru_entry.value
 171         }
 172     }
 173
 174     /// Return the maximum number of key-value pairs the cache can hold.
 175     pub fn capacity(&self) -> uint {
 176         self.max_size
 177     }
 178
 179     /// Change the number of key-value pairs the cache can hold. Remove
 180     /// least-recently-used key-value pairs if necessary.
 181     pub fn change_capacity(&mut self, capacity: uint) {
 182         for _ in range(capacity, self.len()) {
 183             self.remove_lru();
 184         }
 185         self.max_size = capacity;
 186     }
 187
 188     #[inline]
 189     fn remove_lru(&mut self) {
 190         if self.len() > 0 {
 191             let lru = unsafe { (*self.tail).prev };
 192             self.detach(lru);
 193             unsafe {
 194                 match (*lru).key {
 195                     None => (),
 196                     Some(ref k) => { self.map.pop(&KeyRef{k: k}); }
 197                 }
 198             }
 199         }
 200     }
 201
 202     #[inline]
 203     fn detach(&mut self, node: *mut LruEntry<K, V>) {
 204         unsafe {
 205             (*(*node).prev).next = (*node).next;
 206             (*(*node).next).prev = (*node).prev;
 207         }
 208     }
 209
 210     #[inline]
 211     fn attach(&mut self, node: *mut LruEntry<K, V>) {
 212         unsafe {
 213             (*node).next = (*self.head).next;
 214             (*node).prev = self.head;
 215             (*self.head).next = node;
 216             (*(*node).next).prev = node;
 217         }
 218     }
 219 }
 220
 221 impl<A: fmt::Show + Hash + Eq, B: fmt::Show> fmt::Show for LruCache<A, B> {
 222     /// Return a string that lists the key-value pairs from most-recently
 223     /// used to least-recently used.
 224     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 225         try!(write!(f.buf, r"\{"));
 226         let mut cur = self.head;
 227         for i in range(0, self.len()) {
 228             if i > 0 { try!(write!(f.buf, ", ")) }
 229             unsafe {
 230                 cur = (*cur).next;
 231                 match (*cur).key {
 232                     // should never print nil
 233                     None => try!(write!(f.buf, "nil")),
 234                     Some(ref k) => try!(write!(f.buf, "{}", *k)),
 235                 }
 236             }
 237             try!(write!(f.buf, ": "));
 238             unsafe {
 239                 match (*cur).value {
 240                     // should never print nil
 241                     None => try!(write!(f.buf, "nil")),
 242                     Some(ref value) => try!(write!(f.buf, "{}", *value)),
 243                 }
 244             }
 245         }
 246         write!(f.buf, r"\}")
 247     }
 248 }
 249
 250 impl<K: Hash + Eq, V> Container for LruCache<K, V> {
 251     /// Return the number of key-value pairs in the cache.
 252     fn len(&self) -> uint {
 253         self.map.len()
 254     }
 255 }
 256
 257 impl<K: Hash + Eq, V> Mutable for LruCache<K, V> {
 258     /// Clear the cache of all key-value pairs.
 259     fn clear(&mut self) {
 260         self.map.clear();
 261     }
 262 }
 263
 264 #[unsafe_destructor]
 265 impl<K, V> Drop for LruCache<K, V> {
 266     fn drop(&mut self) {
 267         unsafe {
 268             let _: ~LruEntry<K, V> = cast::transmute(self.head);
 269             let _: ~LruEntry<K, V> = cast::transmute(self.tail);
 270         }
 271     }
 272 }
 273
 274 #[cfg(test)]
 275 mod tests {
 276     use super::LruCache;
 277
 278     fn assert_opt_eq<V: Eq>(opt: Option<&V>, v: V) {
 279         assert!(opt.is_some());
 280         assert!(opt.unwrap() == &v);
 281     }
 282
 283     #[test]
 284     fn test_put_and_get() {
 285         let mut cache: LruCache<int, int> = LruCache::new(2);
 286         cache.put(1, 10);
 287         cache.put(2, 20);
 288         assert_opt_eq(cache.get(&1), 10);
 289         assert_opt_eq(cache.get(&2), 20);
 290         assert_eq!(cache.len(), 2);
 291     }
 292
 293     #[test]
 294     fn test_put_update() {
 295         let mut cache: LruCache<~str, ~[u8]> = LruCache::new(1);
 296         cache.put(~"1", ~[10, 10]);
 297         cache.put(~"1", ~[10, 19]);
 298         assert_opt_eq(cache.get(&~"1"), ~[10, 19]);
 299         assert_eq!(cache.len(), 1);
 300     }
 301
 302     #[test]
 303     fn test_expire_lru() {
 304         let mut cache: LruCache<~str, ~str> = LruCache::new(2);
 305         cache.put(~"foo1", ~"bar1");
 306         cache.put(~"foo2", ~"bar2");
 307         cache.put(~"foo3", ~"bar3");
 308         assert!(cache.get(&~"foo1").is_none());
 309         cache.put(~"foo2", ~"bar2update");
 310         cache.put(~"foo4", ~"bar4");
 311         assert!(cache.get(&~"foo3").is_none());
 312     }
 313
 314     #[test]
 315     fn test_pop() {
 316         let mut cache: LruCache<int, int> = LruCache::new(2);
 317         cache.put(1, 10);
 318         cache.put(2, 20);
 319         assert_eq!(cache.len(), 2);
 320         let opt1 = cache.pop(&1);
 321         assert!(opt1.is_some());
 322         assert_eq!(opt1.unwrap(), 10);
 323         assert!(cache.get(&1).is_none());
 324         assert_eq!(cache.len(), 1);
 325     }
 326
 327     #[test]
 328     fn test_change_capacity() {
 329         let mut cache: LruCache<int, int> = LruCache::new(2);
 330         assert_eq!(cache.capacity(), 2);
 331         cache.put(1, 10);
 332         cache.put(2, 20);
 333         cache.change_capacity(1);
 334         assert!(cache.get(&1).is_none());
 335         assert_eq!(cache.capacity(), 1);
 336     }
 337
 338     #[test]
 339     fn test_to_str() {
 340         let mut cache: LruCache<int, int> = LruCache::new(3);
 341         cache.put(1, 10);
 342         cache.put(2, 20);
 343         cache.put(3, 30);
 344         assert_eq!(cache.to_str(), ~"{3: 30, 2: 20, 1: 10}");
 345         cache.put(2, 22);
 346         assert_eq!(cache.to_str(), ~"{2: 22, 3: 30, 1: 10}");
 347         cache.put(6, 60);
 348         assert_eq!(cache.to_str(), ~"{6: 60, 2: 22, 3: 30}");
 349         cache.get(&3);
 350         assert_eq!(cache.to_str(), ~"{3: 30, 6: 60, 2: 22}");
 351         cache.change_capacity(2);
 352         assert_eq!(cache.to_str(), ~"{3: 30, 6: 60}");
 353     }
 354
 355     #[test]
 356     fn test_clear() {
 357         let mut cache: LruCache<int, int> = LruCache::new(2);
 358         cache.put(1, 10);
 359         cache.put(2, 20);
 360         cache.clear();
 361         assert!(cache.get(&1).is_none());
 362         assert!(cache.get(&2).is_none());
 363         assert_eq!(cache.to_str(), ~"{}");
 364     }
 365 }