chore(trie): revamp inner in-memory trie cursor representation

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<Self::StorageTrieCursor, DatabaseError> {
         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<Nibbles>,
+    /// Last key returned by the cursor.
     last_key: Option<Nibbles>,
 }
 
 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<ForwardInMemoryCursor<'a, Nibbles, BranchNodeCompact>>,
+    /// Reference to the set of removed storage node keys.
+    removed_nodes: Option<&'a HashSet<Nibbles>>,
+    /// The flag indicating whether the storage trie was cleared.
+    storage_trie_cleared: bool,
+    /// Last key returned by the cursor.
     last_key: Option<Nibbles>,
 }
 
 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,
+        }
     }
 }
 

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<Nibbles>,
-    pub(crate) storage_tries: Vec<(B256, StorageTrieUpdatesSorted)>,
+    pub(crate) storage_tries: HashMap<B256, StorageTrieUpdatesSorted>,
 }
 
 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<B256, StorageTrieUpdatesSorted> {
         &self.storage_tries
     }
 }