reth_network/

cache.rs

//! Network cache support

use alloy_primitives::map::DefaultHashBuilder;
use core::hash::BuildHasher;
use itertools::Itertools;
use std::{fmt, hash::Hash};

pub use reth_lru::LruMap;

/// A minimal LRU cache based on a [`LruMap`] with limited capacity.
///
/// If the length exceeds the set capacity, the oldest element will be removed
/// In the limit, for each element inserted the oldest existing element will be removed.
pub struct LruCache<T: Hash + Eq + fmt::Debug, S = DefaultHashBuilder>
where
    S: BuildHasher,
{
    inner: LruMap<T, (), S>,
}

impl<T: Hash + Eq + fmt::Debug> LruCache<T> {
    /// Creates a new [`LruCache`] using the given limit
    pub fn new(limit: u32) -> Self {
        Self::with_hasher(limit, Default::default())
    }
}

impl<T: Hash + Eq + fmt::Debug, S: BuildHasher> LruCache<T, S> {
    /// Creates a new [`LruCache`] using the given limit and hash builder.
    pub const fn with_hasher(limit: u32, hash_builder: S) -> Self {
        Self { inner: LruMap::with_hasher(limit, hash_builder) }
    }

    /// Insert an element into the set.
    ///
    /// Automatically evicts the least recently used item if capacity is exceeded.
    ///
    /// This method is more efficient than [`insert_and_get_evicted`](Self::insert_and_get_evicted)
    /// as it performs fewer checks. Use this method when you don't need information about
    /// evicted values.
    ///
    /// If the set did not have this value present, true is returned.
    /// If the set did have this value present, false is returned.
    pub fn insert(&mut self, entry: T) -> bool {
        let mut is_new = false;
        self.inner.get_or_insert(entry, || {
            is_new = true;
        });
        is_new
    }

    /// Same as [`insert`](Self::insert) but returns a tuple, where the second index is the evicted
    /// value, if one was evicted.
    pub fn insert_and_get_evicted(&mut self, entry: T) -> (bool, Option<T>) {
        if self.inner.get(&entry).is_some() {
            return (false, None)
        }
        (true, self.inner.insert(entry, ()).map(|(evicted, ())| evicted))
    }

    /// Gets the given element, if exists, and promotes to lru.
    pub fn get(&mut self, entry: &T) -> Option<&T> {
        self.inner.get(entry)?;
        self.inner.peek_newest().map(|(key, ())| key)
    }

    /// Iterates through entries and returns a reference to the given entry, if exists, without
    /// promoting to lru.
    ///
    /// NOTE: Use this for type that have custom impl of [`PartialEq`] and [`Eq`], that aren't
    /// unique by all fields. If `PartialEq` and `Eq` are derived for a type, it's more efficient to
    /// call [`contains`](Self::contains).
    pub fn find(&self, entry: &T) -> Option<&T> {
        self.iter().find(|key| *key == entry)
    }

    /// Expels the given value. Returns true if the value existed.
    pub fn remove(&mut self, value: &T) -> bool {
        self.inner.remove(value).is_some()
    }

    /// Returns `true` if the set contains a value.
    pub fn contains(&self, value: &T) -> bool {
        self.inner.contains(value)
    }

    /// Returns an iterator over all cached entries in lru order
    pub fn iter(&self) -> impl Iterator<Item = &T> + '_ {
        self.inner.iter().map(|(k, ())| k)
    }

    /// Returns number of elements currently in cache.
    pub const fn len(&self) -> usize {
        self.inner.len()
    }

    /// Returns `true` if there are currently no elements in the cache.
    pub const fn is_empty(&self) -> bool {
        self.inner.is_empty()
    }
}

impl<T, S> Extend<T> for LruCache<T, S>
where
    T: Eq + Hash + fmt::Debug,
    S: BuildHasher,
{
    fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
        for item in iter {
            _ = self.insert(item);
        }
    }
}

