feat(trie): optimize prefix set (#2417)

crates/trie/benches/prefix_set.rs

@@ -1,25 +1,13 @@
-#![allow(dead_code, unused_imports, non_snake_case)]
-
 use criterion::{
     black_box, criterion_group, criterion_main, measurement::WallTime, BenchmarkGroup, Criterion,
 };
 use proptest::{
-    arbitrary::Arbitrary,
     prelude::*,
     strategy::{Strategy, ValueTree},
     test_runner::{basic_result_cache, TestRunner},
 };
-use reth_db::{
-    cursor::{DbCursorRW, DbDupCursorRO, DbDupCursorRW},
-    mdbx::Env,
-    TxHashNumber,
-};
-use reth_primitives::H256;
 use reth_trie::{prefix_set::PrefixSet, Nibbles};
-use std::{
-    collections::{BTreeSet, HashSet},
-    time::Instant,
-};
+use std::collections::BTreeSet;
 
 pub trait PrefixSetAbstraction: Default {
     fn insert(&mut self, key: Nibbles);
@@ -32,7 +20,7 @@ impl PrefixSetAbstraction for PrefixSet {
     }
 
     fn contains(&mut self, key: Nibbles) -> bool {
-        PrefixSet::contains(&self, key)
+        PrefixSet::contains(self, key)
     }
 }
 
@@ -123,10 +111,7 @@ criterion_main!(prefix_set);
 
 mod implementations {
     use super::*;
-    use std::{
-        collections::btree_set::Range, iter::Peekable, marker::PhantomPinned, ops::Bound, pin::Pin,
-        ptr::NonNull,
-    };
+    use std::ops::Bound;
 
     #[derive(Default)]
     pub struct BTreeAnyPrefixSet {

crates/trie/src/prefix_set/mod.rs

@@ -1,5 +1,4 @@
 use crate::Nibbles;
-use std::collections::BTreeSet;
 
 mod loader;
 pub use loader::PrefixSetLoader;
@@ -25,20 +24,46 @@ pub use loader::PrefixSetLoader;
 /// ```
 #[derive(Debug, Default, Clone)]
 pub struct PrefixSet {
-    keys: BTreeSet<Nibbles>,
+    keys: Vec<Nibbles>,
+    sorted: bool,
+    index: usize,
 }
 
 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 {
+    pub fn contains<T: Into<Nibbles>>(&mut self, prefix: T) -> bool {
+        if !self.sorted {
+            self.keys.sort();
+            self.keys.dedup();
+            self.sorted = true;
+        }
+
         let prefix = prefix.into();
-        self.keys.iter().any(|key| key.has_prefix(&prefix))
+
+        while self.index > 0 && self.keys[self.index] > prefix {
+            self.index -= 1;
+        }
+
+        for (idx, key) in self.keys[self.index..].iter().enumerate() {
+            if key.has_prefix(&prefix) {
+                self.index += idx;
+                return true
+            }
+
+            if key > &prefix {
+                self.index += idx;
+                return false
+            }
+        }
+
+        false
     }
 
     /// Inserts the given `nibbles` into the set.
     pub fn insert<T: Into<Nibbles>>(&mut self, nibbles: T) {
-        self.keys.insert(nibbles.into());
+        self.sorted = false;
+        self.keys.push(nibbles.into());
     }
 
     /// Returns the number of elements in the set.