From 0373c5875ac038e4ea71de3f5777f4052db76911 Mon Sep 17 00:00:00 2001 From: Roman Krasiuk Date: Thu, 4 Jul 2024 02:08:24 -0700 Subject: [PATCH] chore(trie): revamp inner in-memory trie cursor representation (#9287) --- crates/trie/trie/src/trie_cursor/in_memory.rs | 56 ++++++++++++++++--- crates/trie/trie/src/updates.rs | 15 +++-- 2 files changed, 54 insertions(+), 17 deletions(-) diff --git a/crates/trie/trie/src/trie_cursor/in_memory.rs b/crates/trie/trie/src/trie_cursor/in_memory.rs index 83ef6e5d6c..983974da38 100644 --- a/crates/trie/trie/src/trie_cursor/in_memory.rs +++ b/crates/trie/trie/src/trie_cursor/in_memory.rs @@ -1,13 +1,19 @@ use super::{TrieCursor, TrieCursorFactory}; -use crate::updates::TrieUpdatesSorted; +use crate::{ + forward_cursor::ForwardInMemoryCursor, + updates::{StorageTrieUpdatesSorted, TrieUpdatesSorted}, +}; use reth_db::DatabaseError; use reth_primitives::B256; use reth_trie_common::{BranchNodeCompact, Nibbles}; +use std::collections::HashSet; /// The trie cursor factory for the trie updates. #[derive(Debug, Clone)] pub struct InMemoryTrieCursorFactory<'a, CF> { + /// Underlying trie cursor factory. cursor_factory: CF, + /// Reference to sorted trie updates. trie_updates: &'a TrieUpdatesSorted, } @@ -32,7 +38,11 @@ impl<'a, CF: TrieCursorFactory> TrieCursorFactory for InMemoryTrieCursorFactory< hashed_address: B256, ) -> Result { let cursor = self.cursor_factory.storage_trie_cursor(hashed_address)?; - Ok(InMemoryStorageTrieCursor::new(cursor, hashed_address, self.trie_updates)) + Ok(InMemoryStorageTrieCursor::new( + hashed_address, + cursor, + self.trie_updates.storage_tries.get(&hashed_address), + )) } } @@ -41,14 +51,25 @@ impl<'a, CF: TrieCursorFactory> TrieCursorFactory for InMemoryTrieCursorFactory< #[derive(Debug)] #[allow(dead_code)] pub struct InMemoryAccountTrieCursor<'a, C> { + /// The database cursor. cursor: C, - trie_updates: &'a TrieUpdatesSorted, + /// Forward-only in-memory cursor over storage trie nodes. + in_memory_cursor: ForwardInMemoryCursor<'a, Nibbles, BranchNodeCompact>, + /// Collection of removed trie nodes. + removed_nodes: &'a HashSet, + /// Last key returned by the cursor. last_key: Option, } impl<'a, C> InMemoryAccountTrieCursor<'a, C> { const fn new(cursor: C, trie_updates: &'a TrieUpdatesSorted) -> Self { - Self { cursor, trie_updates, last_key: None } + let in_memory_cursor = ForwardInMemoryCursor::new(&trie_updates.account_nodes); + Self { + cursor, + in_memory_cursor, + removed_nodes: &trie_updates.removed_nodes, + last_key: None, + } } } @@ -77,16 +98,33 @@ impl<'a, C: TrieCursor> TrieCursor for InMemoryAccountTrieCursor<'a, C> { #[derive(Debug)] #[allow(dead_code)] pub struct InMemoryStorageTrieCursor<'a, C> { - cursor: C, - trie_update_index: usize, - trie_updates: &'a TrieUpdatesSorted, + /// The hashed address of the account that trie belongs to. hashed_address: B256, + /// The database cursor. + cursor: C, + /// Forward-only in-memory cursor over storage trie nodes. + in_memory_cursor: Option>, + /// Reference to the set of removed storage node keys. + removed_nodes: Option<&'a HashSet>, + /// The flag indicating whether the storage trie was cleared. + storage_trie_cleared: bool, + /// Last key returned by the cursor. last_key: Option, } impl<'a, C> InMemoryStorageTrieCursor<'a, C> { - const fn new(cursor: C, hashed_address: B256, trie_updates: &'a TrieUpdatesSorted) -> Self { - Self { cursor, trie_updates, trie_update_index: 0, hashed_address, last_key: None } + fn new(hashed_address: B256, cursor: C, updates: Option<&'a StorageTrieUpdatesSorted>) -> Self { + let in_memory_cursor = updates.map(|u| ForwardInMemoryCursor::new(&u.storage_nodes)); + let removed_nodes = updates.map(|u| &u.removed_nodes); + let storage_trie_cleared = updates.map_or(false, |u| u.is_deleted); + Self { + hashed_address, + cursor, + in_memory_cursor, + removed_nodes, + storage_trie_cleared, + last_key: None, + } } } diff --git a/crates/trie/trie/src/updates.rs b/crates/trie/trie/src/updates.rs index 86119a673c..eba5d1963d 100644 --- a/crates/trie/trie/src/updates.rs +++ b/crates/trie/trie/src/updates.rs @@ -77,12 +77,11 @@ impl TrieUpdates { pub fn into_sorted(self) -> TrieUpdatesSorted { let mut account_nodes = Vec::from_iter(self.account_nodes); account_nodes.sort_unstable_by(|a, b| a.0.cmp(&b.0)); - let mut storage_tries = Vec::from_iter( - self.storage_tries - .into_iter() - .map(|(hashed_address, updates)| (hashed_address, updates.into_sorted())), - ); - storage_tries.sort_unstable_by(|a, b| a.0.cmp(&b.0)); + let storage_tries = self + .storage_tries + .into_iter() + .map(|(hashed_address, updates)| (hashed_address, updates.into_sorted())) + .collect(); TrieUpdatesSorted { removed_nodes: self.removed_nodes, account_nodes, storage_tries } } @@ -296,7 +295,7 @@ impl StorageTrieUpdates { pub struct TrieUpdatesSorted { pub(crate) account_nodes: Vec<(Nibbles, BranchNodeCompact)>, pub(crate) removed_nodes: HashSet, - pub(crate) storage_tries: Vec<(B256, StorageTrieUpdatesSorted)>, + pub(crate) storage_tries: HashMap, } impl TrieUpdatesSorted { @@ -311,7 +310,7 @@ impl TrieUpdatesSorted { } /// Returns reference to updated storage tries. - pub fn storage_tries_ref(&self) -> &[(B256, StorageTrieUpdatesSorted)] { + pub const fn storage_tries_ref(&self) -> &HashMap { &self.storage_tries } }