diff --git a/crates/trie/db/tests/fuzz_in_memory_nodes.rs b/crates/trie/db/tests/fuzz_in_memory_nodes.rs index ebaf5ad84..495289c7f 100644 --- a/crates/trie/db/tests/fuzz_in_memory_nodes.rs +++ b/crates/trie/db/tests/fuzz_in_memory_nodes.rs @@ -18,7 +18,7 @@ proptest! { })] #[test] - fn fuzz_in_memory_nodes(mut init_state: BTreeMap, mut updated_state: BTreeMap) { + fn fuzz_in_memory_nodes(mut init_state: BTreeMap, state_updates: [BTreeMap; 10]) { let factory = create_test_provider_factory(); let provider = factory.provider_rw().unwrap(); let mut hashed_account_cursor = provider.tx_ref().cursor_write::().unwrap(); @@ -29,35 +29,38 @@ proptest! { } // Compute initial root and updates - let (_, trie_updates) = StateRoot::from_tx(provider.tx_ref()) - .root_with_updates() - .unwrap(); - - // Insert state updates into database - let mut changes = PrefixSetMut::default(); - for (hashed_address, balance) in updated_state.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()) - .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()) - ) + let (_, mut trie_nodes) = StateRoot::from_tx(provider.tx_ref()) .root_with_updates() .unwrap(); - // Verify the result - let mut state = BTreeMap::default(); - state.append(&mut init_state); - state.append(&mut updated_state); - 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); + 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 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, 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_nodes.clone().into_sorted()) + ) + .root_with_updates() + .unwrap(); + + trie_nodes.extend(trie_updates); + + // Verify the result + 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); + } + + } } diff --git a/crates/trie/trie/src/updates.rs b/crates/trie/trie/src/updates.rs index 2d35dbf48..3cd802aa6 100644 --- a/crates/trie/trie/src/updates.rs +++ b/crates/trie/trie/src/updates.rs @@ -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(&mut self, walker: TrieWalker, hash_builder: HashBuilder) { // Retrieve deleted keys from trie walker.