feat(trie): multiproof (#9804)

crates/evm/execution-errors/Cargo.toml

@@ -17,6 +17,7 @@ reth-storage-errors.workspace = true
 reth-prune-types.workspace = true
 
 alloy-primitives.workspace = true
+alloy-rlp.workspace = true
 alloy-eips.workspace = true
 revm-primitives.workspace = true
 

crates/evm/execution-errors/src/lib.rs

@@ -23,7 +23,7 @@ use revm_primitives::EVMError;
 use alloc::{boxed::Box, string::String};
 
 pub mod trie;
-pub use trie::{StateRootError, StorageRootError};
+pub use trie::*;
 
 /// Transaction validation errors
 #[derive(thiserror_no_std::Error, Debug, Clone, PartialEq, Eq)]

crates/evm/execution-errors/src/trie.rs

@@ -1,14 +1,34 @@
 //! Errors when computing the state root.
 
-use reth_storage_errors::db::DatabaseError;
+use reth_storage_errors::{db::DatabaseError, provider::ProviderError};
 use thiserror_no_std::Error;
 
+/// State root errors.
+#[derive(Error, Debug, PartialEq, Eq, Clone)]
+pub enum StateProofError {
+    /// Internal database error.
+    #[error(transparent)]
+    Database(#[from] DatabaseError),
+    /// RLP decoding error.
+    #[error(transparent)]
+    Rlp(#[from] alloy_rlp::Error),
+}
+
+impl From<StateProofError> for ProviderError {
+    fn from(value: StateProofError) -> Self {
+        match value {
+            StateProofError::Database(error) => Self::Database(error),
+            StateProofError::Rlp(error) => Self::Rlp(error),
+        }
+    }
+}
+
 /// State root errors.
 #[derive(Error, Debug, PartialEq, Eq, Clone)]
 pub enum StateRootError {
     /// Internal database error.
     #[error(transparent)]
-    DB(#[from] DatabaseError),
+    Database(#[from] DatabaseError),
     /// Storage root error.
     #[error(transparent)]
     StorageRootError(#[from] StorageRootError),
@@ -17,8 +37,8 @@ pub enum StateRootError {
 impl From<StateRootError> for DatabaseError {
     fn from(err: StateRootError) -> Self {
         match err {
-            StateRootError::DB(err) |
-            StateRootError::StorageRootError(StorageRootError::DB(err)) => err,
+            StateRootError::Database(err) |
+            StateRootError::StorageRootError(StorageRootError::Database(err)) => err,
         }
     }
 }
@@ -28,5 +48,5 @@ impl From<StateRootError> for DatabaseError {
 pub enum StorageRootError {
     /// Internal database error.
     #[error(transparent)]
-    DB(#[from] DatabaseError),
+    Database(#[from] DatabaseError),
 }

crates/storage/errors/Cargo.toml

@@ -11,6 +11,7 @@ repository.workspace = true
 workspace = true
 
 [dependencies]
+alloy-rlp.workspace = true
 reth-primitives.workspace = true
 reth-fs-util.workspace = true
 

crates/storage/errors/src/provider.rs

@@ -21,6 +21,9 @@ pub enum ProviderError {
     /// Database error.
     #[error(transparent)]
     Database(#[from] crate::db::DatabaseError),
+    /// RLP error.
+    #[error(transparent)]
+    Rlp(#[from] alloy_rlp::Error),
     /// Filesystem path error.
     #[error("{0}")]
     FsPathError(String),

crates/storage/provider/src/providers/state/historical.rs

@@ -286,7 +286,7 @@ impl<'b, TX: DbTx> StateProofProvider for HistoricalStateProviderRef<'b, TX> {
         let mut revert_state = self.revert_state()?;
         revert_state.extend(hashed_state.clone());
         Proof::overlay_account_proof(self.tx, revert_state, address, slots)
-            .map_err(|err| ProviderError::Database(err.into()))
+            .map_err(Into::<ProviderError>::into)
     }
 }
 

crates/storage/provider/src/providers/state/latest.rs

@@ -96,8 +96,8 @@ impl<'b, TX: DbTx> StateProofProvider for LatestStateProviderRef<'b, TX> {
         address: Address,
         slots: &[B256],
     ) -> ProviderResult<AccountProof> {
-        Ok(Proof::overlay_account_proof(self.tx, hashed_state.clone(), address, slots)
-            .map_err(Into::<reth_db::DatabaseError>::into)?)
+        Proof::overlay_account_proof(self.tx, hashed_state.clone(), address, slots)
+            .map_err(Into::<ProviderError>::into)
     }
 }
 

crates/trie/common/src/lib.rs

@@ -26,7 +26,7 @@ pub use subnode::StoredSubNode;
 mod proofs;
 #[cfg(any(test, feature = "test-utils"))]
 pub use proofs::triehash;
-pub use proofs::{AccountProof, StorageProof};
+pub use proofs::*;
 
 pub mod root;
 

crates/trie/common/src/proofs.rs

@@ -2,12 +2,121 @@
 
 use crate::{Nibbles, TrieAccount};
 use alloy_primitives::{keccak256, Address, Bytes, B256, U256};
-use alloy_rlp::encode_fixed_size;
+use alloy_rlp::{encode_fixed_size, Decodable};
 use alloy_trie::{
+    nodes::TrieNode,
     proof::{verify_proof, ProofVerificationError},
     EMPTY_ROOT_HASH,
 };
-use reth_primitives_traits::Account;
+use reth_primitives_traits::{constants::KECCAK_EMPTY, Account};
+use std::collections::{BTreeMap, HashMap};
+
+/// The state multiproof of target accounts and multiproofs of their storage tries.
+#[derive(Clone, Default, Debug)]
+pub struct MultiProof {
+    /// State trie multiproof for requested accounts.
+    pub account_subtree: BTreeMap<Nibbles, Bytes>,
+    /// Storage trie multiproofs.
+    pub storage_multiproofs: HashMap<B256, StorageMultiProof>,
+}
+
+impl MultiProof {
+    /// Construct the account proof from the multiproof.
+    pub fn account_proof(
+        &self,
+        address: Address,
+        slots: &[B256],
+    ) -> Result<AccountProof, alloy_rlp::Error> {
+        let hashed_address = keccak256(address);
+        let nibbles = Nibbles::unpack(hashed_address);
+
+        // Retrieve the account proof.
+        let proof = self
+            .account_subtree
+            .iter()
+            .filter(|(path, _)| nibbles.starts_with(path))
+            .map(|(_, node)| node.clone())
+            .collect::<Vec<_>>();
+
+        // 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(last) = proof.last() {
+                if let TrieNode::Leaf(leaf) = TrieNode::decode(&mut &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.storage_multiproofs.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 {
+                StorageProof::new(*slot)
+            };
+            storage_proofs.push(proof);
+        }
+        Ok(AccountProof { address, info, proof, storage_root, storage_proofs })
+    }
+}
+
+/// The merkle multiproof of storage trie.
+#[derive(Clone, Debug)]
+pub struct StorageMultiProof {
+    /// Storage trie root.
+    pub root: B256,
+    /// Storage multiproof for requested slots.
+    pub subtree: BTreeMap<Nibbles, Bytes>,
+}
+
+impl Default for StorageMultiProof {
+    fn default() -> Self {
+        Self { root: EMPTY_ROOT_HASH, subtree: BTreeMap::default() }
+    }
+}
+
+impl StorageMultiProof {
+    /// Return storage proofs for the target storage slot (unhashed).
+    pub fn storage_proof(&self, slot: B256) -> Result<StorageProof, alloy_rlp::Error> {
+        let nibbles = Nibbles::unpack(keccak256(slot));
+
+        // Retrieve the storage proof.
+        let proof = self
+            .subtree
+            .iter()
+            .filter(|(path, _)| nibbles.starts_with(path))
+            .map(|(_, node)| node.clone())
+            .collect::<Vec<_>>();
+
+        // Inspect the last node in the proof. If it's a leaf node with matching suffix,
+        // then the node contains the encoded slot value.
+        let value = 'value: {
+            if let Some(last) = proof.last() {
+                if let TrieNode::Leaf(leaf) = TrieNode::decode(&mut &last[..])? {
+                    if nibbles.ends_with(&leaf.key) {
+                        break 'value U256::decode(&mut &leaf.value[..])?
+                    }
+                }
+            }
+            U256::ZERO
+        };
+
+        Ok(StorageProof { key: slot, nibbles, value, proof })
+    }
+}
 
 /// The merkle proof with the relevant account info.
 #[derive(PartialEq, Eq, Debug)]
@@ -37,23 +146,6 @@ impl AccountProof {
         }
     }
 
-    /// 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;
-    }
-
     /// Verify the storage proofs and account proof against the provided state root.
     pub fn verify(&self, root: B256) -> Result<(), ProofVerificationError> {
         // Verify storage proofs.
@@ -106,16 +198,6 @@ impl StorageProof {
         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;
-    }
-
     /// Verify the proof against the provided storage root.
     pub fn verify(&self, root: B256) -> Result<(), ProofVerificationError> {
         let expected =

crates/trie/db/src/proof.rs

@@ -1,5 +1,5 @@
 use reth_db_api::transaction::DbTx;
-use reth_execution_errors::StateRootError;
+use reth_execution_errors::StateProofError;
 use reth_primitives::{Address, B256};
 use reth_trie::{
     hashed_cursor::{DatabaseHashedCursorFactory, HashedPostStateCursorFactory},
@@ -19,7 +19,7 @@ pub trait DatabaseProof<'a, TX> {
         post_state: HashedPostState,
         address: Address,
         slots: &[B256],
-    ) -> Result<AccountProof, StateRootError>;
+    ) -> Result<AccountProof, StateProofError>;
 }
 
 impl<'a, TX: DbTx> DatabaseProof<'a, TX> for Proof<&'a TX, DatabaseHashedCursorFactory<'a, TX>> {
@@ -33,7 +33,7 @@ impl<'a, TX: DbTx> DatabaseProof<'a, TX> for Proof<&'a TX, DatabaseHashedCursorF
         post_state: HashedPostState,
         address: Address,
         slots: &[B256],
-    ) -> Result<AccountProof, StateRootError> {
+    ) -> Result<AccountProof, StateProofError> {
         let prefix_sets = post_state.construct_prefix_sets();
         let sorted = post_state.into_sorted();
         let hashed_cursor_factory =

crates/trie/parallel/src/parallel_root.rs

@@ -210,7 +210,7 @@ impl From<ParallelStateRootError> for ProviderError {
     fn from(error: ParallelStateRootError) -> Self {
         match error {
             ParallelStateRootError::Provider(error) => error,
-            ParallelStateRootError::StorageRoot(StorageRootError::DB(error)) => {
+            ParallelStateRootError::StorageRoot(StorageRootError::Database(error)) => {
                 Self::Database(error)
             }
         }