memo_cache/

lib.rs

#![no_std]

use core::{borrow::Borrow, cell::Cell, fmt, mem};

/// Key equivalence trait, to support `Borrow` types as keys.
trait Equivalent<K: ?Sized> {
    /// Returns `true` if two values are equivalent, `false` otherwise.
    fn equivalent(&self, k: &K) -> bool;
}

impl<Q: ?Sized, K: ?Sized> Equivalent<K> for Q
where
    Q: Eq,
    K: Borrow<Q>,
{
    fn equivalent(&self, k: &K) -> bool {
        self == k.borrow()
    }
}

/// Hit count threshold at which all hit counts are decayed (75% of `u8::MAX`).
const DECAY_THRESHOLD: u8 = 192;

/// A single key/value slot used in the cache.
#[derive(Clone, PartialEq)]
enum KeyValueSlot<K, V> {
    Used { key: K, value: V, hits: Cell<u8> },
    Empty,
}

impl<K, V> KeyValueSlot<K, V> {
    /// Check a used slot key for equivalence.
    ///
    /// Returns `true` for used slots with key equivalence, `false` if otherwise.
    ///
    #[cfg_attr(feature = "inline-more", inline)]
    fn is_key<Q>(&self, k: &Q) -> bool
    where
        Q: Equivalent<K> + ?Sized,
    {
        if let KeyValueSlot::Used { key, .. } = self {
            k.equivalent(key)
        } else {
            false
        }
    }

    /// Get the value of a used slot, counting the access as a hit and raising the decay
    /// flag when the hit count crosses the decay threshold.
    ///
    /// NOTE: Hit counts increment by one, so crossing the threshold passes through it
    ///       exactly once; comparing for equality keeps the hot path free of stores.
    #[cfg_attr(feature = "inline-more", inline)]
    fn get_value(&self, decay_pending: &Cell<bool>) -> Option<&V> {
        if let KeyValueSlot::Used { value, hits, .. } = self {
            let h = hits.get().saturating_add(1);
            hits.set(h);
            if h == DECAY_THRESHOLD {
                decay_pending.set(true);
            }
            Some(value)
        } else {
            None
        }
    }

    /// Get the value of a used slot (for mutation), counting the access as a hit and
    /// raising the decay flag when the hit count crosses the decay threshold.
    #[cfg_attr(feature = "inline-more", inline)]
    fn get_value_mut(&mut self, decay_pending: &Cell<bool>) -> Option<&mut V> {
        if let KeyValueSlot::Used { value, hits, .. } = self {
            let h = hits.get().saturating_add(1);
            hits.set(h);
            if h == DECAY_THRESHOLD {
                decay_pending.set(true);
            }
            Some(value)
        } else {
            None
        }
    }

    /// Get the number of cache hits for this slot.
    fn hits(&self) -> u8 {
        if let KeyValueSlot::Used { hits, .. } = self {
            hits.get()
        } else {
            0
        }
    }

    /// Update the value of a used slot, returning the previous value (will be a no-op
    /// returning `None` for empty slots).
    #[cfg_attr(feature = "inline-more", inline)]
    fn update_value(&mut self, v: V) -> Option<V> {
        if let KeyValueSlot::Used { value, .. } = self {
            Some(mem::replace(value, v))
        } else {
            None
        }
    }
}

/// A small, fixed-size key/value cache with retention management.
///
/// The key/value slots are stored inline (as an array inside the struct), so the cache
/// lives wherever you place it; no heap allocation is required or performed.
///
/// # Thread Safety
///
/// `MemoCache` is `Send` but not `Sync`. It can be moved between threads but cannot
/// be shared across threads via shared references (`&MemoCache`).
///
/// This is because the cache uses interior mutability (via [`Cell`]) for tracking hit counts
/// and random number generation, which is not thread-safe.
///
/// For concurrent access, wrap it in a synchronization primitive. E.g.:
///
/// ```
/// use std::sync::Mutex;
/// use memo_cache::MemoCache;
///
/// let cache = Mutex::new(MemoCache::<u32, String, 8>::new());
///
/// // In thread 1:
/// cache.lock().unwrap().insert(42, "value".to_string());
///
/// // In thread 2:
/// let value = cache.lock().unwrap().get(&42);
/// ```
pub struct MemoCache<K, V, const SIZE: usize> {
    buffer: [KeyValueSlot<K, V>; SIZE],
    rng_state: Cell<u32>,
    /// Number of occupied slots.
    used: usize,
    /// Set when some slot's hit count crossed the decay threshold. May overestimate
    /// (eviction and removal reset or drop slot hit counts without lowering it);
    /// `decay_hits` verifies and corrects it.
    decay_pending: Cell<bool>,
}

