From cb11ab04751cebea8c39cb8b44ee0559695c9e44 Mon Sep 17 00:00:00 2001
From: rotcan <50956594+rotcan@users.noreply.github.com>
Date: Wed, 18 Jun 2025 02:01:08 +0530
Subject: [PATCH] feat(engine): Compare sorted trie updates in witness invalid
 block hook#15689 (#16481)

---
 .../engine/invalid-block-hooks/src/witness.rs |  6 +-
 crates/trie/common/src/updates.rs             | 58 ++++++++++++++++++-
 2 files changed, 60 insertions(+), 4 deletions(-)

diff --git a/crates/engine/invalid-block-hooks/src/witness.rs b/crates/engine/invalid-block-hooks/src/witness.rs
index 7ddf593d33..b3b5428112 100644
--- a/crates/engine/invalid-block-hooks/src/witness.rs
+++ b/crates/engine/invalid-block-hooks/src/witness.rs
@@ -317,13 +317,15 @@ where
 
             if &trie_output != original_updates {
                 // Trie updates are too big to diff, so we just save the original and re-executed
+                let trie_output_sorted = &trie_output.into_sorted_ref();
+                let original_updates_sorted = &original_updates.into_sorted_ref();
                 let original_path = self.save_file(
                     format!("{}_{}.trie_updates.original.json", block.number(), block.hash()),
-                    original_updates,
+                    original_updates_sorted,
                 )?;
                 let re_executed_path = self.save_file(
                     format!("{}_{}.trie_updates.re_executed.json", block.number(), block.hash()),
-                    &trie_output,
+                    trie_output_sorted,
                 )?;
                 warn!(
                     target: "engine::invalid_block_hooks::witness",
diff --git a/crates/trie/common/src/updates.rs b/crates/trie/common/src/updates.rs
index d4362542f0..477e35b2c7 100644
--- a/crates/trie/common/src/updates.rs
+++ b/crates/trie/common/src/updates.rs
@@ -1,8 +1,11 @@
 use crate::{BranchNodeCompact, HashBuilder, Nibbles};
-use alloc::vec::Vec;
+use alloc::{
+    collections::{btree_map::BTreeMap, btree_set::BTreeSet},
+    vec::Vec,
+};
 use alloy_primitives::{
     map::{B256Map, B256Set, HashMap, HashSet},
-    B256,
+    FixedBytes, B256,
 };
 
 /// The aggregation of trie updates.
@@ -114,6 +117,22 @@ impl TrieUpdates {
             .collect();
         TrieUpdatesSorted { removed_nodes: self.removed_nodes, account_nodes, storage_tries }
     }
+
+    /// Converts trie updates into [`TrieUpdatesSortedRef`].
+    pub fn into_sorted_ref<'a>(&'a self) -> TrieUpdatesSortedRef<'a> {
+        let mut account_nodes = self.account_nodes.iter().collect::<Vec<_>>();
+        account_nodes.sort_unstable_by(|a, b| a.0.cmp(b.0));
+
+        TrieUpdatesSortedRef {
+            removed_nodes: self.removed_nodes.iter().collect::<BTreeSet<_>>(),
+            account_nodes,
+            storage_tries: self
+                .storage_tries
+                .iter()
+                .map(|m| (*m.0, m.1.into_sorted_ref().clone()))
+                .collect(),
+        }
+    }
 }
 
 /// Trie updates for storage trie of a single account.
@@ -225,6 +244,15 @@ impl StorageTrieUpdates {
             storage_nodes,
         }
     }
+
+    /// Convert storage trie updates into [`StorageTrieUpdatesSortedRef`].
+    pub fn into_sorted_ref(&self) -> StorageTrieUpdatesSortedRef<'_> {
+        StorageTrieUpdatesSortedRef {
+            is_deleted: self.is_deleted,
+            removed_nodes: self.removed_nodes.iter().collect::<BTreeSet<_>>(),
+            storage_nodes: self.storage_nodes.iter().collect::<BTreeMap<_, _>>(),
+        }
+    }
 }
 
 /// Serializes and deserializes any [`HashSet`] that includes [`Nibbles`] elements, by using the
@@ -350,8 +378,21 @@ mod serde_nibbles_map {
     }
 }
 
+/// Sorted trie updates reference used for serializing trie to file.
+#[derive(PartialEq, Eq, Clone, Default, Debug)]
+#[cfg_attr(any(test, feature = "serde"), derive(serde::Serialize))]
+pub struct TrieUpdatesSortedRef<'a> {
+    /// Sorted collection of updated state nodes with corresponding paths.
+    pub account_nodes: Vec<(&'a Nibbles, &'a BranchNodeCompact)>,
+    /// The set of removed state node keys.
+    pub removed_nodes: BTreeSet<&'a Nibbles>,
+    /// Storage tries stored by hashed address of the account the trie belongs to.
+    pub storage_tries: BTreeMap<FixedBytes<32>, StorageTrieUpdatesSortedRef<'a>>,
+}
+
 /// Sorted trie updates used for lookups and insertions.
 #[derive(PartialEq, Eq, Clone, Default, Debug)]
+#[cfg_attr(any(test, feature = "serde"), derive(serde::Serialize, serde::Deserialize))]
 pub struct TrieUpdatesSorted {
     /// Sorted collection of updated state nodes with corresponding paths.
     pub account_nodes: Vec<(Nibbles, BranchNodeCompact)>,
@@ -378,8 +419,21 @@ impl TrieUpdatesSorted {
     }
 }
 
+/// Sorted storage trie updates reference used for serializing to file.
+#[derive(PartialEq, Eq, Clone, Default, Debug)]
+#[cfg_attr(any(test, feature = "serde"), derive(serde::Serialize))]
+pub struct StorageTrieUpdatesSortedRef<'a> {
+    /// Flag indicating whether the trie has been deleted/wiped.
+    pub is_deleted: bool,
+    /// Sorted collection of updated storage nodes with corresponding paths.
+    pub storage_nodes: BTreeMap<&'a Nibbles, &'a BranchNodeCompact>,
+    /// The set of removed storage node keys.
+    pub removed_nodes: BTreeSet<&'a Nibbles>,
+}
+
 /// Sorted trie updates used for lookups and insertions.
 #[derive(PartialEq, Eq, Clone, Default, Debug)]
+#[cfg_attr(any(test, feature = "serde"), derive(serde::Serialize, serde::Deserialize))]
 pub struct StorageTrieUpdatesSorted {
     /// Flag indicating whether the trie has been deleted/wiped.
     pub is_deleted: bool,