fixed_cache/

lib.rs

#![doc = include_str!("../README.md")]
#![cfg_attr(docsrs, feature(doc_cfg))]
#![allow(clippy::new_without_default)]

use core::{
    cell::UnsafeCell,
    hash::{BuildHasher, Hash},
    mem::MaybeUninit,
    sync::atomic::{AtomicUsize, Ordering},
};
use equivalent::Equivalent;

const LOCKED_BIT: usize = 0x0000_8000;

#[cfg(feature = "rapidhash")]
type DefaultBuildHasher = std::hash::BuildHasherDefault<rapidhash::fast::RapidHasher<'static>>;
#[cfg(not(feature = "rapidhash"))]
type DefaultBuildHasher = std::hash::RandomState;

/// A concurrent, fixed-size, set-associative cache.
///
/// This cache maps keys to values using a fixed number of buckets. Each key hashes to exactly
/// one bucket, and collisions are resolved by eviction (the new value replaces the old one).
///
/// # Thread Safety
///
/// The cache is safe to share across threads (`Send + Sync`). All operations use atomic
/// instructions and never block, making it suitable for high-contention scenarios.
///
/// # Limitations
///
/// - **No `Drop` support**: Key and value types must not implement `Drop`. Use `Copy` types,
///   primitives, or `&'static` references.
/// - **Eviction on collision**: When two keys hash to the same bucket, the older entry is lost.
/// - **No iteration or removal**: Individual entries cannot be enumerated or explicitly removed.
///
/// # Type Parameters
///
/// - `K`: The key type. Must implement [`Hash`] + [`Eq`] and must not implement [`Drop`].
/// - `V`: The value type. Must implement [`Clone`] and must not implement [`Drop`].
/// - `S`: The hash builder type. Must implement [`BuildHasher`]. Defaults to [`RandomState`] or
///   [`rapidhash`] if the `rapidhash` feature is enabled.
///
/// # Example
///
/// ```
/// use fixed_cache::Cache;
///
/// let cache: Cache<u64, u64> = Cache::new(256, Default::default());
///
/// // Insert a value
/// cache.insert(42, 100);
/// assert_eq!(cache.get(&42), Some(100));
///
/// // Get or compute a value
/// let value = cache.get_or_insert_with(123, |&k| k * 2);
/// assert_eq!(value, 246);
/// ```
///
/// [`Hash`]: core::hash::Hash
/// [`Eq`]: core::cmp::Eq
/// [`Clone`]: core::clone::Clone
/// [`Drop`]: core::ops::Drop
/// [`BuildHasher`]: core::hash::BuildHasher
/// [`RandomState`]: std::hash::RandomState
/// [`rapidhash`]: https://crates.io/crates/rapidhash
pub struct Cache<K, V, S = DefaultBuildHasher> {
    entries: *const [Bucket<(K, V)>],
    build_hasher: S,
    drop: bool,
}

impl<K, V, S> core::fmt::Debug for Cache<K, V, S> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("Cache").finish_non_exhaustive()
    }
}

// SAFETY: `Cache` is safe to share across threads because `Bucket` uses atomic operations.
unsafe impl<K: Send, V: Send, S: Send> Send for Cache<K, V, S> {}
unsafe impl<K: Send, V: Send, S: Sync> Sync for Cache<K, V, S> {}

impl<K, V, S> Cache<K, V, S>
where
    K: Hash + Eq,
    S: BuildHasher,
{
    /// Create a new cache with the specified number of entries and hasher.
    ///
    /// Dynamically allocates memory for the cache entries.
    ///
    /// # Panics
    ///
    /// Panics if `num` is not a power of two.
    pub fn new(num: usize, build_hasher: S) -> Self {
        assert!(num.is_power_of_two(), "capacity must be a power of two");
        let entries =
            Box::into_raw((0..num).map(|_| Bucket::new()).collect::<Vec<_>>().into_boxed_slice());
        Self::new_inner(entries, build_hasher, true)
    }

    /// Creates a new cache with the specified entries and hasher.
    ///
    /// # Panics
    ///
    /// Panics if `entries.len()` is not a power of two.
    #[inline]
    pub const fn new_static(entries: &'static [Bucket<(K, V)>], build_hasher: S) -> Self {
        Self::new_inner(entries, build_hasher, false)
    }

    #[inline]
    const fn new_inner(entries: *const [Bucket<(K, V)>], build_hasher: S, drop: bool) -> Self {
        const {
            assert!(!std::mem::needs_drop::<K>(), "dropping keys is not supported yet");
            assert!(!std::mem::needs_drop::<V>(), "dropping values is not supported yet");
        }
        assert!(entries.len().is_power_of_two());
        Self { entries, build_hasher, drop }
    }

    #[inline]
    const fn index_mask(&self) -> usize {
        let n = self.capacity();
        unsafe { core::hint::assert_unchecked(n.is_power_of_two()) };
        n - 1
    }

    #[inline]
    const fn tag_mask(&self) -> usize {
        !self.index_mask()
    }

    /// Returns the hash builder used by this cache.
    #[inline]
    pub const fn hasher(&self) -> &S {
        &self.build_hasher
    }

    /// Returns the number of entries in this cache.
    #[inline]
    pub const fn capacity(&self) -> usize {
        self.entries.len()
    }
}

