diff --git a/crates/common/types/account.rs b/crates/common/types/account.rs index 26c7db96ea6..d211a8f3745 100644 --- a/crates/common/types/account.rs +++ b/crates/common/types/account.rs @@ -1,6 +1,6 @@ use std::collections::HashMap; -use bytes::Bytes; +use bytes::{BufMut, Bytes}; use ethereum_types::{H256, U256}; use ethrex_crypto::keccak::keccak_hash; use ethrex_trie::Trie; @@ -117,7 +117,7 @@ pub struct AccountInfo { pub nonce: u64, } -#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)] pub struct AccountState { pub nonce: u64, pub balance: U256, @@ -125,6 +125,16 @@ pub struct AccountState { pub code_hash: H256, } +/// A slim codec for an [`AccountState`]. +/// +/// The slim codec will optimize both the [storage root](AccountState::storage_root) and the +/// [code hash](AccountState::code_hash)'s encoding so that it does not take space when empty. +/// +/// The correct way to use it is to wrap the [`AccountState`] and encode it using this codec, and +/// not to store the codec as a field in a struct. +#[derive(Clone, Copy, Debug, Default, PartialEq)] +pub struct AccountStateSlimCodec(pub AccountState); + impl Default for AccountInfo { fn default() -> Self { Self { @@ -232,6 +242,101 @@ impl RLPDecode for AccountState { } } +impl RLPEncode for AccountStateSlimCodec { + fn encode(&self, buf: &mut dyn BufMut) { + struct StorageRootCodec<'a>(&'a H256); + impl RLPEncode for StorageRootCodec<'_> { + fn encode(&self, buf: &mut dyn BufMut) { + let data = if *self.0 != *EMPTY_TRIE_HASH { + self.0.as_bytes() + } else { + &[] + }; + + data.encode(buf); + } + } + + struct CodeHashCodec<'a>(&'a H256); + impl RLPEncode for CodeHashCodec<'_> { + fn encode(&self, buf: &mut dyn BufMut) { + let data = if *self.0 != *EMPTY_KECCACK_HASH { + self.0.as_bytes() + } else { + &[] + }; + + data.encode(buf); + } + } + + Encoder::new(buf) + .encode_field(&self.0.nonce) + .encode_field(&self.0.balance) + .encode_field(&StorageRootCodec(&self.0.storage_root)) + .encode_field(&CodeHashCodec(&self.0.code_hash)) + .finish(); + } +} + +impl RLPDecode for AccountStateSlimCodec { + fn decode_unfinished(rlp: &[u8]) -> Result<(Self, &[u8]), RLPDecodeError> { + struct StorageRootCodec(H256); + impl RLPDecode for StorageRootCodec { + fn decode_unfinished(mut rlp: &[u8]) -> Result<(Self, &[u8]), RLPDecodeError> { + let value = match rlp.split_off_first() { + Some(0x80) => *EMPTY_TRIE_HASH, + Some(0xA0) => { + let data; + (data, rlp) = rlp + .split_first_chunk::<32>() + .ok_or(RLPDecodeError::InvalidLength)?; + H256(*data) + } + _ => return Err(RLPDecodeError::InvalidLength), + }; + + Ok((Self(value), rlp)) + } + } + + struct CodeHashCodec(H256); + impl RLPDecode for CodeHashCodec { + fn decode_unfinished(mut rlp: &[u8]) -> Result<(Self, &[u8]), RLPDecodeError> { + let value = match rlp.split_off_first() { + Some(0x80) => *EMPTY_KECCACK_HASH, + Some(0xA0) => { + let data; + (data, rlp) = rlp + .split_first_chunk::<32>() + .ok_or(RLPDecodeError::InvalidLength)?; + H256(*data) + } + _ => return Err(RLPDecodeError::InvalidLength), + }; + + Ok((Self(value), rlp)) + } + } + + let decoder = Decoder::new(rlp)?; + let (nonce, decoder) = decoder.decode_field("nonce")?; + let (balance, decoder) = decoder.decode_field("balance")?; + let (StorageRootCodec(storage_root), decoder) = decoder.decode_field("storage_root")?; + let (CodeHashCodec(code_hash), decoder) = decoder.decode_field("code_hash")?; + + Ok(( + Self(AccountState { + nonce, + balance, + storage_root, + code_hash, + }), + decoder.finish()?, + )) + } +} + pub fn compute_storage_root(storage: &HashMap) -> H256 { let iter = storage.iter().filter_map(|(k, v)| { (!v.is_zero()).then_some((keccak_hash(k.to_big_endian()).to_vec(), v.encode_to_vec())) diff --git a/crates/networking/p2p/peer_handler.rs b/crates/networking/p2p/peer_handler.rs index dc94500f44a..709d66dc066 100644 --- a/crates/networking/p2p/peer_handler.rs +++ b/crates/networking/p2p/peer_handler.rs @@ -733,11 +733,7 @@ impl PeerHandler { peer_id ); all_account_hashes.extend(accounts.iter().map(|unit| unit.hash)); - all_accounts_state.extend( - accounts - .iter() - .map(|unit| AccountState::from(unit.account.clone())), - ); + all_accounts_state.extend(accounts.iter().map(|unit| unit.account)); } let Some((peer_id, connection)) = self @@ -881,7 +877,7 @@ impl PeerHandler { let (account_hashes, account_states): (Vec<_>, Vec<_>) = accounts .clone() .into_iter() - .map(|unit| (unit.hash, AccountState::from(unit.account))) + .map(|unit| (unit.hash, unit.account)) .unzip(); let encoded_accounts = account_states .iter() diff --git a/crates/networking/p2p/rlpx/snap.rs b/crates/networking/p2p/rlpx/snap.rs index 3f1c2f81d76..a249fc1aa57 100644 --- a/crates/networking/p2p/rlpx/snap.rs +++ b/crates/networking/p2p/rlpx/snap.rs @@ -5,8 +5,7 @@ use super::{ use bytes::{BufMut, Bytes}; use ethrex_common::{ H256, U256, - constants::{EMPTY_KECCACK_HASH, EMPTY_TRIE_HASH}, - types::AccountState, + types::{AccountState, AccountStateSlimCodec}, }; use ethrex_rlp::{ decode::RLPDecode, @@ -342,15 +341,7 @@ impl RLPxMessage for TrieNodes { #[derive(Debug, Clone)] pub struct AccountRangeUnit { pub hash: H256, - pub account: AccountStateSlim, -} - -#[derive(Debug, Clone)] -pub struct AccountStateSlim { - pub nonce: u64, - pub balance: U256, - pub storage_root: Bytes, - pub code_hash: Bytes, + pub account: AccountState, } #[derive(Debug, Clone)] @@ -363,7 +354,7 @@ impl RLPEncode for AccountRangeUnit { fn encode(&self, buf: &mut dyn BufMut) { Encoder::new(buf) .encode_field(&self.hash) - .encode_field(&self.account) + .encode_field(&AccountStateSlimCodec(self.account)) .finish(); } } @@ -372,83 +363,12 @@ impl RLPDecode for AccountRangeUnit { fn decode_unfinished(rlp: &[u8]) -> Result<(Self, &[u8]), RLPDecodeError> { let decoder = Decoder::new(rlp)?; let (hash, decoder) = decoder.decode_field("hash")?; - let (account, decoder) = decoder.decode_field("account")?; + let (AccountStateSlimCodec(account), decoder) = + decoder.decode_field::("account")?; Ok((Self { hash, account }, decoder.finish()?)) } } -impl RLPEncode for AccountStateSlim { - fn encode(&self, buf: &mut dyn BufMut) { - Encoder::new(buf) - .encode_field(&self.nonce) - .encode_field(&self.balance) - .encode_field(&self.storage_root) - .encode_field(&self.code_hash) - .finish(); - } -} - -impl RLPDecode for AccountStateSlim { - fn decode_unfinished(rlp: &[u8]) -> Result<(Self, &[u8]), RLPDecodeError> { - let decoder = Decoder::new(rlp)?; - let (nonce, decoder) = decoder.decode_field("nonce")?; - let (balance, decoder) = decoder.decode_field("balance")?; - let (storage_root, decoder) = decoder.decode_field("storage_root")?; - let (code_hash, decoder) = decoder.decode_field("code_hash")?; - Ok(( - Self { - nonce, - balance, - storage_root, - code_hash, - }, - decoder.finish()?, - )) - } -} - -impl From for AccountStateSlim { - fn from(value: AccountState) -> Self { - let storage_root = if value.storage_root == *EMPTY_TRIE_HASH { - Bytes::new() - } else { - Bytes::copy_from_slice(value.storage_root.as_bytes()) - }; - let code_hash = if value.code_hash == *EMPTY_KECCACK_HASH { - Bytes::new() - } else { - Bytes::copy_from_slice(value.code_hash.as_bytes()) - }; - Self { - nonce: value.nonce, - balance: value.balance, - storage_root, - code_hash, - } - } -} - -impl From for AccountState { - fn from(value: AccountStateSlim) -> Self { - let storage_root = if value.storage_root.is_empty() { - *EMPTY_TRIE_HASH - } else { - H256::from_slice(value.storage_root.as_ref()) - }; - let code_hash = if value.code_hash.is_empty() { - *EMPTY_KECCACK_HASH - } else { - H256::from_slice(value.code_hash.as_ref()) - }; - Self { - nonce: value.nonce, - balance: value.balance, - storage_root, - code_hash, - } - } -} - impl RLPEncode for StorageSlot { fn encode(&self, buf: &mut dyn BufMut) { Encoder::new(buf) diff --git a/crates/networking/p2p/snap.rs b/crates/networking/p2p/snap.rs index b5c73105173..9bfaafe1f55 100644 --- a/crates/networking/p2p/snap.rs +++ b/crates/networking/p2p/snap.rs @@ -5,10 +5,11 @@ use ethrex_storage::{Store, error::StoreError}; use crate::rlpx::{ error::PeerConnectionError, snap::{ - AccountRange, AccountRangeUnit, AccountStateSlim, ByteCodes, GetAccountRange, GetByteCodes, - GetStorageRanges, GetTrieNodes, StorageRanges, StorageSlot, TrieNodes, + AccountRange, AccountRangeUnit, ByteCodes, GetAccountRange, GetByteCodes, GetStorageRanges, + GetTrieNodes, StorageRanges, StorageSlot, TrieNodes, }, }; +use ethrex_common::types::AccountStateSlimCodec; // Request Processing @@ -21,8 +22,7 @@ pub async fn process_account_range_request( let mut bytes_used = 0; for (hash, account) in store.iter_accounts_from(request.root_hash, request.starting_hash)? { debug_assert!(hash >= request.starting_hash); - let account = AccountStateSlim::from(account); - bytes_used += 32 + account.length() as u64; + bytes_used += 32 + AccountStateSlimCodec(account).length() as u64; accounts.push(AccountRangeUnit { hash, account }); if hash >= request.limit_hash || bytes_used >= request.response_bytes { break; @@ -177,13 +177,11 @@ pub(crate) fn encodable_to_proof(proof: &[Bytes]) -> Vec> { mod tests { use std::str::FromStr; - use ethrex_common::{BigEndianHash, H256, types::AccountState}; + use ethrex_common::{BigEndianHash, H256, types::AccountStateSlimCodec}; use ethrex_rlp::{decode::RLPDecode, encode::RLPEncode}; use ethrex_storage::EngineType; use ethrex_trie::EMPTY_TRIE_HASH; - use crate::rlpx::snap::AccountStateSlim; - use super::*; // Hive `AccounRange` Tests @@ -1000,7 +998,7 @@ mod tests { let mut state_trie = store.open_direct_state_trie(*EMPTY_TRIE_HASH)?; for (address, account) in accounts { let hashed_address = H256::from_str(address).unwrap().as_bytes().to_vec(); - let account = AccountState::from(AccountStateSlim::decode(&account).unwrap()); + let AccountStateSlimCodec(account) = RLPDecode::decode(&account).unwrap(); state_trie .insert(hashed_address, account.encode_to_vec()) .unwrap();