feat(trie): account & storage proofs (#5041)

crates/primitives/src/trie/hash_builder/mod.rs

@@ -2,8 +2,11 @@ use super::{
     nodes::{rlp_hash, BranchNode, ExtensionNode, LeafNode},
     BranchNodeCompact, Nibbles, TrieMask,
 };
-use crate::{keccak256, proofs::EMPTY_ROOT, B256};
-use std::{collections::HashMap, fmt::Debug};
+use crate::{keccak256, proofs::EMPTY_ROOT, Bytes, B256};
+use std::{
+    collections::{BTreeMap, HashMap},
+    fmt::Debug,
+};
 
 mod state;
 pub use state::HashBuilderState;
@@ -11,6 +14,9 @@ pub use state::HashBuilderState;
 mod value;
 pub use value::HashBuilderValue;
 
+mod proof_retainer;
+pub use proof_retainer::ProofRetainer;
+
 /// A component used to construct the root hash of the trie. The primary purpose of a Hash Builder
 /// is to build the Merkle proof that is essential for verifying the integrity and authenticity of
 /// the trie's contents. It achieves this by constructing the root hash from the hashes of child
@@ -47,6 +53,7 @@ pub struct HashBuilder {
     stored_in_database: bool,
 
     updated_branch_nodes: Option<HashMap<Nibbles, BranchNodeCompact>>,
+    proof_retainer: Option<ProofRetainer>,
 
     rlp_buf: Vec<u8>,
 }
@@ -62,6 +69,7 @@ impl From<HashBuilderState> for HashBuilder {
             hash_masks: state.hash_masks,
             stored_in_database: state.stored_in_database,
             updated_branch_nodes: None,
+            proof_retainer: None,
             rlp_buf: Vec::with_capacity(32),
         }
     }
@@ -90,6 +98,12 @@ impl HashBuilder {
         self
     }
 
+    /// Enable proof retainer for the specified target nibbles.
+    pub fn with_proof_retainer(mut self, targets: Vec<Nibbles>) -> Self {
+        self.proof_retainer = Some(ProofRetainer::new(targets));
+        self
+    }
+
     /// Enables the Hash Builder to store updated branch nodes.
     ///
     /// Call [HashBuilder::split] to get the updates to branch nodes.
@@ -105,6 +119,11 @@ impl HashBuilder {
         (self, updates.unwrap_or_default())
     }
 
