chore(trie): remove database-related types from trie keys (#9175)

paradigmxyz · Jun 28, 2024 · 6e564cd · 6e564cd
1 parent b24ca76
commit 6e564cd
Show file tree

Hide file tree

Showing 5 changed files with 39 additions and 40 deletions.
diff --git a/bin/reth/src/commands/debug_cmd/in_memory_merkle.rs b/bin/reth/src/commands/debug_cmd/in_memory_merkle.rs
@@ -196,7 +196,7 @@ impl Command {
                 (Some(in_mem), Some(incr)) => {
                     similar_asserts::assert_eq!(in_mem.0, incr.0, "Nibbles don't match");
                     if in_mem.1 != incr.1 &&
-                        matches!(in_mem.0, TrieKey::AccountNode(ref nibbles) if nibbles.0.len() > self.skip_node_depth.unwrap_or_default())
+                        matches!(in_mem.0, TrieKey::AccountNode(ref nibbles) if nibbles.len() > self.skip_node_depth.unwrap_or_default())
                     {
                         in_mem_mismatched.push(in_mem);
                         incremental_mismatched.push(incr);

diff --git a/crates/trie/trie/src/trie.rs b/crates/trie/trie/src/trie.rs
@@ -1208,7 +1208,7 @@ mod tests {
             .iter()
             .filter_map(|entry| match entry {
                 (TrieKey::AccountNode(nibbles), TrieOp::Update(node)) => {
-                    Some((nibbles.0.clone(), node.clone()))
+                    Some((nibbles.clone(), node.clone()))
                 }
                 _ => None,
             })
@@ -1295,7 +1295,7 @@ mod tests {
             .iter()
             .filter_map(|entry| match entry {
                 (TrieKey::StorageNode(_, nibbles), TrieOp::Update(node)) => {
-                    Some((nibbles.0.clone(), node.clone()))
+                    Some((nibbles.clone(), node.clone()))
                 }
                 _ => None,
             })

diff --git a/crates/trie/trie/src/trie_cursor/database_cursors.rs b/crates/trie/trie/src/trie_cursor/database_cursors.rs
@@ -61,7 +61,7 @@ where
 
     /// Retrieves the current key in the cursor.
     fn current(&mut self) -> Result<Option<TrieKey>, DatabaseError> {
-        Ok(self.0.current()?.map(|(k, _)| TrieKey::AccountNode(k)))
+        Ok(self.0.current()?.map(|(k, _)| TrieKey::AccountNode(k.0)))
     }
 }
 
@@ -110,7 +110,7 @@ where
 
     /// Retrieves the current value in the storage trie cursor.
     fn current(&mut self) -> Result<Option<TrieKey>, DatabaseError> {
-        Ok(self.cursor.current()?.map(|(k, v)| TrieKey::StorageNode(k, v.nibbles)))
+        Ok(self.cursor.current()?.map(|(k, v)| TrieKey::StorageNode(k, v.nibbles.0)))
     }
 }
 

diff --git a/crates/trie/trie/src/trie_cursor/update.rs b/crates/trie/trie/src/trie_cursor/update.rs
@@ -2,7 +2,7 @@ use super::{TrieCursor, TrieCursorFactory};
 use crate::updates::{TrieKey, TrieOp, TrieUpdatesSorted};
 use reth_db::DatabaseError;
 use reth_primitives::B256;
-use reth_trie_common::{BranchNodeCompact, Nibbles, StoredNibbles, StoredNibblesSubKey};
+use reth_trie_common::{BranchNodeCompact, Nibbles};
 
 /// The trie cursor factory for the trie updates.
 #[derive(Debug, Clone)]
@@ -64,8 +64,7 @@ impl<'a, C: TrieCursor> TrieCursor for TrieUpdatesAccountTrieCursor<'a, C> {
             }
         } else {
             let result = self.cursor.seek_exact(key)?;
-            self.last_key =
-                result.as_ref().map(|(k, _)| TrieKey::AccountNode(StoredNibbles(k.clone())));
+            self.last_key = result.as_ref().map(|(k, _)| TrieKey::AccountNode(k.clone()));
             Ok(result)
         }
     }
@@ -74,12 +73,11 @@ impl<'a, C: TrieCursor> TrieCursor for TrieUpdatesAccountTrieCursor<'a, C> {
         &mut self,
         key: Nibbles,
     ) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
-        let stored_nibbles = StoredNibbles(key.clone());
         let trie_update_entry = self
             .trie_updates
             .trie_operations
             .iter()
-            .find(|(k, _)| matches!(k, TrieKey::AccountNode(nibbles) if nibbles <= &stored_nibbles))
+            .find(|(k, _)| matches!(k, TrieKey::AccountNode(nibbles) if nibbles <= &key))
             .cloned();
 
         if let Some((trie_key, trie_op)) = trie_update_entry {
@@ -89,14 +87,13 @@ impl<'a, C: TrieCursor> TrieCursor for TrieUpdatesAccountTrieCursor<'a, C> {
             };
             self.last_key = Some(trie_key);
             match trie_op {
-                TrieOp::Update(node) => return Ok(Some((nibbles.0, node))),
+                TrieOp::Update(node) => return Ok(Some((nibbles, node))),
                 TrieOp::Delete => return Ok(None),
             }
         }
 
         let result = self.cursor.seek(key)?;
-        self.last_key =
-            result.as_ref().map(|(k, _)| TrieKey::AccountNode(StoredNibbles(k.clone())));
+        self.last_key = result.as_ref().map(|(k, _)| TrieKey::AccountNode(k.clone()));
         Ok(result)
     }
 
@@ -141,9 +138,8 @@ impl<'a, C: TrieCursor> TrieCursor for TrieUpdatesStorageTrieCursor<'a, C> {
             }
         } else {
             let result = self.cursor.seek_exact(key)?;
-            self.last_key = result.as_ref().map(|(k, _)| {
-                TrieKey::StorageNode(self.hashed_address, StoredNibblesSubKey(k.clone()))
-            });
+            self.last_key =
+                result.as_ref().map(|(k, _)| TrieKey::StorageNode(self.hashed_address, k.clone()));
             Ok(result)
         }
     }
