github/reth
1
1
Fork 0
mirror of https://github.com/paradigmxyz/reth.git synced 2026-01-26 23:58:46 -05:00
Code Issues Packages Projects Releases Wiki Activity

feat(trie): extract StorageProof (#11269)

Browse Source
This commit is contained in:
Roman Krasiuk
2024-09-27 17:57:24 +02:00
committed by GitHub
parent dbd9a2bb35
commit 8dfab3354b
3 changed files with 96 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
crates/trie/db/src/proof.rs
View File

@@ -78,7 +78,6 @@ impl<'a, TX: DbTx> DatabaseProof<'a, TX>
&state_sorted,
))
.with_prefix_sets_mut(input.prefix_sets)
.with_targets(targets)
.multiproof()
.multiproof(targets)
}
}

116
crates/trie/trie/src/proof.rs
View File

@@ -1,7 +1,7 @@
use crate::{
hashed_cursor::{HashedCursorFactory, HashedStorageCursor},
node_iter::{TrieElement, TrieNodeIter},
prefix_set::TriePrefixSetsMut,
prefix_set::{PrefixSetMut, TriePrefixSetsMut},
trie_cursor::TrieCursorFactory,
walker::TrieWalker,
HashBuilder, Nibbles,
@@ -30,18 +30,15 @@ pub struct Proof<T, H> {
hashed_cursor_factory: H,
/// A set of prefix sets that have changes.
prefix_sets: TriePrefixSetsMut,
/// Proof targets.
targets: HashMap<B256, HashSet<B256>>,
}
impl<T, H> Proof<T, H> {
/// Create a new [Proof] instance.
/// Create a new [`Proof`] instance.
pub fn new(t: T, h: H) -> Self {
Self {
trie_cursor_factory: t,
hashed_cursor_factory: h,
prefix_sets: TriePrefixSetsMut::default(),
targets: HashMap::default(),
}
}
@@ -51,7 +48,6 @@ impl<T, H> Proof<T, H> {
trie_cursor_factory,
hashed_cursor_factory: self.hashed_cursor_factory,
prefix_sets: self.prefix_sets,
targets: self.targets,
}
}
@@ -61,7 +57,6 @@ impl<T, H> Proof<T, H> {
trie_cursor_factory: self.trie_cursor_factory,
hashed_cursor_factory,
prefix_sets: self.prefix_sets,
targets: self.targets,
}
}
@@ -70,22 +65,11 @@ impl<T, H> Proof<T, H> {
self.prefix_sets = prefix_sets;
self
}
/// Set the target account and slots.
pub fn with_target(self, target: (B256, HashSet<B256>)) -> Self {
self.with_targets(HashMap::from_iter([target]))
}
/// Set the target accounts and slots.
pub fn with_targets(mut self, targets: HashMap<B256, HashSet<B256>>) -> Self {
self.targets = targets;
self
}
}
impl<T, H> Proof<T, H>
where
T: TrieCursorFactory,
T: TrieCursorFactory + Clone,
H: HashedCursorFactory + Clone,
{
/// Generate an account proof from intermediate nodes.
@@ -95,23 +79,28 @@ where
slots: &[B256],
) -> Result<AccountProof, StateProofError> {
Ok(self
.with_target((keccak256(address), slots.iter().map(keccak256).collect()))
.multiproof()?
.multiproof(HashMap::from_iter([(
keccak256(address),
slots.iter().map(keccak256).collect(),
)]))?
.account_proof(address, slots)?)
}
/// Generate a state multiproof according to specified targets.
pub fn multiproof(&self) -> Result<MultiProof, StateProofError> {
pub fn multiproof(
mut self,
mut targets: HashMap<B256, HashSet<B256>>,
) -> 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.extend(self.targets.keys().map(Nibbles::unpack));
prefix_set.extend(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(self.targets.keys().map(Nibbles::unpack));
let retainer = ProofRetainer::from_iter(targets.keys().map(Nibbles::unpack));
let mut hash_builder = HashBuilder::default().with_proof_retainer(retainer);
let mut storages = HashMap::default();
@@ -123,7 +112,19 @@ where
hash_builder.add_branch(node.key, node.value, node.children_are_in_trie);
}
TrieElement::Leaf(hashed_address, account) => {
let storage_multiproof = self.storage_multiproof(hashed_address)?;
let storage_prefix_set = self
.prefix_sets
.storage_prefix_sets
.remove(&hashed_address)
.unwrap_or_default();
let proof_targets = targets.remove(&hashed_address).unwrap_or_default();
let storage_multiproof = StorageProof::new_hashed(
self.trie_cursor_factory.clone(),
self.hashed_cursor_factory.clone(),
hashed_address,
)
.with_prefix_set_mut(storage_prefix_set)
.storage_proof(proof_targets)?;
// Encode account
account_rlp.clear();
@@ -138,30 +139,67 @@ where
let _ = hash_builder.root();
Ok(MultiProof { account_subtree: hash_builder.take_proof_nodes(), storages })
}
}
/// Generate a storage multiproof according to specified targets.
pub fn storage_multiproof(
&self,
hashed_address: B256,
/// Generates storage merkle proofs.
#[derive(Debug)]
pub struct StorageProof<T, H> {
/// The factory for traversing trie nodes.
trie_cursor_factory: T,
/// The factory for hashed cursors.
hashed_cursor_factory: H,
/// The hashed address of an account.
hashed_address: B256,
/// The set of storage slot prefixes that have changed.
prefix_set: PrefixSetMut,
}
impl<T, H> StorageProof<T, H> {
/// Create a new [`StorageProof`] instance.
pub fn new(t: T, h: H, address: Address) -> Self {
Self::new_hashed(t, h, keccak256(address))
}
/// Create a new [`StorageProof`] instance with hashed address.
pub fn new_hashed(t: T, h: H, hashed_address: B256) -> Self {
Self {
trie_cursor_factory: t,
hashed_cursor_factory: h,
hashed_address,
prefix_set: PrefixSetMut::default(),
}
}
/// Set the changed prefixes.
pub fn with_prefix_set_mut(mut self, prefix_set: PrefixSetMut) -> Self {
self.prefix_set = prefix_set;
self
}
}
impl<T, H> StorageProof<T, H>
where
T: TrieCursorFactory,
H: HashedCursorFactory,
{
/// Generate storage proof.
pub fn storage_proof(
mut self,
targets: HashSet<B256>,
) -> Result<StorageMultiProof, StateProofError> {
let mut hashed_storage_cursor =
self.hashed_cursor_factory.hashed_storage_cursor(hashed_address)?;
self.hashed_cursor_factory.hashed_storage_cursor(self.hashed_address)?;
// short circuit on empty storage
if hashed_storage_cursor.is_storage_empty()? {
return Ok(StorageMultiProof::default())
}
let target_nibbles = self
.targets
.get(&hashed_address)
.map_or(Vec::new(), |slots| slots.iter().map(Nibbles::unpack).collect());
let target_nibbles = targets.into_iter().map(Nibbles::unpack).collect::<Vec<_>>();
self.prefix_set.extend(target_nibbles.clone());
let mut prefix_set =
self.prefix_sets.storage_prefix_sets.get(&hashed_address).cloned().unwrap_or_default();
prefix_set.extend(target_nibbles.clone());
let trie_cursor = self.trie_cursor_factory.storage_trie_cursor(hashed_address)?;
let walker = TrieWalker::new(trie_cursor, prefix_set.freeze());
let trie_cursor = self.trie_cursor_factory.storage_trie_cursor(self.hashed_address)?;
let walker = TrieWalker::new(trie_cursor, self.prefix_set.freeze());
let retainer = ProofRetainer::from_iter(target_nibbles);
let mut hash_builder = HashBuilder::default().with_proof_retainer(retainer);

28
crates/trie/trie/src/witness.rs
View File

@@ -1,8 +1,11 @@
use std::collections::BTreeMap;
use crate::{
hashed_cursor::HashedCursorFactory, prefix_set::TriePrefixSetsMut, proof::Proof,
trie_cursor::TrieCursorFactory, HashedPostState,
hashed_cursor::HashedCursorFactory,
prefix_set::TriePrefixSetsMut,
proof::{Proof, StorageProof},
trie_cursor::TrieCursorFactory,
HashedPostState,
};
use alloy_primitives::{
keccak256,
@@ -100,8 +103,7 @@ where
let mut account_multiproof =
Proof::new(self.trie_cursor_factory.clone(), self.hashed_cursor_factory.clone())
.with_prefix_sets_mut(self.prefix_sets.clone())
.with_targets(proof_targets.clone())
.multiproof()?;
.multiproof(proof_targets.clone())?;
// Attempt to compute state root from proofs and gather additional
// information for the witness.
@@ -162,13 +164,19 @@ where
let mut padded_key = key.pack();
padded_key.resize(32, 0);
let target_key = B256::from_slice(&padded_key);
let proof = Proof::new(
let storage_prefix_set = self
.prefix_sets
.storage_prefix_sets
.get(&hashed_address)
.cloned()
.unwrap_or_default();
let proof = StorageProof::new_hashed(
self.trie_cursor_factory.clone(),
self.hashed_cursor_factory.clone(),
hashed_address,
)
.with_prefix_sets_mut(self.prefix_sets.clone())
.with_target((hashed_address, HashSet::from_iter([target_key])))
.storage_multiproof(hashed_address)?;
.with_prefix_set_mut(storage_prefix_set)
.storage_proof(HashSet::from_iter([target_key]))?;
// The subtree only contains the proof for a single target.
let node =
@@ -182,11 +190,11 @@ where
// Right pad the target with 0s.
let mut padded_key = key.pack();
padded_key.resize(32, 0);
let targets = HashMap::from_iter([(B256::from_slice(&padded_key), HashSet::default())]);
let proof =
Proof::new(self.trie_cursor_factory.clone(), self.hashed_cursor_factory.clone())
.with_prefix_sets_mut(self.prefix_sets.clone())
.with_target((B256::from_slice(&padded_key), HashSet::default()))
.multiproof()?;
.multiproof(targets)?;
// The subtree only contains the proof for a single target.
let node =