diff --git a/Cargo.lock b/Cargo.lock index 4b4fffa78f7..77c97a2191c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9628,6 +9628,7 @@ dependencies = [ "reth-revm", "reth-trie", "serde", + "tempfile", "test-case", "thiserror 2.0.16", "tokio", diff --git a/crates/optimism/trie/Cargo.toml b/crates/optimism/trie/Cargo.toml index ef6067d7da7..be6fff9782a 100644 --- a/crates/optimism/trie/Cargo.toml +++ b/crates/optimism/trie/Cargo.toml @@ -45,5 +45,5 @@ reth-codecs = { workspace = true, features = ["test-utils"] } tokio = { workspace = true, features = ["test-util", "rt-multi-thread", "macros"] } test-case.workspace = true reth-db = { workspace = true, features = ["test-utils"] } -reth-db-api = { workspace = true, features = ["test-utils"] } reth-trie = { workspace = true, features = ["test-utils"] } +tempfile.workspace = true diff --git a/crates/optimism/trie/src/api.rs b/crates/optimism/trie/src/api.rs index 89e734bf6c1..81bbe1c3246 100644 --- a/crates/optimism/trie/src/api.rs +++ b/crates/optimism/trie/src/api.rs @@ -2,6 +2,7 @@ use alloy_primitives::{map::HashMap, B256, U256}; use auto_impl::auto_impl; +use reth_db_api::DatabaseError; use reth_primitives_traits::Account; use reth_trie::{updates::TrieUpdates, BranchNodeCompact, HashedPostState, Nibbles}; use std::fmt::Debug; @@ -11,6 +12,10 @@ use thiserror::Error; #[derive(Debug, Error)] pub enum OpProofsStorageError { // TODO: add more errors once we know what they are + /// Error occurred while interacting with the database. + #[error(transparent)] + DatabaseError(#[from] DatabaseError), + /// Other error #[error("Other error: {0}")] Other(eyre::Error), diff --git a/crates/optimism/trie/src/db/models/mod.rs b/crates/optimism/trie/src/db/models/mod.rs index d5ea48e7f31..604d545ac52 100644 --- a/crates/optimism/trie/src/db/models/mod.rs +++ b/crates/optimism/trie/src/db/models/mod.rs @@ -37,7 +37,7 @@ tables! { /// Each entry is identified by a composite key combining the account’s hashed address and the /// compact-encoded trie path. Versions are tracked using block numbers as subkeys. table StorageTrieHistory { - type Key = StorageTrieSubKey; + type Key = StorageTrieKey; type Value = VersionedValue; type SubKey = u64; // block number } @@ -48,7 +48,7 @@ tables! { /// code hash, storage root). table HashedAccountHistory { type Key = B256; - type Value = VersionedValue>; + type Value = VersionedValue; type SubKey = u64; // block number } @@ -57,8 +57,8 @@ tables! { /// Each entry maps a composite key of (hashed address, storage key) to its stored value. /// Used for reconstructing contract storage at any historical block height. table HashedStorageHistory { - type Key = HashedStorageSubKey; - type Value = VersionedValue; + type Key = HashedStorageKey; + type Value = VersionedValue; type SubKey = u64; // block number } diff --git a/crates/optimism/trie/src/db/models/storage.rs b/crates/optimism/trie/src/db/models/storage.rs index 136b474b273..d573681d42c 100644 --- a/crates/optimism/trie/src/db/models/storage.rs +++ b/crates/optimism/trie/src/db/models/storage.rs @@ -1,6 +1,7 @@ -use alloy_primitives::B256; +use alloy_primitives::{B256, U256}; +use derive_more::{Constructor, From, Into}; use reth_db::{ - table::{Decode, Encode}, + table::{Compress, Decode, Decompress, Encode}, DatabaseError, }; use reth_trie::StoredNibbles; @@ -11,21 +12,21 @@ use serde::{Deserialize, Serialize}; /// Used to efficiently index storage branches by both account address and trie path. /// The encoding ensures lexicographic ordering: first by address, then by path. #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] -pub struct StorageTrieSubKey { +pub struct StorageTrieKey { /// Hashed account address pub hashed_address: B256, /// Trie path as nibbles pub path: StoredNibbles, } -impl StorageTrieSubKey { +impl StorageTrieKey { /// Create a new storage branch key pub const fn new(hashed_address: B256, path: StoredNibbles) -> Self { Self { hashed_address, path } } } -impl Encode for StorageTrieSubKey { +impl Encode for StorageTrieKey { type Encoded = Vec; fn encode(self) -> Self::Encoded { @@ -38,7 +39,7 @@ impl Encode for StorageTrieSubKey { } } -impl Decode for StorageTrieSubKey { +impl Decode for StorageTrieKey { fn decode(value: &[u8]) -> Result { if value.len() < 32 { return Err(DatabaseError::Decode); @@ -59,21 +60,21 @@ impl Decode for StorageTrieSubKey { /// Used to efficiently index storage values by both account address and storage key. /// The encoding ensures lexicographic ordering: first by address, then by storage key. #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] -pub struct HashedStorageSubKey { +pub struct HashedStorageKey { /// Hashed account address pub hashed_address: B256, /// Hashed storage key pub hashed_storage_key: B256, } -impl HashedStorageSubKey { +impl HashedStorageKey { /// Create a new hashed storage key pub const fn new(hashed_address: B256, hashed_storage_key: B256) -> Self { Self { hashed_address, hashed_storage_key } } } -impl Encode for HashedStorageSubKey { +impl Encode for HashedStorageKey { type Encoded = [u8; 64]; fn encode(self) -> Self::Encoded { @@ -86,7 +87,7 @@ impl Encode for HashedStorageSubKey { } } -impl Decode for HashedStorageSubKey { +impl Decode for HashedStorageKey { fn decode(value: &[u8]) -> Result { if value.len() != 64 { return Err(DatabaseError::Decode); @@ -99,6 +100,29 @@ impl Decode for HashedStorageSubKey { } } +/// Storage value wrapper for U256 values +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, From, Into, Constructor)] +pub struct StorageValue(pub U256); + +impl Compress for StorageValue { + type Compressed = Vec; + + fn compress_to_buf>(&self, buf: &mut B) { + let be: [u8; 32] = self.0.to_be_bytes::<32>(); + buf.put_slice(&be); + } +} + +impl Decompress for StorageValue { + fn decompress(value: &[u8]) -> Result { + if value.len() != 32 { + return Err(DatabaseError::Decode); + } + let bytes: [u8; 32] = value.try_into().map_err(|_| DatabaseError::Decode)?; + Ok(Self(U256::from_be_bytes(bytes))) + } +} + /// Proof Window key for tracking active proof window bounds /// /// Used to store earliest and latest block numbers in the external storage. @@ -138,10 +162,10 @@ mod tests { fn test_storage_branch_subkey_encode_decode() { let addr = B256::from([1u8; 32]); let path = StoredNibbles(Nibbles::from_nibbles_unchecked([1, 2, 3, 4])); - let key = StorageTrieSubKey::new(addr, path.clone()); + let key = StorageTrieKey::new(addr, path.clone()); let encoded = key.clone().encode(); - let decoded = StorageTrieSubKey::decode(&encoded).unwrap(); + let decoded = StorageTrieKey::decode(&encoded).unwrap(); assert_eq!(key, decoded); assert_eq!(decoded.hashed_address, addr); @@ -155,9 +179,9 @@ mod tests { let path1 = StoredNibbles(Nibbles::from_nibbles_unchecked([1, 2])); let path2 = StoredNibbles(Nibbles::from_nibbles_unchecked([1, 3])); - let key1 = StorageTrieSubKey::new(addr1, path1.clone()); - let key2 = StorageTrieSubKey::new(addr1, path2); - let key3 = StorageTrieSubKey::new(addr2, path1); + let key1 = StorageTrieKey::new(addr1, path1.clone()); + let key2 = StorageTrieKey::new(addr1, path2); + let key3 = StorageTrieKey::new(addr2, path1); // Encoded bytes should be sortable: first by address, then by path let enc1 = key1.encode(); @@ -173,10 +197,10 @@ mod tests { fn test_hashed_storage_subkey_encode_decode() { let addr = B256::from([1u8; 32]); let storage_key = B256::from([2u8; 32]); - let key = HashedStorageSubKey::new(addr, storage_key); + let key = HashedStorageKey::new(addr, storage_key); let encoded = key.clone().encode(); - let decoded = HashedStorageSubKey::decode(&encoded).unwrap(); + let decoded = HashedStorageKey::decode(&encoded).unwrap(); assert_eq!(key, decoded); assert_eq!(decoded.hashed_address, addr); @@ -190,9 +214,9 @@ mod tests { let storage1 = B256::from([10u8; 32]); let storage2 = B256::from([20u8; 32]); - let key1 = HashedStorageSubKey::new(addr1, storage1); - let key2 = HashedStorageSubKey::new(addr1, storage2); - let key3 = HashedStorageSubKey::new(addr2, storage1); + let key1 = HashedStorageKey::new(addr1, storage1); + let key2 = HashedStorageKey::new(addr1, storage2); + let key3 = HashedStorageKey::new(addr2, storage1); // Encoded bytes should be sortable: first by address, then by storage key let enc1 = key1.encode(); @@ -208,7 +232,7 @@ mod tests { fn test_hashed_storage_subkey_size() { let addr = B256::from([1u8; 32]); let storage_key = B256::from([2u8; 32]); - let key = HashedStorageSubKey::new(addr, storage_key); + let key = HashedStorageKey::new(addr, storage_key); let encoded = key.encode(); assert_eq!(encoded.len(), 64, "Encoded size should be exactly 64 bytes"); diff --git a/crates/optimism/trie/src/db/store.rs b/crates/optimism/trie/src/db/store.rs index 980c58e861d..da0c6cda445 100644 --- a/crates/optimism/trie/src/db/store.rs +++ b/crates/optimism/trie/src/db/store.rs @@ -1,11 +1,18 @@ use crate::{ - db::{MdbxAccountCursor, MdbxStorageCursor, MdbxTrieCursor}, + db::{ + models::{ + HashedAccountHistory, HashedStorageHistory, HashedStorageKey, MaybeDeleted, + StorageValue, VersionedValue, + }, + MdbxAccountCursor, MdbxStorageCursor, MdbxTrieCursor, + }, BlockStateDiff, OpProofsStorage, OpProofsStorageError, OpProofsStorageResult, }; use alloy_primitives::{map::HashMap, B256, U256}; use reth_db::{ + cursor::DbDupCursorRW, mdbx::{init_db_for, DatabaseArguments}, - DatabaseEnv, + Database, DatabaseEnv, }; use reth_primitives_traits::Account; use reth_trie::{BranchNodeCompact, Nibbles}; @@ -14,7 +21,7 @@ use std::path::Path; /// MDBX implementation of `OpProofsStorage`. #[derive(Debug)] pub struct MdbxProofsStorage { - _env: DatabaseEnv, + env: DatabaseEnv, } impl MdbxProofsStorage { @@ -22,7 +29,7 @@ impl MdbxProofsStorage { pub fn new(path: &Path) -> Result { let env = init_db_for::<_, super::models::Tables>(path, DatabaseArguments::default()) .map_err(OpProofsStorageError::Other)?; - Ok(Self { _env: env }) + Ok(Self { env }) } } @@ -51,19 +58,51 @@ impl OpProofsStorage for MdbxProofsStorage { async fn store_hashed_accounts( &self, - _accounts: Vec<(B256, Option)>, - _block_number: u64, + accounts: Vec<(B256, Option)>, + block_number: u64, ) -> OpProofsStorageResult<()> { - unimplemented!() + let mut accounts = accounts; + if accounts.is_empty() { + return Ok(()); + } + + // sort the accounts by key to ensure insertion is efficient + accounts.sort_by_key(|(key, _)| *key); + + self.env.update(|tx| { + let mut cursor = tx.new_cursor::()?; + for (key, account) in accounts { + let vv = VersionedValue { block_number, value: MaybeDeleted(account) }; + cursor.append_dup(key, vv)?; + } + Ok(()) + })? } async fn store_hashed_storages( &self, - _hashed_address: B256, - _storages: Vec<(B256, U256)>, - _block_number: u64, + hashed_address: B256, + storages: Vec<(B256, U256)>, + block_number: u64, ) -> OpProofsStorageResult<()> { - unimplemented!() + let mut storages = storages; + if storages.is_empty() { + return Ok(()); + } + + // sort the storages by key to ensure insertion is efficient + storages.sort_by_key(|(key, _)| *key); + + self.env.update(|tx| { + let mut cursor = tx.new_cursor::()?; + for (key, value) in storages { + let vv = + VersionedValue { block_number, value: MaybeDeleted(Some(StorageValue(value))) }; + let storage_key = HashedStorageKey::new(hashed_address, key); + cursor.append_dup(storage_key, vv)?; + } + Ok(()) + })? } async fn get_earliest_block_number(&self) -> OpProofsStorageResult> { @@ -143,3 +182,225 @@ impl OpProofsStorage for MdbxProofsStorage { unimplemented!() } } + +#[cfg(test)] +mod tests { + use super::*; + use reth_db::cursor::DbDupCursorRO; + use tempfile::TempDir; + + const B0: u64 = 0; + + #[tokio::test] + async fn store_hashed_accounts_writes_versioned_values() { + let dir = TempDir::new().unwrap(); + let store = MdbxProofsStorage::new(dir.path()).expect("env"); + + let addr = B256::from([0xAA; 32]); + let account = Account::default(); + store.store_hashed_accounts(vec![(addr, Some(account))], B0).await.expect("write accounts"); + + let tx = store.env.tx().expect("ro tx"); + let mut cur = tx.new_cursor::().expect("cursor"); + let vv = cur.seek_by_key_subkey(addr, B0).expect("seek"); + let vv = vv.expect("entry exists"); + + assert_eq!(vv.block_number, B0); + assert_eq!(vv.value.0, Some(account)); + } + + #[tokio::test] + async fn store_hashed_accounts_multiple_items_unsorted() { + let dir = TempDir::new().unwrap(); + let store = MdbxProofsStorage::new(dir.path()).expect("env"); + + // Unsorted input, mixed Some/None (deletion) + let a1 = B256::from([0x01; 32]); + let a2 = B256::from([0x02; 32]); + let a3 = B256::from([0x03; 32]); + let acc1 = Account { nonce: 2, balance: U256::from(1000u64), ..Default::default() }; + let acc3 = Account { nonce: 1, balance: U256::from(10000u64), ..Default::default() }; + + store + .store_hashed_accounts(vec![(a2, None), (a1, Some(acc1)), (a3, Some(acc3))], B0) + .await + .expect("write"); + + let tx = store.env.tx().expect("ro tx"); + let mut cur = tx.new_cursor::().expect("cursor"); + + let v1 = cur.seek_by_key_subkey(a1, B0).expect("seek a1").expect("exists a1"); + assert_eq!(v1.block_number, B0); + assert_eq!(v1.value.0, Some(acc1)); + + let v2 = cur.seek_by_key_subkey(a2, B0).expect("seek a2").expect("exists a2"); + assert_eq!(v2.block_number, B0); + assert!(v2.value.0.is_none(), "a2 is none"); + + let v3 = cur.seek_by_key_subkey(a3, B0).expect("seek a3").expect("exists a3"); + assert_eq!(v3.block_number, B0); + assert_eq!(v3.value.0, Some(acc3)); + } + + #[tokio::test] + async fn store_hashed_accounts_multiple_calls() { + let dir = TempDir::new().unwrap(); + let store = MdbxProofsStorage::new(dir.path()).expect("env"); + + // Unsorted input, mixed Some/None (deletion) + let a1 = B256::from([0x01; 32]); + let a2 = B256::from([0x02; 32]); + let a3 = B256::from([0x03; 32]); + let a4 = B256::from([0x04; 32]); + let a5 = B256::from([0x05; 32]); + let acc1 = Account { nonce: 2, balance: U256::from(1000u64), ..Default::default() }; + let acc3 = Account { nonce: 1, balance: U256::from(10000u64), ..Default::default() }; + let acc4 = Account { nonce: 5, balance: U256::from(5000u64), ..Default::default() }; + let acc5 = Account { nonce: 10, balance: U256::from(20000u64), ..Default::default() }; + + { + store + .store_hashed_accounts(vec![(a2, None), (a1, Some(acc1)), (a4, Some(acc4))], B0) + .await + .expect("write"); + + let tx = store.env.tx().expect("ro tx"); + let mut cur = tx.new_cursor::().expect("cursor"); + + let v1 = cur.seek_by_key_subkey(a1, B0).expect("seek a1").expect("exists a1"); + assert_eq!(v1.block_number, B0); + assert_eq!(v1.value.0, Some(acc1)); + + let v2 = cur.seek_by_key_subkey(a2, B0).expect("seek a2").expect("exists a2"); + assert_eq!(v2.block_number, B0); + assert!(v2.value.0.is_none(), "a2 is none"); + + let v4 = cur.seek_by_key_subkey(a4, B0).expect("seek a4").expect("exists a4"); + assert_eq!(v4.block_number, B0); + assert_eq!(v4.value.0, Some(acc4)); + } + + { + // Second call + store + .store_hashed_accounts(vec![(a5, Some(acc5)), (a3, Some(acc3))], B0) + .await + .expect("write"); + + let tx = store.env.tx().expect("ro tx"); + let mut cur = tx.new_cursor::().expect("cursor"); + + let v3 = cur.seek_by_key_subkey(a3, B0).expect("seek a3").expect("exists a3"); + assert_eq!(v3.block_number, B0); + assert_eq!(v3.value.0, Some(acc3)); + + let v5 = cur.seek_by_key_subkey(a5, B0).expect("seek a5").expect("exists a5"); + assert_eq!(v5.block_number, B0); + assert_eq!(v5.value.0, Some(acc5)); + } + } + + #[tokio::test] + async fn store_hashed_storages_writes_versioned_values() { + let dir = TempDir::new().unwrap(); + let store = MdbxProofsStorage::new(dir.path()).expect("env"); + + let addr = B256::from([0x11; 32]); + let slot = B256::from([0x22; 32]); + let val = U256::from(0x1234u64); + + store.store_hashed_storages(addr, vec![(slot, val)], B0).await.expect("write storage"); + + let tx = store.env.tx().expect("ro tx"); + let mut cur = tx.new_cursor::().expect("cursor"); + let key = HashedStorageKey::new(addr, slot); + let vv = cur.seek_by_key_subkey(key, B0).expect("seek"); + let vv = vv.expect("entry exists"); + + assert_eq!(vv.block_number, B0); + let inner = vv.value.0.as_ref().expect("Some(StorageValue)"); + assert_eq!(inner.0, val); + } + + #[tokio::test] + async fn store_hashed_storages_multiple_slots_unsorted() { + let dir = TempDir::new().unwrap(); + let store = MdbxProofsStorage::new(dir.path()).expect("env"); + + let addr = B256::from([0x11; 32]); + let s1 = B256::from([0x01; 32]); + let v1 = U256::from(1u64); + let s2 = B256::from([0x02; 32]); + let v2 = U256::from(2u64); + let s3 = B256::from([0x03; 32]); + let v3 = U256::from(3u64); + + store + .store_hashed_storages(addr, vec![(s2, v2), (s1, v1), (s3, v3)], B0) + .await + .expect("write"); + + let tx = store.env.tx().expect("ro tx"); + let mut cur = tx.new_cursor::().expect("cursor"); + + for (slot, expected) in [(s1, v1), (s2, v2), (s3, v3)] { + let key = HashedStorageKey::new(addr, slot); + let vv = cur.seek_by_key_subkey(key, B0).expect("seek").expect("exists"); + assert_eq!(vv.block_number, B0); + let inner = vv.value.0.as_ref().expect("Some(StorageValue)"); + assert_eq!(inner.0, expected); + } + } + + #[tokio::test] + async fn store_hashed_storages_multiple_calls() { + let dir = TempDir::new().unwrap(); + let store = MdbxProofsStorage::new(dir.path()).expect("env"); + + let addr = B256::from([0x11; 32]); + let s1 = B256::from([0x01; 32]); + let v1 = U256::from(1u64); + let s2 = B256::from([0x02; 32]); + let v2 = U256::from(2u64); + let s3 = B256::from([0x03; 32]); + let v3 = U256::from(3u64); + let s4 = B256::from([0x04; 32]); + let v4 = U256::from(4u64); + let s5 = B256::from([0x05; 32]); + let v5 = U256::from(5u64); + + { + store + .store_hashed_storages(addr, vec![(s2, v2), (s1, v1), (s5, v5)], B0) + .await + .expect("write"); + + let tx = store.env.tx().expect("ro tx"); + let mut cur = tx.new_cursor::().expect("cursor"); + + for (slot, expected) in [(s1, v1), (s2, v2), (s5, v5)] { + let key = HashedStorageKey::new(addr, slot); + let vv = cur.seek_by_key_subkey(key, B0).expect("seek").expect("exists"); + assert_eq!(vv.block_number, B0); + let inner = vv.value.0.as_ref().expect("Some(StorageValue)"); + assert_eq!(inner.0, expected); + } + } + + { + // Second call + store.store_hashed_storages(addr, vec![(s4, v4), (s3, v3)], B0).await.expect("write"); + + let tx = store.env.tx().expect("ro tx"); + let mut cur = tx.new_cursor::().expect("cursor"); + + for (slot, expected) in [(s4, v4), (s3, v3)] { + let key = HashedStorageKey::new(addr, slot); + let vv = cur.seek_by_key_subkey(key, B0).expect("seek").expect("exists"); + assert_eq!(vv.block_number, B0); + let inner = vv.value.0.as_ref().expect("Some(StorageValue)"); + assert_eq!(inner.0, expected); + } + } + } +}