improve usability of schnellru::LruMap and LruCache (LinkedHashSet cache) (#6016)

Cargo.lock

@@ -6264,6 +6264,7 @@ dependencies = [
  "async-trait",
  "auto_impl",
  "criterion",
+ "derive_more",
  "enr",
  "ethers-core",
  "ethers-middleware",
@@ -6295,6 +6296,7 @@ dependencies = [
  "reth-tokio-util",
  "reth-tracing",
  "reth-transaction-pool",
+ "schnellru",
  "secp256k1 0.27.0",
  "serde",
  "serde_json",

Cargo.toml

@@ -188,6 +188,7 @@ aquamarine = "0.5"
 bytes = "1.5"
 bitflags = "2.4"
 clap = "4"
+derive_more = "0.99.17"
 eyre = "0.6"
 tracing = "0.1.0"
 tracing-appender = "0.2"

crates/net/network/Cargo.toml

@@ -63,6 +63,8 @@ linked_hash_set = "0.1"
 linked-hash-map = "0.5.6"
 rand.workspace = true
 secp256k1 = { workspace = true, features = ["global-context", "rand-std", "recovery"] }
+derive_more.workspace = true
+schnellru.workspace = true
 
 enr = { workspace = true, features = ["rust-secp256k1"], optional = true }
 tempfile = { workspace = true, optional = true }

crates/net/network/src/cache.rs

@@ -1,5 +1,13 @@
+use core::hash::BuildHasher;
+use derive_more::{Deref, DerefMut};
 use linked_hash_set::LinkedHashSet;
-use std::{borrow::Borrow, hash::Hash, num::NonZeroUsize};
+use schnellru::{self, ByLength, Limiter, RandomState, Unlimited};
+use std::{
+    borrow::Borrow,
+    fmt::{self, Write},
+    hash::Hash,
+    num::NonZeroUsize,
+};
 
 /// A minimal LRU cache based on a `LinkedHashSet` with limited capacity.
 ///
@@ -12,7 +20,7 @@ pub struct LruCache<T: Hash + Eq> {
 }
 
 impl<T: Hash + Eq> LruCache<T> {
-    /// Creates a new `LruCache` using the given limit
+    /// Creates a new [`LruCache`] using the given limit
     pub fn new(limit: NonZeroUsize) -> Self {
         Self { inner: LinkedHashSet::new(), limit }
     }
@@ -26,14 +34,21 @@ impl<T: Hash + Eq> LruCache<T> {
     /// 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 (new_entry, _evicted_val) = self.insert_and_get_evicted(entry);
+        new_entry
+    }
+
+    /// Same as [`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.insert(entry) {
             if self.limit.get() == self.inner.len() {
                 // remove the oldest element in the set
-                _ = self.remove_lru();
+                return (true, self.remove_lru())
             }
-            return true
+            return (true, None)
         }
-        false
+        (false, None)
     }
 
     /// Remove the least recently used entry and return it.
@@ -44,6 +59,12 @@ impl<T: Hash + Eq> LruCache<T> {
         self.inner.pop_front()
     }
 
+    #[allow(dead_code)]
+    /// Expels the given value. Returns true if the value existed.
+    pub fn remove(&mut self, value: &T) -> bool {
+        self.inner.remove(value)
+    }
+
     /// Returns `true` if the set contains a value.
     pub fn contains<Q: ?Sized>(&self, value: &Q) -> bool
     where
@@ -57,6 +78,20 @@ impl<T: Hash + Eq> LruCache<T> {
     pub fn iter(&self) -> impl Iterator<Item = &T> + '_ {
         self.inner.iter()
     }
+
+    /// Returns number of elements currently in cache.
+    #[cfg(test)]
+    #[allow(dead_code)]
+    pub fn len(&self) -> usize {
+        self.inner.len()
+    }
+
+    /// Returns `true` if there are currently no elements in the cache.
+    #[cfg(test)]
+    #[allow(dead_code)]
+    pub fn is_empty(&self) -> bool {
+        self.inner.is_empty()
+    }
 }
 
 impl<T> Extend<T> for LruCache<T>
@@ -70,8 +105,56 @@ where
     }
 }
 
+/// Wrapper of [`schnellru::LruMap`] that implements [`fmt::Debug`].
+#[derive(Deref, DerefMut)]
+pub struct LruMap<K, V, L = ByLength, S = RandomState>(schnellru::LruMap<K, V, L, S>)
+where
+    K: Hash + PartialEq,
+    L: Limiter<K, V>,
+    S: BuildHasher;
+
+impl<K, V, L, S> fmt::Debug for LruMap<K, V, L, S>
+where
+    K: Hash + PartialEq + fmt::Display,
+    V: fmt::Debug,
+    L: Limiter<K, V>,
+    S: BuildHasher,
+{
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let mut debug_struct = f.debug_struct("LruMap");
+        for (k, v) in self.0.iter() {
+            let mut key_str = String::new();
+            write!(&mut key_str, "{k}")?;
+            debug_struct.field(&key_str, &v);
+        }
+        debug_struct.finish()
+    }
+}
+
+impl<K, V> LruMap<K, V>
+where
+    K: Hash + PartialEq,
+{
+    /// Returns a new cache with default limiter and hash builder.
+    #[allow(dead_code)]
+    pub fn new(max_length: u32) -> Self {
+        LruMap(schnellru::LruMap::new(ByLength::new(max_length)))
+    }
+}
+
+impl<K, V> LruMap<K, V, Unlimited>
+where
+    K: Hash + PartialEq,
+{
+    /// Returns a new cache with [`Unlimited`] limiter and default hash builder.
+    #[allow(dead_code)]
+    pub fn new_unlimited() -> Self {
+        LruMap(schnellru::LruMap::new(Unlimited))
+    }
+}
+
 #[cfg(test)]
-mod tests {
+mod test {
     use super::*;
 
     #[test]
@@ -115,4 +198,26 @@ mod tests {
             assert!(cache.contains(e));
         }
     }
+
+    #[test]
+    #[allow(dead_code)]
+    fn test_debug_impl_lru_map() {
+        use derive_more::Display;
+
+        #[derive(Debug, Hash, PartialEq, Eq, Display)]
+        struct Key(i8);
+
+        #[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 { 2: Value(22), 1: Value(11) }", format!("{cache:?}"))
+    }
 }