From 1948e0f79cd7af50f484d7a6972e8232cd504bbd Mon Sep 17 00:00:00 2001
From: Roman Krasiuk <rokrassyuk@gmail.com>
Date: Thu, 16 Jan 2025 10:30:53 +0100
Subject: [PATCH] feat(trie): sparse trie accessors (#13815)

---
 crates/trie/sparse/src/state.rs | 13 +++++++++++++
 crates/trie/sparse/src/trie.rs  |  8 ++++++--
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/crates/trie/sparse/src/state.rs b/crates/trie/sparse/src/state.rs
index 4b47217f9..2def53e57 100644
--- a/crates/trie/sparse/src/state.rs
+++ b/crates/trie/sparse/src/state.rs
@@ -105,6 +105,19 @@ impl<F: BlindedProviderFactory> SparseStateTrie<F> {
         self.storages.get(account)?.as_revealed_ref()?.get_leaf_value(&Nibbles::unpack(slot))
     }
 
+    /// Returns reference to state trie if it was revealed.
+    pub const fn state_trie_ref(&self) -> Option<&RevealedSparseTrie<F::AccountNodeProvider>> {
+        self.state.as_revealed_ref()
+    }
+
+    /// Returns reference to storage trie if it was revealed.
+    pub fn storage_trie_ref(
+        &self,
+        address: &B256,
+    ) -> Option<&RevealedSparseTrie<F::StorageNodeProvider>> {
+        self.storages.get(address).and_then(|e| e.as_revealed_ref())
+    }
+
     /// Returns mutable reference to storage sparse trie if it was revealed.
     pub fn storage_trie_mut(
         &mut self,
diff --git a/crates/trie/sparse/src/trie.rs b/crates/trie/sparse/src/trie.rs
index b5bed53b5..bb1f247a2 100644
--- a/crates/trie/sparse/src/trie.rs
+++ b/crates/trie/sparse/src/trie.rs
@@ -288,6 +288,11 @@ impl<P> RevealedSparseTrie<P> {
         self.updates.as_ref().map_or(Cow::Owned(SparseTrieUpdates::default()), Cow::Borrowed)
     }
 
+    /// Returns reference to all trie nodes.
+    pub const fn nodes_ref(&self) -> &HashMap<Nibbles, SparseNode> {
+        &self.nodes
+    }
+
     /// Returns a reference to the leaf value if present.
     pub fn get_leaf_value(&self, path: &Nibbles) -> Option<&Vec<u8>> {
         self.values.get(path)
@@ -1130,8 +1135,7 @@ impl<P: BlindedProvider> RevealedSparseTrie<P> {
 
                         new_node
                     }
-                    // If more than one child is left set in the branch, we just re-insert it
-                    // as-is.
+                    // If more than one child is left set in the branch, we just re-insert it as-is.
                     else {
                         SparseNode::new_branch(state_mask)
                     }