perf(trie): reduce key copying (#6548)

crates/trie/src/trie_cursor/database_cursors.rs

@@ -7,22 +7,20 @@ use reth_db::{
     DatabaseError,
 };
 use reth_primitives::{
-    trie::{BranchNodeCompact, StoredNibbles, StoredNibblesSubKey},
+    trie::{BranchNodeCompact, Nibbles, StoredNibbles, StoredNibblesSubKey},
     B256,
 };
 
 /// Implementation of the trie cursor factory for a database transaction.
 impl<'a, TX: DbTx> TrieCursorFactory for &'a TX {
-    fn account_trie_cursor(
-        &self,
-    ) -> Result<Box<dyn TrieCursor<Key = StoredNibbles> + '_>, DatabaseError> {
+    fn account_trie_cursor(&self) -> Result<Box<dyn TrieCursor + '_>, DatabaseError> {
         Ok(Box::new(DatabaseAccountTrieCursor::new(self.cursor_read::<tables::AccountsTrie>()?)))
     }
 
     fn storage_tries_cursor(
         &self,
         hashed_address: B256,
-    ) -> Result<Box<dyn TrieCursor<Key = StoredNibblesSubKey> + '_>, DatabaseError> {
+    ) -> Result<Box<dyn TrieCursor + '_>, DatabaseError> {
         Ok(Box::new(DatabaseStorageTrieCursor::new(
             self.cursor_dup_read::<tables::StoragesTrie>()?,
             hashed_address,
@@ -45,23 +43,20 @@ impl<C> TrieCursor for DatabaseAccountTrieCursor<C>
 where
     C: DbCursorRO<tables::AccountsTrie>,
 {
-    /// The type of key used by this cursor.
-    type Key = StoredNibbles;
-
     /// Seeks an exact match for the provided key in the account trie.
     fn seek_exact(
         &mut self,
-        key: Self::Key,
-    ) -> Result<Option<(Vec<u8>, BranchNodeCompact)>, DatabaseError> {
-        Ok(self.0.seek_exact(key)?.map(|value| (value.0 .0.to_vec(), value.1 .0)))
+        key: Nibbles,
+    ) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
+        Ok(self.0.seek_exact(StoredNibbles(key))?.map(|value| (value.0 .0, value.1 .0)))
     }
 
     /// Seeks a key in the account trie that matches or is greater than the provided key.
     fn seek(
         &mut self,
-        key: Self::Key,
-    ) -> Result<Option<(Vec<u8>, BranchNodeCompact)>, DatabaseError> {
-        Ok(self.0.seek(key)?.map(|value| (value.0 .0.to_vec(), value.1 .0)))
+        key: Nibbles,
+    ) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
+        Ok(self.0.seek(StoredNibbles(key))?.map(|value| (value.0 .0, value.1 .0)))
     }
 
     /// Retrieves the current key in the cursor.
@@ -90,30 +85,27 @@ impl<C> TrieCursor for DatabaseStorageTrieCursor<C>
 where
     C: DbDupCursorRO<tables::StoragesTrie> + DbCursorRO<tables::StoragesTrie>,
 {
-    /// Defines the type for keys used in the storage trie cursor.
-    type Key = StoredNibblesSubKey;
-
     /// Seeks an exact match for the given key in the storage trie.
     fn seek_exact(
         &mut self,
-        key: Self::Key,
-    ) -> Result<Option<(Vec<u8>, BranchNodeCompact)>, DatabaseError> {
+        key: Nibbles,
+    ) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
         Ok(self
             .cursor
-            .seek_by_key_subkey(self.hashed_address, key.clone())?
-            .filter(|e| e.nibbles == key)
-            .map(|value| (value.nibbles.to_vec(), value.node)))
+            .seek_by_key_subkey(self.hashed_address, StoredNibblesSubKey(key.clone()))?
+            .filter(|e| e.nibbles == StoredNibblesSubKey(key))
+            .map(|value| (value.nibbles.0, value.node)))
     }
 
     /// Seeks the given key in the storage trie.
     fn seek(
         &mut self,
-        key: Self::Key,
-    ) -> Result<Option<(Vec<u8>, BranchNodeCompact)>, DatabaseError> {
+        key: Nibbles,
+    ) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
         Ok(self
             .cursor
-            .seek_by_key_subkey(self.hashed_address, key)?
-            .map(|value| (value.nibbles.to_vec(), value.node)))
+            .seek_by_key_subkey(self.hashed_address, StoredNibblesSubKey(key))?
+            .map(|value| (value.nibbles.0, value.node)))
     }
 
     /// Retrieves the current value in the storage trie cursor.

crates/trie/src/trie_cursor/mod.rs

@@ -1,7 +1,7 @@
 use crate::updates::TrieKey;
 use reth_db::DatabaseError;
 use reth_primitives::{
-    trie::{BranchNodeCompact, StoredNibbles, StoredNibblesSubKey},
+    trie::{BranchNodeCompact, Nibbles},
     B256,
 };
 
