diff --git a/crates/optimism/trie/src/db/store.rs b/crates/optimism/trie/src/db/store.rs index da0c6cda445..c37e314acc2 100644 --- a/crates/optimism/trie/src/db/store.rs +++ b/crates/optimism/trie/src/db/store.rs @@ -1,8 +1,8 @@ use crate::{ db::{ models::{ - HashedAccountHistory, HashedStorageHistory, HashedStorageKey, MaybeDeleted, - StorageValue, VersionedValue, + AccountTrieHistory, HashedAccountHistory, HashedStorageHistory, HashedStorageKey, + MaybeDeleted, StorageTrieHistory, StorageTrieKey, StorageValue, VersionedValue, }, MdbxAccountCursor, MdbxStorageCursor, MdbxTrieCursor, }, @@ -15,7 +15,7 @@ use reth_db::{ Database, DatabaseEnv, }; use reth_primitives_traits::Account; -use reth_trie::{BranchNodeCompact, Nibbles}; +use reth_trie::{BranchNodeCompact, Nibbles, StoredNibbles}; use std::path::Path; /// MDBX implementation of `OpProofsStorage`. @@ -41,19 +41,48 @@ impl OpProofsStorage for MdbxProofsStorage { async fn store_account_branches( &self, - _block_number: u64, - _updates: Vec<(Nibbles, Option)>, + block_number: u64, + updates: Vec<(Nibbles, Option)>, ) -> OpProofsStorageResult<()> { - unimplemented!() + let mut updates = updates; + if updates.is_empty() { + return Ok(()); + } + + updates.sort_by_key(|(key, _)| *key); + + self.env.update(|tx| { + let mut cursor = tx.new_cursor::()?; + for (nibble, branch_node) in updates { + let vv = VersionedValue { block_number, value: MaybeDeleted(branch_node) }; + cursor.append_dup(StoredNibbles::from(nibble), vv)?; + } + Ok(()) + })? } async fn store_storage_branches( &self, - _block_number: u64, - _hashed_address: B256, - _items: Vec<(Nibbles, Option)>, + block_number: u64, + hashed_address: B256, + items: Vec<(Nibbles, Option)>, ) -> OpProofsStorageResult<()> { - unimplemented!() + let mut items = items; + if items.is_empty() { + return Ok(()); + } + + items.sort_by_key(|(key, _)| *key); + + self.env.update(|tx| { + let mut cursor = tx.new_cursor::()?; + for (nibble, branch_node) in items { + let key = StorageTrieKey::new(hashed_address, StoredNibbles::from(nibble)); + let vv = VersionedValue { block_number, value: MaybeDeleted(branch_node) }; + cursor.append_dup(key, vv)?; + } + Ok(()) + })? } async fn store_hashed_accounts( @@ -186,7 +215,13 @@ impl OpProofsStorage for MdbxProofsStorage { #[cfg(test)] mod tests { use super::*; - use reth_db::cursor::DbDupCursorRO; + use crate::db::{ + models::{AccountTrieHistory, StorageTrieHistory}, + StorageTrieKey, + }; + use alloy_primitives::B256; + use reth_db::{cursor::DbDupCursorRO, transaction::DbTx}; + use reth_trie::{BranchNodeCompact, Nibbles, StoredNibbles}; use tempfile::TempDir; const B0: u64 = 0; @@ -403,4 +438,201 @@ mod tests { } } } + + #[tokio::test] + async fn test_store_account_branches_writes_versioned_values() { + let dir = TempDir::new().unwrap(); + let store = MdbxProofsStorage::new(dir.path()).expect("env"); + + let nibble = Nibbles::from_nibbles_unchecked([0x12, 0x34]); + let branch_node = BranchNodeCompact::new(0b1, 0, 0, vec![], Some(B256::random())); + let updates = vec![(nibble, Some(branch_node.clone()))]; + + store.store_account_branches(B0, updates).await.expect("write"); + + let tx = store.env.tx().expect("ro tx"); + let mut cur = tx.cursor_dup_read::().expect("cursor"); + + let vv = cur + .seek_by_key_subkey(StoredNibbles::from(nibble), B0) + .expect("seek") + .expect("entry exists"); + + assert_eq!(vv.block_number, B0); + assert_eq!(vv.value.0, Some(branch_node)); + } + + #[tokio::test] + async fn test_store_account_branches_multiple_items_unsorted() { + let dir = TempDir::new().unwrap(); + let store = MdbxProofsStorage::new(dir.path()).expect("env"); + + let n1 = Nibbles::from_nibbles_unchecked([0x01]); + let b1 = BranchNodeCompact::new(0b1, 0, 0, vec![], Some(B256::random())); + let n2 = Nibbles::from_nibbles_unchecked([0x02]); + let n3 = Nibbles::from_nibbles_unchecked([0x03]); + let b3 = BranchNodeCompact::new(0b1, 0, 0, vec![], Some(B256::random())); + + let updates = vec![(n2, None), (n1, Some(b1.clone())), (n3, Some(b3.clone()))]; + store.store_account_branches(B0, updates.clone()).await.expect("write"); + + let tx = store.env.tx().expect("ro tx"); + let mut cur = tx.cursor_dup_read::().expect("cursor"); + + for (nibble, branch) in updates { + let v = cur + .seek_by_key_subkey(StoredNibbles::from(nibble), B0) + .expect("seek") + .expect("exists"); + assert_eq!(v.block_number, B0); + assert_eq!(v.value.0, branch); + } + } + + #[tokio::test] + async fn store_account_branches_multiple_calls() { + let dir = TempDir::new().unwrap(); + let store = MdbxProofsStorage::new(dir.path()).expect("env"); + + let n1 = Nibbles::from_nibbles_unchecked([0x01]); + let b1 = BranchNodeCompact::new(0b1, 0, 0, vec![], Some(B256::random())); + let n2 = Nibbles::from_nibbles_unchecked([0x02]); + let n3 = Nibbles::from_nibbles_unchecked([0x03]); + let b3 = BranchNodeCompact::new(0b1, 0, 0, vec![], Some(B256::random())); + let n4 = Nibbles::from_nibbles_unchecked([0x04]); + let b4 = BranchNodeCompact::new(0b1, 0, 0, vec![], Some(B256::random())); + let n5 = Nibbles::from_nibbles_unchecked([0x05]); + let b5 = BranchNodeCompact::new(0b1, 0, 0, vec![], Some(B256::random())); + + { + let updates1 = vec![(n2, None), (n1, Some(b1.clone())), (n4, Some(b4.clone()))]; + store.store_account_branches(B0, updates1.clone()).await.expect("write"); + + let tx = store.env.tx().expect("ro tx"); + let mut cur = tx.cursor_dup_read::().expect("cursor"); + + for (nibble, branch) in updates1 { + let v = cur + .seek_by_key_subkey(StoredNibbles::from(nibble), B0) + .expect("seek") + .expect("exists"); + assert_eq!(v.block_number, B0); + assert_eq!(v.value.0, branch); + } + } + + { + // Second call + let updates2 = vec![(n5, Some(b5.clone())), (n3, Some(b3.clone()))]; + store.store_account_branches(B0, updates2.clone()).await.expect("write"); + + let tx = store.env.tx().expect("ro tx"); + let mut cur = tx.cursor_dup_read::().expect("cursor"); + + for (nibble, branch) in updates2 { + let v = cur + .seek_by_key_subkey(StoredNibbles::from(nibble), B0) + .expect("seek") + .expect("exists"); + assert_eq!(v.block_number, B0); + assert_eq!(v.value.0, branch); + } + } + } + + #[tokio::test] + async fn test_store_storage_branches_writes_versioned_values() { + let dir = TempDir::new().unwrap(); + let store = MdbxProofsStorage::new(dir.path()).expect("env"); + + let hashed_address = B256::random(); + let nibble = Nibbles::from_nibbles_unchecked([0x12, 0x34]); + let branch_node = BranchNodeCompact::new(0b1, 0, 0, vec![], Some(B256::random())); + let items = vec![(nibble, Some(branch_node.clone()))]; + + store.store_storage_branches(B0, hashed_address, items).await.expect("write"); + + let tx = store.env.tx().expect("ro tx"); + let mut cur = tx.cursor_dup_read::().expect("cursor"); + + let key = StorageTrieKey::new(hashed_address, StoredNibbles::from(nibble)); + let vv = cur.seek_by_key_subkey(key, B0).expect("seek").expect("entry exists"); + + assert_eq!(vv.block_number, B0); + assert_eq!(vv.value.0, Some(branch_node)); + } + + #[tokio::test] + async fn store_storage_branches_multiple_items_unsorted() { + let dir = TempDir::new().unwrap(); + let store = MdbxProofsStorage::new(dir.path()).expect("env"); + + let hashed_address = B256::random(); + let n1 = Nibbles::from_nibbles_unchecked([0x01]); + let b1 = BranchNodeCompact::new(0b1, 0, 0, vec![], Some(B256::random())); + let n2 = Nibbles::from_nibbles_unchecked([0x02]); + let n3 = Nibbles::from_nibbles_unchecked([0x03]); + let b3 = BranchNodeCompact::new(0b1, 0, 0, vec![], Some(B256::random())); + + let items = vec![(n2, None), (n1, Some(b1.clone())), (n3, Some(b3.clone()))]; + store.store_storage_branches(B0, hashed_address, items.clone()).await.expect("write"); + + let tx = store.env.tx().expect("ro tx"); + let mut cur = tx.cursor_dup_read::().expect("cursor"); + + for (nibble, branch) in items { + let key = StorageTrieKey::new(hashed_address, StoredNibbles::from(nibble)); + let v = cur.seek_by_key_subkey(key, B0).expect("seek").expect("exists"); + assert_eq!(v.block_number, B0); + assert_eq!(v.value.0, branch); + } + } + + #[tokio::test] + async fn store_storage_branches_multiple_calls() { + let dir = TempDir::new().unwrap(); + let store = MdbxProofsStorage::new(dir.path()).expect("env"); + + let hashed_address = B256::random(); + let n1 = Nibbles::from_nibbles_unchecked([0x01]); + let b1 = BranchNodeCompact::new(0b1, 0, 0, vec![], Some(B256::random())); + let n2 = Nibbles::from_nibbles_unchecked([0x02]); + let n3 = Nibbles::from_nibbles_unchecked([0x03]); + let b3 = BranchNodeCompact::new(0b1, 0, 0, vec![], Some(B256::random())); + let n4 = Nibbles::from_nibbles_unchecked([0x04]); + let b4 = BranchNodeCompact::new(0b1, 0, 0, vec![], Some(B256::random())); + let n5 = Nibbles::from_nibbles_unchecked([0x05]); + let b5 = BranchNodeCompact::new(0b1, 0, 0, vec![], Some(B256::random())); + + { + let items1 = vec![(n2, None), (n1, Some(b1.clone())), (n5, Some(b5.clone()))]; + store.store_storage_branches(B0, hashed_address, items1.clone()).await.expect("write"); + + let tx = store.env.tx().expect("ro tx"); + let mut cur = tx.cursor_dup_read::().expect("cursor"); + + for (nibble, branch) in items1 { + let key = StorageTrieKey::new(hashed_address, StoredNibbles::from(nibble)); + let v = cur.seek_by_key_subkey(key, B0).expect("seek").expect("exists"); + assert_eq!(v.block_number, B0); + assert_eq!(v.value.0, branch); + } + } + + { + // Second call + let items2 = vec![(n4, Some(b4.clone())), (n3, Some(b3.clone()))]; + store.store_storage_branches(B0, hashed_address, items2.clone()).await.expect("write"); + + let tx = store.env.tx().expect("ro tx"); + let mut cur = tx.cursor_dup_read::().expect("cursor"); + + for (nibble, branch) in items2 { + let key = StorageTrieKey::new(hashed_address, StoredNibbles::from(nibble)); + let v = cur.seek_by_key_subkey(key, B0).expect("seek").expect("exists"); + assert_eq!(v.block_number, B0); + assert_eq!(v.value.0, branch); + } + } + } }