Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 26 additions & 25 deletions crates/storage/provider/src/providers/database/provider.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2077,7 +2077,7 @@ impl<TX: DbTxMut + DbTx + 'static, N: NodeTypes> TrieWriter for DatabaseProvider
let mut account_trie_cursor = tx.cursor_write::<tables::AccountsTrie>()?;

// 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) => {
Expand Down Expand Up @@ -2144,7 +2144,7 @@ impl<TX: DbTxMut + DbTx + 'static, N: NodeTypes> TrieWriter for DatabaseProvider
)?;
}

let mut storage_updates = trie_updates.storage_tries.iter().collect::<Vec<_>>();
let mut storage_updates = trie_updates.storage_tries_ref().iter().collect::<Vec<_>>();
storage_updates.sort_unstable_by(|a, b| a.0.cmp(b.0));

num_entries += self.write_storage_trie_changesets(
Expand Down Expand Up @@ -2194,7 +2194,7 @@ impl<TX: DbTx + 'static, N: NodeTypes> TrieReader for DatabaseProvider<TX, N> {
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::<tables::AccountsTrieChangeSets>()?;
Expand All @@ -2207,8 +2207,11 @@ impl<TX: DbTx + 'static, N: NodeTypes> TrieReader for DatabaseProvider<TX, N> {
}
}

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::<Vec<_>>::default();
let mut seen_storage_keys = HashSet::new();
let mut storages_cursor = tx.cursor_dup_read::<tables::StoragesTrieChangeSets>()?;
Expand All @@ -2231,12 +2234,13 @@ impl<TX: DbTx + 'static, N: NodeTypes> TrieReader for DatabaseProvider<TX, N> {
// 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(
Expand All @@ -2254,7 +2258,7 @@ impl<TX: DbTx + 'static, N: NodeTypes> TrieReader for DatabaseProvider<TX, N> {
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::<tables::AccountsTrieChangeSets>()?;
Expand All @@ -2264,10 +2268,11 @@ impl<TX: DbTx + 'static, N: NodeTypes> TrieReader for DatabaseProvider<TX, N> {
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::<tables::StoragesTrieChangeSets>()?;
let storage_range_start = BlockNumberHashedAddress((block_number, B256::ZERO));
let storage_range_end = BlockNumberHashedAddress((block_number + 1, B256::ZERO));
Expand All @@ -2291,8 +2296,7 @@ impl<TX: DbTx + 'static, N: NodeTypes> TrieReader for DatabaseProvider<TX, N> {
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(),
Expand All @@ -2302,7 +2306,7 @@ impl<TX: DbTx + 'static, N: NodeTypes> TrieReader for DatabaseProvider<TX, N> {
.push((nibbles.0, node_value));
}

Ok(trie_updates)
Ok(TrieUpdatesSorted::new(account_nodes, storage_tries))
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This case does not require explicit sorting, as the vecs are constructed in sorted order inherently thanks to the DB

}
}

Expand Down Expand Up @@ -2379,7 +2383,7 @@ impl<TX: DbTxMut + DbTx + 'static, N: NodeTypes> 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);

Expand Down Expand Up @@ -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 =
Expand Down Expand Up @@ -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 {
Expand All @@ -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 =
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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");
Expand All @@ -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");
Expand All @@ -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");

Expand Down
27 changes: 25 additions & 2 deletions crates/trie/common/src/updates.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<BranchNodeCompact>)>,
account_nodes: Vec<(Nibbles, Option<BranchNodeCompact>)>,
/// Storage tries stored by hashed address of the account the trie belongs to.
pub storage_tries: B256Map<StorageTrieUpdatesSorted>,
storage_tries: B256Map<StorageTrieUpdatesSorted>,
}

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<BranchNodeCompact>)>,
storage_tries: B256Map<StorageTrieUpdatesSorted>,
) -> 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()
Expand Down
3 changes: 2 additions & 1 deletion crates/trie/trie/src/trie_cursor/in_memory.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<BranchNodeCompact>)> = 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));
Expand Down
Loading