@@ -19,34 +19,27 @@ pub use self::{
 /// Factory for creating trie cursors.
 pub trait TrieCursorFactory {
     /// Create an account trie cursor.
-    fn account_trie_cursor(
-        &self,
-    ) -> Result<Box<dyn TrieCursor<Key = StoredNibbles> + '_>, DatabaseError>;
+    fn account_trie_cursor(&self) -> Result<Box<dyn TrieCursor + '_>, DatabaseError>;
 
     /// Create a storage tries cursor.
     fn storage_tries_cursor(
         &self,
         hashed_address: B256,
-    ) -> Result<Box<dyn TrieCursor<Key = StoredNibblesSubKey> + '_>, DatabaseError>;
+    ) -> Result<Box<dyn TrieCursor + '_>, DatabaseError>;
 }
 
 /// A cursor for navigating a trie that works with both Tables and DupSort tables.
 #[auto_impl::auto_impl(&mut, Box)]
 pub trait TrieCursor {
-    /// The key type of the cursor.
-    type Key: From<Vec<u8>>;
-
     /// Move the cursor to the key and return if it is an exact match.
     fn seek_exact(
         &mut self,
-        key: Self::Key,
-    ) -> Result<Option<(Vec<u8>, BranchNodeCompact)>, DatabaseError>;
+        key: Nibbles,
+    ) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError>;
 
     /// Move the cursor to the key and return a value matching of greater than the key.
-    fn seek(
-        &mut self,
-        key: Self::Key,
-    ) -> Result<Option<(Vec<u8>, BranchNodeCompact)>, DatabaseError>;
+    fn seek(&mut self, key: Nibbles)
+        -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError>;
 
     /// Get the current entry.
     fn current(&mut self) -> Result<Option<TrieKey>, DatabaseError>;

crates/trie/src/trie_cursor/noop.rs

@@ -1,7 +1,7 @@
 use super::{TrieCursor, TrieCursorFactory};
 use crate::updates::TrieKey;
 use reth_db::DatabaseError;
-use reth_primitives::trie::{BranchNodeCompact, StoredNibbles, StoredNibblesSubKey};
+use reth_primitives::trie::{BranchNodeCompact, Nibbles};
 
 /// Noop trie cursor factory.
 #[derive(Default, Debug)]
@@ -10,9 +10,7 @@ pub struct NoopTrieCursorFactory;
 
 impl TrieCursorFactory for NoopTrieCursorFactory {
     /// Generates a Noop account trie cursor.
-    fn account_trie_cursor(
-        &self,
-    ) -> Result<Box<dyn TrieCursor<Key = StoredNibbles> + '_>, DatabaseError> {
+    fn account_trie_cursor(&self) -> Result<Box<dyn TrieCursor + '_>, DatabaseError> {
         Ok(Box::<NoopAccountTrieCursor>::default())
     }
 
@@ -20,7 +18,7 @@ impl TrieCursorFactory for NoopTrieCursorFactory {
     fn storage_tries_cursor(
         &self,
         _hashed_address: reth_primitives::B256,
-    ) -> Result<Box<dyn TrieCursor<Key = StoredNibblesSubKey> + '_>, DatabaseError> {
+    ) -> Result<Box<dyn TrieCursor + '_>, DatabaseError> {
         Ok(Box::<NoopStorageTrieCursor>::default())
     }
 }
@@ -31,21 +29,19 @@ impl TrieCursorFactory for NoopTrieCursorFactory {
 pub struct NoopAccountTrieCursor;
 
 impl TrieCursor for NoopAccountTrieCursor {
-    type Key = StoredNibbles;
-
     /// Seeks within the account trie.
     fn seek(
         &mut self,
-        _key: Self::Key,
-    ) -> Result<Option<(Vec<u8>, BranchNodeCompact)>, DatabaseError> {
+        _key: Nibbles,
+    ) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
         Ok(None)
     }
 
     /// Seeks an exact match within the account trie.
     fn seek_exact(
         &mut self,
-        _key: Self::Key,
-    ) -> Result<Option<(Vec<u8>, BranchNodeCompact)>, DatabaseError> {
+        _key: Nibbles,
+    ) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
         Ok(None)
     }
 
@@ -61,21 +57,19 @@ impl TrieCursor for NoopAccountTrieCursor {
 pub struct NoopStorageTrieCursor;
 
 impl TrieCursor for NoopStorageTrieCursor {
-    type Key = StoredNibblesSubKey;
-
     /// Seeks a key in storage tries.
     fn seek(
         &mut self,
-        _key: Self::Key,
-    ) -> Result<Option<(Vec<u8>, BranchNodeCompact)>, DatabaseError> {
+        _key: Nibbles,
+    ) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
         Ok(None)
     }
 
     /// Seeks an exact match in storage tries.
     fn seek_exact(
         &mut self,
-        _key: Self::Key,
-    ) -> Result<Option<(Vec<u8>, BranchNodeCompact)>, DatabaseError> {
+        _key: Nibbles,
+    ) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
         Ok(None)
     }
 

crates/trie/src/walker.rs

@@ -120,18 +120,14 @@ impl<C: TrieCursor> TrieWalker<C> {
 
     /// Retrieves the current root node from the DB, seeking either the exact node or the next one.
     fn node(&mut self, exact: bool) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
-        let key = self.key().expect("key must exist");
-        let entry = if exact {
-            self.cursor.seek_exact(key.to_vec().into())?
-        } else {
-            self.cursor.seek(key.to_vec().into())?
-        };
+        let key = self.key().expect("key must exist").clone();
+        let entry = if exact { self.cursor.seek_exact(key)? } else { self.cursor.seek(key)? };
 
         if let Some((_, node)) = &entry {
             assert!(!node.state_mask.is_empty());
         }
 
-        Ok(entry.map(|(k, v)| (Nibbles::from_nibbles_unchecked(k), v)))
+        Ok(entry)
     }
 
     /// Consumes the next node in the trie, updating the stack.