impl<K, V, S> Cache<K, V, S>
where
    K: Hash + Eq,
    V: Clone,
    S: BuildHasher,
{
    /// Get an entry from the cache.
    pub fn get<Q: ?Sized + Hash + Equivalent<K>>(&self, key: &Q) -> Option<V> {
        let (bucket, tag) = self.calc(key);
        self.get_inner(key, bucket, tag)
    }

    #[inline]
    fn get_inner<Q: ?Sized + Hash + Equivalent<K>>(
        &self,
        key: &Q,
        bucket: &Bucket<(K, V)>,
        tag: usize,
    ) -> Option<V> {
        if bucket.try_lock(Some(tag)) {
            // SAFETY: We hold the lock, so we have exclusive access.
            let (ck, v) = unsafe { (*bucket.data.get()).assume_init_ref() };
            if key.equivalent(ck) {
                let v = v.clone();
                bucket.unlock(tag);
                return Some(v);
            }
            bucket.unlock(tag);
            // Hash collision: same hash but different key.
        }

        None
    }

    /// Insert an entry into the cache.
    pub fn insert(&self, key: K, value: V) {
        let (bucket, tag) = self.calc(&key);
        self.insert_inner(|| key, || value, bucket, tag);
    }

    #[inline]
    fn insert_inner(
        &self,
        make_key: impl FnOnce() -> K,
        make_value: impl FnOnce() -> V,
        bucket: &Bucket<(K, V)>,
        tag: usize,
    ) {
        if bucket.try_lock(None) {
            // SAFETY: We hold the lock, so we have exclusive access.
            unsafe {
                let data = (&mut *bucket.data.get()).as_mut_ptr();
                (&raw mut (*data).0).write(make_key());
                (&raw mut (*data).1).write(make_value());
            }
            bucket.unlock(tag);
        }
    }

    /// Gets a value from the cache, or inserts one computed by `f` if not present.
    ///
    /// If the key is found in the cache, returns a clone of the cached value.
    /// Otherwise, calls `f` to compute the value, attempts to insert it, and returns it.
    #[inline]
    pub fn get_or_insert_with<F>(&self, key: K, f: F) -> V
    where
        F: FnOnce(&K) -> V,
    {
        let mut key = std::mem::ManuallyDrop::new(key);
        let mut read = false;
        let r = self.get_or_insert_with_ref(&*key, f, |k| {
            read = true;
            unsafe { std::ptr::read(k) }
        });
        if !read {
            unsafe { std::mem::ManuallyDrop::drop(&mut key) }
        }
        r
    }

    /// Gets a value from the cache, or inserts one computed by `f` if not present.
    ///
    /// If the key is found in the cache, returns a clone of the cached value.
    /// Otherwise, calls `f` to compute the value, attempts to insert it, and returns it.
    ///
    /// This is the same as [`get_or_insert_with`], but takes a reference to the key, and a function
    /// to get the key reference to an owned key.
    ///
    /// [`get_or_insert_with`]: Self::get_or_insert_with
    #[inline]
    pub fn get_or_insert_with_ref<'a, Q, F, Cvt>(&self, key: &'a Q, f: F, cvt: Cvt) -> V
    where
        Q: ?Sized + Hash + Equivalent<K>,
        F: FnOnce(&'a Q) -> V,
        Cvt: FnOnce(&'a Q) -> K,
    {
        let (bucket, tag) = self.calc(key);
        if let Some(v) = self.get_inner(key, bucket, tag) {
            return v;
        }
        let value = f(key);
        self.insert_inner(|| cvt(key), || value.clone(), bucket, tag);
        value
    }

    #[inline]
    fn calc<Q: ?Sized + Hash + Equivalent<K>>(&self, key: &Q) -> (&Bucket<(K, V)>, usize) {
        let hash = self.hash_key(key);
        // SAFETY: index is masked to be within bounds.
        let bucket = unsafe { (&*self.entries).get_unchecked(hash & self.index_mask()) };
        let tag = hash & self.tag_mask();
        (bucket, tag)
    }

    #[inline]
    fn hash_key<Q: ?Sized + Hash + Equivalent<K>>(&self, key: &Q) -> usize {
        let hash = self.build_hasher.hash_one(key);

        if cfg!(target_pointer_width = "32") {
            ((hash >> 32) as usize) ^ (hash as usize)
        } else {
            hash as usize
        }
    }
}

impl<K, V, S> Drop for Cache<K, V, S> {
    fn drop(&mut self) {
        if self.drop {
            drop(unsafe { Box::from_raw(self.entries.cast_mut()) });
        }
    }
}

/// A single cache bucket that holds one key-value pair.
///
/// Buckets are aligned to 128 bytes to avoid false sharing between cache lines.
/// Each bucket contains an atomic tag for lock-free synchronization and uninitialized
/// storage for the data.
///
/// This type is public to allow use with the [`static_cache!`] macro for compile-time
/// cache initialization. You typically don't need to interact with it directly.
#[repr(C, align(128))]
#[doc(hidden)]
pub struct Bucket<T> {
    tag: AtomicUsize,
    data: UnsafeCell<MaybeUninit<T>>,
}

impl<T> Bucket<T> {
    /// Creates a new zeroed bucket.
    #[inline]
    pub const fn new() -> Self {
        Self { tag: AtomicUsize::new(0), data: UnsafeCell::new(MaybeUninit::zeroed()) }
    }

    #[inline]
    fn try_lock(&self, expected: Option<usize>) -> bool {
        let state = self.tag.load(Ordering::Relaxed);
        if let Some(expected) = expected {
            if state != expected {
                return false;
            }
        } else if state & LOCKED_BIT != 0 {
            return false;
        }
        self.tag
            .compare_exchange(state, state | LOCKED_BIT, Ordering::Acquire, Ordering::Relaxed)
            .is_ok()
    }

    #[inline]
    fn unlock(&self, tag: usize) {
        self.tag.store(tag, Ordering::Release);
    }
}

// SAFETY: `Bucket` is a specialized `Mutex<T>` that never blocks.
unsafe impl<T: Send> Send for Bucket<T> {}
unsafe impl<T: Send> Sync for Bucket<T> {}

/// Declares a static cache with the given name, key type, value type, and size.
///
/// The size must be a power of two.
///
/// # Example
///
/// ```
/// # #[cfg(feature = "rapidhash")] {
/// use fixed_cache::{Cache, static_cache};
///
/// type BuildHasher = std::hash::BuildHasherDefault<rapidhash::fast::RapidHasher<'static>>;
///
/// static MY_CACHE: Cache<u64, &'static str, BuildHasher> =
///     static_cache!(u64, &'static str, 1024, BuildHasher::new());
///
/// let value = MY_CACHE.get_or_insert_with(&42, |_k| "hi");
/// assert_eq!(value, "hi");
///
/// let new_value = MY_CACHE.get_or_insert_with(&42, |_k| "not hi");
/// assert_eq!(new_value, "not hi");
/// # }
/// ```
#[macro_export]
macro_rules! static_cache {
    ($K:ty, $V:ty, $size:expr) => {
        $crate::static_cache!($K, $V, $size, Default::default())
    };
    ($K:ty, $V:ty, $size:expr, $hasher:expr) => {{
        static ENTRIES: [$crate::Bucket<($K, $V)>; $size] =
            [const { $crate::Bucket::new() }; $size];
        $crate::Cache::new_static(&ENTRIES, $hasher)
    }};
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::thread;

    const fn iters(n: usize) -> usize {
        if cfg!(miri) { n / 10 } else { n }
    }

    type BH = std::hash::BuildHasherDefault<rapidhash::fast::RapidHasher<'static>>;
    type Cache<K, V> = super::Cache<K, V, BH>;

    fn new_cache<K: Hash + Eq, V: Clone>(size: usize) -> Cache<K, V> {
        Cache::new(size, Default::default())
    }

    #[test]
    fn test_basic_get_or_insert() {
        let cache = new_cache(1024);

        let mut computed = false;
        let value = cache.get_or_insert_with(42, |&k| {
            computed = true;
            k * 2
        });
        assert!(computed);
        assert_eq!(value, 84);

        computed = false;
        let value = cache.get_or_insert_with(42, |&k| {
            computed = true;
            k * 2
        });
        assert!(!computed);
        assert_eq!(value, 84);
    }

    #[test]
    fn test_different_keys() {
        let cache: Cache<&'static str, usize> = static_cache!(&'static str, usize, 1024);

        let v1 = cache.get_or_insert_with("hello", |s| s.len());
        let v2 = cache.get_or_insert_with("world!", |s| s.len());

        assert_eq!(v1, 5);
        assert_eq!(v2, 6);
    }

    #[test]
    fn test_new_dynamic_allocation() {
        let cache: Cache<u32, u32> = new_cache(64);
        assert_eq!(cache.capacity(), 64);

        cache.insert(1, 100);
        assert_eq!(cache.get(&1), Some(100));
    }

    #[test]
    fn test_get_miss() {
        let cache = new_cache::<u64, u64>(64);
        assert_eq!(cache.get(&999), None);
    }

    #[test]
    fn test_insert_and_get() {
        let cache: Cache<u64, &'static str> = new_cache(64);

        cache.insert(1, "one");
        cache.insert(2, "two");
        cache.insert(3, "three");

        assert_eq!(cache.get(&1), Some("one"));
        assert_eq!(cache.get(&2), Some("two"));
        assert_eq!(cache.get(&3), Some("three"));
        assert_eq!(cache.get(&4), None);
    }

    #[test]
    fn test_insert_twice() {
        let cache = new_cache(64);

        cache.insert(42, 1);
        assert_eq!(cache.get(&42), Some(1));

        cache.insert(42, 2);
        let v = cache.get(&42);
        assert!(v == Some(1) || v == Some(2));
    }

    #[test]
    fn test_get_or_insert_with_ref() {
        let cache: Cache<&'static str, usize> = new_cache(64);

        let key = "hello";
        let value = cache.get_or_insert_with_ref(key, |s| s.len(), |s| s);
        assert_eq!(value, 5);

        let value2 = cache.get_or_insert_with_ref(key, |_| 999, |s| s);
        assert_eq!(value2, 5);
    }

    #[test]
    fn test_get_or_insert_with_ref_different_keys() {
        let cache: Cache<&'static str, usize> = new_cache(1024);

        let v1 = cache.get_or_insert_with_ref("foo", |s| s.len(), |s| s);
        let v2 = cache.get_or_insert_with_ref("barbaz", |s| s.len(), |s| s);

        assert_eq!(v1, 3);
        assert_eq!(v2, 6);
    }

    #[test]
    fn test_capacity() {
        let cache = new_cache::<u64, u64>(256);
        assert_eq!(cache.capacity(), 256);

        let cache2 = new_cache::<u64, u64>(128);
        assert_eq!(cache2.capacity(), 128);
    }

    #[test]
    fn test_hasher() {
        let cache = new_cache::<u64, u64>(64);
        let _ = cache.hasher();
    }

    #[test]
    fn test_debug_impl() {
        let cache = new_cache::<u64, u64>(64);
        let debug_str = format!("{:?}", cache);
        assert!(debug_str.contains("Cache"));
    }

    #[test]
    fn test_bucket_new() {
        let bucket: Bucket<(u64, u64)> = Bucket::new();
        assert_eq!(bucket.tag.load(Ordering::Relaxed), 0);
    }

    #[test]
    fn test_many_entries() {
        let cache: Cache<u64, u64> = new_cache(1024);
        let n = iters(500);

        for i in 0..n as u64 {
            cache.insert(i, i * 2);
        }

        let mut hits = 0;
        for i in 0..n as u64 {
            if cache.get(&i) == Some(i * 2) {
                hits += 1;
            }
        }
        assert!(hits > 0);
    }

    #[test]
    fn test_string_keys() {
        let cache: Cache<&'static str, i32> = new_cache(1024);

        cache.insert("alpha", 1);
        cache.insert("beta", 2);
        cache.insert("gamma", 3);

        assert_eq!(cache.get(&"alpha"), Some(1));
        assert_eq!(cache.get(&"beta"), Some(2));
        assert_eq!(cache.get(&"gamma"), Some(3));
    }

    #[test]
    fn test_zero_values() {
        let cache: Cache<u64, u64> = new_cache(64);

        cache.insert(0, 0);
        assert_eq!(cache.get(&0), Some(0));

        cache.insert(1, 0);
        assert_eq!(cache.get(&1), Some(0));
    }

    #[test]
    fn test_clone_value() {
        #[derive(Clone, PartialEq, Debug)]
        struct MyValue(u64);

        let cache: Cache<u64, MyValue> = new_cache(64);

        cache.insert(1, MyValue(123));
        let v = cache.get(&1);
        assert_eq!(v, Some(MyValue(123)));
    }

    fn run_concurrent<F>(num_threads: usize, f: F)
    where
        F: Fn(usize) + Send + Sync,
    {
        thread::scope(|s| {
            for t in 0..num_threads {
                let f = &f;
                s.spawn(move || f(t));
            }
        });
    }

    #[test]
    fn test_concurrent_reads() {
        let cache: Cache<u64, u64> = new_cache(1024);
        let n = iters(100);

        for i in 0..n as u64 {
            cache.insert(i, i * 10);
        }

        run_concurrent(4, |_| {
            for i in 0..n as u64 {
                let _ = cache.get(&i);
            }
        });
    }

    #[test]
    fn test_concurrent_writes() {
        let cache: Cache<u64, u64> = new_cache(1024);
        let n = iters(100);

        run_concurrent(4, |t| {
            for i in 0..n {
                cache.insert((t * 1000 + i) as u64, i as u64);
            }
        });
    }

    #[test]
    fn test_concurrent_read_write() {
        let cache: Cache<u64, u64> = new_cache(256);
        let n = iters(1000);

        run_concurrent(2, |t| {
            for i in 0..n as u64 {
                if t == 0 {
                    cache.insert(i % 100, i);
                } else {
                    let _ = cache.get(&(i % 100));
                }
            }
        });
    }

    #[test]
    fn test_concurrent_get_or_insert() {
        let cache: Cache<u64, u64> = new_cache(1024);
        let n = iters(100);

        run_concurrent(8, |_| {
            for i in 0..n as u64 {
                let _ = cache.get_or_insert_with(i, |&k| k * 2);
            }
        });

        for i in 0..n as u64 {
            if let Some(v) = cache.get(&i) {
                assert_eq!(v, i * 2);
            }
        }
    }

    #[test]
    #[should_panic]
    fn test_non_power_of_two_panics() {
        let _ = new_cache::<u64, u64>(100);
    }

    #[test]
    fn test_power_of_two_sizes() {
        for shift in 1..10 {
            let size = 1 << shift;
            let cache = new_cache::<u64, u64>(size);
            assert_eq!(cache.capacity(), size);
        }
    }

    #[test]
    fn test_small_cache() {
        let cache = new_cache(2);
        assert_eq!(cache.capacity(), 2);

        cache.insert(1, 10);
        cache.insert(2, 20);
        cache.insert(3, 30);

        let count = [1, 2, 3].iter().filter(|&&k| cache.get(&k).is_some()).count();
        assert!(count <= 2);
    }

    #[test]
    fn test_equivalent_key_lookup() {
        let cache = new_cache(64);

        cache.insert("hello", 42);

        assert_eq!(cache.get(&"hello"), Some(42));
    }

    #[test]
    fn test_large_values() {
        let cache: Cache<u64, [u8; 1000]> = new_cache(64);

        let large_value = [42u8; 1000];
        cache.insert(1, large_value);

        assert_eq!(cache.get(&1), Some(large_value));
    }

    #[test]
    fn test_send_sync() {
        fn assert_send<T: Send>() {}
        fn assert_sync<T: Sync>() {}

        assert_send::<Cache<u64, u64>>();
        assert_sync::<Cache<u64, u64>>();
        assert_send::<Bucket<(u64, u64)>>();
        assert_sync::<Bucket<(u64, u64)>>();
    }
}