feat(tree): when comparing trie updates, check the database (#13765)

crates/engine/tree/Cargo.toml

@@ -15,6 +15,7 @@ workspace = true
 reth-chain-state.workspace = true
 reth-chainspec = { workspace = true, optional = true }
 reth-consensus.workspace = true
+reth-db.workspace = true
 reth-engine-primitives.workspace = true
 reth-errors.workspace = true
 reth-evm.workspace = true
@@ -22,8 +23,8 @@ reth-network-p2p.workspace = true
 reth-payload-builder-primitives.workspace = true
 reth-payload-builder.workspace = true
 reth-payload-primitives.workspace = true
-reth-primitives.workspace = true
 reth-primitives-traits.workspace = true
+reth-primitives.workspace = true
 reth-provider.workspace = true
 reth-prune.workspace = true
 reth-revm.workspace = true
@@ -71,7 +72,6 @@ reth-tracing = { workspace = true, optional = true }
 # reth
 reth-chain-state = { workspace = true, features = ["test-utils"] }
 reth-chainspec.workspace = true
-reth-db = { workspace = true, features = ["test-utils"] }
 reth-ethereum-engine-primitives.workspace = true
 reth-ethereum-consensus.workspace = true
 reth-evm = { workspace = true, features = ["test-utils"] }

crates/engine/tree/src/tree/mod.rs

@@ -2383,7 +2383,13 @@ where
                                     state_provider.state_root_with_updates(hashed_state.clone())?;
 
                                 if regular_root == block.header().state_root() {
-                                    compare_trie_updates(&task_trie_updates, &regular_updates);
+                                    let provider = self.provider.database_provider_ro()?;
+                                    compare_trie_updates(
+                                        provider.tx_ref(),
+                                        task_trie_updates.clone(),
+                                        regular_updates,
+                                    )
+                                    .map_err(ProviderError::from)?;
                                 } else {
                                     debug!(target: "engine::tree", "Regular state root does not match block state root");
                                 }

crates/engine/tree/src/tree/trie_updates.rs

@@ -2,16 +2,26 @@ use alloy_primitives::{
     map::{HashMap, HashSet},
     B256,
 };
+use reth_db::{transaction::DbTx, DatabaseError};
 use reth_trie::{
+    trie_cursor::{TrieCursor, TrieCursorFactory},
     updates::{StorageTrieUpdates, TrieUpdates},
     BranchNodeCompact, Nibbles,
 };
+use reth_trie_db::DatabaseTrieCursorFactory;
 use tracing::debug;
 
+#[derive(Debug)]
+struct EntryDiff<T> {
+    task: T,
+    regular: T,
+    database: T,
+}
+
 #[derive(Debug, Default)]
 struct TrieUpdatesDiff {
-    account_nodes: HashMap<Nibbles, (Option<BranchNodeCompact>, Option<BranchNodeCompact>)>,
+    account_nodes: HashMap<Nibbles, EntryDiff<Option<BranchNodeCompact>>>,
-    removed_nodes: HashMap<Nibbles, (bool, bool)>,
+    removed_nodes: HashMap<Nibbles, EntryDiff<bool>>,
     storage_tries: HashMap<B256, StorageTrieDiffEntry>,
 }
 
@@ -24,12 +34,12 @@ impl TrieUpdatesDiff {
 
     pub(super) fn log_differences(mut self) {
         if self.has_differences() {
-            for (path, (task, regular)) in &mut self.account_nodes {
+            for (path, EntryDiff { task, regular, database }) in &mut self.account_nodes {
-                debug!(target: "engine::tree", ?path, ?task, ?regular, "Difference in account trie updates");
+                debug!(target: "engine::tree", ?path, ?task, ?regular, ?database, "Difference in account trie updates");
             }
 
-            for (path, (task, regular)) in &self.removed_nodes {
+            for (path, EntryDiff { task, regular, database }) in &self.removed_nodes {
-                debug!(target: "engine::tree", ?path, ?task, ?regular, "Difference in removed account trie nodes");
+                debug!(target: "engine::tree", ?path, ?task, ?regular, ?database, "Difference in removed account trie nodes");
             }
 
             for (address, storage_diff) in self.storage_tries {
@@ -55,16 +65,17 @@ impl StorageTrieDiffEntry {
                 debug!(target: "engine::tree", ?address, ?task, ?regular, "Difference in storage trie existence");
             }
             Self::Value(mut storage_diff) => {
-                if let Some((task, regular)) = storage_diff.is_deleted {
+                if let Some(EntryDiff { task, regular, database }) = storage_diff.is_deleted {
-                    debug!(target: "engine::tree", ?address, ?task, ?regular, "Difference in storage trie deletion");
+                    debug!(target: "engine::tree", ?address, ?task, ?regular, ?database, "Difference in storage trie deletion");
                 }
 
-                for (path, (task, regular)) in &mut storage_diff.storage_nodes {
+                for (path, EntryDiff { task, regular, database }) in &mut storage_diff.storage_nodes
-                    debug!(target: "engine::tree", ?address, ?path, ?task, ?regular, "Difference in storage trie updates");
+                {
+                    debug!(target: "engine::tree", ?address, ?path, ?task, ?regular, ?database, "Difference in storage trie updates");
                 }
 
-                for (path, (task, regular)) in &storage_diff.removed_nodes {
+                for (path, EntryDiff { task, regular, database }) in &storage_diff.removed_nodes {
-                    debug!(target: "engine::tree", ?address, ?path, ?task, ?regular, "Difference in removed account trie nodes");
+                    debug!(target: "engine::tree", ?address, ?path, ?task, ?regular, ?database, "Difference in removed account trie nodes");
                 }
             }
         }
@@ -73,9 +84,9 @@ impl StorageTrieDiffEntry {
 
 #[derive(Debug, Default)]
 struct StorageTrieUpdatesDiff {
-    is_deleted: Option<(bool, bool)>,
+    is_deleted: Option<EntryDiff<bool>>,
-    storage_nodes: HashMap<Nibbles, (Option<BranchNodeCompact>, Option<BranchNodeCompact>)>,
+    storage_nodes: HashMap<Nibbles, EntryDiff<Option<BranchNodeCompact>>>,
-    removed_nodes: HashMap<Nibbles, (bool, bool)>,
+    removed_nodes: HashMap<Nibbles, EntryDiff<bool>>,
 }
 
 impl StorageTrieUpdatesDiff {
@@ -88,10 +99,20 @@ impl StorageTrieUpdatesDiff {
 
 /// Compares the trie updates from state root task and regular state root calculation, and logs
 /// the differences if there's any.
-pub(super) fn compare_trie_updates(task: &TrieUpdates, regular: &TrieUpdates) {
+pub(super) fn compare_trie_updates(
+    tx: &impl DbTx,
+    task: TrieUpdates,
+    regular: TrieUpdates,
+) -> Result<(), DatabaseError> {
+    let trie_cursor_factory = DatabaseTrieCursorFactory::new(tx);
+
+    let mut task = adjust_trie_updates(task);
+    let mut regular = adjust_trie_updates(regular);
+
     let mut diff = TrieUpdatesDiff::default();
 
     // compare account nodes
+    let mut account_trie_cursor = trie_cursor_factory.account_trie_cursor()?;
     for key in task
         .account_nodes
         .keys()
@@ -99,10 +120,11 @@ pub(super) fn compare_trie_updates(task: &TrieUpdates, regular: &TrieUpdates) {
         .cloned()
         .collect::<HashSet<_>>()
     {
-        let (left, right) = (task.account_nodes.get(&key), regular.account_nodes.get(&key));
+        let (task, regular) = (task.account_nodes.remove(&key), regular.account_nodes.remove(&key));
+        let database = account_trie_cursor.seek_exact(key.clone())?.map(|x| x.1);
 
-        if !branch_nodes_equal(left, right) {
+        if !branch_nodes_equal(task.as_ref(), regular.as_ref(), database.as_ref())? {
-            diff.account_nodes.insert(key, (left.cloned(), right.cloned()));
+            diff.account_nodes.insert(key, EntryDiff { task, regular, database });
         }
     }
 
@@ -114,10 +136,11 @@ pub(super) fn compare_trie_updates(task: &TrieUpdates, regular: &TrieUpdates) {
         .cloned()
         .collect::<HashSet<_>>()
     {
-        let (left, right) =
+        let (task, regular) =
             (task.removed_nodes.contains(&key), regular.removed_nodes.contains(&key));
-        if left != right {
+        let database = account_trie_cursor.seek_exact(key.clone())?.is_none();
-            diff.removed_nodes.insert(key, (left, right));
+        if task != regular {
+            diff.removed_nodes.insert(key, EntryDiff { task, regular, database });
         }
     }
 
@@ -129,31 +152,43 @@ pub(super) fn compare_trie_updates(task: &TrieUpdates, regular: &TrieUpdates) {
         .copied()
         .collect::<HashSet<_>>()
     {
-        let (left, right) = (task.storage_tries.get(&key), regular.storage_tries.get(&key));
+        let (mut task, mut regular) =
-        if left != right {
+            (task.storage_tries.remove(&key), regular.storage_tries.remove(&key));
-            if let Some((left, right)) = left.zip(right) {
+        if task != regular {
-                let storage_diff = compare_storage_trie_updates(left, right);
+            let mut storage_trie_cursor = trie_cursor_factory.storage_trie_cursor(key)?;
+            if let Some((task, regular)) = task.as_mut().zip(regular.as_mut()) {
+                let storage_diff =
+                    compare_storage_trie_updates(&mut storage_trie_cursor, task, regular)?;
                 if storage_diff.has_differences() {
                     diff.storage_tries.insert(key, StorageTrieDiffEntry::Value(storage_diff));
                 }
             } else {
-                diff.storage_tries
+                diff.storage_tries.insert(
-                    .insert(key, StorageTrieDiffEntry::Existence(left.is_some(), right.is_some()));
+                    key,
+                    StorageTrieDiffEntry::Existence(task.is_some(), regular.is_some()),
+                );
             }
         }
     }
 
     // log differences
     diff.log_differences();
+
+    Ok(())
 }
 
 fn compare_storage_trie_updates(
-    task: &StorageTrieUpdates,
+    trie_cursor: &mut impl TrieCursor,
-    regular: &StorageTrieUpdates,
+    task: &mut StorageTrieUpdates,
-) -> StorageTrieUpdatesDiff {
+    regular: &mut StorageTrieUpdates,
+) -> Result<StorageTrieUpdatesDiff, DatabaseError> {
+    let database_deleted = trie_cursor.next()?.is_none();
     let mut diff = StorageTrieUpdatesDiff {
-        is_deleted: (task.is_deleted != regular.is_deleted)
+        is_deleted: (task.is_deleted != regular.is_deleted).then_some(EntryDiff {
-            .then_some((task.is_deleted, regular.is_deleted)),
+            task: task.is_deleted,
+            regular: regular.is_deleted,
+            database: database_deleted,
+        }),
         ..Default::default()
     };
 
@@ -165,9 +200,10 @@ fn compare_storage_trie_updates(
         .cloned()
         .collect::<HashSet<_>>()
     {
-        let (left, right) = (task.storage_nodes.get(&key), regular.storage_nodes.get(&key));
+        let (task, regular) = (task.storage_nodes.remove(&key), regular.storage_nodes.remove(&key));
-        if !branch_nodes_equal(left, right) {
+        let database = trie_cursor.seek_exact(key.clone())?.map(|x| x.1);
-            diff.storage_nodes.insert(key, (left.cloned(), right.cloned()));
+        if !branch_nodes_equal(task.as_ref(), regular.as_ref(), database.as_ref())? {
+            diff.storage_nodes.insert(key, EntryDiff { task, regular, database });
         }
     }
 
@@ -179,30 +215,73 @@ fn compare_storage_trie_updates(
         .cloned()
         .collect::<HashSet<_>>()
     {
-        let (left, right) =
+        let (task, regular) =
             (task.removed_nodes.contains(&key), regular.removed_nodes.contains(&key));
-        if left != right {
+        let database = trie_cursor.seek_exact(key.clone())?.map(|x| x.1).is_none();
-            diff.removed_nodes.insert(key, (left, right));
+        if task != regular {
+            diff.removed_nodes.insert(key, EntryDiff { task, regular, database });
         }
     }
 
-    diff
+    Ok(diff)
+}
+
+/// Filters the removed nodes of both account trie updates and storage trie updates, so that they
+/// don't include those nodes that were also updated.
+fn adjust_trie_updates(trie_updates: TrieUpdates) -> TrieUpdates {
+    TrieUpdates {
+        removed_nodes: trie_updates
+            .removed_nodes
+            .into_iter()
+            .filter(|key| !trie_updates.account_nodes.contains_key(key))
+            .collect(),
+        storage_tries: trie_updates
+            .storage_tries
+            .into_iter()
+            .map(|(address, updates)| {
+                (
+                    address,
+                    StorageTrieUpdates {
+                        removed_nodes: updates
+                            .removed_nodes
+                            .into_iter()
+                            .filter(|key| !updates.storage_nodes.contains_key(key))
+                            .collect(),
+                        ..updates
+                    },
+                )
+            })
+            .collect(),
+        ..trie_updates
+    }
 }
 
 /// Compares the branch nodes from state root task and regular state root calculation.
 ///
+/// If one of the branch nodes is [`None`], it means it's not updated and the other is compared to
+/// the branch node from the database.
+///
 /// Returns `true` if they are equal.
 fn branch_nodes_equal(
     task: Option<&BranchNodeCompact>,
     regular: Option<&BranchNodeCompact>,
-) -> bool {
+    database: Option<&BranchNodeCompact>,
-    if let (Some(task), Some(regular)) = (task.as_ref(), regular.as_ref()) {
+) -> Result<bool, DatabaseError> {
-        task.state_mask == regular.state_mask &&
+    Ok(match (task, regular) {
-            task.tree_mask == regular.tree_mask &&
+        (Some(task), Some(regular)) => {
-            task.hash_mask == regular.hash_mask &&
+            task.state_mask == regular.state_mask &&
-            task.hashes == regular.hashes &&
+                task.tree_mask == regular.tree_mask &&
-            task.root_hash == regular.root_hash
+                task.hash_mask == regular.hash_mask &&
-    } else {
+                task.hashes == regular.hashes &&
-        task == regular
+                task.root_hash == regular.root_hash
-    }
+        }
+        (None, None) => true,
+        _ => {
+            if task.is_some() {
+                task == database
+            } else {
+                regular == database
+            }
+        }
+    })
 }