1use crate::{
2 linked_slab::Token,
3 options::*,
4 shard::{self, CacheShard, InsertStrategy},
5 DefaultHashBuilder, Equivalent, Lifecycle, MemoryUsed, UnitWeighter, Weighter,
6};
7use std::hash::{BuildHasher, Hash};
8
9#[derive(Clone)]
11pub struct Cache<
12 Key,
13 Val,
14 We = UnitWeighter,
15 B = DefaultHashBuilder,
16 L = DefaultLifecycle<Key, Val>,
17> {
18 shard: CacheShard<Key, Val, We, B, L, SharedPlaceholder>,
19}
20
21impl<Key: Eq + Hash, Val> Cache<Key, Val> {
22 pub fn new(items_capacity: usize) -> Self {
24 Self::with(
25 items_capacity,
26 items_capacity as u64,
27 Default::default(),
28 Default::default(),
29 Default::default(),
30 )
31 }
32}
33
34impl<Key: Eq + Hash, Val, We: Weighter<Key, Val>> Cache<Key, Val, We> {
35 pub fn with_weighter(
36 estimated_items_capacity: usize,
37 weight_capacity: u64,
38 weighter: We,
39 ) -> Self {
40 Self::with(
41 estimated_items_capacity,
42 weight_capacity,
43 weighter,
44 Default::default(),
45 Default::default(),
46 )
47 }
48}
49
50impl<Key: Eq + Hash, Val, We: Weighter<Key, Val>, B: BuildHasher, L: Lifecycle<Key, Val>>
51 Cache<Key, Val, We, B, L>
52{
53 pub fn with(
57 estimated_items_capacity: usize,
58 weight_capacity: u64,
59 weighter: We,
60 hash_builder: B,
61 lifecycle: L,
62 ) -> Self {
63 Self::with_options(
64 OptionsBuilder::new()
65 .estimated_items_capacity(estimated_items_capacity)
66 .weight_capacity(weight_capacity)
67 .build()
68 .unwrap(),
69 weighter,
70 hash_builder,
71 lifecycle,
72 )
73 }
74
75 pub fn with_options(options: Options, weighter: We, hash_builder: B, lifecycle: L) -> Self {
94 let shard = CacheShard::new(
95 options.hot_allocation,
96 options.ghost_allocation,
97 options.estimated_items_capacity,
98 options.weight_capacity,
99 weighter,
100 hash_builder,
101 lifecycle,
102 );
103 Self { shard }
104 }
105
106 pub fn is_empty(&self) -> bool {
108 self.shard.len() == 0
109 }
110
111 pub fn len(&self) -> usize {
113 self.shard.len()
114 }
115
116 pub fn weight(&self) -> u64 {
118 self.shard.weight()
119 }
120
121 pub fn capacity(&self) -> u64 {
123 self.shard.capacity()
124 }
125
126 #[cfg(feature = "stats")]
128 pub fn misses(&self) -> u64 {
129 self.shard.misses()
130 }
131
132 #[cfg(feature = "stats")]
134 pub fn hits(&self) -> u64 {
135 self.shard.hits()
136 }
137
138 pub fn reserve(&mut self, additional: usize) {
141 self.shard.reserve(additional);
142 }
143
144 pub fn contains_key<Q>(&self, key: &Q) -> bool
146 where
147 Q: Hash + Equivalent<Key> + ?Sized,
148 {
149 self.shard.contains(self.shard.hash(key), key)
150 }
151
152 pub fn get<Q>(&self, key: &Q) -> Option<&Val>
154 where
155 Q: Hash + Equivalent<Key> + ?Sized,
156 {
157 self.shard.get(self.shard.hash(key), key)
158 }
159
160 pub fn get_mut<Q>(&mut self, key: &Q) -> Option<RefMut<'_, Key, Val, We, B, L>>
164 where
165 Q: Hash + Equivalent<Key> + ?Sized,
166 {
167 self.shard.get_mut(self.shard.hash(key), key).map(RefMut)
168 }
169
170 pub fn peek<Q>(&self, key: &Q) -> Option<&Val>
172 where
173 Q: Hash + Equivalent<Key> + ?Sized,
174 {
175 self.shard.peek(self.shard.hash(key), key)
176 }
177
178 #[cfg(feature = "stats")]
181 pub fn item_stats<Q>(&self, key: &Q) -> Option<crate::ItemStats>
182 where
183 Q: Hash + Equivalent<Key> + ?Sized,
184 {
185 self.shard.item_stats(self.shard.hash(key), key)
186 }
187
188 pub fn peek_mut<Q>(&mut self, key: &Q) -> Option<RefMut<'_, Key, Val, We, B, L>>
192 where
193 Q: Hash + Equivalent<Key> + ?Sized,
194 {
195 self.shard.peek_mut(self.shard.hash(key), key).map(RefMut)
196 }
197
198 pub fn remove<Q>(&mut self, key: &Q) -> Option<(Key, Val)>
201 where
202 Q: Hash + Equivalent<Key> + ?Sized,
203 {
204 self.shard.remove(self.shard.hash(key), key)
205 }
206
207 pub fn remove_if<Q, F>(&mut self, key: &Q, f: F) -> Option<(Key, Val)>
212 where
213 Q: Hash + Equivalent<Key> + ?Sized,
214 F: FnOnce(&Val) -> bool,
215 {
216 self.shard.remove_if(self.shard.hash(key), key, f)
217 }
218
219 pub fn replace(&mut self, key: Key, value: Val, soft: bool) -> Result<(), (Key, Val)> {
225 let mut lcs = Default::default();
226 self.replace_with_lifecycle(key, value, soft, &mut lcs)
227 }
228
229 pub fn replace_with_lifecycle(
240 &mut self,
241 key: Key,
242 value: Val,
243 soft: bool,
244 lcs: &mut L::RequestState,
245 ) -> Result<(), (Key, Val)> {
246 self.shard.insert(
247 lcs,
248 self.shard.hash(&key),
249 key,
250 value,
251 InsertStrategy::Replace { soft },
252 )?;
253 Ok(())
254 }
255
256 pub fn retain<F>(&mut self, f: F)
260 where
261 F: Fn(&Key, &Val) -> bool,
262 {
263 self.shard.retain(f);
264 }
265
266 pub fn get_or_insert_with<Q, E>(
271 &mut self,
272 key: &Q,
273 with: impl FnOnce() -> Result<Val, E>,
274 ) -> Result<Option<&Val>, E>
275 where
276 Q: Hash + Equivalent<Key> + ToOwned<Owned = Key> + ?Sized,
277 {
278 let idx = match self.shard.get_or_placeholder(self.shard.hash(key), key) {
279 Ok((idx, _)) => idx,
280 Err((plh, _)) => {
281 let v = with()?;
282 let mut lcs = Default::default();
283 let replaced = self.shard.replace_placeholder(&mut lcs, &plh, false, v);
284 debug_assert!(replaced.is_ok(), "unsync replace_placeholder can't fail");
285 plh.idx
286 }
287 };
288 Ok(self.shard.peek_token(idx))
289 }
290
291 pub fn get_mut_or_insert_with<'a, Q, E>(
296 &'a mut self,
297 key: &Q,
298 with: impl FnOnce() -> Result<Val, E>,
299 ) -> Result<Option<RefMut<'a, Key, Val, We, B, L>>, E>
300 where
301 Q: Hash + Equivalent<Key> + ToOwned<Owned = Key> + ?Sized,
302 {
303 let idx = match self.shard.get_or_placeholder(self.shard.hash(key), key) {
304 Ok((idx, _)) => idx,
305 Err((plh, _)) => {
306 let v = with()?;
307 let mut lcs = Default::default();
308 let replaced = self.shard.replace_placeholder(&mut lcs, &plh, false, v);
309 debug_assert!(replaced.is_ok(), "unsync replace_placeholder can't fail");
310 plh.idx
311 }
312 };
313 Ok(self.shard.peek_token_mut(idx).map(RefMut))
314 }
315
316 pub fn get_ref_or_guard<Q>(&mut self, key: &Q) -> Result<&Val, Guard<'_, Key, Val, We, B, L>>
320 where
321 Q: Hash + Equivalent<Key> + ToOwned<Owned = Key> + ?Sized,
322 {
323 match self.shard.get_or_placeholder(self.shard.hash(key), key) {
325 Ok((_, v)) => unsafe {
326 let v: *const Val = v;
329 Ok(&*v)
330 },
331 Err((placeholder, _)) => Err(Guard {
332 cache: self,
333 placeholder,
334 inserted: false,
335 }),
336 }
337 }
338
339 pub fn get_mut_or_guard<'a, Q>(
345 &'a mut self,
346 key: &Q,
347 ) -> Result<Option<RefMut<'a, Key, Val, We, B, L>>, Guard<'a, Key, Val, We, B, L>>
348 where
349 Q: Hash + Equivalent<Key> + ToOwned<Owned = Key> + ?Sized,
350 {
351 match self.shard.get_or_placeholder(self.shard.hash(key), key) {
353 Ok((idx, _)) => Ok(self.shard.peek_token_mut(idx).map(RefMut)),
354 Err((placeholder, _)) => Err(Guard {
355 cache: self,
356 placeholder,
357 inserted: false,
358 }),
359 }
360 }
361
362 pub fn insert(&mut self, key: Key, value: Val) {
364 let mut lcs = Default::default();
365 self.insert_with_lifecycle(key, value, &mut lcs);
366 }
367
368 pub fn insert_with_lifecycle(&mut self, key: Key, value: Val, lcs: &mut L::RequestState) {
375 let result = self.shard.insert(
376 lcs,
377 self.shard.hash(&key),
378 key,
379 value,
380 InsertStrategy::Insert,
381 );
382 debug_assert!(result.is_ok());
384 }
385
386 pub fn clear(&mut self) {
388 self.shard.clear();
389 }
390
391 pub fn iter(&self) -> impl Iterator<Item = (&'_ Key, &'_ Val)> + '_ {
393 self.shard.iter()
395 }
396
397 pub fn drain(&mut self) -> impl Iterator<Item = (Key, Val)> + '_ {
401 self.shard.drain()
403 }
404
405 pub fn set_capacity(&mut self, new_weight_capacity: u64) {
410 let mut lcs = Default::default();
411 self.shard.set_capacity(new_weight_capacity, &mut lcs);
412 }
413
414 #[cfg(any(fuzzing, test))]
415 pub fn validate(&self, accept_overweight: bool) {
416 self.shard.validate(accept_overweight);
417 }
418
419 pub fn memory_used(&self) -> MemoryUsed {
424 self.shard.memory_used()
425 }
426}
427
428impl<Key, Val, We, B, L> std::fmt::Debug for Cache<Key, Val, We, B, L> {
429 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
430 f.debug_struct("Cache").finish_non_exhaustive()
431 }
432}
433
434pub struct DefaultLifecycle<Key, Val>(std::marker::PhantomData<(Key, Val)>);
436
437impl<Key, Val> std::fmt::Debug for DefaultLifecycle<Key, Val> {
438 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
439 f.debug_tuple("DefaultLifecycle").finish()
440 }
441}
442
443impl<Key, Val> Default for DefaultLifecycle<Key, Val> {
444 #[inline]
445 fn default() -> Self {
446 Self(Default::default())
447 }
448}
449
450impl<Key, Val> Clone for DefaultLifecycle<Key, Val> {
451 #[inline]
452 fn clone(&self) -> Self {
453 Self(Default::default())
454 }
455}
456
457impl<Key, Val> Lifecycle<Key, Val> for DefaultLifecycle<Key, Val> {
458 type RequestState = ();
459}
460
461#[derive(Debug, Clone)]
462pub(crate) struct SharedPlaceholder {
463 hash: u64,
464 idx: Token,
465}
466
467pub struct Guard<'a, Key, Val, We, B, L> {
468 cache: &'a mut Cache<Key, Val, We, B, L>,
469 placeholder: SharedPlaceholder,
470 inserted: bool,
471}
472
473impl<Key: Eq + Hash, Val, We: Weighter<Key, Val>, B: BuildHasher, L: Lifecycle<Key, Val>>
474 Guard<'_, Key, Val, We, B, L>
475{
476 pub fn insert(self, value: Val) {
478 let mut lcs = Default::default();
479 self.insert_with_lifecycle(value, &mut lcs);
480 }
481
482 pub fn insert_with_lifecycle(mut self, value: Val, lcs: &mut L::RequestState) {
488 let replaced = self
489 .cache
490 .shard
491 .replace_placeholder(lcs, &self.placeholder, false, value);
492 debug_assert!(replaced.is_ok(), "unsync replace_placeholder can't fail");
493 self.inserted = true;
494 }
495}
496
497impl<Key, Val, We, B, L> Drop for Guard<'_, Key, Val, We, B, L> {
498 #[inline]
499 fn drop(&mut self) {
500 #[cold]
501 fn drop_slow<Key, Val, We, B, L>(this: &mut Guard<'_, Key, Val, We, B, L>) {
502 this.cache.shard.remove_placeholder(&this.placeholder);
503 }
504 if !self.inserted {
505 drop_slow(self);
506 }
507 }
508}
509
510pub struct RefMut<'cache, Key, Val, We: Weighter<Key, Val>, B, L>(
511 crate::shard::RefMut<'cache, Key, Val, We, B, L, SharedPlaceholder>,
512);
513
514impl<Key, Val, We: Weighter<Key, Val>, B, L> std::ops::Deref for RefMut<'_, Key, Val, We, B, L> {
515 type Target = Val;
516
517 #[inline]
518 fn deref(&self) -> &Self::Target {
519 self.0.pair().1
520 }
521}
522
523impl<Key, Val, We: Weighter<Key, Val>, B, L> std::ops::DerefMut for RefMut<'_, Key, Val, We, B, L> {
524 #[inline]
525 fn deref_mut(&mut self) -> &mut Self::Target {
526 self.0.value_mut()
527 }
528}
529
530impl shard::SharedPlaceholder for SharedPlaceholder {
531 #[inline]
532 fn new(hash: u64, idx: Token) -> Self {
533 Self { hash, idx }
534 }
535
536 #[inline]
537 fn same_as(&self, _other: &Self) -> bool {
538 true
539 }
540
541 #[inline]
542 fn hash(&self) -> u64 {
543 self.hash
544 }
545
546 #[inline]
547 fn idx(&self) -> Token {
548 self.idx
549 }
550}
551
552#[cfg(test)]
553mod tests {
554 use super::*;
555
556 struct Weighter;
557
558 impl crate::Weighter<u32, u32> for Weighter {
559 fn weight(&self, _key: &u32, val: &u32) -> u64 {
560 *val as u64
561 }
562 }
563
564 #[test]
565 fn test_zero_weights() {
566 let mut cache = Cache::with_weighter(100, 100, Weighter);
567 cache.insert(0, 0);
568 assert_eq!(cache.weight(), 0);
569 for i in 1..100 {
570 cache.insert(i, i);
571 cache.insert(i, i);
572 }
573 assert_eq!(cache.get(&0).copied(), Some(0));
574 assert!(cache.contains_key(&0));
575 let a = cache.weight();
576 *cache.get_mut(&0).unwrap() += 1;
577 assert_eq!(cache.weight(), a + 1);
578 for i in 1..100 {
579 cache.insert(i, i);
580 cache.insert(i, i);
581 }
582 assert_eq!(cache.get(&0), None);
583 assert!(!cache.contains_key(&0));
584
585 cache.insert(0, 1);
586 let a = cache.weight();
587 *cache.get_mut(&0).unwrap() -= 1;
588 assert_eq!(cache.weight(), a - 1);
589 for i in 1..100 {
590 cache.insert(i, i);
591 cache.insert(i, i);
592 }
593 assert_eq!(cache.get(&0).copied(), Some(0));
594 assert!(cache.contains_key(&0));
595 }
596
597 #[test]
598 fn test_set_capacity() {
599 let mut cache = Cache::new(100);
600 for i in 0..80 {
601 cache.insert(i, i);
602 }
603 let initial_len = cache.len();
604 assert!(initial_len <= 80);
605
606 cache.set_capacity(50);
608 assert!(cache.len() <= 50);
609 assert!(cache.weight() <= 50);
610 cache.validate(false);
611
612 cache.set_capacity(200);
614 assert_eq!(cache.capacity(), 200);
615 cache.validate(false);
616
617 for i in 100..180 {
619 cache.insert(i, i);
620 }
621 assert!(cache.len() <= 180);
622 assert!(cache.weight() <= 200);
623 cache.validate(false);
624 }
625
626 #[test]
627 fn test_set_capacity_with_ghosts() {
628 let mut cache = Cache::new(50);
630
631 for i in 0..100 {
633 cache.insert(i, i);
634 }
635 cache.validate(false);
636
637 cache.set_capacity(25);
639 assert!(cache.weight() <= 25);
640 cache.validate(false);
641
642 cache.set_capacity(100);
644 assert_eq!(cache.capacity(), 100);
645 cache.validate(false);
646
647 for i in 100..150 {
649 cache.insert(i, i);
650 }
651 cache.validate(false);
652 }
653
654 #[test]
655 fn test_remove_if() {
656 let mut cache = Cache::new(100);
657
658 cache.insert(1, 10);
660 cache.insert(2, 20);
661 cache.insert(3, 30);
662
663 let removed = cache.remove_if(&2, |v| *v == 20);
665 assert_eq!(removed, Some((2, 20)));
666 assert_eq!(cache.get(&2), None);
667
668 let not_removed = cache.remove_if(&3, |v| *v == 999);
670 assert_eq!(not_removed, None);
671 assert_eq!(cache.get(&3), Some(&30));
672
673 let not_found = cache.remove_if(&999, |_| true);
675 assert_eq!(not_found, None);
676
677 cache.validate(false);
678 }
679}