do not include updated nodes in removed

paradigmxyz · Jan 10, 2025 · 8d4e441 · 8d4e441
1 parent 9da36e7
commit 8d4e441
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 4 deletions.
diff --git 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
@@ -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