feat: refactored store

crates/optimism/trie/src/api.rs

@@ -5,7 +5,7 @@ use crate::{
     OpProofsStorageResult,
 };
 use alloy_eips::{eip1898::BlockWithParent, BlockNumHash};
-use alloy_primitives::{map::HashMap, B256, U256};
+use alloy_primitives::{B256, U256};
 use auto_impl::auto_impl;
 use derive_more::{AddAssign, Constructor};
 use reth_primitives_traits::Account;
@@ -183,11 +183,11 @@ pub trait OpProofsStore: Send + Sync + Debug {
         to: BlockWithParent,
     ) -> impl Future<Output = OpProofsStorageResult<()>> + Send;
 
-    /// Deletes all updates > `latest_common_block_number` and replaces them with the new updates.
+    /// Deletes all updates > `latest_common_block` and replaces them with the new updates.
     fn replace_updates(
         &self,
-        latest_common_block_number: u64,
-        blocks_to_add: HashMap<BlockWithParent, BlockStateDiff>,
+        latest_common_block: BlockNumHash,
+        blocks_to_add: Vec<(BlockWithParent, BlockStateDiff)>,
     ) -> impl Future<Output = OpProofsStorageResult<()>> + Send;
 
     /// Set the earliest block number and hash that has been stored

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

@@ -135,6 +135,17 @@ impl MdbxProofsStorage {
         Ok(())
     }
 
+    /// Internal helper to set latest block number hash within an existing transaction
+    fn inner_set_latest_block_number(
+        tx: &(impl DbTxMut + DbTx),
+        block_number: u64,
+        hash: B256,
+    ) -> OpProofsStorageResult<()> {
+        let mut cursor = tx.cursor_write::<ProofWindow>()?;
+        cursor.upsert(ProofWindowKey::LatestBlock, &BlockNumberHash::new(block_number, hash))?;
+        Ok(())
+    }
+
     /// Persist a batch of versioned history entries to a dup-sorted table.
     ///
     /// # Parameters
