diff --git a/Cargo.lock b/Cargo.lock index 10c3d8dfc97..c2605e4caa2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9617,6 +9617,7 @@ dependencies = [ "bytes", "derive_more", "eyre", + "itertools 0.14.0", "metrics", "reth-codecs", "reth-db", diff --git a/crates/optimism/trie/Cargo.toml b/crates/optimism/trie/Cargo.toml index d73ea537dfa..e4a51cb9e28 100644 --- a/crates/optimism/trie/Cargo.toml +++ b/crates/optimism/trie/Cargo.toml @@ -44,6 +44,7 @@ eyre.workspace = true strum.workspace = true tracing.workspace = true derive_more.workspace = true +itertools.workspace = true [dev-dependencies] reth-codecs = { workspace = true, features = ["test-utils"] } diff --git a/crates/optimism/trie/src/api.rs b/crates/optimism/trie/src/api.rs index 14ff6a7a27b..cdb7e692aab 100644 --- a/crates/optimism/trie/src/api.rs +++ b/crates/optimism/trie/src/api.rs @@ -76,7 +76,7 @@ pub trait OpProofsHashedCursor: Send + Sync { } /// Diff of trie updates and post state for a block. -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Default)] pub struct BlockStateDiff { /// Trie updates for branch nodes pub trie_updates: TrieUpdates, diff --git a/crates/optimism/trie/src/db/store.rs b/crates/optimism/trie/src/db/store.rs index 5f267b38d39..32b2b6a94c8 100644 --- a/crates/optimism/trie/src/db/store.rs +++ b/crates/optimism/trie/src/db/store.rs @@ -9,6 +9,7 @@ use crate::{ BlockStateDiff, OpProofsStorage, OpProofsStorageError, OpProofsStorageResult, }; use alloy_primitives::{map::HashMap, B256, U256}; +use itertools::Itertools; use reth_db::{ cursor::{DbCursorRO, DbCursorRW, DbDupCursorRW}, mdbx::{init_db_for, DatabaseArguments}, @@ -204,10 +205,66 @@ impl OpProofsStorage for MdbxProofsStorage { async fn store_trie_updates( &self, - _block_number: u64, - _block_state_diff: BlockStateDiff, + block_number: u64, + block_state_diff: BlockStateDiff, ) -> OpProofsStorageResult<()> { - unimplemented!() + let sorted_trie_updates = block_state_diff.trie_updates.into_sorted(); + let sorted_account_nodes = sorted_trie_updates.account_nodes; + + let sorted_storage_nodes = sorted_trie_updates + .storage_tries + .into_iter() + .sorted_by_key(|(hashed_address, _)| *hashed_address) + .collect::>(); + + let sorted_post_state = block_state_diff.post_state.into_sorted(); + let sorted_accounts = sorted_post_state.accounts().accounts_sorted(); + + let sorted_storage = sorted_post_state + .account_storages() + .iter() + .sorted_by_key(|(hashed_address, _)| *hashed_address) + .collect::>(); + + self.env.update(|tx| { + let mut account_trie_cursor = tx.new_cursor::()?; + for (path, node) in sorted_account_nodes { + let vv = VersionedValue { block_number, value: MaybeDeleted(node) }; + account_trie_cursor.append_dup(path.into(), vv)?; + } + + let mut storage_trie_cursor = tx.new_cursor::()?; + for (hashed_address, nodes) in sorted_storage_nodes { + // todo: handle is_deleted scenario + for (path, node) in nodes.storage_nodes { + let key = StorageTrieKey::new(hashed_address, path.into()); + let vv = VersionedValue { block_number, value: MaybeDeleted(node) }; + storage_trie_cursor.append_dup(key, vv)?; + } + } + + let mut account_cursor = tx.new_cursor::()?; + for (hashed_address, account) in sorted_accounts { + let vv = VersionedValue { block_number, value: MaybeDeleted(account) }; + account_cursor.append_dup(hashed_address, vv)?; + } + + let mut storage_cursor = tx.new_cursor::()?; + for (hashed_address, storage) in sorted_storage { + // todo: handle wiped storage scenario + let storage_items = storage.storage_slots_sorted().collect::>(); + for (storage_key, storage_value) in storage_items { + let vv = VersionedValue { + block_number, + value: MaybeDeleted(Some(StorageValue(storage_value))), + }; + let key = HashedStorageKey::new(*hashed_address, storage_key); + storage_cursor.append_dup(key, vv)?; + } + } + + Ok(()) + })? } async fn fetch_trie_updates( @@ -251,7 +308,9 @@ mod tests { }; use alloy_primitives::B256; use reth_db::{cursor::DbDupCursorRO, transaction::DbTx}; - use reth_trie::{BranchNodeCompact, Nibbles, StoredNibbles}; + use reth_trie::{ + updates::StorageTrieUpdates, BranchNodeCompact, HashedStorage, Nibbles, StoredNibbles, + }; use tempfile::TempDir; const B0: u64 = 0; @@ -663,6 +722,184 @@ mod tests { } } + #[tokio::test] + async fn test_store_trie_updates_comprehensive() { + let dir = TempDir::new().unwrap(); + let store = MdbxProofsStorage::new(dir.path()).expect("env"); + + // Sample block number + const BLOCK: u64 = 42; + + // Sample addresses and keys + let addr1 = B256::from([0x11; 32]); + let addr2 = B256::from([0x22; 32]); + let slot1 = B256::from([0xA1; 32]); + let slot2 = B256::from([0xA2; 32]); + + // Sample accounts + let acc1 = Account { nonce: 1, balance: U256::from(100), ..Default::default() }; + + // Sample storage values + let val1 = U256::from(1234u64); + let val2 = U256::from(5678u64); + + // Sample trie paths + let account_path1 = Nibbles::from_nibbles_unchecked(vec![0, 1, 2, 3]); + let account_path2 = Nibbles::from_nibbles_unchecked(vec![4, 5, 6, 7]); + let removed_account_path = Nibbles::from_nibbles_unchecked(vec![7, 8, 9]); + + let account_node1 = BranchNodeCompact::default(); + let account_node2 = BranchNodeCompact::default(); + + let storage_path1 = Nibbles::from_nibbles_unchecked(vec![1, 2, 3, 4]); + let storage_path2 = Nibbles::from_nibbles_unchecked(vec![8, 9, 0, 1]); + + let storage_node1 = BranchNodeCompact::default(); + let storage_node2 = BranchNodeCompact::default(); + + // Construct test BlockStateDiff + let mut block_state_diff = BlockStateDiff::default(); + + // Add account trie nodes + block_state_diff.trie_updates.account_nodes.insert(account_path1, account_node1.clone()); + block_state_diff.trie_updates.account_nodes.insert(account_path2, account_node2.clone()); + block_state_diff.trie_updates.removed_nodes.insert(removed_account_path); + + // Add storage trie nodes for two addresses + let mut storage_nodes1 = StorageTrieUpdates::default(); + storage_nodes1.storage_nodes.insert(storage_path1, storage_node1.clone()); + block_state_diff.trie_updates.storage_tries.insert(addr1, storage_nodes1); + + let mut storage_nodes2 = StorageTrieUpdates::default(); + storage_nodes2.storage_nodes.insert(storage_path2, storage_node2.clone()); + block_state_diff.trie_updates.storage_tries.insert(addr2, storage_nodes2); + + // Add hashed accounts (one Some, one None) + block_state_diff.post_state.accounts.insert(addr1, Some(acc1)); + block_state_diff.post_state.accounts.insert(addr2, None); // Deletion + + // Add storage slots for both addresses + let mut storage1 = HashedStorage::default(); + storage1.storage.insert(slot1, val1); + block_state_diff.post_state.storages.insert(addr1, storage1); + + let mut storage2 = HashedStorage::default(); + storage2.storage.insert(slot2, val2); + block_state_diff.post_state.storages.insert(addr2, storage2); + + // Store everything + store.store_trie_updates(BLOCK, block_state_diff).await.expect("store"); + + // Verify account trie nodes + { + let tx = store.env.tx().expect("tx"); + let mut cur = tx.new_cursor::().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); + 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); + assert!(vv2.value.0.is_some()); + + // Check removed node + let vv3 = cur + .seek_by_key_subkey(removed_account_path.into(), BLOCK) + .expect("seek") + .expect("exists"); + assert_eq!(vv3.block_number, BLOCK); + assert!(vv3.value.0.is_none(), "Expected node deletion"); + } + + // Verify storage trie nodes + { + let tx = store.env.tx().expect("tx"); + let mut cur = tx.new_cursor::().expect("cursor"); + + // 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); + 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); + assert!(vv2.value.0.is_some()); + } + + // Verify hashed accounts + { + let tx = store.env.tx().expect("tx"); + let mut cur = tx.new_cursor::().expect("cursor"); + + // Check account1 (exists) + let vv1 = cur.seek_by_key_subkey(addr1, BLOCK).expect("seek").expect("exists"); + assert_eq!(vv1.block_number, BLOCK); + 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); + assert!(vv2.value.0.is_none(), "Expected account deletion"); + } + + // Verify hashed storages + { + let tx = store.env.tx().expect("tx"); + let mut cur = tx.new_cursor::().expect("cursor"); + + // 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 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 inner2 = vv2.value.0.as_ref().expect("Some(StorageValue)"); + assert_eq!(inner2.0, val2); + } + } + + #[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; + + // Create BlockStateDiff with empty collections + let block_state_diff = BlockStateDiff::default(); + + // This should work without errors + store.store_trie_updates(BLOCK, block_state_diff).await.expect("store"); + + // Verify nothing was written (should be empty) + let tx = store.env.tx().expect("tx"); + + let mut cur1 = tx.new_cursor::().expect("cursor"); + assert!(cur1.next_dup_val().expect("first").is_none(), "Account trie should be empty"); + + let mut cur2 = tx.new_cursor::().expect("cursor"); + assert!(cur2.next_dup_val().expect("first").is_none(), "Storage trie should be empty"); + + let mut cur3 = tx.new_cursor::().expect("cursor"); + assert!(cur3.next_dup_val().expect("first").is_none(), "Hashed accounts should be empty"); + + let mut cur4 = tx.new_cursor::().expect("cursor"); + assert!(cur4.next_dup_val().expect("first").is_none(), "Hashed storage should be empty"); + } + #[tokio::test] async fn test_proof_window() { let dir = TempDir::new().unwrap();