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