1use std::collections::HashMap;
11use std::hash::Hash;
12
13pub struct LruCache<K: Eq + Hash + Clone, V: Clone> {
17 capacity: usize,
18 map: HashMap<K, (V, u64)>,
19 tick: u64,
21}
22
23impl<K: Eq + Hash + Clone, V: Clone> LruCache<K, V> {
24 pub fn new(capacity: usize) -> Self {
25 assert!(capacity >= 1, "capacity must be >= 1");
26 Self { capacity, map: HashMap::with_capacity(capacity), tick: 0 }
27 }
28
29 pub fn capacity(&self) -> usize {
30 self.capacity
31 }
32
33 pub fn len(&self) -> usize {
34 self.map.len()
35 }
36
37 pub fn is_empty(&self) -> bool {
38 self.map.is_empty()
39 }
40
41 pub fn contains(&self, key: &K) -> bool {
42 self.map.contains_key(key)
43 }
44
45 pub fn get(&mut self, key: &K) -> Option<V> {
47 let v = self.map.get_mut(key).map(|(v, last)| {
48 self.tick += 1;
49 *last = self.tick;
50 v.clone()
51 });
52 v
53 }
54
55 pub fn put(&mut self, key: K, value: V) -> Option<V> {
58 self.tick += 1;
59 if self.map.contains_key(&key) {
60 let (slot, last) = self.map.get_mut(&key).expect("checked above");
61 *slot = value;
62 *last = self.tick;
63 return None;
64 }
65 if self.map.len() >= self.capacity {
66 if let Some((evict_k, _)) =
68 self.map.iter().min_by_key(|(_, (_, last))| *last).map(|(k, _)| (k.clone(), ()))
69 {
70 let (evicted, _) = self.map.remove(&evict_k).expect("just found");
71 self.map.insert(key, (value, self.tick));
72 return Some(evicted);
73 }
74 }
75 self.map.insert(key, (value, self.tick));
76 None
77 }
78
79 pub fn remove(&mut self, key: &K) -> Option<V> {
80 self.map.remove(key).map(|(v, _)| v)
81 }
82
83 pub fn clear(&mut self) {
84 self.map.clear();
85 self.tick = 0;
86 }
87
88 pub fn peek(&self, key: &K) -> Option<V> {
91 self.map.get(key).map(|(v, _)| v.clone())
92 }
93
94 pub fn iter(&self) -> impl Iterator<Item = (&K, &V)> {
97 self.map.iter().map(|(k, (v, _))| (k, v))
98 }
99}
100
101#[cfg(test)]
102mod tests {
103 use super::*;
104
105 #[test]
106 fn put_and_get() {
107 let mut c = LruCache::<&'static str, i32>::new(3);
108 assert!(c.put("a", 1).is_none());
109 assert!(c.put("b", 2).is_none());
110 assert_eq!(c.get(&"a"), Some(1));
111 assert_eq!(c.get(&"b"), Some(2));
112 assert_eq!(c.get(&"c"), None);
113 assert_eq!(c.len(), 2);
114 }
115
116 #[test]
117 fn lru_eviction_drops_least_recently_used() {
118 let mut c = LruCache::<&'static str, i32>::new(2);
119 c.put("a", 1);
120 c.put("b", 2);
121 let _ = c.get(&"a"); let evicted = c.put("c", 3); assert_eq!(evicted, Some(2));
124 assert!(!c.contains(&"b"));
125 assert!(c.contains(&"a"));
126 assert!(c.contains(&"c"));
127 }
128
129 #[test]
130 fn put_existing_key_updates_value_no_evict() {
131 let mut c = LruCache::<&'static str, i32>::new(2);
132 c.put("a", 1);
133 c.put("b", 2);
134 let evicted = c.put("a", 99);
135 assert!(evicted.is_none());
136 assert_eq!(c.get(&"a"), Some(99));
137 assert_eq!(c.len(), 2);
138 }
139
140 #[test]
141 fn remove_drops_entry() {
142 let mut c = LruCache::<&'static str, i32>::new(2);
143 c.put("a", 1);
144 let r = c.remove(&"a");
145 assert_eq!(r, Some(1));
146 assert!(c.is_empty());
147 }
148
149 #[test]
150 fn clear_resets_state() {
151 let mut c = LruCache::<&'static str, i32>::new(2);
152 c.put("a", 1);
153 c.put("b", 2);
154 c.clear();
155 assert!(c.is_empty());
156 }
157
158 #[test]
159 #[should_panic]
160 fn zero_capacity_panics() {
161 let _: LruCache<&'static str, i32> = LruCache::new(0);
162 }
163
164 #[test]
165 fn peek_does_not_change_recency() {
166 let mut c = LruCache::<&'static str, i32>::new(2);
167 c.put("a", 1);
168 c.put("b", 2);
169 assert_eq!(c.peek(&"a"), Some(1));
170 let evicted = c.put("c", 3);
173 assert_eq!(evicted, Some(1));
174 }
175
176 #[test]
177 fn iter_visits_all_entries() {
178 let mut c = LruCache::<&'static str, i32>::new(3);
179 c.put("a", 1);
180 c.put("b", 2);
181 c.put("c", 3);
182 let mut keys: Vec<&&str> = c.iter().map(|(k, _)| k).collect();
183 keys.sort();
184 assert_eq!(keys, vec![&"a", &"b", &"c"]);
185 }
186}