1use alloy_primitives::map::DefaultHashBuilder;
4use core::hash::BuildHasher;
5use derive_more::{Deref, DerefMut};
6use itertools::Itertools;
7use schnellru::{ByLength, Limiter, Unlimited};
8use std::{fmt, hash::Hash};
9
10pub struct LruCache<T: Hash + Eq + fmt::Debug> {
15 limit: u32,
16 inner: LruMap<T, ()>,
17}
18
19impl<T: Hash + Eq + fmt::Debug> LruCache<T> {
20 pub fn new(limit: u32) -> Self {
22 Self { inner: LruMap::new(limit + 1), limit }
25 }
26
27 pub fn insert(&mut self, entry: T) -> bool {
39 let mut is_new = false;
40 self.inner.get_or_insert(entry, || {
41 is_new = true;
42 });
43 is_new
44 }
45
46 pub fn insert_and_get_evicted(&mut self, entry: T) -> (bool, Option<T>) {
49 let new = self.inner.peek(&entry).is_none();
50 let evicted =
51 (new && (self.limit as usize) <= self.inner.len()).then(|| self.remove_lru()).flatten();
52 _ = self.inner.get_or_insert(entry, || ());
53
54 (new, evicted)
55 }
56
57 pub fn get(&mut self, entry: &T) -> Option<&T> {
59 let _ = self.inner.get(entry)?;
60 self.iter().next()
61 }
62
63 pub fn find(&self, entry: &T) -> Option<&T> {
70 self.iter().find(|key| *key == entry)
71 }
72
73 #[inline]
78 fn remove_lru(&mut self) -> Option<T> {
79 self.inner.pop_oldest().map(|(k, ())| k)
80 }
81
82 pub fn remove(&mut self, value: &T) -> bool {
84 self.inner.remove(value).is_some()
85 }
86
87 pub fn contains(&self, value: &T) -> bool {
89 self.inner.peek(value).is_some()
90 }
91
92 pub fn iter(&self) -> impl Iterator<Item = &T> + '_ {
94 self.inner.iter().map(|(k, ())| k)
95 }
96
97 pub fn len(&self) -> usize {
99 self.inner.len()
100 }
101
102 pub fn is_empty(&self) -> bool {
104 self.inner.is_empty()
105 }
106}
107
108impl<T> Extend<T> for LruCache<T>
109where
110 T: Eq + Hash + fmt::Debug,
111{
112 fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
113 for item in iter {
114 _ = self.insert(item);
115 }
116 }
117}
118
119impl<T> fmt::Debug for LruCache<T>
120where
121 T: fmt::Debug + Hash + Eq,
122{
123 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
124 let mut debug_struct = f.debug_struct("LruCache");
125
126 debug_struct.field("limit", &self.limit);
127
128 debug_struct.field(
129 "ret %iter",
130 &format_args!("Iter: {{{} }}", self.iter().map(|k| format!(" {k:?}")).format(",")),
131 );
132
133 debug_struct.finish()
134 }
135}
136
137#[derive(Deref, DerefMut, Default)]
140pub struct LruMap<K, V, L = ByLength, S = DefaultHashBuilder>(schnellru::LruMap<K, V, L, S>)
141where
142 K: Hash + PartialEq,
143 L: Limiter<K, V>,
144 S: BuildHasher;
145
146impl<K, V, L, S> fmt::Debug for LruMap<K, V, L, S>
147where
148 K: Hash + PartialEq + fmt::Display,
149 V: fmt::Debug,
150 L: Limiter<K, V> + fmt::Debug,
151 S: BuildHasher,
152{
153 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
154 let mut debug_struct = f.debug_struct("LruMap");
155
156 debug_struct.field("limiter", self.limiter());
157
158 debug_struct.field(
159 "ret %iter",
160 &format_args!(
161 "Iter: {{{} }}",
162 self.iter().map(|(k, v)| format!(" {k}: {v:?}")).format(",")
163 ),
164 );
165
166 debug_struct.finish()
167 }
168}
169
170impl<K, V> LruMap<K, V>
171where
172 K: Hash + PartialEq,
173{
174 pub fn new(max_length: u32) -> Self {
176 Self(schnellru::LruMap::with_hasher(ByLength::new(max_length), Default::default()))
177 }
178}
179
180impl<K, V> LruMap<K, V, Unlimited>
181where
182 K: Hash + PartialEq,
183{
184 pub fn new_unlimited() -> Self {
186 Self(schnellru::LruMap::with_hasher(Unlimited, Default::default()))
187 }
188}
189
190#[cfg(test)]
191mod test {
192 use super::*;
193 use derive_more::{Constructor, Display};
194 use std::hash::Hasher;
195
196 #[derive(Debug, Hash, PartialEq, Eq, Display, Clone, Copy)]
197 struct Key(i8);
198
199 #[derive(Debug, Eq, Constructor, Clone, Copy)]
200 struct CompoundKey {
201 id: i8,
203 other: i8,
204 }
205
206 impl PartialEq for CompoundKey {
207 fn eq(&self, other: &Self) -> bool {
208 self.id == other.id
209 }
210 }
211
212 impl Hash for CompoundKey {
213 fn hash<H: Hasher>(&self, state: &mut H) {
214 self.id.hash(state)
215 }
216 }
217
218 #[test]
219 fn test_cache_should_insert_into_empty_set() {
220 let mut cache = LruCache::new(5);
221 let entry = "entry";
222 assert!(cache.insert(entry));
223 assert!(cache.contains(&entry));
224 }
225
226 #[test]
227 fn test_cache_should_not_insert_same_value_twice() {
228 let mut cache = LruCache::new(5);
229 let entry = "entry";
230 assert!(cache.insert(entry));
231 assert!(!cache.insert(entry));
232 }
233
234 #[test]
235 fn test_cache_should_remove_oldest_element_when_exceeding_limit() {
236 let mut cache = LruCache::new(1); let old_entry = "old_entry";
238 let new_entry = "new_entry";
239 cache.insert(old_entry);
240 cache.insert("entry");
241 cache.insert(new_entry);
242 assert!(cache.contains(&new_entry));
243 assert!(!cache.contains(&old_entry));
244 }
245
246 #[test]
247 fn test_cache_should_extend_an_array() {
248 let mut cache = LruCache::new(5);
249 let entries = ["some_entry", "another_entry"];
250 cache.extend(entries);
251 for e in entries {
252 assert!(cache.contains(&e));
253 }
254 }
255
256 #[test]
257 #[expect(dead_code)]
258 fn test_debug_impl_lru_map() {
259 #[derive(Debug)]
260 struct Value(i8);
261
262 let mut cache = LruMap::new(2);
263 let key_1 = Key(1);
264 let value_1 = Value(11);
265 cache.insert(key_1, value_1);
266 let key_2 = Key(2);
267 let value_2 = Value(22);
268 cache.insert(key_2, value_2);
269
270 assert_eq!(
271 "LruMap { limiter: ByLength { max_length: 2 }, ret %iter: Iter: { 2: Value(22), 1: Value(11) } }",
272 format!("{cache:?}")
273 )
274 }
275
276 #[test]
277 fn test_debug_impl_lru_cache() {
278 let mut cache = LruCache::new(2);
279 let key_1 = Key(1);
280 cache.insert(key_1);
281 let key_2 = Key(2);
282 cache.insert(key_2);
283
284 assert_eq!(
285 "LruCache { limit: 2, ret %iter: Iter: { Key(2), Key(1) } }",
286 format!("{cache:?}")
287 )
288 }
289
290 #[test]
291 fn get() {
292 let mut cache = LruCache::new(2);
293 let key_1 = Key(1);
294 cache.insert(key_1);
295 let key_2 = Key(2);
296 cache.insert(key_2);
297
298 _ = cache.get(&key_1);
300
301 assert_eq!(
302 "LruCache { limit: 2, ret %iter: Iter: { Key(1), Key(2) } }",
303 format!("{cache:?}")
304 )
305 }
306
307 #[test]
308 fn get_ty_custom_eq_impl() {
309 let mut cache = LruCache::new(2);
310 let key_1 = CompoundKey::new(1, 11);
311 cache.insert(key_1);
312 let key_2 = CompoundKey::new(2, 22);
313 cache.insert(key_2);
314
315 let key = cache.get(&key_1);
316
317 assert_eq!(key_1.other, key.unwrap().other)
318 }
319
320 #[test]
321 fn peek() {
322 let mut cache = LruCache::new(2);
323 let key_1 = Key(1);
324 cache.insert(key_1);
325 let key_2 = Key(2);
326 cache.insert(key_2);
327
328 _ = cache.find(&key_1);
330
331 assert_eq!(
332 "LruCache { limit: 2, ret %iter: Iter: { Key(2), Key(1) } }",
333 format!("{cache:?}")
334 )
335 }
336
337 #[test]
338 fn peek_ty_custom_eq_impl() {
339 let mut cache = LruCache::new(2);
340 let key_1 = CompoundKey::new(1, 11);
341 cache.insert(key_1);
342 let key_2 = CompoundKey::new(2, 22);
343 cache.insert(key_2);
344
345 let key = cache.find(&key_1);
346
347 assert_eq!(key_1.other, key.unwrap().other)
348 }
349
350 #[test]
351 fn test_insert_methods() {
352 let mut cache = LruCache::new(2);
353
354 assert!(cache.insert("first")); assert!(!cache.insert("first")); assert!(cache.insert("second")); let (is_new, evicted) = cache.insert_and_get_evicted("third");
361 assert!(is_new); assert_eq!(evicted, Some("first")); assert!(cache.contains(&"second"));
365 assert!(cache.contains(&"third"));
366 assert!(!cache.contains(&"first"));
367
368 let (is_new, evicted) = cache.insert_and_get_evicted("second");
370 assert!(!is_new); assert_eq!(evicted, None); }
373}