chore(l1): remove AccountStateSlim

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,14 +117,24 @@ 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,
     pub storage_root: H256,
     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<U256, U256>) -> 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()))

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()

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::<AccountStateSlimCodec>("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<AccountState> 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<AccountStateSlim> 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)

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<Vec<u8>> {
 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();