From d902e7148bbc2559453d22e7d00f05cee9a3df95 Mon Sep 17 00:00:00 2001 From: Kiet Tran Date: Thu, 31 Oct 2013 20:58:15 -0400 Subject: [PATCH] Add lru_cache to libextra --- src/libextra/lib.rs | 1 + src/libextra/lru_cache.rs | 365 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 366 insertions(+) create mode 100644 src/libextra/lru_cache.rs diff --git a/src/libextra/lib.rs b/src/libextra/lib.rs index 106dc5bf5c9da..733795fc1f796 100644 --- a/src/libextra/lib.rs +++ b/src/libextra/lib.rs @@ -71,6 +71,7 @@ pub mod sort; pub mod dlist; pub mod treemap; pub mod btree; +pub mod lru_cache; // And ... other stuff diff --git a/src/libextra/lru_cache.rs b/src/libextra/lru_cache.rs new file mode 100644 index 0000000000000..ff8748e1946c1 --- /dev/null +++ b/src/libextra/lru_cache.rs @@ -0,0 +1,365 @@ +// Copyright 2013 The Rust Project Developers. See the COPYRIGHT +// file at the top-level directory of this distribution and at +// http://rust-lang.org/COPYRIGHT. +// +// Licensed under the Apache License, Version 2.0 or the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. + + +//! A cache that holds a limited number of key-value pairs. When the +//! capacity of the cache is exceeded, the least-recently-used +//! (where "used" means a look-up or putting the pair into the cache) +//! pair is automatically removed. +//! +//! # Example +//! +//! ```rust +//! let mut cache: LruCache = LruCache::new(2); +//! cache.put(1, 10); +//! cache.put(2, 20); +//! cache.put(3, 30); +//! assert!(cache.get(&1).is_none()); +//! assert_eq!(*cache.get(&2).unwrap(), 20); +//! assert_eq!(*cache.get(&3).unwrap(), 30); +//! +//! cache.put(2, 22); +//! assert_eq!(*cache.get(&2).unwrap(), 22); +//! +//! cache.put(6, 60); +//! assert!(cache.get(&3).is_none()); +//! +//! cache.change_capacity(1); +//! assert!(cache.get(&2).is_none()); +//! ``` + +use std::container::Container; +use std::hashmap::HashMap; +use std::to_bytes::Cb; +use std::ptr; +use std::cast; + +struct KeyRef { priv k: *K } + +struct LruEntry { + priv key: Option, + priv value: Option, + priv next: *mut LruEntry, + priv prev: *mut LruEntry, +} + +/// An LRU Cache. +pub struct LruCache { + priv map: HashMap, ~LruEntry>, + priv max_size: uint, + priv head: *mut LruEntry, + priv tail: *mut LruEntry, +} + +impl IterBytes for KeyRef { + fn iter_bytes(&self, lsb0: bool, f: Cb) -> bool { + unsafe{ (*self.k).iter_bytes(lsb0, f) } + } +} + +impl Eq for KeyRef { + fn eq(&self, other: &KeyRef) -> bool { + unsafe{ (*self.k).eq(&*other.k) } + } +} + +impl LruEntry { + fn new() -> LruEntry { + LruEntry { + key: None, + value: None, + next: ptr::mut_null(), + prev: ptr::mut_null(), + } + } + + fn with_key_value(k: K, v: V) -> LruEntry { + LruEntry { + key: Some(k), + value: Some(v), + next: ptr::mut_null(), + prev: ptr::mut_null(), + } + } +} + +impl LruCache { + /// Create an LRU Cache that holds at most `capacity` items. + pub fn new(capacity: uint) -> LruCache { + let cache = LruCache { + map: HashMap::new(), + max_size: capacity, + head: unsafe{ cast::transmute(~LruEntry::::new()) }, + tail: unsafe{ cast::transmute(~LruEntry::::new()) }, + }; + unsafe { + (*cache.head).next = cache.tail; + (*cache.tail).prev = cache.head; + } + return cache; + } + + /// Put a key-value pair into cache. + pub fn put(&mut self, k: K, v: V) { + let mut key_existed = false; + let (node_ptr, node_opt) = match self.map.find_mut(&KeyRef{k: &k}) { + Some(node) => { + key_existed = true; + node.value = Some(v); + let node_ptr: *mut LruEntry = &mut **node; + (node_ptr, None) + } + None => { + let mut node = ~LruEntry::with_key_value(k, v); + let node_ptr: *mut LruEntry = &mut *node; + (node_ptr, Some(node)) + } + }; + if key_existed { + self.detach(node_ptr); + self.attach(node_ptr); + } else { + let keyref = unsafe { (*node_ptr).key.as_ref().unwrap() }; + self.map.swap(KeyRef{k: keyref}, node_opt.unwrap()); + self.attach(node_ptr); + if self.len() > self.capacity() { + self.remove_lru(); + } + } + } + + /// Return a value corresponding to the key in the cache. + pub fn get<'a>(&'a mut self, k: &K) -> Option<&'a V> { + let (value, node_ptr_opt) = match self.map.find_mut(&KeyRef{k: k}) { + None => (None, None), + Some(node) => { + let node_ptr: *mut LruEntry = &mut **node; + unsafe { + match (*node_ptr).value { + None => (None, None), + Some(ref value) => (Some(value), Some(node_ptr)) + } + } + } + }; + match node_ptr_opt { + None => (), + Some(node_ptr) => { + self.detach(node_ptr); + self.attach(node_ptr); + } + } + return value; + } + + /// Remove and return a value corresponding to the key from the cache. + pub fn pop(&mut self, k: &K) -> Option { + match self.map.pop(&KeyRef{k: k}) { + None => None, + Some(lru_entry) => lru_entry.value + } + } + + /// Return the maximum number of key-value pairs the cache can hold. + pub fn capacity(&self) -> uint { + self.max_size + } + + /// Change the number of key-value pairs the cache can hold. Remove + /// least-recently-used key-value pairs if necessary. + pub fn change_capacity(&mut self, capacity: uint) { + for _ in range(capacity, self.len()) { + self.remove_lru(); + } + self.max_size = capacity; + } + + #[inline] + fn remove_lru(&mut self) { + if self.len() > 0 { + let lru = unsafe { (*self.tail).prev }; + self.detach(lru); + unsafe { + match (*lru).key { + None => (), + Some(ref k) => { self.map.pop(&KeyRef{k: k}); } + } + } + } + } + + #[inline] + fn detach(&mut self, node: *mut LruEntry) { + unsafe { + (*(*node).prev).next = (*node).next; + (*(*node).next).prev = (*node).prev; + } + } + + #[inline] + fn attach(&mut self, node: *mut LruEntry) { + unsafe { + (*node).next = (*self.head).next; + (*node).prev = self.head; + (*self.head).next = node; + (*(*node).next).prev = node; + } + } +} + +impl ToStr for LruCache { + /// Return a string that lists the key-value pairs from most-recently + /// used to least-recently used. + #[inline] + fn to_str(&self) -> ~str { + let mut acc = ~"{"; + let mut cur = self.head; + for i in range(0, self.len()) { + if i > 0 { + acc.push_str(", "); + } + unsafe { + cur = (*cur).next; + match (*cur).key { + // should never print nil + None => acc.push_str("nil"), + Some(ref k) => acc.push_str(k.to_str()) + } + } + acc.push_str(": "); + unsafe { + match (*cur).value { + // should never print nil + None => acc.push_str("nil"), + Some(ref value) => acc.push_str(value.to_str()) + } + } + } + acc.push_char('}'); + acc + } +} + +impl Container for LruCache { + /// Return the number of key-value pairs in the cache. + fn len(&self) -> uint { + self.map.len() + } +} + +impl Mutable for LruCache { + /// Clear the cache of all key-value pairs. + fn clear(&mut self) { + self.map.clear(); + } +} + +#[unsafe_destructor] +impl Drop for LruCache { + fn drop(&mut self) { + unsafe { + let _: ~LruEntry = cast::transmute(self.head); + let _: ~LruEntry = cast::transmute(self.tail); + } + } +} + +#[cfg(test)] +mod tests { + use super::LruCache; + + fn assert_opt_eq(opt: Option<&V>, v: V) { + assert!(opt.is_some()); + assert_eq!(opt.unwrap(), &v); + } + + #[test] + fn test_put_and_get() { + let mut cache: LruCache = LruCache::new(2); + cache.put(1, 10); + cache.put(2, 20); + assert_opt_eq(cache.get(&1), 10); + assert_opt_eq(cache.get(&2), 20); + assert_eq!(cache.len(), 2); + } + + #[test] + fn test_put_update() { + let mut cache: LruCache<~str, ~[u8]> = LruCache::new(1); + cache.put(~"1", ~[10, 10]); + cache.put(~"1", ~[10, 19]); + assert_opt_eq(cache.get(&~"1"), ~[10, 19]); + assert_eq!(cache.len(), 1); + } + + #[test] + fn test_expire_lru() { + let mut cache: LruCache<~str, ~str> = LruCache::new(2); + cache.put(~"foo1", ~"bar1"); + cache.put(~"foo2", ~"bar2"); + cache.put(~"foo3", ~"bar3"); + assert!(cache.get(&~"foo1").is_none()); + cache.put(~"foo2", ~"bar2update"); + cache.put(~"foo4", ~"bar4"); + assert!(cache.get(&~"foo3").is_none()); + } + + #[test] + fn test_pop() { + let mut cache: LruCache = LruCache::new(2); + cache.put(1, 10); + cache.put(2, 20); + assert_eq!(cache.len(), 2); + let opt1 = cache.pop(&1); + assert!(opt1.is_some()); + assert_eq!(opt1.unwrap(), 10); + assert!(cache.get(&1).is_none()); + assert_eq!(cache.len(), 1); + } + + #[test] + fn test_change_capacity() { + let mut cache: LruCache = LruCache::new(2); + assert_eq!(cache.capacity(), 2); + cache.put(1, 10); + cache.put(2, 20); + cache.change_capacity(1); + assert!(cache.get(&1).is_none()); + assert_eq!(cache.capacity(), 1); + } + + #[test] + fn test_to_str() { + let mut cache: LruCache = LruCache::new(3); + cache.put(1, 10); + cache.put(2, 20); + cache.put(3, 30); + assert_eq!(cache.to_str(), ~"{3: 30, 2: 20, 1: 10}"); + cache.put(2, 22); + assert_eq!(cache.to_str(), ~"{2: 22, 3: 30, 1: 10}"); + cache.put(6, 60); + assert_eq!(cache.to_str(), ~"{6: 60, 2: 22, 3: 30}"); + cache.get(&3); + assert_eq!(cache.to_str(), ~"{3: 30, 6: 60, 2: 22}"); + cache.change_capacity(2); + assert_eq!(cache.to_str(), ~"{3: 30, 6: 60}"); + } + + #[test] + fn test_clear() { + let mut cache: LruCache = LruCache::new(2); + cache.put(1, 10); + cache.put(2, 20); + cache.clear(); + assert!(cache.get(&1).is_none()); + assert!(cache.get(&2).is_none()); + assert_eq!(cache.to_str(), ~"{}"); + } +}