From 190a1d8bb4df777cf19a90ce072e9c4435a6bcfb Mon Sep 17 00:00:00 2001 From: Alexey Shekhirin Date: Fri, 8 Nov 2024 00:21:53 +0900 Subject: [PATCH] feat(trie): reveal storage slots and calculate storage root in sparse trie (#12145) --- crates/trie/sparse/src/state.rs | 129 ++++++++++++++++++++++++-------- 1 file changed, 96 insertions(+), 33 deletions(-) diff --git a/crates/trie/sparse/src/state.rs b/crates/trie/sparse/src/state.rs index cfb17ef36..126e05e85 100644 --- a/crates/trie/sparse/src/state.rs +++ b/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, /// Collection of revealed account and storage keys. - #[allow(dead_code)] pub(crate) revealed: HashMap>, } @@ -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, + ) -> 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>( + &self, + proof: &mut Peekable, + ) -> SparseStateTrieResult> { + 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) -> SparseStateTrieResult<()> { self.state.update_leaf(path, value)?; @@ -81,6 +120,11 @@ impl SparseStateTrie { pub fn root(&mut self) -> Option { self.state.root() } + + /// Returns storage sparse trie root if the trie has been revealed. + pub fn storage_root(&mut self, account: B256) -> Option { + 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())]) ); } }