1extern crate tetsy_util_mem;
22extern crate lru_cache;
23
24use tetsy_util_mem::{MallocSizeOf, MallocSizeOfExt};
25use lru_cache::LruCache;
26
27use std::hash::Hash;
28
29const INITIAL_CAPACITY: usize = 4;
30
31pub struct MemoryLruCache<K: Eq + Hash, V> {
33 inner: LruCache<K, V>,
34 cur_size: usize,
35 max_size: usize,
36}
37
38fn heap_size_of<T: MallocSizeOf>(val: &T) -> usize {
40 ::std::mem::size_of::<T>() + val.malloc_size_of()
41}
42
43impl<K: Eq + Hash, V: MallocSizeOf> MemoryLruCache<K, V> {
44 pub fn new(max_size: usize) -> Self {
46 MemoryLruCache {
47 inner: LruCache::new(INITIAL_CAPACITY),
48 max_size: max_size,
49 cur_size: 0,
50 }
51 }
52
53 pub fn insert(&mut self, key: K, val: V) {
55 let cap = self.inner.capacity();
56
57 if self.inner.len() == cap && self.cur_size < self.max_size {
60 self.inner.set_capacity(cap * 2);
61 }
62
63 self.cur_size += heap_size_of(&val);
64
65 if let Some(lru) = self.inner.insert(key, val) {
67 self.cur_size -= heap_size_of(&lru);
68 }
69
70 while self.cur_size > self.max_size {
72 match self.inner.remove_lru() {
73 Some((_, v)) => self.cur_size -= heap_size_of(&v),
74 _ => break,
75 }
76 }
77 }
78
79 pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
82 self.inner.get_mut(key)
83 }
84
85 pub fn current_size(&self) -> usize {
87 self.cur_size
88 }
89
90 pub fn backstore(&self) -> &LruCache<K, V> {
92 &self.inner
93 }
94}
95
96#[cfg(test)]
97mod tests {
98 use super::*;
99
100 #[test]
101 fn it_works() {
102 let mut cache = MemoryLruCache::new(256);
103 let val1 = vec![0u8; 100];
104 let size1 = heap_size_of(&val1);
105 cache.insert("hello", val1);
106
107 assert_eq!(cache.current_size(), size1);
108
109 let val2 = vec![0u8; 210];
110 let size2 = heap_size_of(&val2);
111 cache.insert("world", val2);
112
113 assert!(cache.get_mut(&"hello").is_none());
114 assert!(cache.get_mut(&"world").is_some());
115
116 assert_eq!(cache.current_size(), size2);
117 }
118}