diff --git a/crates/optimism/trie/src/api.rs b/crates/optimism/trie/src/api.rs index f9ce8bdc039..9e3ed55ebd9 100644 --- a/crates/optimism/trie/src/api.rs +++ b/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> + 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, + latest_common_block: BlockNumHash, + blocks_to_add: Vec<(BlockWithParent, BlockStateDiff)>, ) -> impl Future> + Send; /// Set the earliest block number and hash that has been stored diff --git a/crates/optimism/trie/src/db/store.rs b/crates/optimism/trie/src/db/store.rs index e7e481f11fc..91c0425fb58 100644 --- a/crates/optimism/trie/src/db/store.rs +++ b/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::()?; + 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::()?; - 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,8 +974,13 @@ 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::()?; - 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(()) })? @@ -976,36 +988,30 @@ impl OpProofsStore for MdbxProofsStorage { async fn replace_updates( &self, - latest_common_block_number: u64, - blocks_to_add: HashMap, + 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::()?; - 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 --- diff --git a/crates/optimism/trie/src/in_memory.rs b/crates/optimism/trie/src/in_memory.rs index aa059e0e3db..ed269180255 100644 --- a/crates/optimism/trie/src/in_memory.rs +++ b/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, + 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); diff --git a/crates/optimism/trie/src/live.rs b/crates/optimism/trie/src/live.rs index e213dd76fc6..7f404c9fdf6 100644 --- a/crates/optimism/trie/src/live.rs +++ b/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 = - 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; diff --git a/crates/optimism/trie/src/metrics.rs b/crates/optimism/trie/src/metrics.rs index 4251d4d5876..cb9e4a9264e 100644 --- a/crates/optimism/trie/src/metrics.rs +++ b/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, + 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] diff --git a/crates/optimism/trie/tests/lib.rs b/crates/optimism/trie/tests/lib.rs index 7bd13ccf582..38ecc84c974 100644 --- a/crates/optimism/trie/tests/lib.rs +++ b/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( ); // ========== Call replace_updates to replace blocks after 100 ========== - let mut blocks_to_add: HashMap = 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( 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( 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!(