chore: separate storage and account cursors (#229)

This PR separates storage and account cursors which makes sense because
these two cursors may have very different seek/next implementations. In
the case of in-memory, a single cursor impl can still handle both, but
for MDBX, it makes sense to implement separate cursors.

Closes #236

---------

Co-authored-by: Arun Dhyani <[email protected]>

Author: 2 people

Cargo.lock

@@ -19,6 +19,10 @@ reth-node-api.workspace = true
 reth-provider.workspace = true
 reth-chainspec.workspace = true
 
+# for ethereum types, `serde-bincode-compat` is added by `reth-storage-api`, however this does not work with `op` until
+# `reth-storage-api` is updated to support `op`, so we add it here.
+reth-optimism-primitives = { workspace = true, features = ["reth-codec", "serde-bincode-compat", "serde"] }
+
 # ethereum
 alloy-primitives.workspace = true
 

crates/optimism/exex/Cargo.toml

@@ -70,7 +70,10 @@ pub struct BlockStateDiff {
 #[auto_impl(Arc)]
 pub trait OpProofsStorage: Send + Sync + Debug {
     /// Cursor for iterating over trie branches.
-    type TrieCursor: OpProofsTrieCursor;
+    type StorageTrieCursor: OpProofsTrieCursor;
+
+    /// Cursor for iterating over account trie branches.
+    type AccountTrieCursor: OpProofsTrieCursor;
 
     /// Cursor for iterating over storage leaves.
     type StorageCursor: OpProofsHashedCursor<Value = U256>;
@@ -123,11 +126,17 @@ pub trait OpProofsStorage: Send + Sync + Debug {
     ) -> impl Future<Output = OpProofsStorageResult<Option<(u64, B256)>>> + Send;
 
     /// Get a trie cursor for the storage backend
-    fn trie_cursor(
+    fn storage_trie_cursor(
+        &self,
+        hashed_address: B256,
+        max_block_number: u64,
+    ) -> OpProofsStorageResult<Self::StorageTrieCursor>;
+
+    /// Get a trie cursor for the account backend
+    fn account_trie_cursor(
         &self,
-        hashed_address: Option<B256>,
         max_block_number: u64,
-    ) -> OpProofsStorageResult<Self::TrieCursor>;
+    ) -> OpProofsStorageResult<Self::AccountTrieCursor>;
 
     /// Get a storage cursor for the storage backend
     fn storage_hashed_cursor(

crates/optimism/trie/src/api.rs

@@ -510,7 +510,7 @@ mod tests {
         job.backfill_accounts_trie().await.unwrap();
 
         // Verify data was stored
-        let mut trie_cursor = storage.trie_cursor(None, 100).unwrap();
+        let mut trie_cursor = storage.account_trie_cursor(100).unwrap();
         let mut count = 0;
         while let Some((path, _node)) = trie_cursor.next().unwrap() {
             assert_eq!(path, nodes[count].0 .0);
@@ -568,7 +568,7 @@ mod tests {
         job.backfill_storages_trie().await.unwrap();
 
         // Verify data was stored for addr1
-        let mut trie_cursor = storage.trie_cursor(Some(addr1), 100).unwrap();
+        let mut trie_cursor = storage.storage_trie_cursor(addr1, 100).unwrap();
         let mut found = vec![];
         while let Some((path, _node)) = trie_cursor.next().unwrap() {
             found.push(path);
@@ -578,7 +578,7 @@ mod tests {
         assert_eq!(found[1], nodes[1].1.nibbles.0);
 
         // Verify data was stored for addr2
-        let mut trie_cursor = storage.trie_cursor(Some(addr2), 100).unwrap();
+        let mut trie_cursor = storage.storage_trie_cursor(addr2, 100).unwrap();
         let mut found = vec![];
         while let Some((path, _node)) = trie_cursor.next().unwrap() {
             found.push(path);
@@ -662,10 +662,10 @@ mod tests {
         let mut storage_cursor = storage.storage_hashed_cursor(addr, 100).unwrap();
         assert!(storage_cursor.next().unwrap().is_some());
 
-        let mut trie_cursor = storage.trie_cursor(None, 100).unwrap();
+        let mut trie_cursor = storage.account_trie_cursor(100).unwrap();
         assert!(trie_cursor.next().unwrap().is_some());
 
-        let mut storage_trie_cursor = storage.trie_cursor(Some(addr), 100).unwrap();
+        let mut storage_trie_cursor = storage.storage_trie_cursor(addr, 100).unwrap();
         assert!(storage_trie_cursor.next().unwrap().is_some());
     }
 

crates/optimism/trie/src/backfill.rs

@@ -27,7 +27,8 @@ impl MdbxProofsStorage {
 }
 
 impl OpProofsStorage for MdbxProofsStorage {
-    type TrieCursor = MdbxTrieCursor;
+    type StorageTrieCursor = MdbxTrieCursor;
+    type AccountTrieCursor = MdbxTrieCursor;
     type StorageCursor = MdbxStorageCursor;
     type AccountHashedCursor = MdbxAccountCursor;
 
@@ -73,11 +74,18 @@ impl OpProofsStorage for MdbxProofsStorage {
         unimplemented!()
     }
 
-    fn trie_cursor(
+    fn storage_trie_cursor(
+        &self,
+        _hashed_address: B256,
+        _max_block_number: u64,
+    ) -> OpProofsStorageResult<Self::StorageTrieCursor> {
+        unimplemented!()
+    }
+
+    fn account_trie_cursor(
         &self,
-        _hashed_address: Option<B256>,
         _max_block_number: u64,
-    ) -> OpProofsStorageResult<Self::TrieCursor> {
+    ) -> OpProofsStorageResult<Self::AccountTrieCursor> {
         unimplemented!()
     }
 

crates/optimism/trie/src/db/store.rs

@@ -380,7 +380,8 @@ impl OpProofsHashedCursor for InMemoryAccountCursor {
 }
 
 impl OpProofsStorage for InMemoryProofsStorage {
-    type TrieCursor = InMemoryTrieCursor;
+    type StorageTrieCursor = InMemoryTrieCursor;
+    type AccountTrieCursor = InMemoryTrieCursor;
     type StorageCursor = InMemoryStorageCursor;
     type AccountHashedCursor = InMemoryAccountCursor;
 
@@ -459,17 +460,28 @@ impl OpProofsStorage for InMemoryProofsStorage {
         }
     }
 
-    fn trie_cursor(
+    fn storage_trie_cursor(
         &self,
-        hashed_address: Option<B256>,
+        hashed_address: B256,
         max_block_number: u64,
-    ) -> OpProofsStorageResult<Self::TrieCursor> {
+    ) -> OpProofsStorageResult<Self::StorageTrieCursor> {
         // For synchronous methods, we need to try_read() and handle potential blocking
         let inner = self
             .inner
             .try_read()
             .map_err(|_| OpProofsStorageError::Other(eyre::eyre!("Failed to acquire read lock")))?;
-        Ok(InMemoryTrieCursor::new(&inner, hashed_address, max_block_number))
+        Ok(InMemoryTrieCursor::new(&inner, Some(hashed_address), max_block_number))
+    }
+
+    fn account_trie_cursor(
+        &self,
+        max_block_number: u64,
+    ) -> OpProofsStorageResult<Self::AccountTrieCursor> {
+        let inner = self
+            .inner
+            .try_read()
+            .map_err(|_| OpProofsStorageError::Other(eyre::eyre!("Failed to acquire read lock")))?;
+        Ok(InMemoryTrieCursor::new(&inner, None, max_block_number))
     }
 
     fn storage_hashed_cursor(