@@ -154,7 +150,7 @@ impl<'a, C: TrieCursor> TrieCursor for TrieUpdatesStorageTrieCursor<'a, C> {
     ) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
         let mut trie_update_entry = self.trie_updates.trie_operations.get(self.trie_update_index);
         while trie_update_entry
-            .filter(|(k, _)| matches!(k, TrieKey::StorageNode(address, nibbles) if address == &self.hashed_address && nibbles.0 < key)).is_some()
+            .filter(|(k, _)| matches!(k, TrieKey::StorageNode(address, nibbles) if address == &self.hashed_address && nibbles < &key)).is_some()
         {
             self.trie_update_index += 1;
             trie_update_entry = self.trie_updates.trie_operations.get(self.trie_update_index);
@@ -169,15 +165,14 @@ impl<'a, C: TrieCursor> TrieCursor for TrieUpdatesStorageTrieCursor<'a, C> {
             };
             self.last_key = Some(trie_key.clone());
             match trie_op {
-                TrieOp::Update(node) => return Ok(Some((nibbles.0, node.clone()))),
+                TrieOp::Update(node) => return Ok(Some((nibbles, node.clone()))),
                 TrieOp::Delete => return Ok(None),
             }
         }
 
         let result = self.cursor.seek(key)?;
-        self.last_key = result.as_ref().map(|(k, _)| {
-            TrieKey::StorageNode(self.hashed_address, StoredNibblesSubKey(k.clone()))
-        });
+        self.last_key =
+            result.as_ref().map(|(k, _)| TrieKey::StorageNode(self.hashed_address, k.clone()));
         Ok(result)
     }
 

