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