impl<T, S> fmt::Debug for LruCache<T, S>
where
    T: fmt::Debug + Hash + Eq,
    S: BuildHasher,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut debug_struct = f.debug_struct("LruCache");

        debug_struct.field("limit", &self.inner.max_length());

        debug_struct.field(
            "ret %iter",
            &format_args!("Iter: {{{} }}", self.iter().map(|k| format!(" {k:?}")).format(",")),
        );

        debug_struct.finish()
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use derive_more::Constructor;
    use std::hash::Hasher;

    #[derive(Debug, Hash, PartialEq, Eq, Clone, Copy)]
    struct Key(i8);

    #[derive(Debug, Eq, Constructor, Clone, Copy)]
    struct CompoundKey {
        // type unique for id
        id: i8,
        other: i8,
    }

    impl PartialEq for CompoundKey {
        fn eq(&self, other: &Self) -> bool {
            self.id == other.id
        }
    }

    impl Hash for CompoundKey {
        fn hash<H: Hasher>(&self, state: &mut H) {
            self.id.hash(state)
        }
    }

    #[test]
    fn test_cache_should_insert_into_empty_set() {
        let mut cache = LruCache::new(5);
        let entry = "entry";
        assert!(cache.insert(entry));
        assert!(cache.contains(&entry));
    }

    #[test]
    fn test_cache_should_not_insert_same_value_twice() {
        let mut cache = LruCache::new(5);
        let entry = "entry";
        assert!(cache.insert(entry));
        assert!(!cache.insert(entry));
    }

    #[test]
    fn test_cache_should_remove_oldest_element_when_exceeding_limit() {
        let mut cache = LruCache::new(1);
        let old_entry = "old_entry";
        let new_entry = "new_entry";
        cache.insert(old_entry);
        cache.insert("entry");
        cache.insert(new_entry);
        assert!(cache.contains(&new_entry));
        assert!(!cache.contains(&old_entry));
    }

    #[test]
    fn test_cache_should_extend_an_array() {
        let mut cache = LruCache::new(5);
        let entries = ["some_entry", "another_entry"];
        cache.extend(entries);
        for e in entries {
            assert!(cache.contains(&e));
        }
    }

    #[test]
    #[expect(dead_code)]
    fn test_debug_impl_lru_map() {
        #[derive(Debug)]
        struct Value(i8);

        let mut cache = LruMap::new(2);
        let key_1 = Key(1);
        let value_1 = Value(11);
        cache.insert(key_1, value_1);
        let key_2 = Key(2);
        let value_2 = Value(22);
        cache.insert(key_2, value_2);

        assert_eq!("LruMap { len: 2, max_length: 2, .. }", format!("{cache:?}"))
    }

    #[test]
    fn test_debug_impl_lru_cache() {
        let mut cache = LruCache::new(2);
        let key_1 = Key(1);
        cache.insert(key_1);
        let key_2 = Key(2);
        cache.insert(key_2);

        assert_eq!(
            "LruCache { limit: 2, ret %iter: Iter: { Key(2), Key(1) } }",
            format!("{cache:?}")
        )
    }

    #[test]
    fn get() {
        let mut cache = LruCache::new(2);
        let key_1 = Key(1);
        cache.insert(key_1);
        let key_2 = Key(2);
        cache.insert(key_2);

        // promotes key 1 to lru
        _ = cache.get(&key_1);

        assert_eq!(
            "LruCache { limit: 2, ret %iter: Iter: { Key(1), Key(2) } }",
            format!("{cache:?}")
        )
    }

    #[test]
    fn get_ty_custom_eq_impl() {
        let mut cache = LruCache::new(2);
        let key_1 = CompoundKey::new(1, 11);
        cache.insert(key_1);
        let key_2 = CompoundKey::new(2, 22);
        cache.insert(key_2);

        let key = cache.get(&key_1);

        assert_eq!(key_1.other, key.unwrap().other)
    }

    #[test]
    fn peek() {
        let mut cache = LruCache::new(2);
        let key_1 = Key(1);
        cache.insert(key_1);
        let key_2 = Key(2);
        cache.insert(key_2);

        // doesn't promote key 1 to lru
        _ = cache.find(&key_1);

        assert_eq!(
            "LruCache { limit: 2, ret %iter: Iter: { Key(2), Key(1) } }",
            format!("{cache:?}")
        )
    }

    #[test]
    fn peek_ty_custom_eq_impl() {
        let mut cache = LruCache::new(2);
        let key_1 = CompoundKey::new(1, 11);
        cache.insert(key_1);
        let key_2 = CompoundKey::new(2, 22);
        cache.insert(key_2);

        let key = cache.find(&key_1);

        assert_eq!(key_1.other, key.unwrap().other)
    }

    #[test]
    fn test_insert_methods() {
        let mut cache = LruCache::new(2);

        // Test basic insert
        assert!(cache.insert("first")); // new entry
        assert!(!cache.insert("first")); // existing entry
        assert!(cache.insert("second")); // new entry

        // Test insert_and_get_evicted
        let (is_new, evicted) = cache.insert_and_get_evicted("third");
        assert!(is_new); // should be new entry
        assert_eq!(evicted, Some("first")); // should evict

        assert!(cache.contains(&"second"));
        assert!(cache.contains(&"third"));
        assert!(!cache.contains(&"first"));

        // Test insert_and_get_evicted with existing entry
        let (is_new, evicted) = cache.insert_and_get_evicted("second");
        assert!(!is_new); // should not be new
        assert_eq!(evicted, None); // should not evict anything
    }
}