impl<K, V, const SIZE: usize> MemoCache<K, V, SIZE>
where
    K: Eq,
{
    const SIZE_CHECK: () = assert!(SIZE > 0, "Cache size must be greater than 0");

    /// Create a new cache.
    ///
    /// # Examples
    ///
    /// ```
    /// use memo_cache::MemoCache;
    ///
    /// let c = MemoCache::<u32, String, 4>::new();
    /// ```
    #[cfg_attr(feature = "inline-more", inline)]
    #[must_use]
    #[allow(clippy::cast_possible_truncation)] // We assume cache size is limited to values representable by `u32`.
    pub fn new() -> Self {
        let () = Self::SIZE_CHECK; // Force evaluation of const assertion.

        Self {
            buffer: [const { KeyValueSlot::Empty }; SIZE],
            rng_state: Cell::new(0x9E37_79B9 ^ (SIZE as u32).wrapping_mul(0x85EB_CA6B)), // Mixed primes.
            used: 0,
            decay_pending: Cell::new(false),
        }
    }

    /// Get the (fixed) capacity of the cache in [number of elements].
    ///
    /// # Examples
    ///
    /// ```
    /// use memo_cache::MemoCache;
    ///
    /// let c = MemoCache::<u32, String, 8>::new();
    ///
    /// assert_eq!(c.capacity(), 8);
    /// ```
    #[cfg_attr(feature = "inline-more", inline)]
    pub const fn capacity(&self) -> usize {
        SIZE
    }

    /// Get the number of occupied slots in the cache.
    ///
    /// # Examples
    ///
    /// ```
    /// use memo_cache::MemoCache;
    ///
    /// let mut c = MemoCache::<u32, &str, 4>::new();
    ///
    /// assert_eq!(c.len(), 0);
    ///
    /// c.insert(42, "The Answer");
    ///
    /// assert_eq!(c.len(), 1);
    /// ```
    #[cfg_attr(feature = "inline-more", inline)]
    pub const fn len(&self) -> usize {
        self.used
    }

    /// Returns `true` if the cache contains no elements.
    ///
    /// # Examples
    ///
    /// ```
    /// use memo_cache::MemoCache;
    ///
    /// let mut c = MemoCache::<u32, &str, 4>::new();
    ///
    /// assert!(c.is_empty());
    ///
    /// c.insert(42, "The Answer");
    ///
    /// assert!(!c.is_empty());
    /// ```
    #[cfg_attr(feature = "inline-more", inline)]
    pub const fn is_empty(&self) -> bool {
        self.used == 0
    }

    /// Get the size of the cache slot eviction search window.
    const fn eviction_window_size() -> usize {
        match SIZE {
            0..=4 => SIZE,      // Tiny cache.
            5..=16 => SIZE / 2, // Small-sized cache.
            17..=64 => 8,       // Medium-sized cache.
            _ => 16,            // Large cache.
        }
    }

    /// Generate a pseudo-random value using Xorshift32 (see: <https://en.wikipedia.org/wiki/Xorshift>).
    #[cfg_attr(feature = "inline-more", inline)]
    fn xorshift32(&self) -> u32 {
        let mut x = self.rng_state.get();
        x ^= x << 13;
        x ^= x >> 17;
        x ^= x << 5;
        self.rng_state.set(x);
        x
    }

    /// Find a cache slot index to evict for replacement.
    ///
    /// Returns the index of an empty slot if (and only if) the cache is not fully occupied.
    fn find_eviction_slot(&self) -> usize {
        // Use an empty slot first, if there is one.
        if self.used < SIZE {
            if let Some(idx) = self
                .buffer
                .iter()
                .position(|s| matches!(s, KeyValueSlot::Empty))
            {
                return idx;
            }

            debug_assert!(false, "Occupancy count says an empty slot exists");
        }

        // The cache is fully occupied, find a slot to evict by scanning a window at a random position.
        let window_size = Self::eviction_window_size();
        let start_idx = ((u64::from(self.xorshift32()) * SIZE as u64) >> 32) as usize;

        let mut evict_idx = start_idx;
        let mut min_hits = u8::MAX;

        for i in 0..window_size {
            let idx = (start_idx + i) % SIZE;
            let hits = self.buffer[idx].hits();

            if hits < min_hits {
                min_hits = hits;
                evict_idx = idx;

                // Early-out for unhit cache entries.
                if hits == 0 {
                    break;
                }
            }
        }

        evict_idx
    }

    /// Evict a suitable slot and replace it. Returns a reference to the replaced slot value.
    #[cfg_attr(feature = "inline-more", inline)]
    fn evict_and_replace(&mut self, k: K, v: V) -> &V {
        self.decay_hits();

        let idx = self.find_eviction_slot();

        // `find_eviction_slot` fills an empty slot iff the cache is not fully occupied.
        if self.used < SIZE {
            self.used += 1;
        }

        let s = &mut self.buffer[idx];

        *s = KeyValueSlot::Used {
            key: k,
            value: v,
            hits: Cell::new(0),
        };

        // NOTE: Return the value directly instead of via `get_value`, which would count
        //       the insertion itself as a hit.
        match s {
            KeyValueSlot::Used { value, .. } => value,
            KeyValueSlot::Empty => unreachable!(), // The slot was filled above.
        }
    }

    /// Decay hit values for all occupied slots if any slot reaches the decay threshold.
    fn decay_hits(&mut self) {
        // Fast path: the flag is raised whenever a slot hit count crosses the threshold,
        // so no slot can have reached it while the flag is down.
        if !self.decay_pending.get() {
            return;
        }

        // The flag may overestimate (eviction and removal reset or drop slot hit counts
        // without lowering it), so compute the true maximum before deciding.
        let true_max = self
            .buffer
            .iter()
            .map(KeyValueSlot::hits)
            .max()
            .unwrap_or(0);

        if true_max >= DECAY_THRESHOLD {
            for s in &mut self.buffer {
                if let KeyValueSlot::Used { hits, .. } = s {
                    let current = hits.get();
                    // NOTE: Subtracting 25% is a no-op for hit counts 1..=3 (the shift
                    //       rounds down to zero). This is deliberate: such entries keep a
                    //       slight edge over never-hit entries, and the imprecision is
                    //       irrelevant at the decay threshold scale.
                    hits.set(current - (current >> 2)); // Subtract 25%.
                }
            }

            // Saturated hit counts may still be at the threshold after one decay round.
            self.decay_pending
                .set(true_max - (true_max >> 2) >= DECAY_THRESHOLD);
        } else {
            self.decay_pending.set(false);
        }
    }

    /// Insert a key/value pair.
    ///
    /// If the key was already present, its value is updated and the previous value is
    /// returned. Otherwise `None` is returned.
    ///
    /// # Notes
    ///
    /// Inserting a new key into a full cache evicts another entry to make room; the
    /// evicted key/value pair is *not* returned.
    ///
    /// # Examples
    ///
    /// ```
    /// use memo_cache::MemoCache;
    ///
    /// let mut c = MemoCache::<u32, &str, 4>::new();
    ///
    /// assert_eq!(c.get(&42), None);
    ///
    /// assert_eq!(c.insert(42, "The Answer"), None);
    ///
    /// assert_eq!(c.get(&42), Some(&"The Answer"));
    ///
    /// assert_eq!(c.insert(42, "Another Answer"), Some("The Answer"));
    /// ```
    #[cfg_attr(feature = "inline-more", inline)]
    pub fn insert(&mut self, k: K, v: V) -> Option<V> {
        if let Some(s) = self.buffer.iter_mut().find(|e| e.is_key(&k)) {
            s.update_value(v)
        } else {
            self.evict_and_replace(k, v);
            None
        }
    }

    /// Returns `true` if the cache contains a value for the specified key.
    ///
    /// # Examples
    ///
    /// ```
    /// use memo_cache::MemoCache;
    ///
    /// let mut c = MemoCache::<u32, &str, 4>::new();
    ///
    /// assert_eq!(c.contains_key(&42), false);
    ///
    /// c.insert(42, "The Answer");
    ///
    /// assert_eq!(c.contains_key(&42), true);
    /// ```
    #[cfg_attr(feature = "inline-more", inline)]
    pub fn contains_key<Q>(&self, k: &Q) -> bool
    where
        K: Borrow<Q>,
        Q: Eq + ?Sized,
    {
        self.buffer.iter().any(|e| e.is_key(k))
    }

    /// Lookup a cache entry by key.
    ///
    /// # Examples
    ///
    /// ```
    /// use memo_cache::MemoCache;
    ///
    /// let mut c = MemoCache::<u32, &str, 4>::new();
    ///
    /// assert_eq!(c.get(&42), None);
    ///
    /// c.insert(42, "The Answer");
    ///
    /// assert_eq!(c.get(&42), Some(&"The Answer"));
    /// ```
    #[cfg_attr(feature = "inline-more", inline)]
    pub fn get<Q>(&self, k: &Q) -> Option<&V>
    where
        K: Borrow<Q>,
        Q: Eq + ?Sized,
    {
        // NOTE: `is_key` only matches used slots, so `get_value` returns `Some` here.
        self.buffer
            .iter()
            .find(|e| e.is_key(k))
            .and_then(|e| e.get_value(&self.decay_pending))
    }

    /// Lookup a cache entry by key (for mutation).
    ///
    /// # Examples
    ///
    /// ```
    /// use memo_cache::MemoCache;
    ///
    /// let mut c = MemoCache::<u32, &str, 4>::new();
    ///
    /// c.insert(42, "The Answer");
    ///
    /// if let Some(v) = c.get_mut(&42) {
    ///     *v = "Another Answer";
    /// }
    ///
    /// assert_eq!(c.get(&42), Some(&"Another Answer"));
    /// ```
    #[cfg_attr(feature = "inline-more", inline)]
    pub fn get_mut<Q>(&mut self, k: &Q) -> Option<&mut V>
    where
        K: Borrow<Q>,
        Q: Eq + ?Sized,
    {
        // NOTE: `is_key` only matches used slots, so `get_value_mut` returns `Some` here.
        let decay_pending = &self.decay_pending;
        self.buffer
            .iter_mut()
            .find(|e| e.is_key(k))
            .and_then(|e| e.get_value_mut(decay_pending))
    }

    /// Remove a key from the cache, returning the value if the key was present.
    ///
    /// The freed slot is reused by subsequent insertions.
    ///
    /// # Examples
    ///
    /// ```
    /// use memo_cache::MemoCache;
    ///
    /// let mut c = MemoCache::<u32, &str, 4>::new();
    ///
    /// c.insert(42, "The Answer");
    ///
    /// assert_eq!(c.remove(&42), Some("The Answer"));
    /// assert_eq!(c.remove(&42), None);
    /// assert_eq!(c.get(&42), None);
    /// ```
    #[cfg_attr(feature = "inline-more", inline)]
    pub fn remove<Q>(&mut self, k: &Q) -> Option<V>
    where
        K: Borrow<Q>,
        Q: Eq + ?Sized,
    {
        let used = &mut self.used;
        self.buffer.iter_mut().find(|e| e.is_key(k)).map(|s| {
            *used -= 1;
            match mem::replace(s, KeyValueSlot::Empty) {
                KeyValueSlot::Used { value, .. } => value,
                KeyValueSlot::Empty => unreachable!(), // The slot was found by key.
            }
        })
    }

    /// Get the index for a given key, if found.
    #[cfg_attr(feature = "inline-more", inline)]
    fn get_key_index<Q>(&self, k: &Q) -> Option<usize>
    where
        K: Borrow<Q>,
        Q: Eq + ?Sized,
    {
        self.buffer.iter().position(|e| e.is_key(k))
    }

    /// Get a value, or, if it does not exist in the cache, insert it using the value computed by `f`.
    /// Returns a reference to the found, or newly inserted value associated with the given key.
    /// If a value is inserted, the key is cloned.
    ///
    /// # Examples
    ///
    /// ```
    /// use memo_cache::MemoCache;
    ///
    /// let mut c = MemoCache::<u32, &str, 4>::new();
    ///
    /// assert_eq!(c.get(&42), None);
    ///
    /// let v = c.get_or_insert_with(&42, |_| "The Answer");
    ///
    /// assert_eq!(v, &"The Answer");
    /// assert_eq!(c.get(&42), Some(&"The Answer"));
    /// ```
    ///
    /// # Notes
    ///
    /// Because this crate is `no_std`, we have no access to `std::borrow::ToOwned`, which means we cannot create a
    /// version of `get_or_insert_with` that can create an owned value from a borrowed key.
    ///
    #[cfg_attr(feature = "inline-more", inline)]
    pub fn get_or_insert_with<F>(&mut self, k: &K, f: F) -> &V
    where
        K: Clone,
        F: FnOnce(&K) -> V,
    {
        if let Some(i) = self.get_key_index(k) {
            // NOTE: The index was found by key, so the slot is used and holds a value.
            self.buffer[i]
                .get_value(&self.decay_pending)
                .expect("Slot found by key must be used")
        } else {
            self.evict_and_replace(k.clone(), f(k))
        }
    }

    /// Get a value, or, if it does not exist in the cache, insert it using the value computed by `f`.
    /// Returns a result with a reference to the found, or newly inserted value associated with the given key.
    /// If a value is inserted, the key is cloned.
    ///
    /// # Examples
    ///
    /// ```
    /// use memo_cache::MemoCache;
    ///
    /// let mut c = MemoCache::<u32, &str, 4>::new();
    ///
    /// assert_eq!(c.get(&42), None);
    ///
    /// let answer : Result<_, &str> = Ok("The Answer");
    /// let v = c.get_or_try_insert_with(&42, |_| answer);
    ///
    /// assert_eq!(v, Ok(&"The Answer"));
    /// assert_eq!(c.get(&42), Some(&"The Answer"));
    ///
    /// let v = c.get_or_try_insert_with(&17, |_| Err("Dunno"));
    ///
    /// assert_eq!(v, Err("Dunno"));
    /// assert_eq!(c.get(&17), None);
    /// ```
    ///
    /// # Errors
    ///
    /// If the function `f` fails, the error of type `E` is returned.
    ///
    /// # Notes
    ///
    /// Because this crate is `no_std`, we have no access to `std::borrow::ToOwned`, which means we cannot create a
    /// version of `get_or_try_insert_with` that can create an owned value from a borrowed key.
    ///
    #[cfg_attr(feature = "inline-more", inline)]
    pub fn get_or_try_insert_with<F, E>(&mut self, k: &K, f: F) -> Result<&V, E>
    where
        K: Clone,
        F: FnOnce(&K) -> Result<V, E>,
    {
        if let Some(i) = self.get_key_index(k) {
            // NOTE: The index was found by key, so the slot is used and holds a value.
            Ok(self.buffer[i]
                .get_value(&self.decay_pending)
                .expect("Slot found by key must be used"))
        } else {
            f(k).map(|v| self.evict_and_replace(k.clone(), v))
        }
    }

    /// Get an iterator over the key/value pairs of the cache, in unspecified order.
    ///
    /// Iterating does not affect the hit counts used by the eviction policy.
    ///
    /// # Examples
    ///
    /// ```
    /// use memo_cache::MemoCache;
    ///
    /// let mut c = MemoCache::<u32, &str, 4>::new();
    ///
    /// c.insert(1, "one");
    /// c.insert(2, "two");
    ///
    /// let mut entries = c.iter().collect::<Vec<_>>();
    /// entries.sort();
    ///
    /// assert_eq!(entries, [(&1, &"one"), (&2, &"two")]);
    /// ```
    #[cfg_attr(feature = "inline-more", inline)]
    pub fn iter(&self) -> impl Iterator<Item = (&K, &V)> {
        self.buffer.iter().filter_map(|s| match s {
            KeyValueSlot::Used { key, value, .. } => Some((key, value)),
            KeyValueSlot::Empty => None,
        })
    }

    /// Clear the cache.
    ///
    /// # Examples
    ///
    /// ```
    /// use memo_cache::MemoCache;
    ///
    /// let mut c = MemoCache::<u32, &str, 4>::new();
    ///
    /// assert_eq!(c.get(&42), None);
    ///
    /// c.insert(42, "The Answer");
    ///
    /// assert_eq!(c.get(&42), Some(&"The Answer"));
    ///
    /// c.clear();
    ///
    /// assert_eq!(c.get(&42), None);
    ///
    #[cfg_attr(feature = "inline-more", inline)]
    pub fn clear(&mut self) {
        self.buffer
            .iter_mut()
            .for_each(|e| *e = KeyValueSlot::Empty);
        self.used = 0;
        self.decay_pending.set(false);
    }
}

