1use std::{
2 hash::{BuildHasher, Hash},
3 hint::unreachable_unchecked,
4 mem::{self, MaybeUninit},
5};
6
7use hashbrown::HashTable;
8
9use crate::{
10 linked_slab::{LinkedSlab, Token},
11 options::DEFAULT_HOT_ALLOCATION,
12 shim::sync::atomic::{self, AtomicU16},
13 Equivalent, Lifecycle, MemoryUsed, Weighter,
14};
15
16#[cfg(feature = "stats")]
17use crate::shim::sync::atomic::AtomicU64;
18
19const MAX_F: u16 = 2;
21
22pub trait SharedPlaceholder: Clone {
23 fn new(hash: u64, idx: Token) -> Self;
24 fn same_as(&self, other: &Self) -> bool;
25 fn hash(&self) -> u64;
26 fn idx(&self) -> Token;
27}
28
29pub enum InsertStrategy {
30 Insert,
31 Replace { soft: bool },
32}
33
34pub enum EntryAction<T> {
39 Retain(T),
44 Remove,
48 ReplaceWithGuard,
53}
54
55pub enum EntryOrPlaceholder<Key, Val, Plh, T> {
57 Kept(T),
59 Removed(Key, Val),
61 Replaced(Plh, Val),
64 ExistingPlaceholder(Plh),
66 NewPlaceholder(Plh),
68}
69
70#[derive(Copy, Clone, Debug, PartialEq, Eq)]
71enum ResidentState {
72 Hot,
73 Cold,
74}
75
76#[derive(Debug)]
77pub struct Resident<Key, Val> {
78 key: Key,
79 value: Val,
80 state: ResidentState,
81 referenced: AtomicU16,
82 #[cfg(feature = "stats")]
86 access_count: AtomicU64,
87}
88
89impl<Key: Clone, Val: Clone> Clone for Resident<Key, Val> {
90 #[inline]
91 fn clone(&self) -> Self {
92 Self {
93 key: self.key.clone(),
94 value: self.value.clone(),
95 state: self.state,
96 referenced: self.referenced.load(atomic::Ordering::Relaxed).into(),
97 #[cfg(feature = "stats")]
98 access_count: self.access_count.load(atomic::Ordering::Relaxed).into(),
99 }
100 }
101}
102
103#[derive(Debug, Clone)]
104struct Placeholder<Key, Plh> {
105 key: Key,
106 hot: ResidentState,
107 shared: Plh,
108}
109
110#[derive(Clone)]
111enum Entry<Key, Val, Plh> {
112 Resident(Resident<Key, Val>),
113 Placeholder(Placeholder<Key, Plh>),
114 Ghost(u64),
115}
116
117pub struct CacheShard<Key, Val, We, B, L, Plh> {
121 hash_builder: B,
122 map: HashTable<Token>,
125 entries: LinkedSlab<Entry<Key, Val, Plh>>,
127 cold_head: Option<Token>,
130 hot_head: Option<Token>,
133 ghost_head: Option<Token>,
136 weight_target_hot: u64,
137 weight_capacity: u64,
138 weight_hot: u64,
139 weight_cold: u64,
140 num_hot: usize,
141 num_cold: usize,
142 num_non_resident: usize,
143 capacity_non_resident: usize,
144 #[cfg(feature = "stats")]
145 hits: AtomicU64,
146 #[cfg(feature = "stats")]
147 misses: AtomicU64,
148 weighter: We,
149 pub(crate) lifecycle: L,
150}
151
152impl<Key: Clone, Val: Clone, We: Clone, B: Clone, L: Clone, Plh: Clone> Clone
153 for CacheShard<Key, Val, We, B, L, Plh>
154{
155 fn clone(&self) -> Self {
156 Self {
157 hash_builder: self.hash_builder.clone(),
158 map: self.map.clone(),
159 entries: self.entries.clone(),
160 cold_head: self.cold_head,
161 hot_head: self.hot_head,
162 ghost_head: self.ghost_head,
163 weight_target_hot: self.weight_target_hot,
164 weight_capacity: self.weight_capacity,
165 weight_hot: self.weight_hot,
166 weight_cold: self.weight_cold,
167 num_hot: self.num_hot,
168 num_cold: self.num_cold,
169 num_non_resident: self.num_non_resident,
170 capacity_non_resident: self.capacity_non_resident,
171 #[cfg(feature = "stats")]
172 hits: self.hits.load(atomic::Ordering::Relaxed).into(),
173 #[cfg(feature = "stats")]
174 misses: self.misses.load(atomic::Ordering::Relaxed).into(),
175 weighter: self.weighter.clone(),
176 lifecycle: self.lifecycle.clone(),
177 }
178 }
179}
180
181#[cfg(feature = "stats")]
182macro_rules! record_hit {
183 ($self: expr) => {{
184 $self.hits.fetch_add(1, atomic::Ordering::Relaxed);
185 }};
186 ($self: expr, $resident: expr) => {{
187 $self.hits.fetch_add(1, atomic::Ordering::Relaxed);
188 $resident
189 .access_count
190 .fetch_add(1, atomic::Ordering::Relaxed);
191 }};
192}
193#[cfg(feature = "stats")]
194macro_rules! record_hit_mut {
195 ($self: expr) => {{
196 *$self.hits.get_mut() += 1;
197 }};
198 ($self: expr, $resident: expr) => {{
199 *$self.hits.get_mut() += 1;
200 *$resident.access_count.get_mut() += 1;
201 }};
202}
203#[cfg(feature = "stats")]
204macro_rules! record_miss {
205 ($self: expr) => {{
206 $self.misses.fetch_add(1, atomic::Ordering::Relaxed);
207 }};
208}
209#[cfg(feature = "stats")]
210macro_rules! record_miss_mut {
211 ($self: expr) => {{
212 *$self.misses.get_mut() += 1;
213 }};
214}
215#[cfg(not(feature = "stats"))]
216macro_rules! record_hit {
217 ($self: expr) => {{}};
218 ($self: expr, $resident: expr) => {{}};
219}
220#[cfg(not(feature = "stats"))]
221macro_rules! record_hit_mut {
222 ($self: expr) => {{}};
223 ($self: expr, $resident: expr) => {{}};
224}
225#[cfg(not(feature = "stats"))]
226macro_rules! record_miss {
227 ($self: expr) => {{}};
228}
229#[cfg(not(feature = "stats"))]
230macro_rules! record_miss_mut {
231 ($self: expr) => {{}};
232}
233
234impl<Key, Val, We, B, L, Plh: SharedPlaceholder> CacheShard<Key, Val, We, B, L, Plh> {
235 pub fn remove_placeholder(&mut self, placeholder: &Plh) {
236 if let Ok(entry) = self.map.find_entry(placeholder.hash(), |&idx| {
237 if idx != placeholder.idx() {
238 return false;
239 }
240 let (entry, _) = self.entries.get(idx).unwrap();
241 matches!(entry, Entry::Placeholder(Placeholder { shared, .. }) if shared.same_as(placeholder))
242 }) {
243 entry.remove();
244 self.entries.remove(placeholder.idx());
245 }
246 }
247
248 #[cold]
249 fn cold_change_weight(&mut self, idx: Token, old_weight: u64, new_weight: u64) {
250 let Some((Entry::Resident(resident), _)) = self.entries.get_mut(idx) else {
251 unsafe { unreachable_unchecked() };
252 };
253 let (weight_ptr, target_head) = if resident.state == ResidentState::Hot {
254 (&mut self.weight_hot, &mut self.hot_head)
255 } else {
256 (&mut self.weight_cold, &mut self.cold_head)
257 };
258 *weight_ptr -= old_weight;
259 *weight_ptr += new_weight;
260
261 if old_weight == 0 && new_weight != 0 {
262 *target_head = Some(self.entries.link(idx, *target_head));
263 } else if old_weight != 0 && new_weight == 0 {
264 *target_head = self.entries.unlink(idx);
265 }
266 }
267}
268
269impl<Key, Val, We, B, L, Plh> CacheShard<Key, Val, We, B, L, Plh> {
270 pub fn memory_used(&self) -> MemoryUsed {
271 MemoryUsed {
272 entries: self.entries.memory_used(),
273 map: self.map.allocation_size(),
274 }
275 }
276
277 pub fn weight(&self) -> u64 {
278 self.weight_hot + self.weight_cold
279 }
280
281 pub fn len(&self) -> usize {
282 self.num_hot + self.num_cold
283 }
284
285 pub fn capacity(&self) -> u64 {
286 self.weight_capacity
287 }
288
289 #[cfg(feature = "stats")]
290 pub fn hits(&self) -> u64 {
291 self.hits.load(atomic::Ordering::Relaxed)
292 }
293
294 #[cfg(feature = "stats")]
295 pub fn misses(&self) -> u64 {
296 self.misses.load(atomic::Ordering::Relaxed)
297 }
298
299 pub fn clear(&mut self) {
300 let _ = self.drain();
301 }
302
303 pub fn drain(&mut self) -> impl Iterator<Item = (Key, Val)> + '_ {
304 self.cold_head = None;
305 self.hot_head = None;
306 self.ghost_head = None;
307 self.num_hot = 0;
308 self.num_cold = 0;
309 self.num_non_resident = 0;
310 self.weight_hot = 0;
311 self.weight_cold = 0;
312 self.map.clear();
313 self.entries.drain().filter_map(|i| match i {
314 Entry::Resident(r) => Some((r.key, r.value)),
315 Entry::Placeholder(_) | Entry::Ghost(_) => None,
316 })
317 }
318
319 pub fn iter(&self) -> impl Iterator<Item = (&'_ Key, &'_ Val)> + '_ {
320 self.entries.iter().filter_map(|i| match i {
321 Entry::Resident(r) => Some((&r.key, &r.value)),
322 Entry::Placeholder(_) | Entry::Ghost(_) => None,
323 })
324 }
325
326 pub fn iter_from(
327 &self,
328 continuation: Option<Token>,
329 ) -> impl Iterator<Item = (Token, &'_ Key, &'_ Val)> + '_ {
330 self.entries
331 .iter_from(continuation)
332 .filter_map(|(token, i)| match i {
333 Entry::Resident(r) => Some((token, &r.key, &r.value)),
334 Entry::Placeholder(_) | Entry::Ghost(_) => None,
335 })
336 }
337}
338
339impl<
340 Key: Eq + Hash,
341 Val,
342 We: Weighter<Key, Val>,
343 B: BuildHasher,
344 L: Lifecycle<Key, Val>,
345 Plh: SharedPlaceholder,
346 > CacheShard<Key, Val, We, B, L, Plh>
347{
348 pub fn new(
349 hot_allocation: f64,
350 ghost_allocation: f64,
351 estimated_items_capacity: usize,
352 weight_capacity: u64,
353 weighter: We,
354 hash_builder: B,
355 lifecycle: L,
356 ) -> Self {
357 let weight_target_hot = ((weight_capacity as f64 * hot_allocation) as u64)
360 .clamp(weight_capacity.min(1), weight_capacity);
361 let capacity_non_resident = (estimated_items_capacity as f64 * ghost_allocation) as usize;
362 Self {
363 hash_builder,
364 map: HashTable::with_capacity(0),
365 entries: LinkedSlab::with_capacity(0),
366 weight_capacity,
367 #[cfg(feature = "stats")]
368 hits: Default::default(),
369 #[cfg(feature = "stats")]
370 misses: Default::default(),
371 cold_head: None,
372 hot_head: None,
373 ghost_head: None,
374 capacity_non_resident,
375 weight_target_hot,
376 num_hot: 0,
377 num_cold: 0,
378 num_non_resident: 0,
379 weight_hot: 0,
380 weight_cold: 0,
381 weighter,
382 lifecycle,
383 }
384 }
385
386 #[cfg(any(fuzzing, test))]
387 pub fn validate(&self, accept_overweight: bool) {
388 self.entries.validate();
389 let mut num_hot = 0;
390 let mut num_cold = 0;
391 let mut num_non_resident = 0;
392 let mut weight_hot = 0;
393 let mut weight_hot_pinned = 0;
394 let mut weight_cold = 0;
395 let mut weight_cold_pinned = 0;
396 for e in self.entries.iter_entries() {
397 match e {
398 Entry::Resident(r) if r.state == ResidentState::Cold => {
399 num_cold += 1;
400 let weight = self.weighter.weight(&r.key, &r.value);
401 if self.lifecycle.is_pinned(&r.key, &r.value) {
402 weight_cold_pinned += weight;
403 } else {
404 weight_cold += weight;
405 }
406 }
407 Entry::Resident(r) => {
408 num_hot += 1;
409 let weight = self.weighter.weight(&r.key, &r.value);
410 if self.lifecycle.is_pinned(&r.key, &r.value) {
411 weight_hot_pinned += weight;
412 } else {
413 weight_hot += weight;
414 }
415 }
416 Entry::Ghost(_) => {
417 num_non_resident += 1;
418 }
419 Entry::Placeholder(_) => (),
420 }
421 }
422 assert_eq!(num_hot, self.num_hot);
441 assert_eq!(num_cold, self.num_cold);
442 assert_eq!(num_non_resident, self.num_non_resident);
443 assert_eq!(weight_hot + weight_hot_pinned, self.weight_hot);
444 assert_eq!(weight_cold + weight_cold_pinned, self.weight_cold);
445 if !accept_overweight {
446 assert!(weight_hot + weight_cold <= self.weight_capacity);
447 }
448 assert!(num_non_resident <= self.capacity_non_resident);
449 }
450
451 pub fn reserve(&mut self, additional: usize) {
454 let additional = additional.saturating_add(additional.min(self.capacity_non_resident));
459 self.entries.reserve(additional);
460 self.map.reserve(additional, |&idx| {
461 let (entry, _) = self.entries.get(idx).unwrap();
462 match entry {
463 Entry::Resident(Resident { key, .. })
464 | Entry::Placeholder(Placeholder { key, .. }) => {
465 Self::hash_static(&self.hash_builder, key)
466 }
467 Entry::Ghost(non_resident_hash) => *non_resident_hash,
468 }
469 })
470 }
471
472 pub fn retain<F>(&mut self, f: F)
473 where
474 F: Fn(&Key, &Val) -> bool,
475 {
476 let retained_tokens = self
477 .map
478 .iter()
479 .filter_map(|&idx| match self.entries.get(idx) {
480 Some((entry, _idx)) => match entry {
481 Entry::Resident(r) => {
482 if !f(&r.key, &r.value) {
483 let hash = self.hash(&r.key);
484 Some((idx, hash))
485 } else {
486 None
487 }
488 }
489 Entry::Placeholder(_) | Entry::Ghost(_) => None,
490 },
491 None => None,
492 })
493 .collect::<Vec<_>>();
494 for (idx, hash) in retained_tokens {
495 self.remove_internal(hash, idx);
496 }
497 }
498
499 #[inline]
500 fn hash_static<Q>(hasher: &B, key: &Q) -> u64
501 where
502 Q: Hash + Equivalent<Key> + ?Sized,
503 {
504 hasher.hash_one(key)
505 }
506
507 #[inline]
508 pub fn hash<Q>(&self, key: &Q) -> u64
509 where
510 Q: Hash + Equivalent<Key> + ?Sized,
511 {
512 Self::hash_static(&self.hash_builder, key)
513 }
514
515 #[inline]
516 fn search<Q>(&self, hash: u64, k: &Q) -> Option<Token>
517 where
518 Q: Hash + Equivalent<Key> + ?Sized,
519 {
520 let mut hash_match = None;
521 for bucket in self.map.iter_hash(hash) {
522 let idx = *bucket;
523 let (entry, _) = self.entries.get(idx).unwrap();
524 match entry {
525 Entry::Resident(Resident { key, .. })
526 | Entry::Placeholder(Placeholder { key, .. })
527 if k.equivalent(key) =>
528 {
529 return Some(idx);
530 }
531 Entry::Ghost(non_resident_hash) if *non_resident_hash == hash => {
532 hash_match = Some(idx);
533 }
534 _ => (),
535 }
536 }
537 hash_match
538 }
539
540 #[inline]
541 fn search_resident<Q>(&self, hash: u64, k: &Q) -> Option<(Token, &Resident<Key, Val>)>
542 where
543 Q: Hash + Equivalent<Key> + ?Sized,
544 {
545 let mut resident = MaybeUninit::uninit();
546 self.map
547 .find(hash, |&idx| {
548 let (entry, _) = self.entries.get(idx).unwrap();
549 match entry {
550 Entry::Resident(r) if k.equivalent(&r.key) => {
551 resident.write(r);
552 true
553 }
554 _ => false,
555 }
556 })
557 .map(|idx| {
558 (*idx, unsafe { resident.assume_init() })
561 })
562 }
563
564 pub fn contains<Q>(&self, hash: u64, key: &Q) -> bool
565 where
566 Q: Hash + Equivalent<Key> + ?Sized,
567 {
568 self.map
569 .find(hash, |&idx| {
570 let (entry, _) = self.entries.get(idx).unwrap();
571 matches!(entry, Entry::Resident(r) if key.equivalent(&r.key))
572 })
573 .is_some()
574 }
575
576 pub fn get_key_value<Q>(&self, hash: u64, key: &Q) -> Option<(&Key, &Val)>
577 where
578 Q: Hash + Equivalent<Key> + ?Sized,
579 {
580 if let Some((_, resident)) = self.search_resident(hash, key) {
581 let referenced = resident.referenced.load(atomic::Ordering::Relaxed);
582 if referenced < MAX_F {
584 resident.referenced.fetch_add(1, atomic::Ordering::Relaxed);
587 }
588 record_hit!(self, resident);
589 Some((&resident.key, &resident.value))
590 } else {
591 record_miss!(self);
592 None
593 }
594 }
595
596 #[inline]
597 pub fn get<Q>(&self, hash: u64, key: &Q) -> Option<&Val>
598 where
599 Q: Hash + Equivalent<Key> + ?Sized,
600 {
601 self.get_key_value(hash, key).map(|(_k, v)| v)
602 }
603
604 pub fn get_mut<Q>(&mut self, hash: u64, key: &Q) -> Option<RefMut<'_, Key, Val, We, B, L, Plh>>
605 where
606 Q: Hash + Equivalent<Key> + ?Sized,
607 {
608 let Some((idx, _)) = self.search_resident(hash, key) else {
609 record_miss_mut!(self);
610 return None;
611 };
612 let (Entry::Resident(resident), _) = (unsafe { self.entries.get_mut_unchecked(idx) })
613 else {
614 unsafe { unreachable_unchecked() };
615 };
616 if *resident.referenced.get_mut() < MAX_F {
617 *resident.referenced.get_mut() += 1;
618 }
619 record_hit_mut!(self, resident);
620
621 let old_weight = self.weighter.weight(&resident.key, &resident.value);
622 Some(RefMut {
623 guard: WeightGuard {
624 shard: self as *mut _,
625 idx,
626 old_weight,
627 },
628 _phantom: std::marker::PhantomData,
629 })
630 }
631
632 #[inline]
633 pub fn peek_token(&self, token: Token) -> Option<&Val> {
634 if let Some((Entry::Resident(resident), _)) = self.entries.get(token) {
635 Some(&resident.value)
636 } else {
637 None
638 }
639 }
640
641 #[inline]
642 pub fn peek_token_mut(&mut self, token: Token) -> Option<RefMut<'_, Key, Val, We, B, L, Plh>> {
643 if let Some((Entry::Resident(resident), _)) = self.entries.get_mut(token) {
644 let old_weight = self.weighter.weight(&resident.key, &resident.value);
645 Some(RefMut {
646 guard: WeightGuard {
647 shard: self as *mut _,
648 idx: token,
649 old_weight,
650 },
651 _phantom: std::marker::PhantomData,
652 })
653 } else {
654 None
655 }
656 }
657
658 pub fn peek<Q>(&self, hash: u64, key: &Q) -> Option<&Val>
659 where
660 Q: Hash + Equivalent<Key> + ?Sized,
661 {
662 let (_, resident) = self.search_resident(hash, key)?;
663 Some(&resident.value)
664 }
665
666 #[cfg(feature = "stats")]
669 pub fn item_stats<Q>(&self, hash: u64, key: &Q) -> Option<crate::ItemStats>
670 where
671 Q: Hash + Equivalent<Key> + ?Sized,
672 {
673 let (_, resident) = self.search_resident(hash, key)?;
674 Some(crate::ItemStats {
675 access_count: resident.access_count.load(atomic::Ordering::Relaxed),
676 })
677 }
678
679 pub fn peek_mut<Q>(&mut self, hash: u64, key: &Q) -> Option<RefMut<'_, Key, Val, We, B, L, Plh>>
680 where
681 Q: Hash + Equivalent<Key> + ?Sized,
682 {
683 let (idx, _) = self.search_resident(hash, key)?;
684 self.peek_token_mut(idx)
685 }
686
687 pub fn remove<Q>(&mut self, hash: u64, key: &Q) -> Option<(Key, Val)>
688 where
689 Q: Hash + Equivalent<Key> + ?Sized,
690 {
691 let idx = self.search(hash, key)?;
694 self.remove_internal(hash, idx)
695 }
696
697 pub fn remove_if<Q, F>(&mut self, hash: u64, key: &Q, f: F) -> Option<(Key, Val)>
698 where
699 Q: Hash + Equivalent<Key> + ?Sized,
700 F: FnOnce(&Val) -> bool,
701 {
702 let (idx, resident) = self.search_resident(hash, key)?;
703 if f(&resident.value) {
704 self.remove_internal(hash, idx)
705 } else {
706 None
707 }
708 }
709
710 pub fn remove_token(&mut self, token: Token) -> Option<(Key, Val)> {
711 let Some((Entry::Resident(resident), _)) = self.entries.get(token) else {
712 return None;
713 };
714 let hash = Self::hash_static(&self.hash_builder, &resident.key);
715 self.remove_internal(hash, token)
716 }
717
718 pub fn remove_next(&mut self, continuation: Option<Token>) -> Option<(Token, Key, Val)> {
719 let (token, key, _) = self
720 .entries
721 .iter_from(continuation)
722 .filter_map(|(token, i)| match i {
723 Entry::Resident(r) => Some((token, &r.key, &r.value)),
724 Entry::Placeholder(_) | Entry::Ghost(_) => None,
725 })
726 .next()?;
727 let hash = Self::hash_static(&self.hash_builder, key);
728 self.remove_internal(hash, token)
729 .map(|(k, v)| (token, k, v))
730 }
731
732 fn remove_internal(&mut self, hash: u64, idx: Token) -> Option<(Key, Val)> {
733 self.map_remove(hash, idx);
734 let mut result = None;
735 let (entry, next) = self.entries.remove(idx).unwrap();
736 let list_head = match entry {
737 Entry::Resident(r) => {
738 let weight = self.weighter.weight(&r.key, &r.value);
739 result = Some((r.key, r.value));
740 if r.state == ResidentState::Hot {
741 self.num_hot -= 1;
742 self.weight_hot -= weight;
743 &mut self.hot_head
744 } else {
745 debug_assert!(r.state == ResidentState::Cold);
746 self.num_cold -= 1;
747 self.weight_cold -= weight;
748 &mut self.cold_head
749 }
750 }
751 Entry::Ghost(_) => {
752 self.num_non_resident -= 1;
754 &mut self.ghost_head
755 }
756 Entry::Placeholder(_) => {
757 return None;
759 }
760 };
761 if *list_head == Some(idx) {
762 *list_head = next;
763 }
764 result
765 }
766
767 #[must_use]
769 fn advance_cold(&mut self, lcs: &mut L::RequestState) -> bool {
770 let Some(mut idx) = self.cold_head else {
771 return self.advance_hot(lcs);
772 };
773 loop {
774 let (entry, next) = self.entries.get_mut(idx).unwrap();
775 let Entry::Resident(resident) = entry else {
776 unsafe { unreachable_unchecked() };
777 };
778 debug_assert_eq!(resident.state, ResidentState::Cold);
779 if *resident.referenced.get_mut() != 0 {
780 *resident.referenced.get_mut() -= 1;
781 resident.state = ResidentState::Hot;
782 let weight = self.weighter.weight(&resident.key, &resident.value);
783 self.weight_hot += weight;
784 self.weight_cold -= weight;
785 self.num_hot += 1;
786 self.num_cold -= 1;
787 self.cold_head = self.entries.unlink(idx);
788 self.hot_head = Some(self.entries.link(idx, self.hot_head));
789 while self.weight_hot > self.weight_target_hot && self.advance_hot(lcs) {}
791 return true;
792 }
793
794 if self.lifecycle.is_pinned(&resident.key, &resident.value) {
795 if Some(next) == self.cold_head {
796 return self.advance_hot(lcs);
797 }
798 idx = next;
799 continue;
800 }
801
802 self.weight_cold -= self.weighter.weight(&resident.key, &resident.value);
803 self.lifecycle
804 .before_evict(lcs, &resident.key, &mut resident.value);
805 if self.weighter.weight(&resident.key, &resident.value) == 0 {
806 self.cold_head = self.entries.unlink(idx);
807 return true;
808 }
809 let hash = Self::hash_static(&self.hash_builder, &resident.key);
810 let Entry::Resident(evicted) = mem::replace(entry, Entry::Ghost(hash)) else {
811 unsafe { core::hint::unreachable_unchecked() };
813 };
814 self.cold_head = self.entries.unlink(idx);
815 self.ghost_head = Some(self.entries.link(idx, self.ghost_head));
816 self.num_cold -= 1;
817 self.num_non_resident += 1;
818 if self.num_non_resident > self.capacity_non_resident {
820 self.advance_ghost();
821 }
822 self.lifecycle
823 .on_evict_cold(lcs, evicted.key, evicted.value);
824 return true;
825 }
826 }
827
828 #[must_use]
830 fn advance_hot(&mut self, lcs: &mut L::RequestState) -> bool {
831 let mut unpinned = 0usize;
832 let Some(mut idx) = self.hot_head else {
833 return false;
834 };
835 loop {
836 let (entry, next) = self.entries.get_mut(idx).unwrap();
837 let Entry::Resident(resident) = entry else {
838 unsafe { unreachable_unchecked() };
839 };
840 debug_assert_eq!(resident.state, ResidentState::Hot);
841 if self.lifecycle.is_pinned(&resident.key, &resident.value) {
842 *resident.referenced.get_mut() = (*resident.referenced.get_mut())
843 .min(MAX_F)
844 .saturating_sub(1);
845 if Some(next) == self.hot_head {
846 if unpinned == 0 {
847 return false;
849 }
850 unpinned = 0;
852 }
853 idx = next;
854 continue;
855 }
856 unpinned += 1;
857 if *resident.referenced.get_mut() != 0 {
858 *resident.referenced.get_mut() = (*resident.referenced.get_mut()).min(MAX_F) - 1;
859 idx = next;
860 continue;
861 }
862 self.weight_hot -= self.weighter.weight(&resident.key, &resident.value);
863 self.lifecycle
864 .before_evict(lcs, &resident.key, &mut resident.value);
865 if self.weighter.weight(&resident.key, &resident.value) == 0 {
866 self.hot_head = self.entries.unlink(idx);
867 } else {
868 self.num_hot -= 1;
869 let hash = Self::hash_static(&self.hash_builder, &resident.key);
870 let Some((Entry::Resident(evicted), next)) = self.entries.remove(idx) else {
871 unsafe { core::hint::unreachable_unchecked() };
873 };
874 self.hot_head = next;
875 self.lifecycle.on_evict_hot(lcs, evicted.key, evicted.value);
876 self.map_remove(hash, idx);
877 }
878 return true;
879 }
880 }
881
882 #[inline]
883 fn advance_ghost(&mut self) {
884 debug_assert_ne!(self.num_non_resident, 0);
885 let idx = self.ghost_head.unwrap();
886 let (entry, _) = self.entries.get_mut(idx).unwrap();
887 let Entry::Ghost(hash) = *entry else {
888 unsafe { unreachable_unchecked() };
889 };
890 self.num_non_resident -= 1;
891 self.map_remove(hash, idx);
892 let (_, next) = self.entries.remove(idx).unwrap();
893 self.ghost_head = next;
894 }
895
896 fn insert_existing(
897 &mut self,
898 lcs: &mut L::RequestState,
899 idx: Token,
900 key: Key,
901 value: Val,
902 weight: u64,
903 strategy: InsertStrategy,
904 ) -> Result<(), (Key, Val)> {
905 let (entry, _) = self.entries.get_mut(idx).unwrap();
907 let referenced;
908 let enter_state;
909 match entry {
910 Entry::Resident(resident) => {
911 enter_state = resident.state;
912 referenced = resident
913 .referenced
914 .get_mut()
915 .saturating_add(
916 !matches!(strategy, InsertStrategy::Replace { soft: true }) as u16
917 )
918 .min(MAX_F);
919 }
920 _ if matches!(strategy, InsertStrategy::Replace { .. }) => {
921 return Err((key, value));
922 }
923 Entry::Ghost(_) => {
924 referenced = 0;
925 enter_state = ResidentState::Hot;
926 }
927 Entry::Placeholder(ph) => {
928 referenced = 1; enter_state = ph.hot;
930 }
931 }
932
933 let evicted = mem::replace(
934 entry,
935 Entry::Resident(Resident {
936 key,
937 value,
938 state: enter_state,
939 referenced: referenced.into(),
940 #[cfg(feature = "stats")]
941 access_count: Default::default(),
942 }),
943 );
944 match evicted {
945 Entry::Resident(evicted) => {
946 debug_assert_eq!(evicted.state, enter_state);
947 let evicted_weight = self.weighter.weight(&evicted.key, &evicted.value);
948 let list_head = if enter_state == ResidentState::Hot {
949 self.weight_hot -= evicted_weight;
950 self.weight_hot += weight;
951 &mut self.hot_head
952 } else {
953 self.weight_cold -= evicted_weight;
954 self.weight_cold += weight;
955 &mut self.cold_head
956 };
957 if evicted_weight == 0 && weight != 0 {
958 *list_head = Some(self.entries.link(idx, *list_head));
959 } else if evicted_weight != 0 && weight == 0 {
960 *list_head = self.entries.unlink(idx);
961 }
962 match enter_state {
963 ResidentState::Hot => {
964 self.lifecycle.on_evict_hot(lcs, evicted.key, evicted.value)
965 }
966 ResidentState::Cold => {
967 self.lifecycle
968 .on_evict_cold(lcs, evicted.key, evicted.value)
969 }
970 }
971 }
972 Entry::Ghost(_) => {
973 self.weight_hot += weight;
974 self.num_hot += 1;
975 self.num_non_resident -= 1;
976 let next_ghost = self.entries.unlink(idx);
977 if self.ghost_head == Some(idx) {
978 self.ghost_head = next_ghost;
979 }
980 if weight != 0 {
981 self.hot_head = Some(self.entries.link(idx, self.hot_head));
982 }
983 }
984 Entry::Placeholder(_) => {
985 let list_head = if enter_state == ResidentState::Hot {
986 self.num_hot += 1;
987 self.weight_hot += weight;
988 &mut self.hot_head
989 } else {
990 self.num_cold += 1;
991 self.weight_cold += weight;
992 &mut self.cold_head
993 };
994 if weight != 0 {
995 *list_head = Some(self.entries.link(idx, *list_head));
996 }
997 }
998 }
999
1000 while self.weight_hot + self.weight_cold > self.weight_capacity && self.advance_cold(lcs) {}
1001 Ok(())
1002 }
1003
1004 #[inline]
1005 fn map_insert(&mut self, hash: u64, idx: Token) {
1006 self.map.insert_unique(hash, idx, |&i| {
1007 let (entry, _) = self.entries.get(i).unwrap();
1008 match entry {
1009 Entry::Resident(Resident { key, .. })
1010 | Entry::Placeholder(Placeholder { key, .. }) => {
1011 Self::hash_static(&self.hash_builder, key)
1012 }
1013 Entry::Ghost(hash) => *hash,
1014 }
1015 });
1016 }
1017
1018 #[inline]
1019 fn map_remove(&mut self, hash: u64, idx: Token) {
1020 if let Ok(entry) = self.map.find_entry(hash, |&i| i == idx) {
1021 entry.remove();
1022 return;
1023 }
1024 #[cfg(debug_assertions)]
1025 panic!("key not found");
1026 }
1027
1028 pub fn replace_placeholder(
1029 &mut self,
1030 lcs: &mut L::RequestState,
1031 placeholder: &Plh,
1032 referenced: bool,
1033 mut value: Val,
1034 ) -> Result<(), Val> {
1035 let entry = match self.entries.get_mut(placeholder.idx()) {
1036 Some((entry, _)) if matches!(&*entry, Entry::Placeholder(p) if p.shared.same_as(placeholder)) => {
1037 entry
1038 }
1039 _ => return Err(value),
1040 };
1041 let Entry::Placeholder(Placeholder {
1042 key,
1043 hot: mut placeholder_hot,
1044 ..
1045 }) = mem::replace(entry, Entry::Ghost(0))
1046 else {
1047 unsafe { core::hint::unreachable_unchecked() };
1049 };
1050 let mut weight = self.weighter.weight(&key, &value);
1051 if weight > self.weight_target_hot && !self.lifecycle.is_pinned(&key, &value) {
1053 self.lifecycle.before_evict(lcs, &key, &mut value);
1054 weight = self.weighter.weight(&key, &value);
1055 if weight > self.weight_target_hot {
1057 return self.handle_overweight_replace_placeholder(lcs, placeholder, key, value);
1058 }
1059 }
1060
1061 if self.weight_hot + weight <= self.weight_target_hot {
1063 placeholder_hot = ResidentState::Hot;
1064 }
1065 *entry = Entry::Resident(Resident {
1066 key,
1067 value,
1068 state: placeholder_hot,
1069 referenced: (referenced as u16).into(),
1070 #[cfg(feature = "stats")]
1071 access_count: Default::default(),
1072 });
1073
1074 let list_head = if placeholder_hot == ResidentState::Hot {
1075 self.num_hot += 1;
1076 self.weight_hot += weight;
1077 &mut self.hot_head
1078 } else {
1079 self.num_cold += 1;
1080 self.weight_cold += weight;
1081 &mut self.cold_head
1082 };
1083
1084 if weight != 0 {
1085 *list_head = Some(self.entries.link(placeholder.idx(), *list_head));
1086 while self.weight_hot + self.weight_cold > self.weight_capacity
1087 && self.advance_cold(lcs)
1088 {}
1089 }
1090
1091 Ok(())
1092 }
1093
1094 #[cold]
1095 fn handle_overweight_replace_placeholder(
1096 &mut self,
1097 lcs: &mut L::RequestState,
1098 placeholder: &Plh,
1099 key: Key,
1100 value: Val,
1101 ) -> Result<(), Val> {
1102 self.entries.remove(placeholder.idx());
1103 self.map_remove(placeholder.hash(), placeholder.idx());
1104 self.lifecycle.on_evict_cold(lcs, key, value);
1105 Ok(())
1106 }
1107
1108 pub fn insert(
1109 &mut self,
1110 lcs: &mut L::RequestState,
1111 hash: u64,
1112 key: Key,
1113 mut value: Val,
1114 strategy: InsertStrategy,
1115 ) -> Result<(), (Key, Val)> {
1116 let mut weight = self.weighter.weight(&key, &value);
1117 if weight > self.weight_target_hot && !self.lifecycle.is_pinned(&key, &value) {
1119 self.lifecycle.before_evict(lcs, &key, &mut value);
1120 weight = self.weighter.weight(&key, &value);
1121 if weight > self.weight_target_hot {
1123 return self.handle_insert_overweight(lcs, hash, key, value, strategy);
1124 }
1125 }
1126
1127 if let Some(idx) = self.search(hash, &key) {
1128 return self.insert_existing(lcs, idx, key, value, weight, strategy);
1129 } else if matches!(strategy, InsertStrategy::Replace { .. }) {
1130 return Err((key, value));
1131 }
1132
1133 let enter_hot = self.weight_hot + weight <= self.weight_target_hot;
1135 while self.weight_hot + self.weight_cold + weight > self.weight_capacity
1137 && self.advance_cold(lcs)
1138 {}
1139
1140 let (state, list_head) = if enter_hot {
1141 self.num_hot += 1;
1142 self.weight_hot += weight;
1143 (ResidentState::Hot, &mut self.hot_head)
1144 } else {
1145 self.num_cold += 1;
1146 self.weight_cold += weight;
1147 (ResidentState::Cold, &mut self.cold_head)
1148 };
1149 let idx = self.entries.insert(Entry::Resident(Resident {
1150 key,
1151 value,
1152 state,
1153 referenced: Default::default(),
1154 #[cfg(feature = "stats")]
1155 access_count: Default::default(),
1156 }));
1157 if weight != 0 {
1158 *list_head = Some(self.entries.link(idx, *list_head));
1159 }
1160 self.map_insert(hash, idx);
1161 Ok(())
1162 }
1163
1164 #[cold]
1165 fn handle_insert_overweight(
1166 &mut self,
1167 lcs: &mut L::RequestState,
1168 hash: u64,
1169 key: Key,
1170 value: Val,
1171 strategy: InsertStrategy,
1172 ) -> Result<(), (Key, Val)> {
1173 if let Some((idx, resident)) = self.search_resident(hash, &key) {
1175 let prev_state = resident.state;
1176 if let Some((ek, ev)) = self.remove_internal(hash, idx) {
1177 match prev_state {
1178 ResidentState::Hot => self.lifecycle.on_evict_hot(lcs, ek, ev),
1179 ResidentState::Cold => self.lifecycle.on_evict_cold(lcs, ek, ev),
1180 }
1181 }
1182 }
1183 if matches!(strategy, InsertStrategy::Replace { .. }) {
1184 return Err((key, value));
1185 }
1186 self.lifecycle.on_evict_cold(lcs, key, value);
1187 Ok(())
1188 }
1189
1190 pub fn get_or_placeholder<Q>(
1191 &mut self,
1192 hash: u64,
1193 key: &Q,
1194 ) -> Result<(Token, &Val), (Plh, bool)>
1195 where
1196 Q: Hash + Equivalent<Key> + ToOwned<Owned = Key> + ?Sized,
1197 {
1198 let idx = self.search(hash, key);
1199 if let Some(idx) = idx {
1200 if let Some((Entry::Resident(resident), _)) = self.entries.get_mut(idx) {
1201 if *resident.referenced.get_mut() < MAX_F {
1202 *resident.referenced.get_mut() += 1;
1203 }
1204 record_hit_mut!(self, resident);
1205 unsafe {
1206 let value_ptr: *const Val = &resident.value;
1209 return Ok((idx, &*value_ptr));
1210 }
1211 }
1212 }
1213 let (shared, is_new) = unsafe { self.non_resident_to_placeholder(hash, key, idx) };
1214 Err((shared, is_new))
1215 }
1216
1217 pub fn entry_or_placeholder<Q, T, F>(
1226 &mut self,
1227 hash: u64,
1228 key: &Q,
1229 on_occupied: &mut F,
1230 ) -> EntryOrPlaceholder<Key, Val, Plh, T>
1231 where
1232 Q: Hash + Equivalent<Key> + ToOwned<Owned = Key> + ?Sized,
1233 F: FnMut(&Key, &mut Val) -> EntryAction<T>,
1234 {
1235 let idx = self.search(hash, key);
1236 if let Some(idx) = idx {
1237 let shard = self as *mut _;
1238 if let Some((Entry::Resident(r), _)) = self.entries.get_mut(idx) {
1239 let action = {
1244 let (key_ptr, val_ptr) = (&r.key as *const Key, &mut r.value as *mut Val);
1245 let _guard = WeightGuard::<Key, Val, We, B, L, Plh> {
1246 idx,
1247 old_weight: self.weighter.weight(&r.key, &r.value),
1248 shard,
1249 };
1250 on_occupied(unsafe { &*key_ptr }, unsafe { &mut *val_ptr })
1251 };
1252
1253 return match action {
1254 EntryAction::Retain(t) => {
1255 let Some((Entry::Resident(resident), _)) = self.entries.get_mut(idx) else {
1256 unsafe { unreachable_unchecked() };
1258 };
1259 if *resident.referenced.get_mut() < MAX_F {
1260 *resident.referenced.get_mut() += 1;
1261 }
1262 record_hit_mut!(self, resident);
1263 EntryOrPlaceholder::Kept(t)
1264 }
1265 EntryAction::Remove => {
1266 let (key, val) = self.remove_internal(hash, idx).unwrap();
1267 EntryOrPlaceholder::Removed(key, val)
1268 }
1269 EntryAction::ReplaceWithGuard => {
1270 let Some((Entry::Resident(r), _)) = self.entries.get_mut(idx) else {
1271 unsafe { unreachable_unchecked() };
1273 };
1274 let state = r.state;
1275 let current_weight = self.weighter.weight(&r.key, &r.value);
1276 let list_head = if state == ResidentState::Hot {
1277 self.num_hot -= 1;
1278 self.weight_hot -= current_weight;
1279 &mut self.hot_head
1280 } else {
1281 self.num_cold -= 1;
1282 self.weight_cold -= current_weight;
1283 &mut self.cold_head
1284 };
1285 if current_weight != 0 {
1286 let next = self.entries.unlink(idx);
1287 if *list_head == Some(idx) {
1288 *list_head = next;
1289 }
1290 }
1291 let shared = Plh::new(hash, idx);
1292 let (entry, _) = unsafe { self.entries.get_mut_unchecked(idx) };
1293 let Entry::Resident(r) = mem::replace(entry, Entry::Ghost(0)) else {
1294 unsafe { unreachable_unchecked() }
1295 };
1296 *entry = Entry::Placeholder(Placeholder {
1297 key: r.key,
1298 hot: state,
1299 shared: shared.clone(),
1300 });
1301 EntryOrPlaceholder::Replaced(shared, r.value)
1302 }
1303 };
1304 }
1305 }
1306 let (shared, is_new) = unsafe { self.non_resident_to_placeholder(hash, key, idx) };
1307 if is_new {
1308 EntryOrPlaceholder::NewPlaceholder(shared)
1309 } else {
1310 EntryOrPlaceholder::ExistingPlaceholder(shared)
1311 }
1312 }
1313
1314 unsafe fn non_resident_to_placeholder<Q>(
1318 &mut self,
1319 hash: u64,
1320 key: &Q,
1321 idx: Option<Token>,
1322 ) -> (Plh, bool)
1323 where
1324 Q: Hash + Equivalent<Key> + ToOwned<Owned = Key> + ?Sized,
1325 {
1326 if let Some(idx) = idx {
1327 let (entry, _) = unsafe { self.entries.get_mut_unchecked(idx) };
1328 match entry {
1329 Entry::Placeholder(p) => {
1330 record_hit_mut!(self);
1331 (p.shared.clone(), false)
1332 }
1333 Entry::Ghost(_) => {
1334 let shared = Plh::new(hash, idx);
1335 *entry = Entry::Placeholder(Placeholder {
1336 key: key.to_owned(),
1337 hot: ResidentState::Hot,
1338 shared: shared.clone(),
1339 });
1340 self.num_non_resident -= 1;
1341 let next = self.entries.unlink(idx);
1342 if self.ghost_head == Some(idx) {
1343 self.ghost_head = next;
1344 }
1345 record_miss_mut!(self);
1346 (shared, true)
1347 }
1348 Entry::Resident(_) => unsafe { unreachable_unchecked() },
1349 }
1350 } else {
1351 let idx = self.entries.next_free();
1352 let shared = Plh::new(hash, idx);
1353 let idx_ = self.entries.insert(Entry::Placeholder(Placeholder {
1354 key: key.to_owned(),
1355 hot: ResidentState::Cold,
1356 shared: shared.clone(),
1357 }));
1358 debug_assert_eq!(idx, idx_);
1359 self.map_insert(hash, idx);
1360 record_miss_mut!(self);
1361 (shared, true)
1362 }
1363 }
1364
1365 pub fn set_capacity(&mut self, new_weight_capacity: u64, lcs: &mut L::RequestState) {
1366 if self.weight_capacity == 0 {
1368 self.weight_capacity = new_weight_capacity;
1369 self.weight_target_hot = ((new_weight_capacity as f64 * DEFAULT_HOT_ALLOCATION) as u64)
1370 .clamp(new_weight_capacity.min(1), new_weight_capacity);
1371 } else {
1373 let old_new_ratio = new_weight_capacity as f64 / self.weight_capacity as f64;
1374 let hot_ratio = self.weight_target_hot as f64 / self.weight_capacity as f64;
1375
1376 self.weight_capacity = new_weight_capacity;
1377 self.weight_target_hot = ((new_weight_capacity as f64 * hot_ratio) as u64)
1378 .clamp(new_weight_capacity.min(1), new_weight_capacity);
1379 self.capacity_non_resident =
1380 (self.capacity_non_resident as f64 * old_new_ratio) as usize;
1381 }
1382
1383 while self.weight_hot + self.weight_cold > self.weight_capacity && self.advance_cold(lcs) {}
1385 while self.num_non_resident > self.capacity_non_resident {
1387 self.advance_ghost();
1388 }
1389 }
1390}
1391
1392struct WeightGuard<Key, Val, We: Weighter<Key, Val>, B, L, Plh: SharedPlaceholder> {
1395 shard: *mut CacheShard<Key, Val, We, B, L, Plh>,
1396 idx: Token,
1397 old_weight: u64,
1398}
1399
1400impl<Key, Val, We: Weighter<Key, Val>, B, L, Plh: SharedPlaceholder> Drop
1401 for WeightGuard<Key, Val, We, B, L, Plh>
1402{
1403 fn drop(&mut self) {
1404 unsafe {
1407 let shard = &mut *self.shard;
1408 let (entry, _) = shard.entries.get_unchecked(self.idx);
1409 let Entry::Resident(r) = entry else {
1410 unreachable_unchecked()
1411 };
1412 let new_weight = shard.weighter.weight(&r.key, &r.value);
1413 if self.old_weight != new_weight {
1414 shard.cold_change_weight(self.idx, self.old_weight, new_weight);
1415 }
1416 }
1417 }
1418}
1419
1420pub struct RefMut<'cache, Key, Val, We: Weighter<Key, Val>, B, L, Plh: SharedPlaceholder> {
1423 guard: WeightGuard<Key, Val, We, B, L, Plh>,
1424 _phantom: std::marker::PhantomData<&'cache mut CacheShard<Key, Val, We, B, L, Plh>>,
1425}
1426
1427impl<Key, Val, We: Weighter<Key, Val>, B, L, Plh: SharedPlaceholder>
1428 RefMut<'_, Key, Val, We, B, L, Plh>
1429{
1430 pub(crate) fn pair(&self) -> (&Key, &Val) {
1431 unsafe {
1434 let shard = &*self.guard.shard;
1435 let (entry, _) = shard.entries.get_unchecked(self.guard.idx);
1436 let Entry::Resident(Resident { key, value, .. }) = entry else {
1437 core::hint::unreachable_unchecked()
1438 };
1439 (key, value)
1440 }
1441 }
1442
1443 pub(crate) fn value_mut(&mut self) -> &mut Val {
1444 unsafe {
1446 let shard = &mut *self.guard.shard;
1447 let (entry, _) = shard.entries.get_mut_unchecked(self.guard.idx);
1448 let Entry::Resident(Resident { value, .. }) = entry else {
1449 core::hint::unreachable_unchecked()
1450 };
1451 value
1452 }
1453 }
1454}
1455
1456#[cfg(test)]
1457mod tests {
1458 use super::*;
1459
1460 #[cfg(not(feature = "stats"))]
1461 #[test]
1462 fn reserve_caps_ghost_headroom() {
1463 let mut shard = CacheShard::<
1467 u64,
1468 u64,
1469 crate::UnitWeighter,
1470 crate::DefaultHashBuilder,
1471 crate::sync::DefaultLifecycle<u64, u64>,
1472 crate::sync_placeholder::SharedPlaceholder<u64>,
1473 >::new(
1474 DEFAULT_HOT_ALLOCATION,
1475 0.5, 1_000_000, u64::MAX, crate::UnitWeighter,
1479 crate::DefaultHashBuilder::default(),
1480 crate::sync::DefaultLifecycle::default(),
1481 );
1482 assert_eq!(shard.capacity_non_resident, 500_000);
1483 shard.reserve(100);
1484 assert!(
1487 shard.entries.capacity() < 1_000,
1488 "slab over-allocated: {}",
1489 shard.entries.capacity()
1490 );
1491 }
1492
1493 #[test]
1494 fn entry_overhead() {
1495 use std::mem::size_of;
1496 assert_eq!(
1503 size_of::<Entry<u64, u64, crate::sync_placeholder::SharedPlaceholder<u64>>>()
1504 - size_of::<[u64; 2]>(),
1505 16
1506 );
1507 #[cfg(not(feature = "stats"))]
1508 let unsync_overhead = 16;
1509 #[cfg(feature = "stats")]
1510 let unsync_overhead = 24;
1511 assert_eq!(
1512 size_of::<Entry<u64, u64, crate::unsync::SharedPlaceholder>>() - size_of::<[u64; 2]>(),
1513 unsync_overhead
1514 );
1515 }
1516}