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