impl<K, V, const SIZE: usize> Default for MemoCache<K, V, SIZE>
where
    K: Eq,
{
    fn default() -> Self {
        Self::new()
    }
}

impl<K, V, const SIZE: usize> fmt::Debug for MemoCache<K, V, SIZE>
where
    K: Eq + fmt::Debug,
    V: fmt::Debug,
{
    /// Format the cache entries as a map, in unspecified order.
    ///
    /// # Examples
    ///
    /// ```
    /// use memo_cache::MemoCache;
    ///
    /// let mut c = MemoCache::<u32, &str, 4>::new();
    ///
    /// c.insert(42, "The Answer");
    ///
    /// assert_eq!(format!("{c:?}"), r#"{42: "The Answer"}"#);
    /// ```
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_map().entries(self.iter()).finish()
    }
}

impl<K, V, const SIZE: usize> Clone for MemoCache<K, V, SIZE>
where
    K: Eq + Clone,
    V: Clone,
{
    fn clone(&self) -> Self {
        // Remix the RNG state so the clone does not replay the original's eviction
        // randomness.
        let mut rng_state = self.rng_state.get() ^ 0x5851_F42D; // Mixing prime.
        if rng_state == 0 {
            rng_state = 0x9E37_79B9; // The Xorshift32 state must be nonzero.
        }

        Self {
            buffer: self.buffer.clone(),
            rng_state: Cell::new(rng_state),
            used: self.used,
            decay_pending: self.decay_pending.clone(),
        }
    }
}

#[cfg(test)]
mod tests_internal {
    use super::*;

    #[test]
    fn test_new_state() {
        const SIZE: usize = 8;

        let c = MemoCache::<i32, i32, SIZE>::new();

        // Verify cache size.
        assert_eq!(c.buffer.len(), SIZE);
        assert_eq!(c.capacity(), SIZE);

        // All slots should be empty.
        assert!(c.buffer.iter().all(|s| s == &KeyValueSlot::Empty));
    }

    // Compile-time assertion: MemoCache should be Send (can move between threads).
    const _: fn() = || {
        fn assert_send<T: Send>() {}
        assert_send::<MemoCache<i32, i32, 4>>();
    };

    // MemoCache should NOT be Sync (cannot share across threads).
    //
    // The following compile-time test should fail:
    //
    // const _: fn() = || {
    //     fn assert_sync<T: Sync>() {}
    //     assert_sync::<MemoCache<i32, i32, 4>>();
    // };
}