1#![forbid(unsafe_code)]
2
3use ahash::AHashMap;
30use std::collections::VecDeque;
31use std::hash::Hash;
32
33struct Entry<K, V> {
35 key: K,
36 value: V,
37 freq: u8,
38}
39
40pub struct S3Fifo<K, V> {
42 index: AHashMap<K, Location>,
44 small: VecDeque<Entry<K, V>>,
46 main: VecDeque<Entry<K, V>>,
48 ghost: VecDeque<K>,
50 small_cap: usize,
52 main_cap: usize,
54 ghost_cap: usize,
56 hits: u64,
58 misses: u64,
59}
60
61#[derive(Debug, Clone, Copy, PartialEq, Eq)]
63enum Location {
64 Small,
65 Main,
66}
67
68#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
70pub struct S3FifoStats {
71 pub hits: u64,
73 pub misses: u64,
75 pub small_size: usize,
77 pub main_size: usize,
79 pub ghost_size: usize,
81 pub capacity: usize,
83}
84
85impl<K, V> S3Fifo<K, V>
86where
87 K: Hash + Eq + Clone,
88{
89 pub fn new(capacity: usize) -> Self {
94 let capacity = capacity.max(2);
95 let small_cap = (capacity / 10).max(1);
96 let main_cap = capacity - small_cap;
97 let ghost_cap = small_cap;
98
99 Self {
100 index: AHashMap::with_capacity(capacity),
101 small: VecDeque::with_capacity(small_cap),
102 main: VecDeque::with_capacity(main_cap),
103 ghost: VecDeque::with_capacity(ghost_cap),
104 small_cap,
105 main_cap,
106 ghost_cap,
107 hits: 0,
108 misses: 0,
109 }
110 }
111
112 pub fn get(&mut self, key: &K) -> Option<&V> {
114 match self.index.get(key)? {
115 Location::Small => {
116 self.hits += 1;
117 for entry in self.small.iter_mut() {
119 if entry.key == *key {
120 entry.freq = entry.freq.saturating_add(1).min(3);
121 return Some(&entry.value);
122 }
123 }
124 None
125 }
126 Location::Main => {
127 self.hits += 1;
128 for entry in self.main.iter_mut() {
130 if entry.key == *key {
131 entry.freq = entry.freq.saturating_add(1).min(3);
132 return Some(&entry.value);
133 }
134 }
135 None
136 }
137 }
138 }
139
140 pub fn insert(&mut self, key: K, value: V) -> Option<V> {
143 if let Some(&loc) = self.index.get(&key) {
145 match loc {
146 Location::Small => {
147 for entry in self.small.iter_mut() {
148 if entry.key == key {
149 let old = std::mem::replace(&mut entry.value, value);
150 entry.freq = entry.freq.saturating_add(1).min(3);
151 return Some(old);
152 }
153 }
154 }
155 Location::Main => {
156 for entry in self.main.iter_mut() {
157 if entry.key == key {
158 let old = std::mem::replace(&mut entry.value, value);
159 entry.freq = entry.freq.saturating_add(1).min(3);
160 return Some(old);
161 }
162 }
163 }
164 }
165 }
166
167 self.misses += 1;
168
169 let in_ghost = self.remove_from_ghost(&key);
171
172 if in_ghost {
173 self.evict_main_if_full();
175 self.main.push_back(Entry {
176 key: key.clone(),
177 value,
178 freq: 0,
179 });
180 self.index.insert(key, Location::Main);
181 } else {
182 self.evict_small_if_full();
184 self.small.push_back(Entry {
185 key: key.clone(),
186 value,
187 freq: 0,
188 });
189 self.index.insert(key, Location::Small);
190 }
191
192 None
193 }
194
195 pub fn remove(&mut self, key: &K) -> Option<V> {
197 let loc = self.index.remove(key)?;
198 match loc {
199 Location::Small => {
200 if let Some(pos) = self.small.iter().position(|e| e.key == *key) {
201 return Some(self.small.remove(pos).unwrap().value);
202 }
203 }
204 Location::Main => {
205 if let Some(pos) = self.main.iter().position(|e| e.key == *key) {
206 return Some(self.main.remove(pos).unwrap().value);
207 }
208 }
209 }
210 None
211 }
212
213 pub fn len(&self) -> usize {
215 self.small.len() + self.main.len()
216 }
217
218 pub fn is_empty(&self) -> bool {
220 self.small.is_empty() && self.main.is_empty()
221 }
222
223 pub fn capacity(&self) -> usize {
225 self.small_cap + self.main_cap
226 }
227
228 pub fn stats(&self) -> S3FifoStats {
230 S3FifoStats {
231 hits: self.hits,
232 misses: self.misses,
233 small_size: self.small.len(),
234 main_size: self.main.len(),
235 ghost_size: self.ghost.len(),
236 capacity: self.small_cap + self.main_cap,
237 }
238 }
239
240 pub fn clear(&mut self) {
242 self.index.clear();
243 self.small.clear();
244 self.main.clear();
245 self.ghost.clear();
246 self.hits = 0;
247 self.misses = 0;
248 }
249
250 pub fn contains_key(&self, key: &K) -> bool {
252 self.index.contains_key(key)
253 }
254
255 fn remove_from_ghost(&mut self, key: &K) -> bool {
259 if let Some(pos) = self.ghost.iter().position(|k| k == key) {
260 self.ghost.remove(pos);
261 true
262 } else {
263 false
264 }
265 }
266
267 fn evict_small_if_full(&mut self) {
269 while self.small.len() >= self.small_cap {
270 if let Some(entry) = self.small.pop_front() {
271 self.index.remove(&entry.key);
272 if entry.freq > 0 {
273 self.evict_main_if_full();
275 self.index.insert(entry.key.clone(), Location::Main);
276 self.main.push_back(Entry {
277 key: entry.key,
278 value: entry.value,
279 freq: 0, });
281 } else {
282 if self.ghost.len() >= self.ghost_cap {
284 self.ghost.pop_front();
285 }
286 self.ghost.push_back(entry.key);
287 }
288 }
289 }
290 }
291
292 fn evict_main_if_full(&mut self) {
294 while self.main.len() >= self.main_cap {
295 if let Some(mut entry) = self.main.pop_front() {
296 if entry.freq > 0 {
297 entry.freq -= 1;
299 self.main.push_back(entry);
300 } else {
301 self.index.remove(&entry.key);
303 }
304 }
305 }
306 }
307}
308
309impl<K, V> std::fmt::Debug for S3Fifo<K, V> {
310 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
311 f.debug_struct("S3Fifo")
312 .field("small", &self.small.len())
313 .field("main", &self.main.len())
314 .field("ghost", &self.ghost.len())
315 .field("hits", &self.hits)
316 .field("misses", &self.misses)
317 .finish()
318 }
319}
320
321#[cfg(test)]
322mod tests {
323 use super::*;
324
325 #[test]
326 fn empty_cache() {
327 let cache: S3Fifo<&str, i32> = S3Fifo::new(10);
328 assert!(cache.is_empty());
329 assert_eq!(cache.len(), 0);
330 }
331
332 #[test]
333 fn insert_and_get() {
334 let mut cache = S3Fifo::new(10);
335 cache.insert("key1", 42);
336 assert_eq!(cache.get(&"key1"), Some(&42));
337 assert_eq!(cache.len(), 1);
338 }
339
340 #[test]
341 fn miss_returns_none() {
342 let mut cache: S3Fifo<&str, i32> = S3Fifo::new(10);
343 assert_eq!(cache.get(&"missing"), None);
344 }
345
346 #[test]
347 fn update_existing_key() {
348 let mut cache = S3Fifo::new(10);
349 cache.insert("key1", 1);
350 let old = cache.insert("key1", 2);
351 assert_eq!(old, Some(1));
352 assert_eq!(cache.get(&"key1"), Some(&2));
353 assert_eq!(cache.len(), 1);
354 }
355
356 #[test]
357 fn remove_key() {
358 let mut cache = S3Fifo::new(10);
359 cache.insert("key1", 42);
360 let removed = cache.remove(&"key1");
361 assert_eq!(removed, Some(42));
362 assert!(cache.is_empty());
363 assert_eq!(cache.get(&"key1"), None);
364 }
365
366 #[test]
367 fn remove_nonexistent() {
368 let mut cache: S3Fifo<&str, i32> = S3Fifo::new(10);
369 assert_eq!(cache.remove(&"missing"), None);
370 }
371
372 #[test]
373 fn eviction_at_capacity() {
374 let mut cache = S3Fifo::new(5);
375 for i in 0..10 {
376 cache.insert(i, i * 10);
377 }
378 assert!(cache.len() <= cache.capacity());
380 }
381
382 #[test]
383 fn small_to_main_promotion() {
384 let mut cache = S3Fifo::new(10); cache.insert("keep", 1);
389 cache.get(&"keep"); cache.insert("new", 2);
393
394 assert_eq!(cache.get(&"keep"), Some(&1));
396 }
397
398 #[test]
399 fn ghost_readmission() {
400 let mut cache = S3Fifo::new(10); cache.insert("ghost_key", 1);
406 cache.insert("displacer", 2); assert_eq!(cache.get(&"ghost_key"), None);
410
411 cache.insert("ghost_key", 3);
413 assert_eq!(cache.get(&"ghost_key"), Some(&3));
414 }
415
416 #[test]
417 fn stats_tracking() {
418 let mut cache = S3Fifo::new(20);
420 cache.insert("a", 1);
421 cache.insert("b", 2);
422 cache.get(&"a"); cache.get(&"a"); cache.get(&"c"); let stats = cache.stats();
427 assert_eq!(stats.hits, 2);
428 assert_eq!(stats.misses, 2); }
431
432 #[test]
433 fn clear_resets() {
434 let mut cache = S3Fifo::new(10);
435 cache.insert("a", 1);
436 cache.insert("b", 2);
437 cache.get(&"a");
438 cache.clear();
439
440 assert!(cache.is_empty());
441 assert_eq!(cache.len(), 0);
442 let stats = cache.stats();
443 assert_eq!(stats.hits, 0);
444 assert_eq!(stats.misses, 0);
445 assert_eq!(stats.ghost_size, 0);
446 }
447
448 #[test]
449 fn contains_key() {
450 let mut cache = S3Fifo::new(10);
451 cache.insert("a", 1);
452 assert!(cache.contains_key(&"a"));
453 assert!(!cache.contains_key(&"b"));
454 }
455
456 #[test]
457 fn capacity_split() {
458 let cache: S3Fifo<i32, i32> = S3Fifo::new(100);
459 assert_eq!(cache.capacity(), 100);
460 assert_eq!(cache.small_cap, 10);
461 assert_eq!(cache.main_cap, 90);
462 assert_eq!(cache.ghost_cap, 10);
463 }
464
465 #[test]
466 fn minimum_capacity() {
467 let cache: S3Fifo<i32, i32> = S3Fifo::new(0);
468 assert!(cache.capacity() >= 2);
469 }
470
471 #[test]
472 fn freq_capped_at_3() {
473 let mut cache = S3Fifo::new(10);
474 cache.insert("a", 1);
475 for _ in 0..10 {
476 cache.get(&"a");
477 }
478 assert_eq!(cache.get(&"a"), Some(&1));
480 }
481
482 #[test]
483 fn main_eviction_gives_second_chance() {
484 let mut cache = S3Fifo::new(5); for i in 0..4 {
489 cache.insert(i, i);
490 cache.get(&i);
492 }
493
494 for i in 10..20 {
496 cache.insert(i, i);
497 }
498
499 assert!(cache.len() <= cache.capacity());
501 }
502
503 #[test]
504 fn debug_format() {
505 let cache: S3Fifo<&str, i32> = S3Fifo::new(10);
506 let debug = format!("{cache:?}");
507 assert!(debug.contains("S3Fifo"));
508 assert!(debug.contains("small"));
509 assert!(debug.contains("main"));
510 }
511
512 #[test]
513 fn large_workload() {
514 let mut cache = S3Fifo::new(100);
515
516 for i in 0..200 {
519 cache.insert(i, i * 10);
520 if i >= 50 {
522 for hot in 50..std::cmp::min(i, 100) {
523 cache.get(&hot);
524 }
525 }
526 }
527
528 let mut hot_hits = 0;
530 for i in 50..100 {
531 if cache.get(&i).is_some() {
532 hot_hits += 1;
533 }
534 }
535
536 assert!(hot_hits > 20, "hot set retention: {hot_hits}/50");
538 }
539
540 #[test]
541 fn scan_resistance() {
542 let mut cache = S3Fifo::new(100);
543
544 for i in 0..50 {
546 cache.insert(i, i);
547 cache.get(&i);
548 cache.get(&i);
549 }
550
551 for i in 1000..2000 {
553 cache.insert(i, i);
554 }
555
556 let mut survivors = 0;
558 for i in 0..50 {
559 if cache.get(&i).is_some() {
560 survivors += 1;
561 }
562 }
563
564 assert!(
566 survivors > 10,
567 "scan resistance: {survivors}/50 working set items survived"
568 );
569 }
570
571 #[test]
572 fn ghost_size_bounded() {
573 let mut cache = S3Fifo::new(10);
574
575 for i in 0..100 {
577 cache.insert(i, i);
578 }
579
580 let stats = cache.stats();
581 assert!(
582 stats.ghost_size <= cache.ghost_cap,
583 "ghost should be bounded: {} <= {}",
584 stats.ghost_size,
585 cache.ghost_cap
586 );
587 }
588}