From 1d9687af71ecd4e3a7e7b53a5093205267e1486f Mon Sep 17 00:00:00 2001
From: Dan Cline <6798349+Rjected@users.noreply.github.com>
Date: Mon, 10 Feb 2025 21:02:54 -0500
Subject: [PATCH] chore: add root_with_updates method to sparse trie

---
 crates/engine/tree/src/tree/root.rs |  3 +-
 crates/trie/sparse/src/state.rs     | 50 ++++++++++++++++++++---------
 crates/trie/sparse/src/trie.rs      |  6 ++++
 3 files changed, 42 insertions(+), 17 deletions(-)

diff --git a/crates/engine/tree/src/tree/root.rs b/crates/engine/tree/src/tree/root.rs
index 274e01cdfc0d..9522552c82a2 100644
--- a/crates/engine/tree/src/tree/root.rs
+++ b/crates/engine/tree/src/tree/root.rs
@@ -1011,11 +1011,10 @@ where
     debug!(target: "engine::root", num_iterations, "All proofs processed, ending calculation");
 
     let start = Instant::now();
-    let root = trie.root().expect("sparse trie should be revealed");
+    let (root, trie_updates) = trie.root_with_updates().expect("sparse trie should be revealed");
     let elapsed = start.elapsed();
     metrics.sparse_trie_final_update_duration_histogram.record(elapsed);
 
-    let trie_updates = trie.take_trie_updates().expect("retention must be enabled");
     Ok((root, trie_updates, num_iterations))
 }
 
diff --git a/crates/trie/sparse/src/state.rs b/crates/trie/sparse/src/state.rs
index f7d21741e1c4..cec4c04bb386 100644
--- a/crates/trie/sparse/src/state.rs
+++ b/crates/trie/sparse/src/state.rs
@@ -482,30 +482,50 @@ impl<F: BlindedProviderFactory> SparseStateTrie<F> {
         self.state.root()
     }
 
+    /// Returns sparse trie root and trie updates if the trie has been revealed.
+    pub fn root_with_updates(&mut self) -> Option<(B256, TrieUpdates)> {
+        let storage_tries = self.storage_trie_updates();
+        self.state.root_with_updates().map(|(root, updates)| {
+            let updates = TrieUpdates {
+                account_nodes: updates.updated_nodes,
+                removed_nodes: updates.removed_nodes,
+                storage_tries,
+            };
+            (root, updates)
+        })
+    }
+
+    /// Returns storage trie updates for tries that have been revealed.
+    ///
+    /// Panics if any of the storage tries are not revealed.
+    pub fn storage_trie_updates(&mut self) -> B256HashMap<StorageTrieUpdates> {
+        self.storages
+            .iter_mut()
+            .map(|(address, trie)| {
+                let trie = trie.as_revealed_mut().unwrap();
+                let updates = trie.take_updates();
+                let updates = StorageTrieUpdates {
+                    is_deleted: updates.wiped,
+                    storage_nodes: updates.updated_nodes,
+                    removed_nodes: updates.removed_nodes,
+                };
+                (*address, updates)
+            })
+            .filter(|(_, updates)| !updates.is_empty())
+            .collect()
+    }
+
     /// Returns [`TrieUpdates`] by taking the updates from the revealed sparse tries.
     ///
     /// Returns `None` if the accounts trie is not revealed.
     pub fn take_trie_updates(&mut self) -> Option<TrieUpdates> {
+        let storage_tries = self.storage_trie_updates();
         self.state.as_revealed_mut().map(|state| {
             let updates = state.take_updates();
             TrieUpdates {
                 account_nodes: updates.updated_nodes,
                 removed_nodes: updates.removed_nodes,
-                storage_tries: self
-                    .storages
-                    .iter_mut()
-                    .map(|(address, trie)| {
-                        let trie = trie.as_revealed_mut().unwrap();
-                        let updates = trie.take_updates();
-                        let updates = StorageTrieUpdates {
-                            is_deleted: updates.wiped,
-                            storage_nodes: updates.updated_nodes,
-                            removed_nodes: updates.removed_nodes,
-                        };
-                        (*address, updates)
-                    })
-                    .filter(|(_, updates)| !updates.is_empty())
-                    .collect(),
+                storage_tries,
             }
         })
     }
diff --git a/crates/trie/sparse/src/trie.rs b/crates/trie/sparse/src/trie.rs
index 162131f58a79..948eebb61359 100644
--- a/crates/trie/sparse/src/trie.rs
+++ b/crates/trie/sparse/src/trie.rs
@@ -134,6 +134,12 @@ impl<P> SparseTrie<P> {
         Some(self.as_revealed_mut()?.root())
     }
 
+    /// Returns both the trie root and takes sparse trie updates if the trie has been revealed.
+    pub fn root_with_updates(&mut self) -> Option<(B256, SparseTrieUpdates)> {
+        let revealed = self.as_revealed_mut()?;
+        Some((revealed.root(), revealed.take_updates()))
+    }
+
     /// Calculates the hashes of the nodes below the provided level.
     pub fn calculate_below_level(&mut self, level: usize) {
         self.as_revealed_mut().unwrap().update_rlp_node_level(level);