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

paradigmxyz · Jul 26, 2024 · efcc77e · efcc77e
1 parent 6d036cd
commit efcc77e
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 27 deletions.
diff --git 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<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())
-        .root_with_updates()
-        .unwrap();
+        let (_, mut trie_nodes) = 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));
-        }
+        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, _) = StateRoot::from_tx(provider.tx_ref())
-            .with_prefix_sets(TriePrefixSets { account_prefix_set: changes.freeze(), ..Default::default() })
-            .with_trie_cursor_factory(InMemoryTrieCursorFactory::new(
+            // 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);
+        }
 
-                DatabaseTrieCursorFactory::new(provider.tx_ref()),
-                &trie_updates.into_sorted())
-            )
-            .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);
     }
 }
diff --git 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<C>(&mut self, walker: TrieWalker<C>, hash_builder: HashBuilder) {
         // Retrieve deleted keys from trie walker.