1use alloc::boxed::Box;
40use alloc::sync::Arc;
41use alloc::vec::Vec;
42use core::hash::Hasher;
43use core::sync::atomic::{AtomicUsize, Ordering};
44
45use arc_swap::ArcSwap;
46use rustc_hash::FxHasher;
47
48const LOAD_FACTOR_NUM: usize = 4;
49const LOAD_FACTOR_DEN: usize = 3;
50
51#[inline]
52fn fxhash(key: &[u8]) -> u64 {
53 let mut hasher = FxHasher::default();
54 hasher.write(key);
55 hasher.finish()
56}
57
58enum SlotState<V> {
62 Empty,
64 Tomb,
67 Full {
69 hash: u64,
70 key: Box<[u8]>,
71 bucket: Arc<V>,
72 },
73}
74
75struct Slot<V>(ArcSwap<SlotState<V>>);
78
79pub struct BucketTable<V> {
81 slots: Box<[Slot<V>]>,
82 tomb: Arc<SlotState<V>>,
87 count: AtomicUsize,
89 mask: u64,
90 max_keys: usize,
91}
92
93impl<V> BucketTable<V> {
94 #[must_use]
104 pub fn with_max_keys(max_keys: usize) -> Self {
105 let needed = max_keys
106 .saturating_mul(LOAD_FACTOR_NUM)
107 .div_ceil(LOAD_FACTOR_DEN);
108 let mut capacity = needed.max(1).next_power_of_two().max(2);
109 if capacity <= max_keys {
110 capacity = (max_keys + 1).next_power_of_two();
111 }
112 assert!(
113 capacity > max_keys,
114 "bucket table capacity must exceed max_keys for probe termination"
115 );
116 let empty: Arc<SlotState<V>> = Arc::new(SlotState::Empty);
117 let tomb: Arc<SlotState<V>> = Arc::new(SlotState::Tomb);
118 let mut slots = Vec::with_capacity(capacity);
119 for _ in 0..capacity {
120 slots.push(Slot(ArcSwap::new(empty.clone())));
121 }
122 Self {
123 slots: slots.into_boxed_slice(),
124 tomb,
125 count: AtomicUsize::new(0),
126 mask: (capacity as u64) - 1,
127 max_keys,
128 }
129 }
130
131 #[inline]
132 fn capacity(&self) -> usize {
133 self.slots.len()
134 }
135
136 #[inline]
137 fn home(&self, hash: u64) -> usize {
138 (hash & self.mask) as usize
139 }
140
141 #[inline]
142 fn next(&self, index: usize) -> usize {
143 ((index as u64 + 1) & self.mask) as usize
144 }
145
146 #[must_use]
148 pub fn len(&self) -> usize {
149 self.count.load(Ordering::Relaxed)
150 }
151
152 #[must_use]
153 pub fn is_empty(&self) -> bool {
154 self.len() == 0
155 }
156
157 pub fn get_or_insert(&self, key: &[u8], make: impl FnOnce() -> V) -> Arc<V> {
165 let hash = fxhash(key);
166 if let Some(bucket) = self.get(key, hash) {
168 return bucket;
169 }
170 let bucket = Arc::new(make());
173 let full = Arc::new(SlotState::Full {
174 hash,
175 key: key.into(),
176 bucket: bucket.clone(),
177 });
178 'restart: loop {
179 let mut first_tomb: Option<usize> = None;
180 let mut index = self.home(hash);
181 for _ in 0..self.capacity() {
182 let guard = self.slots[index].0.load();
183 match &**guard {
184 SlotState::Empty => {
185 let (target, current) = match first_tomb {
187 Some(tomb_index) => (tomb_index, self.tomb.clone()),
188 None => (index, arc_swap::Guard::into_inner(guard)),
189 };
190 if self.install(target, ¤t, &full) {
191 return bucket;
192 }
193 continue 'restart;
194 }
195 SlotState::Tomb => {
196 if first_tomb.is_none() {
197 first_tomb = Some(index);
198 }
199 index = self.next(index);
200 }
201 SlotState::Full {
202 hash: stored_hash,
203 key: stored_key,
204 bucket: found,
205 } => {
206 if *stored_hash == hash && stored_key.as_ref() == key {
208 return found.clone();
209 }
210 index = self.next(index);
211 }
212 }
213 }
214 if let Some(tomb_index) = first_tomb {
219 let tomb = self.tomb.clone();
220 if self.install(tomb_index, &tomb, &full) {
221 return bucket;
222 }
223 continue 'restart;
224 }
225 debug_assert!(
226 self.count.load(Ordering::Relaxed) >= self.max_keys,
227 "wrapped with no empty/tomb yet count below max_keys"
228 );
229 self.evict_one_lru(|_| 0);
230 }
231 }
232
233 #[must_use]
238 pub fn get(&self, key: &[u8], hash: u64) -> Option<Arc<V>> {
239 let mut index = self.home(hash);
240 for _ in 0..self.capacity() {
241 let guard = self.slots[index].0.load();
242 match &**guard {
243 SlotState::Empty => return None,
244 SlotState::Full {
245 hash: stored_hash,
246 key: stored_key,
247 bucket,
248 } if *stored_hash == hash && stored_key.as_ref() == key => {
249 return Some(bucket.clone());
250 }
251 _ => index = self.next(index),
252 }
253 }
254 None
255 }
256
257 fn install(
262 &self,
263 target: usize,
264 current: &Arc<SlotState<V>>,
265 full: &Arc<SlotState<V>>,
266 ) -> bool {
267 let prev = self.slots[target].0.compare_and_swap(current, full.clone());
268 let won = Arc::ptr_eq(&arc_swap::Guard::into_inner(prev), current);
269 if won {
270 self.count.fetch_add(1, Ordering::Relaxed);
271 }
272 won
273 }
274
275 fn tombstone(&self, index: usize, expected_full: &Arc<SlotState<V>>) {
280 let prev = self.slots[index]
281 .0
282 .compare_and_swap(expected_full, self.tomb.clone());
283 if Arc::ptr_eq(&arc_swap::Guard::into_inner(prev), expected_full) {
284 self.count.fetch_sub(1, Ordering::Relaxed);
285 }
286 }
287
288 pub fn evict_one_lru(&self, metric: impl Fn(&V) -> u64) {
294 let mut victim: Option<(usize, Arc<SlotState<V>>)> = None;
295 let mut victim_metric: u64 = 0;
296 for index in 0..self.capacity() {
297 let guard = self.slots[index].0.load();
298 if let SlotState::Full { bucket, .. } = &**guard {
299 let measured = metric(bucket);
300 if victim.is_none() || measured < victim_metric {
301 victim = Some((index, arc_swap::Guard::into_inner(guard)));
302 victim_metric = measured;
303 }
304 }
305 }
306 if let Some((index, full)) = victim {
307 self.tombstone(index, &full);
308 }
309 }
310
311 pub fn clear(&self) {
318 for index in 0..self.capacity() {
319 let guard = self.slots[index].0.load();
320 if matches!(&**guard, SlotState::Full { .. }) {
321 self.tombstone(index, &arc_swap::Guard::into_inner(guard));
322 }
323 }
324 }
325
326 pub fn sweep_idle(&self, now: u64, ttl: u64, metric: impl Fn(&V) -> u64) {
329 let cutoff = now.saturating_sub(ttl);
330 for index in 0..self.capacity() {
331 let guard = self.slots[index].0.load();
332 let idle = match &**guard {
333 SlotState::Full { bucket, .. } => metric(bucket) < cutoff,
334 _ => false,
335 };
336 if idle {
337 self.tombstone(index, &arc_swap::Guard::into_inner(guard));
338 }
339 }
340 }
341}
342
343#[cfg(test)]
344#[allow(clippy::unwrap_used, clippy::expect_used)]
345mod tests {
346 use super::*;
347 use core::sync::atomic::AtomicU64 as StdAtomicU64;
348 use core::sync::atomic::Ordering::Relaxed;
349 use std::sync::Barrier;
350
351 struct Timed {
355 last_access: StdAtomicU64,
356 }
357
358 impl Timed {
359 fn at(value: u64) -> Self {
360 Self {
361 last_access: StdAtomicU64::new(value),
362 }
363 }
364 fn set(&self, value: u64) {
365 self.last_access.store(value, Relaxed);
366 }
367 fn metric(value: &Timed) -> u64 {
368 value.last_access.load(Relaxed)
369 }
370 }
371
372 impl<V> BucketTable<V> {
375 fn full_key_at(&self, index: usize) -> Option<Vec<u8>> {
376 match &**self.slots[index].0.load() {
377 SlotState::Full { key, .. } => Some(key.to_vec()),
378 _ => None,
379 }
380 }
381 }
382
383 #[proxima::test]
384 #[case::one(1)]
385 #[case::two(2)]
386 #[case::three(3)]
387 #[case::seven(7)]
388 #[case::hundred(100)]
389 #[case::hundred_k(100_000)]
390 async fn cap_exceeds_max_keys(#[case] max_keys: usize) {
391 let table: BucketTable<u8> = BucketTable::with_max_keys(max_keys);
392 assert!(
393 table.capacity() > max_keys,
394 "cap {} must exceed max_keys {max_keys}",
395 table.capacity()
396 );
397 assert!(table.capacity().is_power_of_two());
398 }
399
400 #[test]
401 fn phase_walk_evict_reuses_tombstone_and_sweeps() {
402 let table: BucketTable<Timed> = BucketTable::with_max_keys(8);
406
407 let arc_a = table.get_or_insert(b"a", || Timed::at(10));
408 let arc_b = table.get_or_insert(b"b", || Timed::at(20));
409 assert_eq!(table.len(), 2);
410
411 let arc_a_again = table.get_or_insert(b"a", || Timed::at(999));
412 assert!(
413 Arc::ptr_eq(&arc_a, &arc_a_again),
414 "same key returns the same Arc"
415 );
416 assert_eq!(Arc::strong_count(&arc_a), 3, "no second make ran for a");
417
418 arc_a.set(30);
420 table.evict_one_lru(Timed::metric);
421 assert_eq!(table.len(), 1);
422 let arc_b_after = table.get_or_insert(b"b", || Timed::at(40));
423 assert!(
424 !Arc::ptr_eq(&arc_b, &arc_b_after),
425 "b was evicted, fresh Arc on reinsert"
426 );
427
428 let arc_c = table.get_or_insert(b"c", || Timed::at(15));
429 assert_eq!(table.len(), 3, "c reuses a tombstone, three live");
430
431 arc_a.set(30);
434 arc_c.set(15);
435 arc_b_after.set(40);
436 table.sweep_idle(50, 20, Timed::metric);
437 assert_eq!(table.len(), 2, "only c (strictly < cutoff) swept");
438 assert!(Arc::ptr_eq(
439 &table.get_or_insert(b"a", || Timed::at(0)),
440 &arc_a
441 ));
442 assert!(Arc::ptr_eq(
443 &table.get_or_insert(b"b", || Timed::at(0)),
444 &arc_b_after
445 ));
446 }
447
448 #[test]
449 fn sweep_strict_less_than_cutoff_spares_equal() {
450 let table: BucketTable<Timed> = BucketTable::with_max_keys(8);
451 let equal = table.get_or_insert(b"equal", || Timed::at(100));
452 let below = table.get_or_insert(b"below", || Timed::at(99));
453 table.sweep_idle(200, 100, Timed::metric);
455 assert_eq!(table.len(), 1);
456 assert!(Arc::ptr_eq(
457 &table.get_or_insert(b"equal", || Timed::at(0)),
458 &equal
459 ));
460 drop(below);
461 }
462
463 #[test]
464 fn lru_tie_breaks_to_lowest_slot_index() {
465 let table: BucketTable<Timed> = BucketTable::with_max_keys(8);
467 let keys: [&[u8]; 3] = [b"k0", b"k1", b"k2"];
468 let mut arcs = Vec::new();
469 for key in keys {
470 arcs.push(table.get_or_insert(key, || Timed::at(7)));
471 }
472 let mut occupied: Vec<(usize, &[u8])> = Vec::new();
474 for slot_index in 0..table.capacity() {
475 if let Some(key_bytes) = table.full_key_at(slot_index)
476 && let Some(found) = keys
477 .iter()
478 .copied()
479 .find(|candidate| *candidate == key_bytes.as_slice())
480 {
481 occupied.push((slot_index, found));
482 }
483 }
484 occupied.sort_by_key(|(slot_index, _)| *slot_index);
485 let lowest_key = occupied[0].1;
486
487 table.evict_one_lru(Timed::metric);
488 assert_eq!(table.len(), 2);
489 let reinserted = table.get_or_insert(lowest_key, || Timed::at(7));
491 let original = arcs
492 .iter()
493 .zip(keys)
494 .find(|(_, key)| *key == lowest_key)
495 .map(|(arc, _)| arc)
496 .unwrap();
497 assert!(
498 !Arc::ptr_eq(&reinserted, original),
499 "lowest-index slot was the victim"
500 );
501 }
502
503 #[test]
508 fn collision_probe_and_concurrent_same_key_share_one_bucket() {
509 let table: Arc<BucketTable<Timed>> = Arc::new(BucketTable::with_max_keys(1));
510 assert_eq!(table.capacity(), 2);
511
512 let make_calls = Arc::new(StdAtomicU64::new(0));
516 let barrier = Arc::new(Barrier::new(2));
517 let mut handles = Vec::new();
518 for _ in 0..2 {
519 let table = table.clone();
520 let make_calls = make_calls.clone();
521 let barrier = barrier.clone();
522 handles.push(std::thread::spawn(move || {
523 barrier.wait();
524 table.get_or_insert(b"shared", || {
525 make_calls.fetch_add(1, Relaxed);
526 Timed::at(1)
527 })
528 }));
529 }
530 let first = handles.pop().unwrap().join().unwrap();
531 let second = handles.pop().unwrap().join().unwrap();
532 assert!(
533 Arc::ptr_eq(&first, &second),
534 "racing same key shares one bucket"
535 );
536 let builds = make_calls.load(Relaxed);
537 assert!(
538 (1..=2).contains(&builds),
539 "make ran at most once per racer, got {builds}"
540 );
541 assert_eq!(table.len(), 1, "exactly one entry published");
542
543 let other = table.get_or_insert(b"other", || Timed::at(2));
545 assert_eq!(table.len(), 2);
546
547 for slot_index in 0..table.capacity() {
549 if table.full_key_at(slot_index).as_deref() == Some(b"shared".as_slice()) {
550 let full = arc_swap::Guard::into_inner(table.slots[slot_index].0.load());
551 table.tombstone(slot_index, &full);
552 }
553 }
554 assert_eq!(table.len(), 1);
555 let other_again = table.get_or_insert(b"other", || Timed::at(99));
556 assert!(
557 Arc::ptr_eq(&other, &other_again),
558 "found across the tombstone"
559 );
560
561 let reclaimed = table.get_or_insert(b"shared", || Timed::at(3));
563 assert!(
564 !Arc::ptr_eq(&reclaimed, &first),
565 "reclaimed slot, fresh Arc"
566 );
567 assert_eq!(table.len(), 2);
568 }
569
570 #[test]
571 fn concurrent_same_key_reclaim_no_duplicate() {
572 let table: Arc<BucketTable<Timed>> = Arc::new(BucketTable::with_max_keys(64));
577 let barrier = Arc::new(Barrier::new(8));
578 let mut handles = Vec::new();
579 for _ in 0..8 {
580 let table = table.clone();
581 let barrier = barrier.clone();
582 handles.push(std::thread::spawn(move || {
583 barrier.wait();
584 for _ in 0..50 {
585 let _ = table.get_or_insert(b"K", || Timed::at(1));
586 }
587 }));
588 }
589 for handle in handles {
590 handle.join().unwrap();
591 }
592 let k_slots = (0..table.capacity())
594 .filter(|index| table.full_key_at(*index).as_deref() == Some(b"K".as_slice()))
595 .count();
596 assert_eq!(
597 k_slots, 1,
598 "K occupies exactly one slot, no duplicate insert"
599 );
600 assert_eq!(table.len(), 1, "count does not double-count K");
601 let first = table.get_or_insert(b"K", || Timed::at(9));
602 let second = table.get_or_insert(b"K", || Timed::at(9));
603 assert!(Arc::ptr_eq(&first, &second));
604 }
605
606 #[test]
607 fn clear_tombstones_every_published_entry() {
608 let table = BucketTable::with_max_keys(8);
609 for key in [b"one".as_slice(), b"two", b"three"] {
610 table.get_or_insert(key, || Timed::at(1));
611 }
612 assert_eq!(table.len(), 3);
613 table.clear();
614 assert!(table.is_empty());
615 for key in [b"one".as_slice(), b"two", b"three"] {
616 assert!(table.get(key, fxhash(key)).is_none());
617 }
618 }
619
620 #[test]
621 fn threaded_stress_returns_correct_key_and_stays_bounded() {
622 let max_keys = 64;
626 let table: Arc<BucketTable<KeyedValue>> = Arc::new(BucketTable::with_max_keys(max_keys));
627 let capacity = table.capacity();
628 let barrier = Arc::new(Barrier::new(8));
629 let mut handles = Vec::new();
630 for thread_id in 0..8u64 {
631 let table = table.clone();
632 let barrier = barrier.clone();
633 handles.push(std::thread::spawn(move || {
634 barrier.wait();
635 for op in 0..10_000u64 {
636 let key = format!("k{}", (thread_id * 7 + op) % 128);
637 let bytes = key.as_bytes();
638 let bucket = table.get_or_insert(bytes, || KeyedValue::new(bytes));
639 assert_eq!(
640 bucket.key.as_slice(),
641 bytes,
642 "returned bucket carries the requested key"
643 );
644 bucket.last_access.store(op, Relaxed);
645 if op % 256 == 0 {
646 table.evict_one_lru(|value| value.last_access.load(Relaxed));
647 }
648 }
649 }));
650 }
651 for handle in handles {
652 handle.join().unwrap();
653 }
654 assert!(
655 table.len() <= capacity,
656 "len {} must stay within capacity {capacity}",
657 table.len()
658 );
659 }
660
661 struct KeyedValue {
664 key: Vec<u8>,
665 last_access: StdAtomicU64,
666 }
667
668 impl KeyedValue {
669 fn new(key: &[u8]) -> Self {
670 Self {
671 key: key.to_vec(),
672 last_access: StdAtomicU64::new(0),
673 }
674 }
675 }
676}