chore: add root_with_updates method to sparse trie (#14393)

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))
 }
 

crates/trie/sparse/src/state.rs

@@ -482,17 +482,24 @@ impl<F: BlindedProviderFactory> SparseStateTrie<F> {
         self.state.root()
     }
 
-    /// 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> {
-        self.state.as_revealed_mut().map(|state| {
-            let updates = state.take_updates();
-            TrieUpdates {
+    /// 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: self
-                    .storages
+                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();
@@ -505,7 +512,20 @@ impl<F: BlindedProviderFactory> SparseStateTrie<F> {
                 (*address, updates)
             })
             .filter(|(_, updates)| !updates.is_empty())
-                    .collect(),
+            .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,
             }
         })
     }

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);