feat: latest block updates in store_trie_updates

crates/optimism/trie/src/api.rs

@@ -1,5 +1,6 @@
 //! Storage API for external storage of intermediary trie nodes.
 
+use alloy_eips::eip1898::BlockWithParent;
 use alloy_primitives::{map::HashMap, B256, U256};
 use auto_impl::auto_impl;
 use reth_db::DatabaseError;
@@ -17,6 +18,16 @@ pub enum OpProofsStorageError {
     /// Parent block number is less than earliest stored block number
     #[error("Parent block number is less than earliest stored block number")]
     UnknownParent,
+    /// Block is out of order
+    #[error("Block {block_number} is out of order (parent: {parent_block_hash}, latest stored block hash: {latest_block_hash})")]
+    OutOfOrder {
+        /// The block number being inserted
+        block_number: u64,
+        /// The parent hash of the block being inserted
+        parent_block_hash: B256,
+        /// block hash of the latest stored block
+        latest_block_hash: B256,
+    },
     /// Block update failed since parent state
     #[error("Cannot execute block updates for block {0} without parent state {1} (latest stored block number: {2})")]
     BlockUpdateFailed(u64, u64, u64),
@@ -187,7 +198,7 @@ pub trait OpProofsStore: Send + Sync + Debug {
     /// so should only happen for legacy reasons.
     fn store_trie_updates(
         &self,
-        block_number: u64,
+        block_ref: BlockWithParent,
         block_state_diff: BlockStateDiff,
     ) -> impl Future<Output = OpProofsStorageResult<()>> + Send;
 

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

@@ -10,6 +10,7 @@ use crate::{
     },
     BlockStateDiff, OpProofsStorageError, OpProofsStorageResult, OpProofsStore,
 };
+use alloy_eips::eip1898::BlockWithParent;
 use alloy_primitives::{map::HashMap, B256, U256};
 use itertools::Itertools;
 use reth_db::{
@@ -238,9 +239,10 @@ impl OpProofsStore for MdbxProofsStorage {
 
     async fn store_trie_updates(
         &self,
-        block_number: u64,
+        block_ref: BlockWithParent,
         block_state_diff: BlockStateDiff,
     ) -> OpProofsStorageResult<()> {
+        let block_number = block_ref.block.number;
         let sorted_trie_updates = block_state_diff.trie_updates.into_sorted();
         let sorted_account_nodes = sorted_trie_updates.account_nodes;
 
@@ -259,6 +261,18 @@ impl OpProofsStore for MdbxProofsStorage {
             .sorted_by_key(|(hashed_address, _)| *hashed_address)
             .collect::<Vec<_>>();
 
+        // check latest stored block is the parent of incoming block
+        // todo: move this check inside the update transaction
+        let latest_hash =
+            self.get_latest_block_number().await?.map(|(_, hash)| hash).unwrap_or(B256::ZERO);
+        if latest_hash != block_ref.parent {
+            return Err(OpProofsStorageError::OutOfOrder {
+                block_number,
+                parent_block_hash: block_ref.parent,
+                latest_block_hash: latest_hash,
+            });
+        }
+
         self.env.update(|tx| {
             let mut account_trie_cursor = tx.new_cursor::<AccountTrieHistory>()?;
             for (path, node) in sorted_account_nodes {
@@ -296,6 +310,12 @@ impl OpProofsStore for MdbxProofsStorage {
                 }
             }
 
+            // update proof window latest block
+            let mut proof_window_cursor = tx.new_cursor::<ProofWindow>()?;
+            proof_window_cursor.append(
+                ProofWindowKey::LatestBlock,
+                &BlockNumberHash::new(block_number, block_ref.block.hash),
+            )?;
             Ok(())
         })?
     }
@@ -339,6 +359,7 @@ mod tests {
         models::{AccountTrieHistory, StorageTrieHistory},
         StorageTrieKey,
     };
+    use alloy_eips::NumHash;
     use alloy_primitives::B256;
     use reth_db::{cursor::DbDupCursorRO, transaction::DbTx};
     use reth_trie::{
@@ -761,7 +782,8 @@ mod tests {
         let store = MdbxProofsStorage::new(dir.path()).expect("env");
 
         // Sample block number
-        const BLOCK: u64 = 42;
+        const BLOCK: BlockWithParent =
+            BlockWithParent::new(B256::ZERO, NumHash::new(42, B256::ZERO));
 
         // Sample addresses and keys
         let addr1 = B256::from([0x11; 32]);
@@ -829,23 +851,27 @@ mod tests {
             let mut cur = tx.new_cursor::<AccountTrieHistory>().expect("cursor");
 
             // Check first node
-            let vv1 =
-                cur.seek_by_key_subkey(account_path1.into(), BLOCK).expect("seek").expect("exists");
-            assert_eq!(vv1.block_number, BLOCK);
+            let vv1 = cur
+                .seek_by_key_subkey(account_path1.into(), BLOCK.block.number)
+                .expect("seek")
+                .expect("exists");
+            assert_eq!(vv1.block_number, BLOCK.block.number);
             assert!(vv1.value.0.is_some());
 
             // Check second node
-            let vv2 =
-                cur.seek_by_key_subkey(account_path2.into(), BLOCK).expect("seek").expect("exists");
-            assert_eq!(vv2.block_number, BLOCK);
+            let vv2 = cur
+                .seek_by_key_subkey(account_path2.into(), BLOCK.block.number)
+                .expect("seek")
+                .expect("exists");
+            assert_eq!(vv2.block_number, BLOCK.block.number);
             assert!(vv2.value.0.is_some());
 
             // Check removed node
             let vv3 = cur
-                .seek_by_key_subkey(removed_account_path.into(), BLOCK)
+                .seek_by_key_subkey(removed_account_path.into(), BLOCK.block.number)
                 .expect("seek")
                 .expect("exists");
-            assert_eq!(vv3.block_number, BLOCK);
+            assert_eq!(vv3.block_number, BLOCK.block.number);
             assert!(vv3.value.0.is_none(), "Expected node deletion");
         }
 
@@ -856,14 +882,16 @@ mod tests {
 
             // Check node for addr1
             let key1 = StorageTrieKey::new(addr1, storage_path1.into());
-            let vv1 = cur.seek_by_key_subkey(key1, BLOCK).expect("seek").expect("exists");
-            assert_eq!(vv1.block_number, BLOCK);
+            let vv1 =
+                cur.seek_by_key_subkey(key1, BLOCK.block.number).expect("seek").expect("exists");
+            assert_eq!(vv1.block_number, BLOCK.block.number);
             assert!(vv1.value.0.is_some());
 
             // Check node for addr2
             let key2 = StorageTrieKey::new(addr2, storage_path2.into());
-            let vv2 = cur.seek_by_key_subkey(key2, BLOCK).expect("seek").expect("exists");
-            assert_eq!(vv2.block_number, BLOCK);
+            let vv2 =
+                cur.seek_by_key_subkey(key2, BLOCK.block.number).expect("seek").expect("exists");
+            assert_eq!(vv2.block_number, BLOCK.block.number);
             assert!(vv2.value.0.is_some());
         }
 
@@ -873,13 +901,15 @@ mod tests {
             let mut cur = tx.new_cursor::<HashedAccountHistory>().expect("cursor");
 
             // Check account1 (exists)
-            let vv1 = cur.seek_by_key_subkey(addr1, BLOCK).expect("seek").expect("exists");
-            assert_eq!(vv1.block_number, BLOCK);
+            let vv1 =
+                cur.seek_by_key_subkey(addr1, BLOCK.block.number).expect("seek").expect("exists");
+            assert_eq!(vv1.block_number, BLOCK.block.number);
             assert_eq!(vv1.value.0, Some(acc1));
 
             // Check account2 (deletion)
-            let vv2 = cur.seek_by_key_subkey(addr2, BLOCK).expect("seek").expect("exists");
-            assert_eq!(vv2.block_number, BLOCK);
+            let vv2 =
+                cur.seek_by_key_subkey(addr2, BLOCK.block.number).expect("seek").expect("exists");
+            assert_eq!(vv2.block_number, BLOCK.block.number);
             assert!(vv2.value.0.is_none(), "Expected account deletion");
         }
 
@@ -890,26 +920,97 @@ mod tests {
 
             // Check storage for addr1
             let key1 = HashedStorageKey::new(addr1, slot1);
-            let vv1 = cur.seek_by_key_subkey(key1, BLOCK).expect("seek").expect("exists");
-            assert_eq!(vv1.block_number, BLOCK);
+            let vv1 =
+                cur.seek_by_key_subkey(key1, BLOCK.block.number).expect("seek").expect("exists");
+            assert_eq!(vv1.block_number, BLOCK.block.number);
             let inner1 = vv1.value.0.as_ref().expect("Some(StorageValue)");
             assert_eq!(inner1.0, val1);
 
             // Check storage for addr2
             let key2 = HashedStorageKey::new(addr2, slot2);
-            let vv2 = cur.seek_by_key_subkey(key2, BLOCK).expect("seek").expect("exists");
-            assert_eq!(vv2.block_number, BLOCK);
+            let vv2 =
+                cur.seek_by_key_subkey(key2, BLOCK.block.number).expect("seek").expect("exists");
+            assert_eq!(vv2.block_number, BLOCK.block.number);
             let inner2 = vv2.value.0.as_ref().expect("Some(StorageValue)");
             assert_eq!(inner2.0, val2);
         }
+
+        // check the latest block number in proof window
+        {
+            let tx = store.env.tx().expect("tx");
+            let mut proof_window_cursor = tx.new_cursor::<ProofWindow>().expect("cursor");
+            let latest_block = proof_window_cursor
+                .seek(ProofWindowKey::LatestBlock)
+                .expect("seek")
+                .expect("exists");
+            assert_eq!(latest_block.1.number(), BLOCK.block.number);
+            assert_eq!(*latest_block.1.hash(), BLOCK.block.hash);
+        }
+    }
+
+    #[tokio::test]
+    async fn store_trie_updates_out_of_order_rejects() {
+        let dir = tempfile::TempDir::new().unwrap();
+        let store = MdbxProofsStorage::new(dir.path()).expect("env");
+
+        // set latest to some hash H1
+        let existing_block = BlockWithParent::new(B256::random(), NumHash::new(1, B256::random()));
+        store
+            .set_earliest_block_number(existing_block.block.number, existing_block.block.hash)
+            .await
+            .expect("set");
+
+        // incoming block whose parent != existing latest
+        let bad_parent = B256::from([0xFF; 32]);
+        let bad_block: BlockWithParent =
+            BlockWithParent::new(bad_parent, NumHash::new(2, B256::ZERO));
+        let diff = BlockStateDiff::default();
+
+        let res = store.store_trie_updates(bad_block, diff).await;
+        assert!(matches!(res, Err(OpProofsStorageError::OutOfOrder { .. })));
+        // verify nothing written: proof window still unchanged
+        let latest = store.get_latest_block_number().await.expect("get latest");
+        assert_eq!(latest.unwrap().1, existing_block.block.hash);
+    }
+
+    #[tokio::test]
+    async fn store_trie_updates_multiple_blocks_append_versions() {
+        let dir = tempfile::TempDir::new().unwrap();
+        let store = MdbxProofsStorage::new(dir.path()).expect("env");
+
+        let addr = B256::from([0x21; 32]);
+        // block A (parent = ZERO)
+        let block_a = BlockWithParent::new(B256::ZERO, NumHash::new(1, B256::random()));
+        let mut diff_a = BlockStateDiff::default();
+        diff_a.post_state.accounts.insert(addr, Some(Account::default()));
+
+        store.store_trie_updates(block_a, diff_a).await.expect("store A");
+
+        // block B (parent = hash of A)
+        let block_b = BlockWithParent::new(block_a.block.hash, NumHash::new(2, B256::random()));
+        let mut diff_b = BlockStateDiff::default();
+        diff_b.post_state.accounts.insert(addr, Some(Account { nonce: 5, ..Default::default() }));
+
+        store.store_trie_updates(block_b, diff_b).await.expect("store B");
+
+        // verify we can retrieve entries for both block numbers
+        let tx = store.env.tx().expect("tx");
+        let mut cur = tx.new_cursor::<HashedAccountHistory>().expect("cursor");
+        let v_a =
+            cur.seek_by_key_subkey(addr, block_a.block.number).expect("seek").expect("exists");
+        let v_b =
+            cur.seek_by_key_subkey(addr, block_b.block.number).expect("seek").expect("exists");
+        assert_eq!(v_a.block_number, block_a.block.number);
+        assert_eq!(v_b.block_number, block_b.block.number);
     }
 
     #[tokio::test]
     async fn test_store_trie_updates_empty_collections() {
         let dir = TempDir::new().unwrap();
         let store = MdbxProofsStorage::new(dir.path()).expect("env");
 
-        const BLOCK: u64 = 42;
+        const BLOCK: BlockWithParent =
+            BlockWithParent::new(B256::ZERO, NumHash::new(42, B256::ZERO));
 
         // Create BlockStateDiff with empty collections
         let block_state_diff = BlockStateDiff::default();

crates/optimism/trie/src/in_memory.rs

@@ -4,6 +4,7 @@ use crate::{
     BlockStateDiff, OpProofsHashedCursorRO, OpProofsStorageError, OpProofsStorageResult,
     OpProofsStore, OpProofsTrieCursorRO,
 };
+use alloy_eips::eip1898::BlockWithParent;
 use alloy_primitives::{map::HashMap, B256, U256};
 use reth_primitives_traits::Account;
 use reth_trie::{updates::TrieUpdates, BranchNodeCompact, HashedPostState, Nibbles};
@@ -506,12 +507,12 @@ impl OpProofsStore for InMemoryProofsStorage {
 
     async fn store_trie_updates(
         &self,
-        block_number: u64,
+        block_ref: BlockWithParent,
         block_state_diff: BlockStateDiff,
     ) -> OpProofsStorageResult<()> {
         let mut inner = self.inner.write().await;
 
-        inner.store_trie_updates(block_number, block_state_diff);
+        inner.store_trie_updates(block_ref.block.number, block_state_diff);
 
         Ok(())
     }
@@ -628,6 +629,7 @@ impl OpProofsStore for InMemoryProofsStorage {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use alloy_eips::NumHash;
     use alloy_primitives::U256;
     use reth_primitives_traits::Account;
 
@@ -662,9 +664,11 @@ mod tests {
         let block_state_diff =
             BlockStateDiff { trie_updates: trie_updates.clone(), post_state: post_state.clone() };
 
-        storage.store_trie_updates(5, block_state_diff).await?;
+        const BLOCK: BlockWithParent =
+            BlockWithParent::new(B256::ZERO, NumHash::new(5, B256::ZERO));
+        storage.store_trie_updates(BLOCK, block_state_diff).await?;
 
-        let retrieved_diff = storage.fetch_trie_updates(5).await?;
+        let retrieved_diff = storage.fetch_trie_updates(BLOCK.block.number).await?;
         assert_eq!(retrieved_diff.trie_updates, trie_updates);
         assert_eq!(retrieved_diff.post_state, post_state);
 

crates/optimism/trie/src/live.rs

@@ -5,6 +5,7 @@ use crate::{
     provider::OpProofsStateProviderRef,
     OpProofsStorage,
 };
+use alloy_eips::{eip1898::BlockWithParent, NumHash};
 use derive_more::Constructor;
 use reth_evm::{execute::Executor, ConfigureEvm};
 use reth_primitives_traits::{AlloyBlockHeader, BlockTy, RecoveredBlock};
@@ -14,7 +15,7 @@ use reth_provider::{
 };
 use reth_revm::database::StateProviderDatabase;
 use std::time::Instant;
-use tracing::debug;
+use tracing::info;
 
 /// Live trie collector for external proofs storage.
 #[derive(Debug, Constructor)]
@@ -64,7 +65,8 @@ where
             .into());
         }
 
-        let block_number = block.number();
+        let block_ref =
+            BlockWithParent::new(block.parent_hash(), NumHash::new(block.number(), block.hash()));
 
         // TODO: should we check block hash here?
 
@@ -101,19 +103,24 @@ where
 
         self.storage
             .store_trie_updates(
-                block_number,
+                block_ref,
                 BlockStateDiff { trie_updates, post_state: hashed_state },
             )
             .await?;
 
         let write_trie_updates_duration = start.elapsed() - calculate_state_root_duration;
-
-        debug!("execute_and_store_block_updates duration: {:?}", start.elapsed());
-        debug!("- fetch_block_duration: {:?}", fetch_block_duration);
-        debug!("- init_provider_duration: {:?}", init_provider_duration);
-        debug!("- execute_block_duration: {:?}", execute_block_duration);
-        debug!("- calculate_state_root_duration: {:?}", calculate_state_root_duration);
-        debug!("- write_trie_updates_duration: {:?}", write_trie_updates_duration);
+        let execute_and_store_total_duration = start.elapsed();
+
+        info!(
+            block_number = block.number(),
+            ?execute_and_store_total_duration,
+            ?fetch_block_duration,
+            ?init_provider_duration,
+            ?execute_block_duration,
+            ?calculate_state_root_duration,
+            ?write_trie_updates_duration,
+            "Stored trie updates",
+        );
 
         Ok(())
     }