memo_cache/

lib.rs

#![no_std]

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

/// 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()
    }
}

/// 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.
    #[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.
    #[cfg_attr(feature = "inline-more", inline)]
    fn get_value(&self) -> Option<&V> {
        if let KeyValueSlot::Used { value, hits, .. } = self {
            hits.set(hits.get().saturating_add(1));
            Some(value)
        } else {
            None
        }
    }

    /// Get the value of a used slot (for mutation).
    #[cfg_attr(feature = "inline-more", inline)]
    fn get_value_mut(&mut self) -> Option<&mut V> {
        if let KeyValueSlot::Used { value, hits, .. } = self {
            hits.set(hits.get().saturating_add(1));
            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.
    #[cfg_attr(feature = "inline-more", inline)]
    fn update_value(&mut self, v: V) {
        if let KeyValueSlot::Used { value, .. } = self {
            *value = v;
        }
    }
}

/// A small, fixed-size, heap-allocated key/value cache with retention management.
pub struct MemoCache<K, V, const SIZE: usize> {
    buffer: [KeyValueSlot<K, V>; SIZE],
    rng_state: Cell<u32>,
}

impl<K, V, const SIZE: usize> MemoCache<K, V, SIZE>
where
    K: Clone + Eq,
    V: Clone,
{
    /// Create a new cache.
    ///
    /// # Examples
    ///
    /// ```
    /// use memo_cache::MemoCache;
    ///
    /// let c = MemoCache::<u32, String, 4>::new();
    /// ```
    #[cfg_attr(feature = "inline-more", inline)]
    #[must_use]
    pub fn new() -> Self {
        assert!(SIZE > 0, "Cache size must be greater than 0");

        Self {
            buffer: [const { KeyValueSlot::Empty }; SIZE],
            rng_state: Cell::new(0x9E37_79B9), // The golden ratio prime.
        }
    }

    /// 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 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.
    fn find_eviction_slot(&self) -> usize {
        // Check for empty slots first.
        for (idx, slot) in self.buffer.iter().enumerate() {
            if let KeyValueSlot::Empty = slot {
                return idx;
            }
        }

        // 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;
            // SAFETY: idx is guaranteed to be < SIZE due to the modulo operation.
            let hits = unsafe { self.buffer.get_unchecked(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 {
        let idx = self.find_eviction_slot();
        let s = unsafe { self.buffer.get_unchecked_mut(idx) };

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

        // SAFETY: The slot was filled with a key/value above.
        unsafe { s.get_value().unwrap_unchecked() }
    }

    /// Decay hit values for all occupied slots if any slot reaches the maximum value.
    fn decay_hits(&mut self) {
        let decay_needed = self.buffer.iter().any(|s| {
            if let KeyValueSlot::Used { hits, .. } = s {
                hits.get() == u8::MAX
            } else {
                false
            }
        });

        if decay_needed {
            for s in &mut self.buffer {
                if let KeyValueSlot::Used { hits, .. } = s {
                    hits.set(hits.get() >> 1); // Divide by two.
                }
            }
        }
    }

    /// Insert a key/value pair.
    ///
    /// # 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 insert(&mut self, k: K, v: V) {
        self.decay_hits();

        match self.buffer.iter_mut().find(|e| e.is_key(&k)) {
            Some(s) => s.update_value(v),
            None => {
                self.evict_and_replace(k, v);
            }
        }
    }

    /// 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,
    {
        self.buffer
            .iter()
            .find(|e| e.is_key(k))
            // SAFETY: We must unwrap the `Option` we get from the cache.
            .map(|e| unsafe { e.get_value().unwrap_unchecked() })
    }

    /// 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,
    {
        self.buffer
            .iter_mut()
            .find(|e| e.is_key(k))
            // SAFETY: We must unwrap the `Option` we get from the cache.
            .map(|e| unsafe { e.get_value_mut().unwrap_unchecked() })
    }

    /// 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
        F: FnOnce(&K) -> V,
    {
        if let Some(i) = self.get_key_index(k) {
            // SAFETY: The key index was retrieved from a found key.
            unsafe { self.buffer.get_unchecked(i).get_value().unwrap_unchecked() }
        } 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
        F: FnOnce(&K) -> Result<V, E>,
    {
        if let Some(i) = self.get_key_index(k) {
            // SAFETY: The key index was retrieved from a found key.
            Ok(unsafe { self.buffer.get_unchecked(i).get_value().unwrap_unchecked() })
        } else {
            f(k).map(|v| self.evict_and_replace(k.clone(), v))
        }
    }

    /// 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);
    }
}

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

#[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));
    }
}