Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add lru_cache to libextra #10211

Closed
wants to merge 1 commit into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions src/libextra/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,7 @@ pub mod sort;
pub mod dlist;
pub mod treemap;
pub mod btree;
pub mod lru_cache;

// And ... other stuff

Expand Down
365 changes: 365 additions & 0 deletions src/libextra/lru_cache.rs
Original file line number Diff line number Diff line change
@@ -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 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, 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<int, int> = 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<K> { priv k: *K }

struct LruEntry<K, V> {
priv key: Option<K>,
priv value: Option<V>,
priv next: *mut LruEntry<K, V>,
priv prev: *mut LruEntry<K, V>,
}

/// An LRU Cache.
pub struct LruCache<K, V> {
priv map: HashMap<KeyRef<K>, ~LruEntry<K, V>>,
priv max_size: uint,
priv head: *mut LruEntry<K, V>,
priv tail: *mut LruEntry<K, V>,
}

impl<K: IterBytes> IterBytes for KeyRef<K> {
fn iter_bytes(&self, lsb0: bool, f: Cb) -> bool {
unsafe{ (*self.k).iter_bytes(lsb0, f) }
}
}

impl<K: Eq> Eq for KeyRef<K> {
fn eq(&self, other: &KeyRef<K>) -> bool {
unsafe{ (*self.k).eq(&*other.k) }
}
}

impl<K, V> LruEntry<K, V> {
fn new() -> LruEntry<K, V> {
LruEntry {
key: None,
value: None,
next: ptr::mut_null(),
prev: ptr::mut_null(),
}
}

fn with_key_value(k: K, v: V) -> LruEntry<K, V> {
LruEntry {
key: Some(k),
value: Some(v),
next: ptr::mut_null(),
prev: ptr::mut_null(),
}
}
}

impl<K: IterBytes + Eq, V> LruCache<K, V> {
/// Create an LRU Cache that holds at most `capacity` items.
pub fn new(capacity: uint) -> LruCache<K, V> {
let cache = LruCache {
map: HashMap::new(),
max_size: capacity,
head: unsafe{ cast::transmute(~LruEntry::<K, V>::new()) },
tail: unsafe{ cast::transmute(~LruEntry::<K, V>::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<K, V> = &mut **node;
(node_ptr, None)
}
None => {
let mut node = ~LruEntry::with_key_value(k, v);
let node_ptr: *mut LruEntry<K, V> = &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<K, V> = &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<V> {
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<K, V>) {
unsafe {
(*(*node).prev).next = (*node).next;
(*(*node).next).prev = (*node).prev;
}
}

#[inline]
fn attach(&mut self, node: *mut LruEntry<K, V>) {
unsafe {
(*node).next = (*self.head).next;
(*node).prev = self.head;
(*self.head).next = node;
(*(*node).next).prev = node;
}
}
}

impl<A: ToStr + IterBytes + Eq, B: ToStr> ToStr for LruCache<A, B> {
/// 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<K: IterBytes + Eq, V> Container for LruCache<K, V> {
/// Return the number of key-value pairs in the cache.
fn len(&self) -> uint {
self.map.len()
}
}

impl<K: IterBytes + Eq, V> Mutable for LruCache<K, V> {
/// Clear the cache of all key-value pairs.
fn clear(&mut self) {
self.map.clear();
}
}

#[unsafe_destructor]
impl<K, V> Drop for LruCache<K, V> {
fn drop(&mut self) {
unsafe {
let _: ~LruEntry<K, V> = cast::transmute(self.head);
let _: ~LruEntry<K, V> = cast::transmute(self.tail);
}
}
}

#[cfg(test)]
mod tests {
use super::LruCache;

fn assert_opt_eq<V: 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<int, int> = 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<int, int> = 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<int, int> = 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<int, int> = 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<int, int> = 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(), ~"{}");
}
}