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