feat(trie): multiproof (#9804)

This commit is contained in:
Roman Krasiuk
2024-07-25 08:33:28 -07:00
committed by GitHub
parent db8706e5d8
commit c1b5410867
13 changed files with 202 additions and 100 deletions

2
Cargo.lock generated
View File

@ -7290,6 +7290,7 @@ version = "1.0.3"
dependencies = [
"alloy-eips",
"alloy-primitives",
"alloy-rlp",
"reth-consensus",
"reth-prune-types",
"reth-storage-errors",
@ -8647,6 +8648,7 @@ dependencies = [
name = "reth-storage-errors"
version = "1.0.3"
dependencies = [
"alloy-rlp",
"reth-fs-util",
"reth-primitives",
"thiserror-no-std",

View File

@ -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

View File

@ -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)]

View File

@ -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),
}

View File

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

View File

@ -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),

View File

@ -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)
}
}

View File

@ -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)
}
}

View File

@ -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;

View File

@ -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 =

View File

@ -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 =

View File

@ -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)
}
}

View File

@ -7,9 +7,12 @@ use crate::{
HashBuilder, Nibbles,
};
use alloy_rlp::{BufMut, Encodable};
use reth_execution_errors::{StateRootError, StorageRootError};
use reth_primitives::{constants::EMPTY_ROOT_HASH, keccak256, Address, B256};
use reth_trie_common::{proof::ProofRetainer, AccountProof, StorageProof, TrieAccount};
use reth_execution_errors::trie::StateProofError;
use reth_primitives::{keccak256, Address, B256};
use reth_trie_common::{
proof::ProofRetainer, AccountProof, MultiProof, StorageMultiProof, TrieAccount,
};
use std::collections::HashMap;
/// A struct for generating merkle proofs.
///
@ -24,6 +27,8 @@ pub struct Proof<T, H> {
trie_cursor_factory: T,
/// A set of prefix sets that have changes.
prefix_sets: TriePrefixSetsMut,
/// Proof targets.
targets: HashMap<B256, Vec<B256>>,
}
impl<T, H> Proof<T, H> {
@ -33,6 +38,7 @@ impl<T, H> Proof<T, H> {
trie_cursor_factory: t,
hashed_cursor_factory: h,
prefix_sets: TriePrefixSetsMut::default(),
targets: HashMap::default(),
}
}
@ -42,6 +48,7 @@ impl<T, H> Proof<T, H> {
trie_cursor_factory: self.trie_cursor_factory,
hashed_cursor_factory,
prefix_sets: self.prefix_sets,
targets: self.targets,
}
}
@ -50,6 +57,12 @@ impl<T, H> Proof<T, H> {
self.prefix_sets = prefix_sets;
self
}
/// Set the target accounts and slots.
pub fn with_targets(mut self, targets: HashMap<B256, Vec<B256>>) -> Self {
self.targets = targets;
self
}
}
impl<T, H> Proof<T, H>
@ -59,26 +72,34 @@ where
{
/// Generate an account proof from intermediate nodes.
pub fn account_proof(
&self,
self,
address: Address,
slots: &[B256],
) -> Result<AccountProof, StateRootError> {
let target_hashed_address = keccak256(address);
let target_nibbles = Nibbles::unpack(target_hashed_address);
let mut account_proof = AccountProof::new(address);
) -> Result<AccountProof, StateProofError> {
Ok(self
.with_targets(HashMap::from([(
keccak256(address),
slots.iter().map(keccak256).collect(),
)]))
.multi_proof()?
.account_proof(address, slots)?)
}
/// Generate a state multiproof according to specified targets.
pub fn multi_proof(&self) -> Result<MultiProof, StateProofError> {
let hashed_account_cursor = self.hashed_cursor_factory.hashed_account_cursor()?;
let trie_cursor = self.trie_cursor_factory.account_trie_cursor()?;
// Create the walker.
let mut prefix_set = self.prefix_sets.account_prefix_set.clone();
prefix_set.insert(target_nibbles.clone());
prefix_set.extend(self.targets.keys().map(Nibbles::unpack));
let walker = TrieWalker::new(trie_cursor, prefix_set.freeze());
// Create a hash builder to rebuild the root node since it is not available in the database.
let retainer = ProofRetainer::from_iter([target_nibbles]);
let retainer = ProofRetainer::from_iter(self.targets.keys().map(Nibbles::unpack));
let mut hash_builder = HashBuilder::default().with_proof_retainer(retainer);
let mut storage_multiproofs = HashMap::default();
let mut account_rlp = Vec::with_capacity(128);
let mut account_node_iter = TrieNodeIter::new(walker, hashed_account_cursor);
while let Some(account_node) = account_node_iter.try_next()? {
@ -87,55 +108,40 @@ where
hash_builder.add_branch(node.key, node.value, node.children_are_in_trie);
}
TrieElement::Leaf(hashed_address, account) => {
let storage_root = if hashed_address == target_hashed_address {
let (storage_root, storage_proofs) =
self.storage_root_with_proofs(hashed_address, slots)?;
account_proof.set_account(account, storage_root, storage_proofs);
storage_root
} else {
self.storage_root(hashed_address)?
};
let storage_multiproof = self.storage_multiproof(hashed_address)?;
// Encode account
account_rlp.clear();
let account = TrieAccount::from((account, storage_root));
let account = TrieAccount::from((account, storage_multiproof.root));
account.encode(&mut account_rlp as &mut dyn BufMut);
hash_builder.add_leaf(Nibbles::unpack(hashed_address), &account_rlp);
storage_multiproofs.insert(hashed_address, storage_multiproof);
}
}
}
let _ = hash_builder.root();
let proofs = hash_builder.take_proofs();
account_proof.set_proof(proofs.values().cloned().collect());
Ok(account_proof)
Ok(MultiProof { account_subtree: hash_builder.take_proofs(), storage_multiproofs })
}
/// Compute storage root.
pub fn storage_root(&self, hashed_address: B256) -> Result<B256, StorageRootError> {
let (storage_root, _) = self.storage_root_with_proofs(hashed_address, &[])?;
Ok(storage_root)
}
/// Compute the storage root and retain proofs for requested slots.
pub fn storage_root_with_proofs(
/// Generate a storage multiproof according to specified targets.
pub fn storage_multiproof(
&self,
hashed_address: B256,
slots: &[B256],
) -> Result<(B256, Vec<StorageProof>), StorageRootError> {
) -> Result<StorageMultiProof, StateProofError> {
let mut hashed_storage_cursor =
self.hashed_cursor_factory.hashed_storage_cursor(hashed_address)?;
let mut proofs = slots.iter().copied().map(StorageProof::new).collect::<Vec<_>>();
// short circuit on empty storage
if hashed_storage_cursor.is_storage_empty()? {
return Ok((EMPTY_ROOT_HASH, proofs))
return Ok(StorageMultiProof::default())
}
let target_nibbles = proofs.iter().map(|p| p.nibbles.clone()).collect::<Vec<_>>();
let target_nibbles = self
.targets
.get(&hashed_address)
.map_or(Vec::new(), |slots| slots.iter().map(Nibbles::unpack).collect());
let mut prefix_set =
self.prefix_sets.storage_prefix_sets.get(&hashed_address).cloned().unwrap_or_default();
prefix_set.extend(target_nibbles.clone());
@ -151,28 +157,15 @@ where
hash_builder.add_branch(node.key, node.value, node.children_are_in_trie);
}
TrieElement::Leaf(hashed_slot, value) => {
let nibbles = Nibbles::unpack(hashed_slot);
if let Some(proof) = proofs.iter_mut().find(|proof| proof.nibbles == nibbles) {
proof.set_value(value);
}
hash_builder.add_leaf(nibbles, alloy_rlp::encode_fixed_size(&value).as_ref());
hash_builder.add_leaf(
Nibbles::unpack(hashed_slot),
alloy_rlp::encode_fixed_size(&value).as_ref(),
);
}
}
}
let root = hash_builder.root();
let all_proof_nodes = hash_builder.take_proofs();
for proof in &mut proofs {
// Iterate over all proof nodes and find the matching ones.
// The filtered results are guaranteed to be in order.
let matching_proof_nodes = all_proof_nodes
.iter()
.filter(|(path, _)| proof.nibbles.starts_with(path))
.map(|(_, node)| node.clone());
proof.set_proof(matching_proof_nodes.collect());
}
Ok((root, proofs))
Ok(StorageMultiProof { root, subtree: hash_builder.take_proofs() })
}
}