diff --git a/crates/trie/trie/src/updates.rs b/crates/trie/trie/src/updates.rs
@@ -16,16 +16,16 @@ use std::collections::{hash_map::IntoIter, HashMap, HashSet};
 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
 pub enum TrieKey {
     /// A node in the account trie.
-    AccountNode(StoredNibbles),
+    AccountNode(Nibbles),
     /// A node in the storage trie.
-    StorageNode(B256, StoredNibblesSubKey),
+    StorageNode(B256, Nibbles),
     /// Storage trie of an account.
     StorageTrie(B256),
 }
 
 impl TrieKey {
     /// Returns reference to account node key if the key is for [`Self::AccountNode`].
-    pub const fn as_account_node_key(&self) -> Option<&StoredNibbles> {
+    pub const fn as_account_node_key(&self) -> Option<&Nibbles> {
         if let Self::AccountNode(nibbles) = &self {
             Some(nibbles)
         } else {
@@ -34,7 +34,7 @@ impl TrieKey {
     }
 
     /// Returns reference to storage node key if the key is for [`Self::StorageNode`].
-    pub const fn as_storage_node_key(&self) -> Option<(&B256, &StoredNibblesSubKey)> {
+    pub const fn as_storage_node_key(&self) -> Option<(&B256, &Nibbles)> {
         if let Self::StorageNode(key, subkey) = &self {
             Some((key, subkey))
         } else {
@@ -121,9 +121,9 @@ impl TrieUpdates {
     /// Extend the updates with account trie updates.
     pub fn extend_with_account_updates(&mut self, updates: HashMap<Nibbles, BranchNodeCompact>) {
         self.extend(
-            updates.into_iter().map(|(nibbles, node)| {
-                (TrieKey::AccountNode(nibbles.into()), TrieOp::Update(node))
-            }),
+            updates
+                .into_iter()
+                .map(|(nibbles, node)| (TrieKey::AccountNode(nibbles), TrieOp::Update(node))),
         );
     }
 
@@ -162,7 +162,7 @@ impl TrieUpdates {
         // Add storage node updates from hash builder.
         let (_, hash_builder_updates) = hash_builder.split();
         self.extend(hash_builder_updates.into_iter().map(|(nibbles, node)| {
-            (TrieKey::StorageNode(hashed_address, nibbles.into()), TrieOp::Update(node))
+            (TrieKey::StorageNode(hashed_address, nibbles), TrieOp::Update(node))
         }));
     }
 
@@ -179,18 +179,21 @@ impl TrieUpdates {
         trie_operations.sort_unstable_by(|a, b| a.0.cmp(&b.0));
         for (key, operation) in trie_operations {
             match key {
-                TrieKey::AccountNode(nibbles) => match operation {
-                    TrieOp::Delete => {
-                        if account_trie_cursor.seek_exact(nibbles)?.is_some() {
-                            account_trie_cursor.delete_current()?;
+                TrieKey::AccountNode(nibbles) => {
+                    let nibbles = StoredNibbles(nibbles);
+                    match operation {
+                        TrieOp::Delete => {
+                            if account_trie_cursor.seek_exact(nibbles)?.is_some() {
+                                account_trie_cursor.delete_current()?;
+                            }
                         }
-                    }
-                    TrieOp::Update(node) => {
-                        if !nibbles.0.is_empty() {
-                            account_trie_cursor.upsert(nibbles, StoredBranchNode(node))?;
+                        TrieOp::Update(node) => {
+                            if !nibbles.0.is_empty() {
+                                account_trie_cursor.upsert(nibbles, StoredBranchNode(node))?;
+                            }
                         }
                     }
-                },
+                }
                 TrieKey::StorageTrie(hashed_address) => match operation {
                     TrieOp::Delete => {
                         if storage_trie_cursor.seek_exact(hashed_address)?.is_some() {
@@ -201,6 +204,7 @@ impl TrieUpdates {
                 },
                 TrieKey::StorageNode(hashed_address, nibbles) => {
                     if !nibbles.is_empty() {
+                        let nibbles = StoredNibblesSubKey(nibbles);
                         // Delete the old entry if it exists.
                         if storage_trie_cursor
                             .seek_by_key_subkey(hashed_address, nibbles.clone())?
@@ -250,7 +254,7 @@ impl TrieUpdatesSorted {
     pub fn find_account_node(&self, key: &Nibbles) -> Option<(TrieKey, TrieOp)> {
         self.trie_operations
             .iter()
-            .find(|(k, _)| matches!(k, TrieKey::AccountNode(nibbles) if &nibbles.0 == key))
+            .find(|(k, _)| matches!(k, TrieKey::AccountNode(nibbles) if nibbles == key))
             .cloned()
     }
 
@@ -261,7 +265,7 @@ impl TrieUpdatesSorted {
         key: &Nibbles,
     ) -> Option<(TrieKey, TrieOp)> {
         self.trie_operations.iter().find(|(k, _)| {
-            matches!(k, TrieKey::StorageNode(address, nibbles) if address == hashed_address && &nibbles.0 == key)
+            matches!(k, TrieKey::StorageNode(address, nibbles) if address == hashed_address && nibbles == key)
         }).cloned()
     }
 }