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 entries: Vec<Option<Entry<K, V>>>,
46 free_indices: Vec<usize>,
48 small: VecDeque<usize>,
50 main: VecDeque<usize>,
52 ghost: VecDeque<K>,
54 small_cap: usize,
56 main_cap: usize,
58 ghost_cap: usize,
60 hits: u64,
62 misses: u64,
63}
64
65#[derive(Debug, Clone, Copy, PartialEq, Eq)]
67enum Location {
68 Small(usize),
69 Main(usize),
70}
71
72impl Location {
73 fn idx(&self) -> usize {
74 match self {
75 Self::Small(i) | Self::Main(i) => *i,
76 }
77 }
78}
79
80#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
82pub struct S3FifoStats {
83 pub hits: u64,
85 pub misses: u64,
87 pub small_size: usize,
89 pub main_size: usize,
91 pub ghost_size: usize,
93 pub capacity: usize,
95}
96
97impl<K, V> S3Fifo<K, V>
98where
99 K: Hash + Eq + Clone,
100{
101 pub fn new(capacity: usize) -> Self {
106 let capacity = capacity.max(2);
107 let small_cap = (capacity / 10).max(1);
108 let main_cap = capacity - small_cap;
109 let ghost_cap = small_cap;
110
111 Self {
112 index: AHashMap::with_capacity(capacity),
113 entries: Vec::with_capacity(capacity),
114 free_indices: Vec::new(),
115 small: VecDeque::with_capacity(small_cap),
116 main: VecDeque::with_capacity(main_cap),
117 ghost: VecDeque::with_capacity(ghost_cap),
118 small_cap,
119 main_cap,
120 ghost_cap,
121 hits: 0,
122 misses: 0,
123 }
124 }
125
126 pub fn get(&mut self, key: &K) -> Option<&V> {
128 if let Some(loc) = self.index.get(key) {
129 self.hits += 1;
130 let idx = loc.idx();
131 let entry = self.entries[idx]
133 .as_mut()
134 .expect("S3Fifo invariant violated: valid index required");
135 entry.freq = entry.freq.saturating_add(1).min(3);
136 Some(&entry.value)
137 } else {
138 None
139 }
140 }
141
142 pub fn insert(&mut self, key: K, value: V) -> Option<V> {
145 if let Some(loc) = self.index.get(&key) {
147 let idx = loc.idx();
148 let entry = self.entries[idx]
149 .as_mut()
150 .expect("S3Fifo invariant violated: valid index required");
151 let old = std::mem::replace(&mut entry.value, value);
152 entry.freq = entry.freq.saturating_add(1).min(3);
153 return Some(old);
154 }
155
156 self.misses += 1;
157
158 let in_ghost = self.remove_from_ghost(&key);
160
161 if in_ghost {
162 self.evict_main_if_full();
164 let idx = self.alloc_entry(key.clone(), value);
165 self.main.push_back(idx);
166 self.index.insert(key, Location::Main(idx));
167 } else {
168 self.evict_small_if_full();
170 let idx = self.alloc_entry(key.clone(), value);
171 self.small.push_back(idx);
172 self.index.insert(key, Location::Small(idx));
173 }
174
175 None
176 }
177
178 pub fn remove(&mut self, key: &K) -> Option<V> {
180 let loc = self.index.remove(key)?;
181 let idx = loc.idx();
182
183 match loc {
187 Location::Small(_) => {
188 if let Some(pos) = self.small.iter().position(|&i| i == idx) {
189 self.small.remove(pos);
190 }
191 }
192 Location::Main(_) => {
193 if let Some(pos) = self.main.iter().position(|&i| i == idx) {
194 self.main.remove(pos);
195 }
196 }
197 }
198
199 self.free_entry(idx)
200 }
201
202 pub fn len(&self) -> usize {
204 self.index.len()
205 }
206
207 pub fn is_empty(&self) -> bool {
209 self.index.is_empty()
210 }
211
212 pub fn capacity(&self) -> usize {
214 self.small_cap + self.main_cap
215 }
216
217 pub fn stats(&self) -> S3FifoStats {
219 S3FifoStats {
220 hits: self.hits,
221 misses: self.misses,
222 small_size: self.small.len(),
223 main_size: self.main.len(),
224 ghost_size: self.ghost.len(),
225 capacity: self.small_cap + self.main_cap,
226 }
227 }
228
229 pub fn clear(&mut self) {
231 self.index.clear();
232 self.entries.clear();
233 self.free_indices.clear();
234 self.small.clear();
235 self.main.clear();
236 self.ghost.clear();
237 self.hits = 0;
238 self.misses = 0;
239 }
240
241 pub fn contains_key(&self, key: &K) -> bool {
243 self.index.contains_key(key)
244 }
245
246 fn alloc_entry(&mut self, key: K, value: V) -> usize {
250 let entry = Some(Entry {
251 key,
252 value,
253 freq: 0,
254 });
255
256 if let Some(idx) = self.free_indices.pop() {
257 self.entries[idx] = entry;
258 idx
259 } else {
260 let idx = self.entries.len();
261 self.entries.push(entry);
262 idx
263 }
264 }
265
266 fn free_entry(&mut self, idx: usize) -> Option<V> {
268 let entry = self.entries[idx].take()?;
269 self.free_indices.push(idx);
270 Some(entry.value)
271 }
272
273 fn remove_from_ghost(&mut self, key: &K) -> bool {
275 if let Some(pos) = self.ghost.iter().position(|k| k == key) {
276 self.ghost.remove(pos);
277 true
278 } else {
279 false
280 }
281 }
282
283 fn evict_small_if_full(&mut self) {
285 while self.small.len() >= self.small_cap {
286 if let Some(idx) = self.small.pop_front() {
287 let freq = self.entries[idx]
289 .as_ref()
290 .expect("S3Fifo invariant violated: valid index required")
291 .freq;
292
293 if freq > 0 {
294 self.entries[idx]
296 .as_mut()
297 .expect("S3Fifo invariant violated: valid index required")
298 .freq = 0;
299
300 let key = self.entries[idx]
302 .as_ref()
303 .expect("S3Fifo invariant violated: valid index required")
304 .key
305 .clone();
306
307 self.evict_main_if_full();
308
309 self.index.insert(key, Location::Main(idx));
310 self.main.push_back(idx);
311 } else {
312 let entry = self.entries[idx]
315 .take()
316 .expect("S3Fifo invariant violated: valid index required");
317 self.free_indices.push(idx);
318
319 self.index.remove(&entry.key);
320
321 if self.ghost.len() >= self.ghost_cap {
322 self.ghost.pop_front();
323 }
324 self.ghost.push_back(entry.key);
325 }
326 }
327 }
328 }
329
330 fn evict_main_if_full(&mut self) {
332 while self.main.len() >= self.main_cap {
333 if let Some(idx) = self.main.pop_front() {
334 let freq = self.entries[idx]
335 .as_ref()
336 .expect("S3Fifo invariant violated: valid index required")
337 .freq;
338
339 if freq > 0 {
340 self.entries[idx]
342 .as_mut()
343 .expect("S3Fifo invariant violated: valid index required")
344 .freq -= 1;
345 self.main.push_back(idx);
346 } else {
347 let entry = self.entries[idx]
349 .take()
350 .expect("S3Fifo invariant violated: valid index required");
351 self.free_indices.push(idx);
352 self.index.remove(&entry.key);
353 }
354 }
355 }
356 }
357}
358
359impl<K, V> std::fmt::Debug for S3Fifo<K, V> {
360 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
361 f.debug_struct("S3Fifo")
362 .field("small", &self.small.len())
363 .field("main", &self.main.len())
364 .field("ghost", &self.ghost.len())
365 .field("hits", &self.hits)
366 .field("misses", &self.misses)
367 .finish()
368 }
369}
370
371#[cfg(test)]
372mod tests {
373 use super::*;
374
375 #[test]
376 fn empty_cache() {
377 let cache: S3Fifo<&str, i32> = S3Fifo::new(10);
378 assert!(cache.is_empty());
379 assert_eq!(cache.len(), 0);
380 }
381
382 #[test]
383 fn insert_and_get() {
384 let mut cache = S3Fifo::new(10);
385 cache.insert("key1", 42);
386 assert_eq!(cache.get(&"key1"), Some(&42));
387 assert_eq!(cache.len(), 1);
388 }
389
390 #[test]
391 fn miss_returns_none() {
392 let mut cache: S3Fifo<&str, i32> = S3Fifo::new(10);
393 assert_eq!(cache.get(&"missing"), None);
394 }
395
396 #[test]
397 fn update_existing_key() {
398 let mut cache = S3Fifo::new(10);
399 cache.insert("key1", 1);
400 let old = cache.insert("key1", 2);
401 assert_eq!(old, Some(1));
402 assert_eq!(cache.get(&"key1"), Some(&2));
403 assert_eq!(cache.len(), 1);
404 }
405
406 #[test]
407 fn remove_key() {
408 let mut cache = S3Fifo::new(10);
409 cache.insert("key1", 42);
410 let removed = cache.remove(&"key1");
411 assert_eq!(removed, Some(42));
412 assert!(cache.is_empty());
413 assert_eq!(cache.get(&"key1"), None);
414 }
415
416 #[test]
417 fn remove_nonexistent() {
418 let mut cache: S3Fifo<&str, i32> = S3Fifo::new(10);
419 assert_eq!(cache.remove(&"missing"), None);
420 }
421
422 #[test]
423 fn eviction_at_capacity() {
424 let mut cache = S3Fifo::new(5);
425 for i in 0..10 {
426 cache.insert(i, i * 10);
427 }
428 assert!(cache.len() <= cache.capacity());
430 }
431
432 #[test]
433 fn small_to_main_promotion() {
434 let mut cache = S3Fifo::new(10); cache.insert("keep", 1);
439 cache.get(&"keep"); cache.insert("new", 2);
443
444 assert_eq!(cache.get(&"keep"), Some(&1));
446 }
447
448 #[test]
449 fn ghost_readmission() {
450 let mut cache = S3Fifo::new(10); cache.insert("ghost_key", 1);
456 cache.insert("displacer", 2); assert_eq!(cache.get(&"ghost_key"), None);
460
461 cache.insert("ghost_key", 3);
463 assert_eq!(cache.get(&"ghost_key"), Some(&3));
464 }
465
466 #[test]
467 fn stats_tracking() {
468 let mut cache = S3Fifo::new(20);
470 cache.insert("a", 1);
471 cache.insert("b", 2);
472 cache.get(&"a"); cache.get(&"a"); cache.get(&"c"); let stats = cache.stats();
477 assert_eq!(stats.hits, 2);
478 assert_eq!(stats.misses, 2); }
481
482 #[test]
483 fn clear_resets() {
484 let mut cache = S3Fifo::new(10);
485 cache.insert("a", 1);
486 cache.insert("b", 2);
487 cache.get(&"a");
488 cache.clear();
489
490 assert!(cache.is_empty());
491 assert_eq!(cache.len(), 0);
492 let stats = cache.stats();
493 assert_eq!(stats.hits, 0);
494 assert_eq!(stats.misses, 0);
495 assert_eq!(stats.ghost_size, 0);
496 }
497
498 #[test]
499 fn contains_key() {
500 let mut cache = S3Fifo::new(10);
501 cache.insert("a", 1);
502 assert!(cache.contains_key(&"a"));
503 assert!(!cache.contains_key(&"b"));
504 }
505
506 #[test]
507 fn capacity_split() {
508 let cache: S3Fifo<i32, i32> = S3Fifo::new(100);
509 assert_eq!(cache.capacity(), 100);
510 assert_eq!(cache.small_cap, 10);
511 assert_eq!(cache.main_cap, 90);
512 assert_eq!(cache.ghost_cap, 10);
513 }
514
515 #[test]
516 fn minimum_capacity() {
517 let cache: S3Fifo<i32, i32> = S3Fifo::new(0);
518 assert!(cache.capacity() >= 2);
519 }
520
521 #[test]
522 fn freq_capped_at_3() {
523 let mut cache = S3Fifo::new(10);
524 cache.insert("a", 1);
525 for _ in 0..10 {
526 cache.get(&"a");
527 }
528 assert_eq!(cache.get(&"a"), Some(&1));
530 }
531
532 #[test]
533 fn main_eviction_gives_second_chance() {
534 let mut cache = S3Fifo::new(5); for i in 0..4 {
539 cache.insert(i, i);
540 cache.get(&i);
542 }
543
544 for i in 10..20 {
546 cache.insert(i, i);
547 }
548
549 assert!(cache.len() <= cache.capacity());
551 }
552
553 #[test]
554 fn debug_format() {
555 let cache: S3Fifo<&str, i32> = S3Fifo::new(10);
556 let debug = format!("{cache:?}");
557 assert!(debug.contains("S3Fifo"));
558 assert!(debug.contains("small"));
559 assert!(debug.contains("main"));
560 }
561
562 #[test]
563 fn large_workload() {
564 let mut cache = S3Fifo::new(100);
565
566 for i in 0..200 {
569 cache.insert(i, i * 10);
570 if i >= 50 {
572 for hot in 50..std::cmp::min(i, 100) {
573 cache.get(&hot);
574 }
575 }
576 }
577
578 let mut hot_hits = 0;
580 for i in 50..100 {
581 if cache.get(&i).is_some() {
582 hot_hits += 1;
583 }
584 }
585
586 assert!(hot_hits > 20, "hot set retention: {hot_hits}/50");
588 }
589
590 #[test]
591 fn scan_resistance() {
592 let mut cache = S3Fifo::new(100);
593
594 for i in 0..50 {
596 cache.insert(i, i);
597 cache.get(&i);
598 cache.get(&i);
599 }
600
601 for i in 1000..2000 {
603 cache.insert(i, i);
604 }
605
606 let mut survivors = 0;
608 for i in 0..50 {
609 if cache.get(&i).is_some() {
610 survivors += 1;
611 }
612 }
613
614 assert!(
616 survivors > 10,
617 "scan resistance: {survivors}/50 working set items survived"
618 );
619 }
620
621 #[test]
622 fn ghost_size_bounded() {
623 let mut cache = S3Fifo::new(10);
624
625 for i in 0..100 {
627 cache.insert(i, i);
628 }
629
630 let stats = cache.stats();
631 assert!(
632 stats.ghost_size <= cache.ghost_cap,
633 "ghost should be bounded: {} <= {}",
634 stats.ghost_size,
635 cache.ghost_cap
636 );
637 }
638}