From 8dfab3354bbbaa255c62e654933835ad96d1de3b Mon Sep 17 00:00:00 2001 From: Roman Krasiuk Date: Fri, 27 Sep 2024 17:57:24 +0200 Subject: [PATCH] feat(trie): extract `StorageProof` (#11269) --- crates/trie/db/src/proof.rs | 3 +- crates/trie/trie/src/proof.rs | 116 +++++++++++++++++++++----------- crates/trie/trie/src/witness.rs | 28 +++++--- 3 files changed, 96 insertions(+), 51 deletions(-) diff --git a/crates/trie/db/src/proof.rs b/crates/trie/db/src/proof.rs index 9416d078090c..1d5fda84cc5b 100644 --- a/crates/trie/db/src/proof.rs +++ b/crates/trie/db/src/proof.rs @@ -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) } } diff --git a/crates/trie/trie/src/proof.rs b/crates/trie/trie/src/proof.rs index e3bdccafefd2..3e9ca5783814 100644 --- a/crates/trie/trie/src/proof.rs +++ b/crates/trie/trie/src/proof.rs @@ -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 { hashed_cursor_factory: H, /// A set of prefix sets that have changes. prefix_sets: TriePrefixSetsMut, - /// Proof targets. - targets: HashMap>, } impl Proof { - /// 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 Proof { trie_cursor_factory, hashed_cursor_factory: self.hashed_cursor_factory, prefix_sets: self.prefix_sets, - targets: self.targets, } } @@ -61,7 +57,6 @@ impl Proof { trie_cursor_factory: self.trie_cursor_factory, hashed_cursor_factory, prefix_sets: self.prefix_sets, - targets: self.targets, } } @@ -70,22 +65,11 @@ impl Proof { self.prefix_sets = prefix_sets; self } - - /// Set the target account and slots. - pub fn with_target(self, target: (B256, HashSet)) -> Self { - self.with_targets(HashMap::from_iter([target])) - } - - /// Set the target accounts and slots. - pub fn with_targets(mut self, targets: HashMap>) -> Self { - self.targets = targets; - self - } } impl Proof where - T: TrieCursorFactory, + T: TrieCursorFactory + Clone, H: HashedCursorFactory + Clone, { /// Generate an account proof from intermediate nodes. @@ -95,23 +79,28 @@ where slots: &[B256], ) -> Result { 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 { + pub fn multiproof( + mut self, + mut targets: HashMap>, + ) -> Result { 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 }) } +} + +/// Generates storage merkle proofs. +#[derive(Debug)] +pub struct StorageProof { + /// 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 StorageProof { + /// 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(), + } + } - /// Generate a storage multiproof according to specified targets. - pub fn storage_multiproof( - &self, - hashed_address: B256, + /// Set the changed prefixes. + pub fn with_prefix_set_mut(mut self, prefix_set: PrefixSetMut) -> Self { + self.prefix_set = prefix_set; + self + } +} + +impl StorageProof +where + T: TrieCursorFactory, + H: HashedCursorFactory, +{ + /// Generate storage proof. + pub fn storage_proof( + mut self, + targets: HashSet, ) -> Result { 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::>(); + 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); diff --git a/crates/trie/trie/src/witness.rs b/crates/trie/trie/src/witness.rs index ef5b358d31cb..972afc10c342 100644 --- a/crates/trie/trie/src/witness.rs +++ b/crates/trie/trie/src/witness.rs @@ -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 =