feat(trie): trie node update extension (#9838)

This commit is contained in:
Roman Krasiuk
2024-07-26 08:03:20 -07:00
committed by GitHub
parent 6d036cd95b
commit efcc77ef44
2 changed files with 47 additions and 28 deletions

View File

@ -18,7 +18,7 @@ proptest! {
})]
#[test]
fn fuzz_in_memory_nodes(mut init_state: BTreeMap<B256, U256>, mut updated_state: BTreeMap<B256, U256>) {
fn fuzz_in_memory_nodes(mut init_state: BTreeMap<B256, U256>, state_updates: [BTreeMap<B256, U256>; 10]) {
let factory = create_test_provider_factory();
let provider = factory.provider_rw().unwrap();
let mut hashed_account_cursor = provider.tx_ref().cursor_write::<tables::HashedAccounts>().unwrap();
@ -29,35 +29,38 @@ proptest! {
}
// Compute initial root and updates
let (_, trie_updates) = StateRoot::from_tx(provider.tx_ref())
let (_, mut trie_nodes) = StateRoot::from_tx(provider.tx_ref())
.root_with_updates()
.unwrap();
let mut state = init_state;
for mut state_update in state_updates {
// Insert state updates into database
let mut changes = PrefixSetMut::default();
for (hashed_address, balance) in updated_state.clone() {
for (hashed_address, balance) in state_update.clone() {
hashed_account_cursor.upsert(hashed_address, Account { balance, ..Default::default() }).unwrap();
changes.insert(Nibbles::unpack(hashed_address));
}
// Compute root with in-memory trie nodes overlay
let (state_root, _) = StateRoot::from_tx(provider.tx_ref())
let (state_root, trie_updates) = StateRoot::from_tx(provider.tx_ref())
.with_prefix_sets(TriePrefixSets { account_prefix_set: changes.freeze(), ..Default::default() })
.with_trie_cursor_factory(InMemoryTrieCursorFactory::new(
DatabaseTrieCursorFactory::new(provider.tx_ref()),
&trie_updates.into_sorted())
DatabaseTrieCursorFactory::new(provider.tx_ref()), &trie_nodes.clone().into_sorted())
)
.root_with_updates()
.unwrap();
trie_nodes.extend(trie_updates);
// Verify the result
let mut state = BTreeMap::default();
state.append(&mut init_state);
state.append(&mut updated_state);
state.append(&mut state_update);
let expected_root = state_root_prehashed(
state.iter().map(|(&key, &balance)| (key, (Account { balance, ..Default::default() }, std::iter::empty())))
);
assert_eq!(expected_root, state_root);
}
}
}

View File

@ -34,6 +34,15 @@ impl TrieUpdates {
&self.storage_tries
}
/// Extends the trie updates.
pub fn extend(&mut self, other: Self) {
self.account_nodes.extend(other.account_nodes);
self.removed_nodes.extend(other.removed_nodes);
for (hashed_address, storage_trie) in other.storage_tries {
self.storage_tries.entry(hashed_address).or_default().extend(storage_trie);
}
}
/// Insert storage updates for a given hashed address.
pub fn insert_storage_updates(
&mut self,
@ -138,6 +147,13 @@ impl StorageTrieUpdates {
self.is_deleted = deleted;
}
/// Extends storage trie updates.
pub fn extend(&mut self, other: Self) {
self.is_deleted |= other.is_deleted;
self.storage_nodes.extend(other.storage_nodes);
self.removed_nodes.extend(other.removed_nodes);
}
/// Finalize storage trie updates for by taking updates from walker and hash builder.
pub fn finalize<C>(&mut self, walker: TrieWalker<C>, hash_builder: HashBuilder) {
// Retrieve deleted keys from trie walker.