Skip to main content

r_tool/cache/
lru_cache.rs

1use std::collections::{HashMap, LinkedList};
2
3pub struct LRUCache<K, V> {
4    capacity: usize,
5    cache: HashMap<K, (V, usize)>,
6    order: LinkedList<K>,
7}
8
9impl<K, V> LRUCache<K, V>
10    where
11        K: Eq + std::hash::Hash + Clone,
12{
13    pub fn new(capacity: usize) -> Self {
14        LRUCache {
15            capacity,
16            cache: HashMap::with_capacity(capacity),
17            order: LinkedList::new(),
18        }
19    }
20
21    pub fn get(&mut self, key: &K) -> Option<&V> {
22        if let Some((value, timestamp)) = self.cache.get_mut(key) {
23            *timestamp = self.order.len();
24            let popped_key = self.order.pop_front().unwrap();
25            self.order.push_back(popped_key);
26            Some(value)
27        } else {
28            None
29        }
30    }
31
32    pub fn put(&mut self, key: K, value: V) {
33        if self.cache.len() >= self.capacity {
34            if let Some(oldest_key) = self.order.pop_front() {
35                self.cache.remove(&oldest_key);
36            }
37        }
38
39        let timestamp = self.order.len();
40        self.cache.insert(key.clone(), (value, timestamp));
41        self.order.push_back(key);
42    }
43}
44
45#[cfg(test)]
46mod tests {
47    use super::*;
48
49    #[test]
50    fn test_lru_cache() {
51        let mut lru_cache = LRUCache::new(3);
52
53        lru_cache.put("one", 1);
54        lru_cache.put("two", 2);
55        lru_cache.put("three", 3);
56
57        assert_eq!(lru_cache.get(&"one"), Some(&1));
58
59        lru_cache.put("four", 4);
60
61        assert_eq!(lru_cache.cache.len(), 3);
62        assert_eq!(lru_cache.get(&"one"), Some(&1));
63        assert_eq!(lru_cache.get(&"two"), None);
64        assert_eq!(lru_cache.get(&"three"), Some(&3));
65        assert_eq!(lru_cache.get(&"four"), Some(&4));
66    }
67}