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     #[cfg(stage0)]
 212     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 213         try!(write!(f, r"\{"));
 214         let mut cur = self.head;
 215         for i in range(0, self.len()) {
 216             if i > 0 { try!(write!(f, ", ")) }
 217             unsafe {
 218                 cur = (*cur).next;
 219                 try!(write!(f, "{}", (*cur).key));
 220             }
 221             try!(write!(f, ": "));
 222             unsafe {
 223                 try!(write!(f, "{}", (*cur).value));
 224             }
 225         }
 226         write!(f, r"\}")
 227     }
 228     /// Return a string that lists the key-value pairs from most-recently
 229     /// used to least-recently used.
 230     #[cfg(not(stage0))]
 231     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 232         try!(write!(f, "{{"));
 233         let mut cur = self.head;
 234         for i in range(0, self.len()) {
 235             if i > 0 { try!(write!(f, ", ")) }
 236             unsafe {
 237                 cur = (*cur).next;
 238                 try!(write!(f, "{}", (*cur).key));
 239             }
 240             try!(write!(f, ": "));
 241             unsafe {
 242                 try!(write!(f, "{}", (*cur).value));
 243             }
 244         }
 245         write!(f, r"}}")
 246     }
 247 }
 248
 249 impl<K: Hash + Eq, V> Collection for LruCache<K, V> {
 250     /// Return the number of key-value pairs in the cache.
 251     fn len(&self) -> uint {
 252         self.map.len()
 253     }
 254 }
 255
 256 impl<K: Hash + Eq, V> Mutable for LruCache<K, V> {
 257     /// Clear the cache of all key-value pairs.
 258     fn clear(&mut self) {
 259         self.map.clear();
 260     }
 261 }
 262
 263 #[unsafe_destructor]
 264 impl<K, V> Drop for LruCache<K, V> {
 265     fn drop(&mut self) {
 266         unsafe {
 267             let node: Box<LruEntry<K, V>> = mem::transmute(self.head);
 268             // Prevent compiler from trying to drop the un-initialized field in the sigil node.
 269             let box LruEntry { key: k, value: v, .. } = node;
 270             mem::forget(k);
 271             mem::forget(v);
 272         }
 273     }
 274 }
 275
 276 #[cfg(test)]
 277 mod tests {
 278     use prelude::*;
 279     use super::LruCache;
 280
 281     fn assert_opt_eq<V: PartialEq>(opt: Option<&V>, v: V) {
 282         assert!(opt.is_some());
 283         assert!(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<String, Vec<u8>> = LruCache::new(1);
 299         cache.put("1".to_string(), vec![10, 10]);
 300         cache.put("1".to_string(), vec![10, 19]);
 301         assert_opt_eq(cache.get(&"1".to_string()), vec![10, 19]);
 302         assert_eq!(cache.len(), 1);
 303     }
 304
 305     #[test]
 306     fn test_expire_lru() {
 307         let mut cache: LruCache<String, String> = LruCache::new(2);
 308         cache.put("foo1".to_string(), "bar1".to_string());
 309         cache.put("foo2".to_string(), "bar2".to_string());
 310         cache.put("foo3".to_string(), "bar3".to_string());
 311         assert!(cache.get(&"foo1".to_string()).is_none());
 312         cache.put("foo2".to_string(), "bar2update".to_string());
 313         cache.put("foo4".to_string(), "bar4".to_string());
 314         assert!(cache.get(&"foo3".to_string()).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}".to_string());
 348         cache.put(2, 22);
 349         assert_eq!(cache.to_str(), "{2: 22, 3: 30, 1: 10}".to_string());
 350         cache.put(6, 60);
 351         assert_eq!(cache.to_str(), "{6: 60, 2: 22, 3: 30}".to_string());
 352         cache.get(&3);
 353         assert_eq!(cache.to_str(), "{3: 30, 6: 60, 2: 22}".to_string());
 354         cache.change_capacity(2);
 355         assert_eq!(cache.to_str(), "{3: 30, 6: 60}".to_string());
 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(), "{}".to_string());
 367     }
 368 }