From ab90993dadcf00ff8e6ea5c76e89a256b134694a Mon Sep 17 00:00:00 2001 From: Roman Krasiuk Date: Thu, 4 Jul 2024 19:32:10 +0200 Subject: [PATCH 1/3] feat(trie): allow supplying prefix sets to `Proof` --- crates/trie/trie/src/prefix_set/mod.rs | 9 +++++++++ crates/trie/trie/src/proof.rs | 27 ++++++++++++++++++-------- 2 files changed, 28 insertions(+), 8 deletions(-) diff --git a/crates/trie/trie/src/prefix_set/mod.rs b/crates/trie/trie/src/prefix_set/mod.rs index 9d08ad36a98d..3474bf74c8e1 100644 --- a/crates/trie/trie/src/prefix_set/mod.rs +++ b/crates/trie/trie/src/prefix_set/mod.rs @@ -131,6 +131,15 @@ impl PrefixSetMut { self.keys.push(nibbles); } + /// Extend prefix set keys with contents of provided iterator. + pub fn extend(&mut self, nibbles_iter: I) + where + I: IntoIterator, + { + self.sorted = false; + self.keys.extend(nibbles_iter); + } + /// Returns the number of elements in the set. pub fn len(&self) -> usize { self.keys.len() diff --git a/crates/trie/trie/src/proof.rs b/crates/trie/trie/src/proof.rs index 2342ece9869a..dcad8bef41da 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::PrefixSetMut, + prefix_set::{PrefixSetMut, TriePrefixSetsMut}, trie_cursor::{DatabaseAccountTrieCursor, DatabaseStorageTrieCursor}, walker::TrieWalker, HashBuilder, Nibbles, @@ -24,23 +24,31 @@ pub struct Proof<'a, TX, H> { tx: &'a TX, /// The factory for hashed cursors. hashed_cursor_factory: H, + /// A set of prefix sets that have changes. + prefix_sets: TriePrefixSetsMut, } impl<'a, TX, H> Proof<'a, TX, H> { /// Creates a new proof generator. - pub const fn new(tx: &'a TX, hashed_cursor_factory: H) -> Self { - Self { tx, hashed_cursor_factory } + pub fn new(tx: &'a TX, hashed_cursor_factory: H) -> Self { + Self { tx, hashed_cursor_factory, prefix_sets: TriePrefixSetsMut::default() } } /// Set the hashed cursor factory. pub fn with_hashed_cursor_factory(self, hashed_cursor_factory: HF) -> Proof<'a, TX, HF> { - Proof { tx: self.tx, hashed_cursor_factory } + Proof { tx: self.tx, hashed_cursor_factory, prefix_sets: self.prefix_sets } + } + + /// Set the prefix sets. They have to be mutable in order to allow extension with proof target. + pub fn with_prefix_sets_mut(mut self, prefix_sets: TriePrefixSetsMut) -> Self { + self.prefix_sets = prefix_sets; + self } } impl<'a, TX> Proof<'a, TX, &'a TX> { /// Create a new [Proof] instance from database transaction. - pub const fn from_tx(tx: &'a TX) -> Self { + pub fn from_tx(tx: &'a TX) -> Self { Self::new(tx, tx) } } @@ -65,7 +73,7 @@ where DatabaseAccountTrieCursor::new(self.tx.cursor_read::()?); // Create the walker. - let mut prefix_set = PrefixSetMut::default(); + let mut prefix_set = self.prefix_sets.account_prefix_set.clone(); prefix_set.insert(target_nibbles.clone()); let walker = TrieWalker::new(trie_cursor, prefix_set.freeze()); @@ -130,12 +138,15 @@ where } let target_nibbles = proofs.iter().map(|p| p.nibbles.clone()).collect::>(); - let prefix_set = PrefixSetMut::from(target_nibbles.clone()).freeze(); + let mut prefix_set = PrefixSetMut::from( + self.prefix_sets.storage_prefix_sets.get(&hashed_address).cloned().unwrap_or_default(), + ); + prefix_set.extend(target_nibbles.clone()); let trie_cursor = DatabaseStorageTrieCursor::new( self.tx.cursor_dup_read::()?, hashed_address, ); - let walker = TrieWalker::new(trie_cursor, prefix_set); + let walker = TrieWalker::new(trie_cursor, prefix_set.freeze()); let retainer = ProofRetainer::from_iter(target_nibbles); let mut hash_builder = HashBuilder::default().with_proof_retainer(retainer); From c4cd32abe8bbe2f15433cd9a7685fb5049e226d3 Mon Sep 17 00:00:00 2001 From: Roman Krasiuk Date: Thu, 4 Jul 2024 19:35:55 +0200 Subject: [PATCH 2/3] fix --- crates/trie/trie/src/proof.rs | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/crates/trie/trie/src/proof.rs b/crates/trie/trie/src/proof.rs index dcad8bef41da..f0774eb48a80 100644 --- a/crates/trie/trie/src/proof.rs +++ b/crates/trie/trie/src/proof.rs @@ -138,9 +138,8 @@ where } let target_nibbles = proofs.iter().map(|p| p.nibbles.clone()).collect::>(); - let mut prefix_set = PrefixSetMut::from( - self.prefix_sets.storage_prefix_sets.get(&hashed_address).cloned().unwrap_or_default(), - ); + 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 = DatabaseStorageTrieCursor::new( self.tx.cursor_dup_read::()?, From 55762c5a2940004027baa83d6ddf28e4a36379a1 Mon Sep 17 00:00:00 2001 From: Roman Krasiuk Date: Thu, 4 Jul 2024 19:55:30 +0200 Subject: [PATCH 3/3] rm unused import --- crates/trie/trie/src/proof.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/trie/trie/src/proof.rs b/crates/trie/trie/src/proof.rs index f0774eb48a80..eb492f81f4f6 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::{PrefixSetMut, TriePrefixSetsMut}, + prefix_set::TriePrefixSetsMut, trie_cursor::{DatabaseAccountTrieCursor, DatabaseStorageTrieCursor}, walker::TrieWalker, HashBuilder, Nibbles,