chore: add root_with_updates method to sparse trie

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

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