@@ -573,11 +584,7 @@ impl MdbxProofsStorage {
         change_set_cursor.append(block_number, change_set)?;
 
         // Update proof window's latest block
-        let mut proof_window_cursor = tx.new_cursor::<ProofWindow>()?;
-        proof_window_cursor.upsert(
-            ProofWindowKey::LatestBlock,
-            &BlockNumberHash::new(block_number, block_ref.block.hash),
-        )?;
+        Self::inner_set_latest_block_number(tx, block_number, block_ref.block.hash)?;
 
         Ok(WriteCounts {
             account_trie_updates_written_total: change_set.account_trie_keys.len() as u64,
@@ -967,45 +974,44 @@ impl OpProofsStore for MdbxProofsStorage {
 
             let new_latest_block =
                 BlockNumberHash::new(to.block.number.saturating_sub(1), to.parent);
-            let mut proof_window_cursor = tx.new_cursor::<ProofWindow>()?;
-            proof_window_cursor.upsert(ProofWindowKey::LatestBlock, &new_latest_block)?;
+
+            // Update proof window's Latest block
+            Self::inner_set_latest_block_number(
+                tx,
+                new_latest_block.number(),
+                *new_latest_block.hash(),
+            )?;
 
             Ok(())
         })?
     }
 
     async fn replace_updates(
         &self,
-        latest_common_block_number: u64,
-        blocks_to_add: HashMap<BlockWithParent, BlockStateDiff>,
+        latest_common_block: BlockNumHash,
+        mut blocks_to_add: Vec<(BlockWithParent, BlockStateDiff)>,
     ) -> OpProofsStorageResult<()> {
+        // Sort the vec list by block number
+        blocks_to_add.sort_unstable_by_key(|(bwp, _)| bwp.block.number);
+
         let history_to_delete = self
             .env
-            .view(|tx| self.collect_history_ranged(tx, latest_common_block_number + 1..))??;
+            .view(|tx| self.collect_history_ranged(tx, latest_common_block.number + 1..))??;
 
         self.env.update(|tx| {
-            self.delete_history_ranged(tx, latest_common_block_number + 1.., history_to_delete)?;
-
-            // Sort by block number: Hashmap does not guarantee order
-            // todo: use a sorted vec instead
-            let mut blocks_to_add_vec: Vec<(BlockWithParent, BlockStateDiff)> =
-                blocks_to_add.into_iter().collect();
-
-            blocks_to_add_vec.sort_unstable_by_key(|(bwp, _)| bwp.block.number);
-
-            // update the proof window
-            // todo: refactor to use block hash from the block to add. We need to pass the
-            // BlockNumHash type for the latest_common_block_number
-            let mut proof_window_cursor = tx.new_cursor::<ProofWindow>()?;
-            proof_window_cursor.upsert(
-                ProofWindowKey::LatestBlock,
-                &BlockNumberHash::new(
-                    latest_common_block_number,
-                    blocks_to_add_vec.first().unwrap().0.parent,
-                ),
+            // Remove the old history
+            self.delete_history_ranged(tx, latest_common_block.number + 1.., history_to_delete)?;
+
+            // Update the ProofWindow Latest Block to latest_common_block so we can perform
+            // `store_trie_updates_append_only`.
+            Self::inner_set_latest_block_number(
+                tx,
+                latest_common_block.number,
+                latest_common_block.hash,
             )?;
 
-            for (block_with_parent, diff) in blocks_to_add_vec {
+            // Apply the new history
+            for (block_with_parent, diff) in blocks_to_add {
                 self.store_trie_updates_append_only(tx, block_with_parent, diff)?;
             }
             Ok(())
@@ -3215,12 +3221,13 @@ mod tests {
         let b3p = BlockWithParent::new(b2.block.hash, NumHash::new(3, B256::random())); // 3'
         let b4p = BlockWithParent::new(b3p.block.hash, NumHash::new(4, B256::random())); // 4'
 
-        // Build blocks_to_add (HashMap). Order is not guaranteed.
-        let mut blocks_to_add = HashMap::default();
-        blocks_to_add.insert(b3p, make_diff(300)); // new value at height 3
-        blocks_to_add.insert(b4p, make_diff(400)); // new value at height 4
+        // Build blocks_to_add vec.
+        let blocks_to_add = vec![(b3p, make_diff(300)), (b4p, make_diff(400))];
 
-        store.replace_updates(2, blocks_to_add).await.expect("replace_updates succeeds");
+        store
+            .replace_updates(BlockNumHash::new(2, b2.block.hash), blocks_to_add)
+            .await
+            .expect("replace_updates succeeds");
 
         // --- Verify post-conditions ---
 

crates/optimism/trie/src/in_memory.rs

@@ -3,8 +3,8 @@
 use crate::{
     api::WriteCounts, BlockStateDiff, OpProofsStorageError, OpProofsStorageResult, OpProofsStore,
 };
-use alloy_eips::eip1898::BlockWithParent;
-use alloy_primitives::{map::HashMap, B256, U256};
+use alloy_eips::{eip1898::BlockWithParent, BlockNumHash};
+use alloy_primitives::{B256, U256};
 use reth_db::DatabaseError;
 use reth_primitives_traits::Account;
 use reth_trie::{
@@ -714,10 +714,11 @@ impl OpProofsStore for InMemoryProofsStorage {
 
     async fn replace_updates(
         &self,
-        latest_common_block_number: u64,
-        blocks_to_add: HashMap<BlockWithParent, BlockStateDiff>,
+        latest_common_block: BlockNumHash,
+        blocks_to_add: Vec<(BlockWithParent, BlockStateDiff)>,
     ) -> OpProofsStorageResult<()> {
         let mut inner = self.inner.write().await;
+        let latest_common_block_number = latest_common_block.number;
 
         // Remove all updates after latest_common_block_number
         inner.trie_updates.retain(|block, _| *block <= latest_common_block_number);

crates/optimism/trie/src/live.rs

@@ -4,8 +4,7 @@ use crate::{
     api::OperationDurations, provider::OpProofsStateProviderRef, BlockStateDiff, OpProofsStorage,
     OpProofsStorageError, OpProofsStore,
 };
-use alloy_eips::{eip1898::BlockWithParent, NumHash};
-use alloy_primitives::map::{DefaultHashBuilder, HashMap};
+use alloy_eips::{eip1898::BlockWithParent, BlockNumHash, NumHash};
 use derive_more::Constructor;
 use reth_evm::{execute::Executor, ConfigureEvm};
 use reth_primitives_traits::{AlloyBlockHeader, BlockTy, RecoveredBlock};
@@ -187,22 +186,23 @@ where
 
         let start = Instant::now();
         let mut operation_durations = OperationDurations::default();
-        let latest_common_block_number = block_updates[0].0.block.number.saturating_sub(1);
-
-        let mut block_trie_updates: HashMap<BlockWithParent, BlockStateDiff> =
-            HashMap::with_capacity_and_hasher(block_updates.len(), DefaultHashBuilder::default());
+        let first = &block_updates[0].0;
+        let latest_common_block =
+            BlockNumHash::new(first.block.number.saturating_sub(1), first.parent);
+        let mut block_trie_updates: Vec<(BlockWithParent, BlockStateDiff)> =
+            Vec::with_capacity(block_updates.len());
 
         for (block, trie_updates, hashed_state) in &block_updates {
-            block_trie_updates.insert(
+            block_trie_updates.push((
                 *block,
                 BlockStateDiff {
                     sorted_trie_updates: (**trie_updates).clone(),
                     sorted_post_state: (**hashed_state).clone(),
                 },
-            );
+            ));
         }
 
-        self.storage.replace_updates(latest_common_block_number, block_trie_updates).await?;
+        self.storage.replace_updates(latest_common_block, block_trie_updates).await?;
         let write_duration = start.elapsed();
         operation_durations.total_duration_seconds = write_duration;
         operation_durations.write_duration_seconds = write_duration;

crates/optimism/trie/src/metrics.rs

@@ -551,10 +551,10 @@ where
     #[inline]
     async fn replace_updates(
         &self,
-        latest_common_block_number: u64,
-        blocks_to_add: HashMap<BlockWithParent, BlockStateDiff>,
+        latest_common_block: BlockNumHash,
+        blocks_to_add: Vec<(BlockWithParent, BlockStateDiff)>,
     ) -> OpProofsStorageResult<()> {
-        self.storage.replace_updates(latest_common_block_number, blocks_to_add).await
+        self.storage.replace_updates(latest_common_block, blocks_to_add).await
     }
 
     #[inline]

crates/optimism/trie/tests/lib.rs

@@ -1,7 +1,7 @@
 //! Common test suite for [`OpProofsStore`] implementations.
 
-use alloy_eips::{eip1898::BlockWithParent, NumHash};
-use alloy_primitives::{map::HashMap, B256, U256};
+use alloy_eips::{eip1898::BlockWithParent, BlockNumHash, NumHash};
+use alloy_primitives::{B256, U256};
 use reth_optimism_trie::{
     db::MdbxProofsStorage, BlockStateDiff, InMemoryProofsStorage, OpProofsStorageError,
     OpProofsStore,
@@ -1580,7 +1580,7 @@ async fn test_replace_updates_applies_all_updates<S: OpProofsStore>(
     );
 
     // ========== Call replace_updates to replace blocks after 100 ==========
-    let mut blocks_to_add: HashMap<BlockWithParent, BlockStateDiff> = HashMap::default();
+    let mut blocks_to_add: Vec<(BlockWithParent, BlockStateDiff)> = Vec::default();
 
     // New data for block 101
     let new_account_addr = B256::repeat_byte(0x40);
@@ -1611,13 +1611,13 @@ async fn test_replace_updates_applies_all_updates<S: OpProofsStore>(
     new_storage.storage.insert(new_storage_slot, new_storage_value);
     new_post_state.storages.insert(new_storage_addr, new_storage);
 
-    blocks_to_add.insert(
+    blocks_to_add.push((
         block_ref_101,
         BlockStateDiff {
             sorted_trie_updates: new_trie_updates.into_sorted(),
             sorted_post_state: new_post_state.into_sorted(),
         },
-    );
+    ));
 
     // New data for block 102
     let block_102_account_addr = B256::repeat_byte(0x70);
@@ -1630,16 +1630,18 @@ async fn test_replace_updates_applies_all_updates<S: OpProofsStore>(
     let mut post_state_102 = HashedPostState::default();
     post_state_102.accounts.insert(block_102_account_addr, Some(block_102_account));
 
-    blocks_to_add.insert(
+    blocks_to_add.push((
         block_ref_102,
         BlockStateDiff {
             sorted_trie_updates: trie_updates_102.into_sorted(),
             sorted_post_state: post_state_102.into_sorted(),
         },
-    );
+    ));
 
     // Execute replace_updates
-    storage.replace_updates(100, blocks_to_add).await?;
+    storage
+        .replace_updates(BlockNumHash::new(100, block_ref_100.block.hash), blocks_to_add)
+        .await?;
     // ========== Verify that data up to block 100 still exists ==========
     let mut cursor_50 = storage.account_trie_cursor(75)?;
     assert!(