+    /// Take and return the proofs retained.
+    pub fn take_proofs(&mut self) -> BTreeMap<Nibbles, Bytes> {
+        self.proof_retainer.take().map(ProofRetainer::into_proofs).unwrap_or_default()
+    }
+
     /// The number of total updates accrued.
     /// Returns `0` if [Self::with_updates] was not called.
     pub fn updates_len(&self) -> usize {
@@ -141,13 +160,6 @@ impl HashBuilder {
         self.stored_in_database = stored_in_database;
     }
 
-    fn set_key_value<T: Into<HashBuilderValue>>(&mut self, key: Nibbles, value: T) {
-        tracing::trace!(target: "trie::hash_builder", key = ?self.key, value = ?self.value, "old key/value");
-        self.key = key;
-        self.value = value.into();
-        tracing::trace!(target: "trie::hash_builder", key = ?self.key, value = ?self.value, "new key/value");
-    }
-
     /// Returns the current root hash of the trie builder.
     pub fn root(&mut self) -> B256 {
         // Clears the internal state
@@ -159,6 +171,13 @@ impl HashBuilder {
         self.current_root()
     }
 
+    fn set_key_value<T: Into<HashBuilderValue>>(&mut self, key: Nibbles, value: T) {
+        tracing::trace!(target: "trie::hash_builder", key = ?self.key, value = ?self.value, "old key/value");
+        self.key = key;
+        self.value = value.into();
+        tracing::trace!(target: "trie::hash_builder", key = ?self.key, value = ?self.value, "new key/value");
+    }
+
     fn current_root(&self) -> B256 {
         if let Some(node_ref) = self.stack.last() {
             if node_ref.len() == B256::len_bytes() + 1 {
@@ -252,6 +271,7 @@ impl HashBuilder {
 
                         self.rlp_buf.clear();
                         self.stack.push(leaf_node.rlp(&mut self.rlp_buf));
+                        self.retain_proof_from_buf(&current);
                     }
                     HashBuilderValue::Hash(hash) => {
                         tracing::debug!(target: "trie::hash_builder", ?hash, "pushing branch node hash");
@@ -281,6 +301,7 @@ impl HashBuilder {
                 }, "extension node rlp");
                 self.rlp_buf.clear();
                 self.stack.push(extension_node.rlp(&mut self.rlp_buf));
+                self.retain_proof_from_buf(&current.slice(0, len_from));
                 self.resize_masks(len_from);
             }
 
@@ -292,7 +313,7 @@ impl HashBuilder {
             // Insert branch nodes in the stack
             if !succeeding.is_empty() || preceding_exists {
                 // Pushes the corresponding branch node to the stack
-                let children = self.push_branch_node(len);
+                let children = self.push_branch_node(&current, len);
                 // Need to store the branch node in an efficient format
                 // outside of the hash builder
                 self.store_branch_node(&current, len, children);
@@ -323,7 +344,7 @@ impl HashBuilder {
     /// Given the size of the longest common prefix, it proceeds to create a branch node
     /// from the state mask and existing stack state, and store its RLP to the top of the stack,
     /// after popping all the relevant elements from the stack.
-    fn push_branch_node(&mut self, len: usize) -> Vec<B256> {
+    fn push_branch_node(&mut self, current: &Nibbles, len: usize) -> Vec<B256> {
         let state_mask = self.groups[len];
         let hash_mask = self.hash_masks[len];
         let branch_node = BranchNode::new(&self.stack);
@@ -331,6 +352,7 @@ impl HashBuilder {
 
         self.rlp_buf.clear();
         let rlp = branch_node.rlp(state_mask, &mut self.rlp_buf);
+        self.retain_proof_from_buf(&current.slice(0, len));
 
         // Clears the stack from the branch node elements
         let first_child_idx = self.stack.len() - state_mask.count_ones() as usize;
@@ -387,6 +409,12 @@ impl HashBuilder {
         }
     }
 
+    fn retain_proof_from_buf(&mut self, prefix: &Nibbles) {
+        if let Some(proof_retainer) = self.proof_retainer.as_mut() {
+            proof_retainer.retain(prefix, &self.rlp_buf)
+        }
+    }
+
     fn update_masks(&mut self, current: &Nibbles, len_from: usize) {
         if len_from > 0 {
             let flag = TrieMask::from_nibble(current[len_from - 1]);

crates/primitives/src/trie/hash_builder/proof_retainer.rs

@@ -0,0 +1,37 @@
+use crate::{trie::Nibbles, Bytes};
+use std::collections::BTreeMap;
+
+/// Proof retainer is used to store proofs during merkle trie construction.
+/// It is intended to be used within the [`HashBuilder`](crate::trie::HashBuilder).
+#[derive(Debug)]
+pub struct ProofRetainer {
+    /// The nibbles of the target trie keys to retain proofs for.
+    targets: Vec<Nibbles>,
+    /// The map of retained proofs (RLP serialized trie nodes)
+    /// with their corresponding key in the trie.
+    proofs: BTreeMap<Nibbles, Bytes>,
+}
+
+impl ProofRetainer {
+    /// Create new retainer with target nibbles.
+    pub fn new(targets: Vec<Nibbles>) -> Self {
+        Self { targets, proofs: Default::default() }
+    }
+
+    /// Returns `true` if the given prefix matches the retainer target.
+    pub fn matches(&self, prefix: &Nibbles) -> bool {
+        self.targets.iter().any(|target| target.starts_with(prefix))
+    }
+
+    /// Returns all collected proofs.
+    pub fn into_proofs(self) -> BTreeMap<Nibbles, Bytes> {
+        self.proofs
+    }
+
+    /// Retain the proof if the key matches any of the targets.
+    pub fn retain(&mut self, prefix: &Nibbles, proof: &[u8]) {
+        if self.matches(prefix) {
+            self.proofs.insert(prefix.clone(), Bytes::from(proof.to_vec()));
+        }
+    }
+}

crates/primitives/src/trie/mod.rs

@@ -8,6 +8,10 @@ pub use nodes::BranchNodeCompact;
 pub mod hash_builder;
 pub use hash_builder::HashBuilder;
 
+/// Merkle trie proofs.
+mod proofs;
+pub use proofs::{AccountProof, StorageProof};
+
 mod mask;
 mod nibbles;
 mod storage;

crates/primitives/src/trie/nibbles.rs

@@ -1,6 +1,6 @@
 use crate::Bytes;
 use alloy_rlp::RlpEncodableWrapper;
-use derive_more::{Deref, DerefMut, From, Index};
+use derive_more::{Deref, From, Index};
 use reth_codecs::{main_codec, Compact};
 use serde::{Deserialize, Serialize};
 
@@ -63,18 +63,7 @@ impl Compact for StoredNibblesSubKey {
 /// `hex_data` has its upper 4 bits set to zero and the lower 4 bits
 /// representing the nibble value.
 #[derive(
-    Default,
-    Clone,
-    Eq,
-    PartialEq,
-    RlpEncodableWrapper,
-    PartialOrd,
-    Ord,
-    Hash,
-    Index,
-    From,
-    Deref,
-    DerefMut,
+    Default, Clone, Eq, PartialEq, RlpEncodableWrapper, PartialOrd, Ord, Hash, Index, From, Deref,
 )]
 pub struct Nibbles {
     /// The inner representation of the nibble sequence.
@@ -276,8 +265,6 @@ impl Nibbles {
 
     /// Extend the current nibbles with another nibbles.
     pub fn extend(&mut self, b: impl AsRef<[u8]>) {
-        // self.hex_data.extend_from_slice(b.as_ref());
-
         let mut bytes = self.hex_data.to_vec();
         bytes.extend_from_slice(b.as_ref());
         self.hex_data = bytes.into();

crates/primitives/src/trie/proofs.rs

@@ -0,0 +1,84 @@
+use super::Nibbles;
+use crate::{keccak256, Account, Address, Bytes, B256, U256};
+
+/// The merkle proof with the relevant account info.
+#[derive(PartialEq, Eq, Default, Debug)]
+pub struct AccountProof {
+    /// The address associated with the account.
+    pub address: Address,
+    /// Account info.
+    pub info: Option<Account>,
+    /// Array of rlp-serialized merkle trie nodes which starting from the root node and
+    /// following the path of the hashed address as key.
+    pub proof: Vec<Bytes>,
+    /// The storage trie root.
+    pub storage_root: B256,
+    /// Array of storage proofs as requested.
+    pub storage_proofs: Vec<StorageProof>,
+}
+
+impl AccountProof {
+    /// Create new account proof entity.
+    pub fn new(address: Address) -> Self {
+        Self { address, ..Default::default() }
+    }
+
+    /// Set account info, storage root and requested storage proofs.
+    pub fn set_account(
+        &mut self,
+        info: Account,
+        storage_root: B256,
+        storage_proofs: Vec<StorageProof>,
+    ) {
+        self.info = Some(info);
+        self.storage_root = storage_root;
+        self.storage_proofs = storage_proofs;
+    }
+
+    /// Set proof path.
+    pub fn set_proof(&mut self, proof: Vec<Bytes>) {
+        self.proof = proof;
+    }
+}
+
+/// The merkle proof of the storage entry.
+#[derive(PartialEq, Eq, Default, Debug)]
+pub struct StorageProof {
+    /// The raw storage key.
+    pub key: B256,
+    /// The hashed storage key nibbles.
+    pub nibbles: Nibbles,
+    /// The storage value.
+    pub value: U256,
+    /// Array of rlp-serialized merkle trie nodes which starting from the storage root node and
+    /// following the path of the hashed storage slot as key.
+    pub proof: Vec<Bytes>,
+}
+
+impl StorageProof {
+    /// Create new storage proof from the storage slot.
+    pub fn new(key: B256) -> Self {
+        let nibbles = Nibbles::unpack(keccak256(key));
+        Self { key, nibbles, ..Default::default() }
+    }
+
+    /// Create new storage proof from the storage slot and its pre-hashed image.
+    pub fn new_with_hashed(key: B256, hashed_key: B256) -> Self {
+        Self { key, nibbles: Nibbles::unpack(hashed_key), ..Default::default() }
+    }
+
+    /// Create new storage proof from the storage slot and its pre-hashed image.
+    pub fn new_with_nibbles(key: B256, nibbles: Nibbles) -> Self {
+        Self { key, nibbles, ..Default::default() }
+    }
+
+    /// Set storage value.
+    pub fn set_value(&mut self, value: U256) {
+        self.value = value;
+    }
+
+    /// Set proof path.
+    pub fn set_proof(&mut self, proof: Vec<Bytes>) {
+        self.proof = proof;
+    }
+}