From 8d4e441cb12bc56395aab996f7adaaa097f2d23c Mon Sep 17 00:00:00 2001
From: Alexey Shekhirin <5773434+shekhirin@users.noreply.github.com>
Date: Fri, 10 Jan 2025 13:02:30 +0000
Subject: [PATCH] do not include updated nodes in removed

---
 crates/engine/tree/src/tree/mod.rs          |  4 +--
 crates/engine/tree/src/tree/trie_updates.rs | 37 +++++++++++++++++++--
 2 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/crates/engine/tree/src/tree/mod.rs b/crates/engine/tree/src/tree/mod.rs
index f211176b32eb..bb1ab030879c 100644
--- a/crates/engine/tree/src/tree/mod.rs
+++ b/crates/engine/tree/src/tree/mod.rs
@@ -2381,8 +2381,8 @@ where
                                     let provider = self.provider.database_provider_ro()?;
                                     compare_trie_updates(
                                         provider.tx_ref(),
-                                        &task_trie_updates,
-                                        &regular_updates,
+                                        task_trie_updates.clone(),
+                                        regular_updates,
                                     )
                                     .map_err(ProviderError::from)?;
                                 } else {
diff --git a/crates/engine/tree/src/tree/trie_updates.rs b/crates/engine/tree/src/tree/trie_updates.rs
index 53c9af2d61a3..e5cea6c4614a 100644
--- a/crates/engine/tree/src/tree/trie_updates.rs
+++ b/crates/engine/tree/src/tree/trie_updates.rs
@@ -93,11 +93,14 @@ impl StorageTrieUpdatesDiff {
 /// the differences if there's any.
 pub(super) fn compare_trie_updates(
     tx: &impl DbTx,
-    task: &TrieUpdates,
-    regular: &TrieUpdates,
+    task: TrieUpdates,
+    regular: TrieUpdates,
 ) -> Result<(), DatabaseError> {
     let trie_cursor_factory = DatabaseTrieCursorFactory::new(tx);
 
+    let task = adjust_trie_updates(task);
+    let regular = adjust_trie_updates(regular);
+
     let mut diff = TrieUpdatesDiff::default();
 
     // compare account nodes
@@ -204,6 +207,36 @@ fn compare_storage_trie_updates(
     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