diff --git a/crates/trie/common/src/proofs.rs b/crates/trie/common/src/proofs.rs index 1f37e639d..9eff5c90e 100644 --- a/crates/trie/common/src/proofs.rs +++ b/crates/trie/common/src/proofs.rs @@ -144,6 +144,119 @@ impl MultiProof { } } +/// This is a type of [`MultiProof`] that uses decoded proofs, meaning these proofs are stored as a +/// collection of [`TrieNode`]s instead of RLP-encoded bytes. +#[derive(Clone, Default, Debug, PartialEq, Eq)] +pub struct DecodedMultiProof { + /// State trie multiproof for requested accounts. + pub account_subtree: DecodedProofNodes, + /// The hash masks of the branch nodes in the account proof. + pub branch_node_hash_masks: HashMap, + /// The tree masks of the branch nodes in the account proof. + pub branch_node_tree_masks: HashMap, + /// Storage trie multiproofs. + pub storages: B256Map, +} + +impl DecodedMultiProof { + /// Return the account proof nodes for the given account path. + pub fn account_proof_nodes(&self, path: &Nibbles) -> Vec<(Nibbles, TrieNode)> { + self.account_subtree.matching_nodes_sorted(path) + } + + /// Return the storage proof nodes for the given storage slots of the account path. + pub fn storage_proof_nodes( + &self, + hashed_address: B256, + slots: impl IntoIterator, + ) -> Vec<(B256, Vec<(Nibbles, TrieNode)>)> { + self.storages + .get(&hashed_address) + .map(|storage_mp| { + slots + .into_iter() + .map(|slot| { + let nibbles = Nibbles::unpack(slot); + (slot, storage_mp.subtree.matching_nodes_sorted(&nibbles)) + }) + .collect() + }) + .unwrap_or_default() + } + + /// Construct the account proof from the multiproof. + pub fn account_proof( + &self, + address: Address, + slots: &[B256], + ) -> Result { + let hashed_address = keccak256(address); + let nibbles = Nibbles::unpack(hashed_address); + + // Retrieve the account proof. + let proof = self + .account_proof_nodes(&nibbles) + .into_iter() + .map(|(_, node)| node) + .collect::>(); + + // Inspect the last node in the proof. If it's a leaf node with matching suffix, + // then the node contains the encoded trie account. + let info = 'info: { + if let Some(TrieNode::Leaf(leaf)) = proof.last() { + if nibbles.ends_with(&leaf.key) { + let account = TrieAccount::decode(&mut &leaf.value[..])?; + break 'info Some(Account { + balance: account.balance, + nonce: account.nonce, + bytecode_hash: (account.code_hash != KECCAK_EMPTY) + .then_some(account.code_hash), + }) + } + } + None + }; + + // Retrieve proofs for requested storage slots. + let storage_multiproof = self.storages.get(&hashed_address); + let storage_root = storage_multiproof.map(|m| m.root).unwrap_or(EMPTY_ROOT_HASH); + let mut storage_proofs = Vec::with_capacity(slots.len()); + for slot in slots { + let proof = if let Some(multiproof) = &storage_multiproof { + multiproof.storage_proof(*slot)? + } else { + DecodedStorageProof::new(*slot) + }; + storage_proofs.push(proof); + } + Ok(DecodedAccountProof { address, info, proof, storage_root, storage_proofs }) + } + + /// Extends this multiproof with another one, merging both account and storage + /// proofs. + pub fn extend(&mut self, other: Self) { + self.account_subtree.extend_from(other.account_subtree); + + self.branch_node_hash_masks.extend(other.branch_node_hash_masks); + self.branch_node_tree_masks.extend(other.branch_node_tree_masks); + + for (hashed_address, storage) in other.storages { + match self.storages.entry(hashed_address) { + hash_map::Entry::Occupied(mut entry) => { + debug_assert_eq!(entry.get().root, storage.root); + let entry = entry.get_mut(); + entry.subtree.extend_from(storage.subtree); + entry.branch_node_hash_masks.extend(storage.branch_node_hash_masks); + entry.branch_node_tree_masks.extend(storage.branch_node_tree_masks); + } + hash_map::Entry::Vacant(entry) => { + entry.insert(storage); + } + } + } + } +} + /// The merkle multiproof of storage trie. #[derive(Clone, Debug, PartialEq, Eq)] pub struct StorageMultiProof { @@ -334,6 +447,41 @@ impl AccountProof { } } +/// The merkle proof with the relevant account info. +#[derive(Clone, PartialEq, Eq, Debug)] +pub struct DecodedAccountProof { + /// The address associated with the account. + pub address: Address, + /// Account info. + pub info: Option, + /// Array of merkle trie nodes which starting from the root node and following the path of the + /// hashed address as key. + pub proof: Vec, + /// The storage trie root. + pub storage_root: B256, + /// Array of storage proofs as requested. + pub storage_proofs: Vec, +} + +impl Default for DecodedAccountProof { + fn default() -> Self { + Self::new(Address::default()) + } +} + +impl DecodedAccountProof { + /// Create new account proof entity. + pub const fn new(address: Address) -> Self { + Self { + address, + info: None, + proof: Vec::new(), + storage_root: EMPTY_ROOT_HASH, + storage_proofs: Vec::new(), + } + } +} + /// The merkle proof of the storage entry. #[derive(Clone, PartialEq, Eq, Default, Debug)] #[cfg_attr(any(test, feature = "serde"), derive(serde::Serialize, serde::Deserialize))]