1use std::alloc::{Allocator, Global};
16use std::borrow::Borrow;
17use std::cell::RefCell;
18use std::hash::{BuildHasher, Hash};
19use std::mem::MaybeUninit;
20use std::ptr::NonNull;
21use std::sync::atomic::Ordering;
22
23pub use ahash::RandomState;
24use hashbrown::HashTable;
25use hashbrown::hash_table::Entry;
26use hashbrown0_14 as _;
27
28use crate::sequence::{AtomicSequence, Sequence, Sequencer};
29
30thread_local! {
31 pub static SEQUENCER: RefCell<Sequencer> = const { RefCell::new(Sequencer::new(Sequencer::DEFAULT_STEP, Sequencer::DEFAULT_LAG)) };
32}
33
34static SEQUENCER_DEFAULT_STEP: AtomicSequence = AtomicSequence::new(Sequencer::DEFAULT_STEP);
35static SEQUENCER_DEFAULT_LAG: AtomicSequence = AtomicSequence::new(Sequencer::DEFAULT_LAG);
36
37pub fn init_global_sequencer_args(step: Sequence, lag: Sequence) {
38 SEQUENCER_DEFAULT_STEP.store(step, Ordering::Relaxed);
39 SEQUENCER_DEFAULT_LAG.store(lag, Ordering::Relaxed);
40}
41
42struct LruEntry<K, V>
43where
44 K: Hash + Eq,
45{
46 prev: Option<NonNull<LruEntry<K, V>>>,
47 next: Option<NonNull<LruEntry<K, V>>>,
48 key: MaybeUninit<K>,
49 value: MaybeUninit<V>,
50 hash: u64,
51 sequence: Sequence,
52}
53
54impl<K, V> LruEntry<K, V>
55where
56 K: Hash + Eq,
57{
58 fn key(&self) -> &K {
59 unsafe { self.key.assume_init_ref() }
60 }
61
62 fn value(&self) -> &V {
63 unsafe { self.value.assume_init_ref() }
64 }
65
66 fn value_mut(&mut self) -> &mut V {
67 unsafe { self.value.assume_init_mut() }
68 }
69}
70
71unsafe impl<K, V> Send for LruEntry<K, V> where K: Hash + Eq {}
72unsafe impl<K, V> Sync for LruEntry<K, V> where K: Hash + Eq {}
73
74pub struct LruCache<K, V, S = RandomState, A = Global>
75where
76 K: Hash + Eq,
77 S: BuildHasher + Send + Sync + 'static,
78 A: Clone + Allocator,
79{
80 map: HashTable<NonNull<LruEntry<K, V>>, A>,
81 dummy: Box<LruEntry<K, V>, A>,
83
84 alloc: A,
85 hash_builder: S,
86}
87
88unsafe impl<K, V, S, A> Send for LruCache<K, V, S, A>
89where
90 K: Hash + Eq,
91 S: BuildHasher + Send + Sync + 'static,
92 A: Clone + Allocator,
93{
94}
95unsafe impl<K, V, S, A> Sync for LruCache<K, V, S, A>
96where
97 K: Hash + Eq,
98 S: BuildHasher + Send + Sync + 'static,
99 A: Clone + Allocator,
100{
101}
102
103impl<K, V> LruCache<K, V>
104where
105 K: Hash + Eq,
106{
107 pub fn unbounded() -> Self {
108 Self::unbounded_with_hasher_in(RandomState::default(), Global)
109 }
110}
111
112impl<K, V, S, A> LruCache<K, V, S, A>
113where
114 K: Hash + Eq,
115 S: BuildHasher + Send + Sync + 'static,
116 A: Clone + Allocator,
117{
118 pub fn unbounded_with_hasher_in(hash_builder: S, alloc: A) -> Self {
119 let map = HashTable::new_in(alloc.clone());
120 let mut dummy = Box::new_in(
121 LruEntry {
122 prev: None,
123 next: None,
124 key: MaybeUninit::uninit(),
125 value: MaybeUninit::uninit(),
126 hash: 0,
127 sequence: Sequence::default(),
128 },
129 alloc.clone(),
130 );
131 let ptr = unsafe { NonNull::new_unchecked(dummy.as_mut() as *mut _) };
132 dummy.next = Some(ptr);
133 dummy.prev = Some(ptr);
134 Self {
135 map,
136 dummy,
137 alloc,
138 hash_builder,
139 }
140 }
141
142 pub fn put(&mut self, key: K, mut value: V) -> Option<V> {
143 unsafe {
144 let hash = self.hash_builder.hash_one(&key);
145
146 match self
147 .map
148 .entry(hash, |p| p.as_ref().key() == &key, |p| p.as_ref().hash)
149 {
150 Entry::Occupied(o) => {
151 let mut ptr = *o.get();
152 let entry = ptr.as_mut();
153 std::mem::swap(&mut value, entry.value_mut());
154 Self::detach(ptr);
155 self.attach(ptr);
156 Some(value)
157 }
158 Entry::Vacant(v) => {
159 let entry = Box::new_in(
160 LruEntry {
161 prev: None,
162 next: None,
163 key: MaybeUninit::new(key),
164 value: MaybeUninit::new(value),
165 hash,
166 sequence: 0,
168 },
169 self.alloc.clone(),
170 );
171 let ptr = NonNull::new_unchecked(Box::into_raw_with_allocator(entry).0);
172 v.insert(ptr);
173 self.attach(ptr);
174 None
175 }
176 }
177 }
178 }
179
180 pub fn remove(&mut self, key: &K) -> Option<V> {
181 unsafe {
182 let hash = self.hash_builder.hash_one(key);
183
184 match self
185 .map
186 .entry(hash, |p| p.as_ref().key() == key, |p| p.as_ref().hash)
187 {
188 Entry::Occupied(o) => {
189 let ptr = *o.get();
190
191 Self::detach(ptr);
193
194 let mut entry = Box::from_raw_in(ptr.as_ptr(), self.alloc.clone());
196 entry.key.assume_init_drop();
197 let value = entry.value.assume_init();
198
199 o.remove();
201
202 Some(value)
203 }
204 Entry::Vacant(_) => None,
205 }
206 }
207 }
208
209 pub fn get<'a, Q>(&'a mut self, key: &Q) -> Option<&'a V>
210 where
211 K: Borrow<Q>,
212 Q: Hash + Eq + ?Sized,
213 {
214 unsafe {
215 let key = key.borrow();
216 let hash = self.hash_builder.hash_one(key);
217 if let Some(ptr) = self.map.find(hash, |p| p.as_ref().key().borrow() == key) {
218 let ptr = *ptr;
219 Self::detach(ptr);
220 self.attach(ptr);
221 Some(ptr.as_ref().value())
222 } else {
223 None
224 }
225 }
226 }
227
228 pub fn get_mut<'a, Q>(&'a mut self, key: &Q) -> Option<&'a mut V>
229 where
230 K: Borrow<Q>,
231 Q: Hash + Eq + ?Sized,
232 {
233 unsafe {
234 let key = key.borrow();
235 let hash = self.hash_builder.hash_one(key);
236 if let Some(ptr) = self
237 .map
238 .find_mut(hash, |p| p.as_ref().key().borrow() == key)
239 {
240 let mut ptr = *ptr;
241 Self::detach(ptr);
242 self.attach(ptr);
243 Some(ptr.as_mut().value_mut())
244 } else {
245 None
246 }
247 }
248 }
249
250 pub fn peek<'a, Q>(&'a self, key: &Q) -> Option<&'a V>
251 where
252 K: Borrow<Q>,
253 Q: Hash + Eq + ?Sized,
254 {
255 unsafe {
256 let key = key.borrow();
257 let hash = self.hash_builder.hash_one(key);
258 self.map
259 .find(hash, |p| p.as_ref().key().borrow() == key)
260 .map(|ptr| ptr.as_ref().value())
261 }
262 }
263
264 pub fn peek_mut<'a, Q>(&'a mut self, key: &Q) -> Option<&'a mut V>
265 where
266 K: Borrow<Q>,
267 Q: Hash + Eq + ?Sized,
268 {
269 unsafe {
270 let key = key.borrow();
271 let hash = self.hash_builder.hash_one(key);
272 self.map
273 .find(hash, |p| p.as_ref().key().borrow() == key)
274 .map(|ptr| ptr.clone().as_mut().value_mut())
275 }
276 }
277
278 pub fn contains<Q>(&self, key: &Q) -> bool
279 where
280 K: Borrow<Q>,
281 Q: Hash + Eq + ?Sized,
282 {
283 unsafe {
284 let key = key.borrow();
285 let hash = self.hash_builder.hash_one(key);
286 self.map
287 .find(hash, |p| p.as_ref().key().borrow() == key)
288 .is_some()
289 }
290 }
291
292 pub fn len(&self) -> usize {
293 self.map.len()
294 }
295
296 pub fn is_empty(&self) -> bool {
297 self.len() == 0
298 }
299
300 pub fn pop_with_sequence(&mut self, sequence: Sequence) -> Option<(K, V, Sequence)> {
302 unsafe {
303 if self.is_empty() {
304 return None;
305 }
306
307 let ptr = self.dummy.next.unwrap_unchecked();
308 if ptr.as_ref().sequence >= sequence {
309 return None;
310 }
311
312 Self::detach(ptr);
313
314 let entry = Box::from_raw_in(ptr.as_ptr(), self.alloc.clone());
315
316 let key = entry.key.assume_init();
317 let value = entry.value.assume_init();
318 let sequence = entry.sequence;
319
320 let hash = self.hash_builder.hash_one(&key);
321
322 match self
323 .map
324 .entry(hash, |p| p.as_ref().key() == &key, |p| p.as_ref().hash)
325 {
326 Entry::Occupied(o) => {
327 o.remove();
328 }
329 Entry::Vacant(_) => {}
330 }
331
332 Some((key, value, sequence))
333 }
334 }
335
336 pub fn clear(&mut self) {
337 unsafe {
338 let mut map = HashTable::new_in(self.alloc.clone());
339 std::mem::swap(&mut map, &mut self.map);
340
341 for ptr in map.drain() {
342 Self::detach(ptr);
343 let mut entry = Box::from_raw_in(ptr.as_ptr(), self.alloc.clone());
344 entry.key.assume_init_drop();
345 entry.value.assume_init_drop();
346 }
347
348 debug_assert!(self.is_empty());
349 debug_assert_eq!(
350 self.dummy.as_mut() as *mut _,
351 self.dummy.next.unwrap_unchecked().as_ptr()
352 )
353 }
354 }
355
356 fn detach(mut ptr: NonNull<LruEntry<K, V>>) {
357 unsafe {
358 let entry = ptr.as_mut();
359
360 debug_assert!(entry.prev.is_some() && entry.next.is_some());
361
362 entry.prev.unwrap_unchecked().as_mut().next = entry.next;
363 entry.next.unwrap_unchecked().as_mut().prev = entry.prev;
364
365 entry.next = None;
366 entry.prev = None;
367 }
368 }
369
370 fn attach(&mut self, mut ptr: NonNull<LruEntry<K, V>>) {
371 unsafe {
372 let entry = ptr.as_mut();
373
374 debug_assert!(entry.prev.is_none() && entry.next.is_none());
375
376 entry.next = Some(NonNull::new_unchecked(self.dummy.as_mut() as *mut _));
377 entry.prev = self.dummy.prev;
378
379 self.dummy.prev.unwrap_unchecked().as_mut().next = Some(ptr);
380 self.dummy.prev = Some(ptr);
381
382 entry.sequence = SEQUENCER.with(|s| s.borrow_mut().alloc());
383 }
384 }
385}
386
387impl<K, V, S, A> Drop for LruCache<K, V, S, A>
388where
389 K: Hash + Eq,
390 S: BuildHasher + Send + Sync + 'static,
391 A: Clone + Allocator,
392{
393 fn drop(&mut self) {
394 self.clear()
395 }
396}
397
398#[cfg(test)]
399mod tests {
400 use super::*;
401
402 #[test]
403 fn test_unbounded() {
404 let cache: LruCache<i32, &str> = LruCache::unbounded();
405 assert!(cache.is_empty());
406 assert_eq!(cache.len(), 0);
407 }
408
409 #[test]
410 fn test_unbounded_with_hasher_in() {
411 let cache: LruCache<i32, &str, RandomState, Global> =
412 LruCache::unbounded_with_hasher_in(RandomState::default(), Global);
413 assert!(cache.is_empty());
414 assert_eq!(cache.len(), 0);
415 }
416
417 #[test]
418 fn test_put() {
419 let mut cache = LruCache::unbounded();
420
421 assert_eq!(cache.put(1, "one"), None);
423 assert_eq!(cache.len(), 1);
424
425 assert_eq!(cache.put(1, "ONE"), Some("one"));
427 assert_eq!(cache.len(), 1);
428
429 assert_eq!(cache.put(2, "two"), None);
431 assert_eq!(cache.len(), 2);
432 }
433
434 #[test]
435 fn test_remove() {
436 let mut cache = LruCache::unbounded();
437
438 assert_eq!(cache.remove(&1), None);
440
441 cache.put(1, "one");
443 assert_eq!(cache.remove(&1), Some("one"));
444 assert!(cache.is_empty());
445
446 assert_eq!(cache.remove(&1), None);
448
449 cache.put(1, "one");
451 cache.put(2, "two");
452 cache.put(3, "three");
453 assert_eq!(cache.remove(&2), Some("two"));
454 assert_eq!(cache.len(), 2);
455 assert!(!cache.contains(&2));
456 }
457
458 #[test]
459 fn test_get() {
460 let mut cache = LruCache::unbounded();
461
462 assert_eq!(cache.get(&1), None);
464
465 cache.put(1, "one");
467 assert_eq!(cache.get(&1), Some(&"one"));
468
469 cache.put(2, "two");
471 let _ = cache.get(&1); let (key, _, _) = cache.pop_with_sequence(u64::MAX).unwrap();
475 assert_eq!(key, 2); }
477
478 #[test]
479 fn test_get_mut() {
480 let mut cache = LruCache::unbounded();
481
482 assert_eq!(cache.get_mut(&1), None);
484
485 cache.put(1, String::from("one"));
487 {
488 let val = cache.get_mut(&1).unwrap();
489 *val = String::from("ONE");
490 }
491 assert_eq!(cache.get(&1), Some(&String::from("ONE")));
492
493 cache.put(2, String::from("two"));
495 let _ = cache.get_mut(&1); let (key, _, _) = cache.pop_with_sequence(u64::MAX).unwrap();
499 assert_eq!(key, 2); }
501
502 #[test]
503 fn test_peek() {
504 let mut cache = LruCache::unbounded();
505
506 assert_eq!(cache.peek(&1), None);
508
509 cache.put(1, "one");
511 cache.put(2, "two");
512 assert_eq!(cache.peek(&1), Some(&"one"));
513
514 let (key, _, _) = cache.pop_with_sequence(u64::MAX).unwrap();
516 assert_eq!(key, 1); }
518
519 #[test]
520 fn test_peek_mut() {
521 let mut cache = LruCache::unbounded();
522
523 assert_eq!(cache.peek_mut(&1), None);
525
526 cache.put(1, String::from("one"));
528 cache.put(2, String::from("two"));
529 {
530 let val = cache.peek_mut(&1).unwrap();
531 *val = String::from("ONE");
532 }
533 assert_eq!(cache.peek(&1), Some(&String::from("ONE")));
534
535 let (key, _, _) = cache.pop_with_sequence(u64::MAX).unwrap();
537 assert_eq!(key, 1); }
539
540 #[test]
541 fn test_contains() {
542 let mut cache = LruCache::unbounded();
543
544 assert!(!cache.contains(&1));
546
547 cache.put(1, "one");
549 assert!(cache.contains(&1));
550
551 cache.remove(&1);
553 assert!(!cache.contains(&1));
554 }
555
556 #[test]
557 fn test_len_and_is_empty() {
558 let mut cache = LruCache::unbounded();
559
560 assert!(cache.is_empty());
562 assert_eq!(cache.len(), 0);
563
564 cache.put(1, "one");
566 assert!(!cache.is_empty());
567 assert_eq!(cache.len(), 1);
568
569 cache.put(2, "two");
571 assert_eq!(cache.len(), 2);
572 cache.remove(&1);
573 assert_eq!(cache.len(), 1);
574 cache.clear();
575 assert!(cache.is_empty());
576 assert_eq!(cache.len(), 0);
577 }
578
579 #[test]
580 fn test_clear() {
581 let mut cache = LruCache::unbounded();
582
583 cache.clear();
585 assert!(cache.is_empty());
586
587 cache.put(1, "one");
589 cache.put(2, "two");
590 cache.put(3, "three");
591 assert_eq!(cache.len(), 3);
592 cache.clear();
593 assert!(cache.is_empty());
594 assert_eq!(cache.len(), 0);
595 assert!(!cache.contains(&1));
596 assert!(!cache.contains(&2));
597 assert!(!cache.contains(&3));
598 }
599
600 #[test]
601 fn test_lru_behavior() {
602 let mut cache = LruCache::unbounded();
603
604 cache.put(1, "one");
606 cache.put(2, "two");
607 cache.put(3, "three");
608
609 let _ = cache.get(&1); let (key, _, _) = cache.pop_with_sequence(u64::MAX).unwrap();
614 assert_eq!(key, 2);
615 let (key, _, _) = cache.pop_with_sequence(u64::MAX).unwrap();
616 assert_eq!(key, 3);
617 let (key, _, _) = cache.pop_with_sequence(u64::MAX).unwrap();
618 assert_eq!(key, 1);
619 assert!(cache.is_empty());
620 }
621}