feat(trie): reveal storage slots and calculate storage root in sparse trie

crates/trie/sparse/src/state.rs

@@ -1,3 +1,5 @@
+use std::iter::Peekable;
+
 use crate::{SparseStateTrieError, SparseStateTrieResult, SparseTrie};
 use alloy_primitives::{
     map::{HashMap, HashSet},
@@ -12,10 +14,8 @@ pub struct SparseStateTrie {
     /// Sparse account trie.
     pub(crate) state: SparseTrie,
     /// Sparse storage tries.
-    #[allow(dead_code)]
     pub(crate) storages: HashMap<B256, SparseTrie>,
     /// Collection of revealed account and storage keys.
-    #[allow(dead_code)]
     pub(crate) revealed: HashMap<B256, HashSet<B256>>,
 }
 
@@ -35,7 +35,7 @@ impl SparseStateTrie {
         self.revealed.get(account).map_or(false, |slots| slots.contains(slot))
     }
 
-    /// Reveal unknown trie paths from provided leaf path and its proof.
+    /// Reveal unknown trie paths from provided leaf path and its proof for the account.
     /// NOTE: This method does not extensively validate the proof.
     pub fn reveal_account(
         &mut self,
@@ -44,22 +44,12 @@ impl SparseStateTrie {
     ) -> SparseStateTrieResult<()> {
         let mut proof = proof.into_iter().peekable();
 
-        // reveal root and initialize the trie if not already
-        let Some((path, node)) = proof.next() else { return Ok(()) };
-        if !path.is_empty() {
-            return Err(SparseStateTrieError::InvalidRootNode { path, node })
-        }
-
-        // Decode root node and perform sanity check.
-        let root_node = TrieNode::decode(&mut &node[..])?;
-        if matches!(root_node, TrieNode::EmptyRoot) && proof.peek().is_some() {
-            return Err(SparseStateTrieError::InvalidRootNode { path, node })
-        }
+        let Some(root_node) = self.validate_proof(&mut proof)? else { return Ok(()) };
 
         // Reveal root node if it wasn't already.
         let trie = self.state.reveal_root(root_node)?;
 
-        // add the remaining proof nodes
+        // Reveal the remaining proof nodes.
         for (path, bytes) in proof {
             let node = TrieNode::decode(&mut &bytes[..])?;
             trie.reveal_node(path, node)?;
@@ -71,6 +61,55 @@ impl SparseStateTrie {
         Ok(())
     }
 
+    /// Reveal unknown trie paths from provided leaf path and its proof for the storage slot.
+    /// NOTE: This method does not extensively validate the proof.
+    pub fn reveal_storage_slot(
+        &mut self,
+        account: B256,
+        slot: B256,
+        proof: impl IntoIterator<Item = (Nibbles, Bytes)>,
+    ) -> SparseStateTrieResult<()> {
+        let mut proof = proof.into_iter().peekable();
+
+        let Some(root_node) = self.validate_proof(&mut proof)? else { return Ok(()) };
+
+        // Reveal root node if it wasn't already.
+        let trie = self.storages.entry(account).or_default().reveal_root(root_node)?;
+
+        // Reveal the remaining proof nodes.
+        for (path, bytes) in proof {
+            let node = TrieNode::decode(&mut &bytes[..])?;
+            trie.reveal_node(path, node)?;
+        }
+
+        // Mark leaf path as revealed.
+        self.revealed.entry(account).or_default().insert(slot);
+
+        Ok(())
+    }
+
+    /// Validates the root node of the proof and returns it if it exists and is valid.
+    fn validate_proof<I: Iterator<Item = (Nibbles, Bytes)>>(
+        &self,
+        proof: &mut Peekable<I>,
+    ) -> SparseStateTrieResult<Option<TrieNode>> {
+        let mut proof = proof.into_iter().peekable();
+
+        // Validate root node.
+        let Some((path, node)) = proof.next() else { return Ok(None) };
+        if !path.is_empty() {
+            return Err(SparseStateTrieError::InvalidRootNode { path, node })
+        }
+
+        // Decode root node and perform sanity check.
+        let root_node = TrieNode::decode(&mut &node[..])?;
+        if matches!(root_node, TrieNode::EmptyRoot) && proof.peek().is_some() {
+            return Err(SparseStateTrieError::InvalidRootNode { path, node })
+        }
+
+        Ok(Some(root_node))
+    }
+
     /// Update the leaf node.
     pub fn update_leaf(&mut self, path: Nibbles, value: Vec<u8>) -> SparseStateTrieResult<()> {
         self.state.update_leaf(path, value)?;
@@ -81,6 +120,11 @@ impl SparseStateTrie {
     pub fn root(&mut self) -> Option<B256> {
         self.state.root()
     }
+
+    /// Returns storage sparse trie root if the trie has been revealed.
+    pub fn storage_root(&mut self, account: B256) -> Option<B256> {
+        self.storages.get_mut(&account).and_then(|trie| trie.root())
+    }
 }
 
 #[cfg(test)]
@@ -93,7 +137,30 @@ mod tests {
     use reth_trie_common::proof::ProofRetainer;
 
     #[test]
-    fn sparse_trie_reveal_empty() {
+    fn validate_proof_first_node_not_root() {
+        let sparse = SparseStateTrie::default();
+        let proof = [(Nibbles::from_nibbles([0x1]), Bytes::from([EMPTY_STRING_CODE]))];
+        assert_matches!(
+            sparse.validate_proof(&mut proof.into_iter().peekable()),
+            Err(SparseStateTrieError::InvalidRootNode { .. })
+        );
+    }
+
+    #[test]
+    fn validate_proof_invalid_proof_with_empty_root() {
+        let sparse = SparseStateTrie::default();
+        let proof = [
+            (Nibbles::default(), Bytes::from([EMPTY_STRING_CODE])),
+            (Nibbles::from_nibbles([0x1]), Bytes::new()),
+        ];
+        assert_matches!(
+            sparse.validate_proof(&mut proof.into_iter().peekable()),
+            Err(SparseStateTrieError::InvalidRootNode { .. })
+        );
+    }
+
+    #[test]
+    fn reveal_account_empty() {
         let retainer = ProofRetainer::from_iter([Nibbles::default()]);
         let mut hash_builder = HashBuilder::default().with_proof_retainer(retainer);
         hash_builder.root();
@@ -107,25 +174,21 @@ mod tests {
     }
 
     #[test]
-    fn reveal_first_node_not_root() {
-        let mut sparse = SparseStateTrie::default();
-        let proof = [(Nibbles::from_nibbles([0x1]), Bytes::from([EMPTY_STRING_CODE]))];
-        assert_matches!(
-            sparse.reveal_account(Default::default(), proof),
-            Err(SparseStateTrieError::InvalidRootNode { .. })
-        );
-    }
+    fn reveal_storage_slot_empty() {
+        let retainer = ProofRetainer::from_iter([Nibbles::default()]);
+        let mut hash_builder = HashBuilder::default().with_proof_retainer(retainer);
+        hash_builder.root();
+        let proofs = hash_builder.take_proof_nodes();
+        assert_eq!(proofs.len(), 1);
 
-    #[test]
-    fn reveal_invalid_proof_with_empty_root() {
         let mut sparse = SparseStateTrie::default();
-        let proof = [
-            (Nibbles::default(), Bytes::from([EMPTY_STRING_CODE])),
-            (Nibbles::from_nibbles([0x1]), Bytes::new()),
-        ];
-        assert_matches!(
-            sparse.reveal_account(Default::default(), proof),
-            Err(SparseStateTrieError::InvalidRootNode { .. })
+        assert!(sparse.storages.is_empty());
+        sparse
+            .reveal_storage_slot(Default::default(), Default::default(), proofs.into_inner())
+            .unwrap();
+        assert_eq!(
+            sparse.storages,
+            HashMap::from_iter([(Default::default(), SparseTrie::revealed_empty())])
         );
     }
 }