From d83c07c13cd738d658d7ea21df70de5d4e6b4640 Mon Sep 17 00:00:00 2001
From: Roman Krasiuk <rokrassyuk@gmail.com>
Date: Tue, 11 Apr 2023 09:41:28 +0300
Subject: [PATCH] feat(trie): prefix set (#2181)

---
 crates/primitives/src/trie/nibbles.rs |  6 +++
 crates/trie/src/lib.rs                |  4 ++
 crates/trie/src/prefix_set.rs         | 70 +++++++++++++++++++++++++++
 3 files changed, 80 insertions(+)
 create mode 100644 crates/trie/src/prefix_set.rs

diff --git a/crates/primitives/src/trie/nibbles.rs b/crates/primitives/src/trie/nibbles.rs
index 82b634fae0..5292c313f8 100644
--- a/crates/primitives/src/trie/nibbles.rs
+++ b/crates/primitives/src/trie/nibbles.rs
@@ -36,6 +36,12 @@ impl From<&[u8]> for Nibbles {
     }
 }
 
+impl<const N: usize> From<&[u8; N]> for Nibbles {
+    fn from(arr: &[u8; N]) -> Self {
+        Nibbles::from_hex(arr.to_vec())
+    }
+}
+
 impl std::fmt::Debug for Nibbles {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("Nibbles").field("hex_data", &hex::encode(&self.hex_data)).finish()
diff --git a/crates/trie/src/lib.rs b/crates/trie/src/lib.rs
index f0ae16d486..c67f22e87d 100644
--- a/crates/trie/src/lib.rs
+++ b/crates/trie/src/lib.rs
@@ -14,3 +14,7 @@ pub mod nodes;
 
 /// The implementation of hash builder.
 pub mod hash_builder;
+
+/// The implementation of a container for storing intermediate changes to a trie.
+/// The container indicates when the trie has been modified.
+pub mod prefix_set;
diff --git a/crates/trie/src/prefix_set.rs b/crates/trie/src/prefix_set.rs
new file mode 100644
index 0000000000..749796db1f
--- /dev/null
+++ b/crates/trie/src/prefix_set.rs
@@ -0,0 +1,70 @@
+use std::collections::BTreeSet;
+
+use reth_primitives::trie::Nibbles;
+
+/// A container for efficiently storing and checking for the presence of key prefixes.
+///
+/// This data structure stores a set of `Nibbles` and provides methods to insert
+/// new elements and check whether any existing element has a given prefix.
+///
+/// Internally, this implementation uses a `BTreeSet` to store the `Nibbles`, which
+/// ensures that they are always sorted and deduplicated.
+///
+/// # Examples
+///
+/// ```
+/// use reth_trie::prefix_set::PrefixSet;
+///
+/// let mut prefix_set = PrefixSet::default();
+/// prefix_set.insert(b"key1");
+/// prefix_set.insert(b"key2");
+///
+/// assert_eq!(prefix_set.contains(b"key"), true);
+/// ```
+#[derive(Debug, Default, Clone)]
+pub struct PrefixSet {
+    keys: BTreeSet<Nibbles>,
+}
+
+impl PrefixSet {
+    /// Returns `true` if any of the keys in the set has the given prefix or
+    /// if the given prefix is a prefix of any key in the set.
+    pub fn contains<T: Into<Nibbles>>(&self, prefix: T) -> bool {
+        let prefix = prefix.into();
+        self.keys.iter().any(|key| key.has_prefix(&prefix))
+    }
+
+    /// Inserts the given `nibbles` into the set.
+    pub fn insert<T: Into<Nibbles>>(&mut self, nibbles: T) {
+        self.keys.insert(nibbles.into());
+    }
+
+    /// Returns the number of elements in the set.
+    pub fn len(&self) -> usize {
+        self.keys.len()
+    }
+
+    /// Returns `true` if the set is empty.
+    pub fn is_empty(&self) -> bool {
+        self.keys.is_empty()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_contains_with_multiple_inserts_and_duplicates() {
+        let mut prefix_set = PrefixSet::default();
+        prefix_set.insert(b"123");
+        prefix_set.insert(b"124");
+        prefix_set.insert(b"456");
+        prefix_set.insert(b"123"); // Duplicate
+
+        assert!(prefix_set.contains(b"12"));
+        assert!(prefix_set.contains(b"45"));
+        assert!(!prefix_set.contains(b"78"));
+        assert_eq!(prefix_set.len(), 3); // Length should be 3 (excluding duplicate)
+    }
+}