fix

bin/reth/src/commands/debug_cmd/in_memory_merkle.rs

@@ -25,7 +25,7 @@ use reth_provider::{
 use reth_revm::database::StateProviderDatabase;
 use reth_stages::StageId;
 use reth_tasks::TaskExecutor;
-use reth_trie::{updates::TrieKey, StateRoot};
+use reth_trie::StateRoot;
 use std::{path::PathBuf, sync::Arc};
 use tracing::*;
 
@@ -188,15 +188,16 @@ impl Command {
         // Compare updates
         let mut in_mem_mismatched = Vec::new();
         let mut incremental_mismatched = Vec::new();
-        let mut in_mem_updates_iter = in_memory_updates.into_iter().peekable();
-        let mut incremental_updates_iter = incremental_trie_updates.into_iter().peekable();
+        let mut in_mem_updates_iter = in_memory_updates.account_nodes_ref().into_iter().peekable();
+        let mut incremental_updates_iter =
+            incremental_trie_updates.account_nodes_ref().into_iter().peekable();
 
         while in_mem_updates_iter.peek().is_some() || incremental_updates_iter.peek().is_some() {
             match (in_mem_updates_iter.next(), incremental_updates_iter.next()) {
                 (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.len() > self.skip_node_depth.unwrap_or_default())
+                        in_mem.0.len() > self.skip_node_depth.unwrap_or_default()
                     {
                         in_mem_mismatched.push(in_mem);
                         incremental_mismatched.push(incr);

crates/stages/stages/benches/setup/mod.rs

@@ -139,7 +139,11 @@ pub(crate) fn txs_testdata(num_blocks: u64) -> TestStageDB {
         let offset = transitions.len() as u64;
 
         db.insert_changesets(transitions, None).unwrap();
-        db.commit(|tx| Ok(updates.write_to_database(tx)?)).unwrap();
+        db.commit(|tx| {
+            updates.write_to_database(tx)?;
+            Ok(())
+        })
+        .unwrap();
 
         let (transitions, final_state) = random_changeset_range(
             &mut rng,

crates/trie/trie/src/updates.rs

@@ -10,82 +10,6 @@ use reth_db_api::{
 use reth_primitives::B256;
 use std::collections::{HashMap, HashSet};
 
-/// The key of a trie node.
-#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
-#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
-pub enum TrieKey {
-    /// A node in the account trie.
-    AccountNode(Nibbles),
-    /// A node in the storage trie.
-    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<&Nibbles> {
-        if let Self::AccountNode(nibbles) = &self {
-            Some(nibbles)
-        } else {
-            None
-        }
-    }
-
-    /// Returns reference to storage node key if the key is for [`Self::StorageNode`].
-    pub const fn as_storage_node_key(&self) -> Option<(&B256, &Nibbles)> {
-        if let Self::StorageNode(key, subkey) = &self {
-            Some((key, subkey))
-        } else {
-            None
-        }
-    }
-
-    /// Returns reference to storage trie key if the key is for [`Self::StorageTrie`].
-    pub const fn as_storage_trie_key(&self) -> Option<&B256> {
-        if let Self::StorageTrie(key) = &self {
-            Some(key)
-        } else {
-            None
-        }
-    }
-}
-
-/// The operation to perform on the trie.
-#[derive(PartialEq, Eq, Debug, Clone)]
-#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
-pub enum TrieOp {
-    /// Delete the node entry.
-    Delete,
-    /// Update the node entry with the provided value.
-    Update(BranchNodeCompact),
-}
-
-impl TrieOp {
-    /// Returns `true` if the operation is an update.
-    pub const fn is_update(&self) -> bool {
-        matches!(self, Self::Update(..))
-    }
-
-    /// Returns reference to updated branch node if operation is [`Self::Update`].
-    pub const fn as_update(&self) -> Option<&BranchNodeCompact> {
-        if let Self::Update(node) = &self {
-            Some(node)
-        } else {
-            None
-        }
-    }
-
-    /// Returns owned updated branch node if operation is [`Self::Update`].
-    pub fn into_update(self) -> Option<BranchNodeCompact> {
-        if let Self::Update(node) = self {
-            Some(node)
-        } else {
-            None
-        }
-    }
-}
-
 /// The aggregation of trie updates.
 #[derive(PartialEq, Eq, Clone, Default, Debug)]
 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
@@ -96,6 +20,11 @@ pub struct TrieUpdates {
 }
 
 impl TrieUpdates {
+    /// Returns reference to updated account nodes.
+    pub fn account_nodes_ref(&self) -> &HashMap<Nibbles, BranchNodeCompact> {
+        &self.account_nodes
+    }
+
     /// Insert storage updates for a given hashed address.
     pub fn insert_storage_updates(
         &mut self,