diff --git a/crates/storage/provider/src/providers/database/provider.rs b/crates/storage/provider/src/providers/database/provider.rs index d5e49d822b2..93baa4309d2 100644 --- a/crates/storage/provider/src/providers/database/provider.rs +++ b/crates/storage/provider/src/providers/database/provider.rs @@ -2077,7 +2077,7 @@ impl TrieWriter for DatabaseProvider let mut account_trie_cursor = tx.cursor_write::()?; // Process sorted account nodes - for (key, updated_node) in &trie_updates.account_nodes { + for (key, updated_node) in trie_updates.account_nodes_ref() { let nibbles = StoredNibbles(*key); match updated_node { Some(node) => { @@ -2144,7 +2144,7 @@ impl TrieWriter for DatabaseProvider )?; } - let mut storage_updates = trie_updates.storage_tries.iter().collect::>(); + let mut storage_updates = trie_updates.storage_tries_ref().iter().collect::>(); storage_updates.sort_unstable_by(|a, b| a.0.cmp(b.0)); num_entries += self.write_storage_trie_changesets( @@ -2194,7 +2194,7 @@ impl TrieReader for DatabaseProvider { let tx = self.tx_ref(); // Read account trie changes directly into a Vec - data is already sorted by nibbles - // within each block, and we want the oldest (first) version of each node + // within each block, and we want the oldest (first) version of each node sorted by path. let mut account_nodes = Vec::new(); let mut seen_account_keys = HashSet::new(); let mut accounts_cursor = tx.cursor_dup_read::()?; @@ -2207,8 +2207,11 @@ impl TrieReader for DatabaseProvider { } } + account_nodes.sort_by_key(|(path, _)| *path); + // Read storage trie changes - data is sorted by (block, hashed_address, nibbles) - // Keep track of seen (address, nibbles) pairs to only keep the oldest version + // Keep track of seen (address, nibbles) pairs to only keep the oldest version per address, + // sorted by path. let mut storage_tries = B256Map::>::default(); let mut seen_storage_keys = HashSet::new(); let mut storages_cursor = tx.cursor_dup_read::()?; @@ -2231,12 +2234,13 @@ impl TrieReader for DatabaseProvider { // Convert to StorageTrieUpdatesSorted let storage_tries = storage_tries .into_iter() - .map(|(address, nodes)| { + .map(|(address, mut nodes)| { + nodes.sort_by_key(|(path, _)| *path); (address, StorageTrieUpdatesSorted { storage_nodes: nodes, is_deleted: false }) }) .collect(); - Ok(TrieUpdatesSorted { account_nodes, storage_tries }) + Ok(TrieUpdatesSorted::new(account_nodes, storage_tries)) } fn get_block_trie_updates( @@ -2254,7 +2258,7 @@ impl TrieReader for DatabaseProvider { let cursor_factory = InMemoryTrieCursorFactory::new(db_cursor_factory, &reverts); // Step 3: Collect all account trie nodes that changed in the target block - let mut trie_updates = TrieUpdatesSorted::default(); + let mut account_nodes = Vec::new(); // Walk through all account trie changes for this block let mut accounts_trie_cursor = tx.cursor_dup_read::()?; @@ -2264,10 +2268,11 @@ impl TrieReader for DatabaseProvider { let (_, TrieChangeSetsEntry { nibbles, .. }) = entry?; // Look up the current value of this trie node using the overlay cursor let node_value = account_cursor.seek_exact(nibbles.0)?.map(|(_, node)| node); - trie_updates.account_nodes.push((nibbles.0, node_value)); + account_nodes.push((nibbles.0, node_value)); } // Step 4: Collect all storage trie nodes that changed in the target block + let mut storage_tries = B256Map::default(); let mut storages_trie_cursor = tx.cursor_dup_read::()?; let storage_range_start = BlockNumberHashedAddress((block_number, B256::ZERO)); let storage_range_end = BlockNumberHashedAddress((block_number + 1, B256::ZERO)); @@ -2291,8 +2296,7 @@ impl TrieReader for DatabaseProvider { let cursor = storage_cursor.as_mut().expect("storage_cursor was just initialized above"); let node_value = cursor.seek_exact(nibbles.0)?.map(|(_, node)| node); - trie_updates - .storage_tries + storage_tries .entry(hashed_address) .or_insert_with(|| StorageTrieUpdatesSorted { storage_nodes: Vec::new(), @@ -2302,7 +2306,7 @@ impl TrieReader for DatabaseProvider { .push((nibbles.0, node_value)); } - Ok(trie_updates) + Ok(TrieUpdatesSorted::new(account_nodes, storage_tries)) } } @@ -2379,7 +2383,7 @@ impl StorageTrieWriter for DatabaseP // Get the overlay updates for this storage trie, or use an empty array let overlay_updates = updates_overlay - .and_then(|overlay| overlay.storage_tries.get(hashed_address)) + .and_then(|overlay| overlay.storage_tries_ref().get(hashed_address)) .map(|updates| updates.storage_nodes_ref()) .unwrap_or(&EMPTY_UPDATES); @@ -3463,7 +3467,7 @@ mod tests { storage_tries.insert(storage_address1, storage_trie1); storage_tries.insert(storage_address2, storage_trie2); - let trie_updates = TrieUpdatesSorted { account_nodes, storage_tries }; + let trie_updates = TrieUpdatesSorted::new(account_nodes, storage_tries); // Write the changesets let num_written = @@ -3679,10 +3683,7 @@ mod tests { overlay_storage_tries.insert(storage_address1, overlay_storage_trie1); overlay_storage_tries.insert(storage_address2, overlay_storage_trie2); - let overlay = TrieUpdatesSorted { - account_nodes: overlay_account_nodes, - storage_tries: overlay_storage_tries, - }; + let overlay = TrieUpdatesSorted::new(overlay_account_nodes, overlay_storage_tries); // Normal storage trie: one Some (update) and one None (new) let storage_trie1 = StorageTrieUpdatesSorted { @@ -3709,7 +3710,7 @@ mod tests { storage_tries.insert(storage_address1, storage_trie1); storage_tries.insert(storage_address2, storage_trie2); - let trie_updates = TrieUpdatesSorted { account_nodes, storage_tries }; + let trie_updates = TrieUpdatesSorted::new(account_nodes, storage_tries); // Write the changesets WITH OVERLAY let num_written = @@ -4275,7 +4276,7 @@ mod tests { storage_tries.insert(storage_address1, storage_trie1); storage_tries.insert(storage_address2, storage_trie2); - let trie_updates = TrieUpdatesSorted { account_nodes, storage_tries }; + let trie_updates = TrieUpdatesSorted::new(account_nodes, storage_tries); // Write the sorted trie updates let num_entries = provider_rw.write_trie_updates_sorted(&trie_updates).unwrap(); @@ -4539,11 +4540,11 @@ mod tests { let result = provider.get_block_trie_updates(target_block).unwrap(); // Verify account trie updates - assert_eq!(result.account_nodes.len(), 2, "Should have 2 account trie updates"); + assert_eq!(result.account_nodes_ref().len(), 2, "Should have 2 account trie updates"); // Check nibbles1 - should have the current value (node1) let nibbles1_update = result - .account_nodes + .account_nodes_ref() .iter() .find(|(n, _)| n == &account_nibbles1) .expect("Should find nibbles1"); @@ -4556,7 +4557,7 @@ mod tests { // Check nibbles2 - should have the current value (node2) let nibbles2_update = result - .account_nodes + .account_nodes_ref() .iter() .find(|(n, _)| n == &account_nibbles2) .expect("Should find nibbles2"); @@ -4569,14 +4570,14 @@ mod tests { // nibbles3 should NOT be in the result (it was changed in next_block, not target_block) assert!( - !result.account_nodes.iter().any(|(n, _)| n == &account_nibbles3), + !result.account_nodes_ref().iter().any(|(n, _)| n == &account_nibbles3), "nibbles3 should not be in target_block updates" ); // Verify storage trie updates - assert_eq!(result.storage_tries.len(), 1, "Should have 1 storage trie"); + assert_eq!(result.storage_tries_ref().len(), 1, "Should have 1 storage trie"); let storage_updates = result - .storage_tries + .storage_tries_ref() .get(&storage_address1) .expect("Should have storage updates for address1"); diff --git a/crates/trie/common/src/updates.rs b/crates/trie/common/src/updates.rs index e3e098ac8e5..b0d178cd1d0 100644 --- a/crates/trie/common/src/updates.rs +++ b/crates/trie/common/src/updates.rs @@ -432,12 +432,35 @@ pub struct TrieUpdatesSortedRef<'a> { pub struct TrieUpdatesSorted { /// Sorted collection of updated state nodes with corresponding paths. None indicates that a /// node was removed. - pub account_nodes: Vec<(Nibbles, Option)>, + account_nodes: Vec<(Nibbles, Option)>, /// Storage tries stored by hashed address of the account the trie belongs to. - pub storage_tries: B256Map, + storage_tries: B256Map, } impl TrieUpdatesSorted { + /// Creates a new `TrieUpdatesSorted` with the given account nodes and storage tries. + /// + /// # Panics + /// + /// In debug mode, panics if `account_nodes` is not sorted by the `Nibbles` key, + /// or if any storage trie's `storage_nodes` is not sorted by its `Nibbles` key. + pub fn new( + account_nodes: Vec<(Nibbles, Option)>, + storage_tries: B256Map, + ) -> Self { + debug_assert!( + account_nodes.is_sorted_by_key(|item| &item.0), + "account_nodes must be sorted by Nibbles key" + ); + debug_assert!( + storage_tries.values().all(|storage_trie| { + storage_trie.storage_nodes.is_sorted_by_key(|item| &item.0) + }), + "all storage_nodes in storage_tries must be sorted by Nibbles key" + ); + Self { account_nodes, storage_tries } + } + /// Returns `true` if the updates are empty. pub fn is_empty(&self) -> bool { self.account_nodes.is_empty() && self.storage_tries.is_empty() diff --git a/crates/trie/trie/src/trie_cursor/in_memory.rs b/crates/trie/trie/src/trie_cursor/in_memory.rs index 1c7f179ad0a..e76bf7b2be3 100644 --- a/crates/trie/trie/src/trie_cursor/in_memory.rs +++ b/crates/trie/trie/src/trie_cursor/in_memory.rs @@ -47,7 +47,8 @@ where // if the storage trie has no updates then we use this as the in-memory overlay. static EMPTY_UPDATES: Vec<(Nibbles, Option)> = Vec::new(); - let storage_trie_updates = self.trie_updates.as_ref().storage_tries.get(&hashed_address); + let storage_trie_updates = + self.trie_updates.as_ref().storage_tries_ref().get(&hashed_address); let (storage_nodes, cleared) = storage_trie_updates .map(|u| (u.storage_nodes_ref(), u.is_deleted())) .unwrap_or((&EMPTY_UPDATES, false));