refactor: reuse alloy_trie::TrieAccount

hoank101 · Dec 7, 2024 · 378245a · 378245a
1 parent 552c623
commit 378245a
Show file tree

Hide file tree

Showing 14 changed files with 50 additions and 127 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/trie/common/Cargo.toml b/crates/trie/common/Cargo.toml
@@ -20,7 +20,6 @@ alloy-consensus.workspace = true
 reth-primitives-traits.workspace = true
 reth-codecs.workspace = true
 revm-primitives.workspace = true
-
 alloy-genesis.workspace = true
 alloy-rpc-types-eth = { workspace = true, optional = true }
 alloy-serde = { workspace = true, optional = true }

diff --git a/crates/trie/common/src/account.rs b/crates/trie/common/src/account.rs
@@ -1,57 +1,14 @@
-use crate::root::storage_root_unhashed;
 use alloy_consensus::constants::KECCAK_EMPTY;
-use alloy_genesis::GenesisAccount;
-use alloy_primitives::{keccak256, B256, U256};
-use alloy_rlp::{RlpDecodable, RlpEncodable};
-use alloy_trie::EMPTY_ROOT_HASH;
+use alloy_primitives::B256;
+use alloy_trie::TrieAccount;
 use reth_primitives_traits::Account;
-use revm_primitives::AccountInfo;
 
-/// An Ethereum account as represented in the trie.
-#[derive(Clone, Copy, Debug, PartialEq, Eq, Default, RlpEncodable, RlpDecodable)]
-pub struct TrieAccount {
-    /// Account nonce.
-    pub nonce: u64,
-    /// Account balance.
-    pub balance: U256,
-    /// Account's storage root.
-    pub storage_root: B256,
-    /// Hash of the account's bytecode.
-    pub code_hash: B256,
-}
-
-impl TrieAccount {
-    /// Get account's storage root.
-    pub const fn storage_root(&self) -> B256 {
-        self.storage_root
-    }
-}
-
-impl From<GenesisAccount> for TrieAccount {
-    fn from(account: GenesisAccount) -> Self {
-        let storage_root = account
-            .storage
-            .map(|storage| {
-                storage_root_unhashed(
-                    storage
-                        .into_iter()
-                        .filter(|(_, value)| !value.is_zero())
-                        .map(|(slot, value)| (slot, U256::from_be_bytes(*value))),
-                )
-            })
-            .unwrap_or(EMPTY_ROOT_HASH);
-
-        Self {
-            nonce: account.nonce.unwrap_or_default(),
-            balance: account.balance,
-            storage_root,
-            code_hash: account.code.map_or(KECCAK_EMPTY, keccak256),
-        }
-    }
-}
+/// Wrapper type to implement foreign trait conversions
+#[derive(Debug)]
+pub struct AccountWithStorageRoot(pub Account, pub B256);
 
-impl From<(Account, B256)> for TrieAccount {
-    fn from((account, storage_root): (Account, B256)) -> Self {
+impl From<AccountWithStorageRoot> for TrieAccount {
+    fn from(AccountWithStorageRoot(account, storage_root): AccountWithStorageRoot) -> Self {
         Self {
             nonce: account.nonce,
             balance: account.balance,
@@ -61,21 +18,13 @@ impl From<(Account, B256)> for TrieAccount {
     }
 }
 
-impl From<(AccountInfo, B256)> for TrieAccount {
-    fn from((account, storage_root): (AccountInfo, B256)) -> Self {
-        Self {
-            nonce: account.nonce,
-            balance: account.balance,
-            storage_root,
-            code_hash: account.code_hash,
-        }
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;
-    use alloy_primitives::Bytes;
+    use crate::root::storage_root_unhashed;
+    use alloy_genesis::GenesisAccount;
+    use alloy_primitives::{keccak256, Bytes, U256};
+    use alloy_trie::{TrieAccount, EMPTY_ROOT_HASH};
     use std::collections::BTreeMap;
 
     #[test]
@@ -88,17 +37,8 @@ mod tests {
         // Check the fields are properly set.
         assert_eq!(trie_account.nonce, 0);
         assert_eq!(trie_account.balance, U256::default());
-        assert_eq!(trie_account.storage_root(), EMPTY_ROOT_HASH);
+        assert_eq!(trie_account.storage_root, EMPTY_ROOT_HASH);
         assert_eq!(trie_account.code_hash, KECCAK_EMPTY);
-
-        // Check that the default Account converts to the same TrieAccount
-        assert_eq!(Into::<TrieAccount>::into((Account::default(), EMPTY_ROOT_HASH)), trie_account);
-
-        // Check that the default AccountInfo converts to the same TrieAccount
-        assert_eq!(
-            Into::<TrieAccount>::into((AccountInfo::default(), EMPTY_ROOT_HASH)),
-            trie_account
-        );
     }
 
     #[test]
@@ -126,35 +66,8 @@ mod tests {
         // Check that the fields are properly set.
         assert_eq!(trie_account.nonce, 10);
         assert_eq!(trie_account.balance, U256::from(1000));
-        assert_eq!(trie_account.storage_root(), expected_storage_root);
+        assert_eq!(trie_account.storage_root, expected_storage_root);
         assert_eq!(trie_account.code_hash, keccak256([0x60, 0x61]));
-
-        // Check that the Account converts to the same TrieAccount
-        assert_eq!(
-            Into::<TrieAccount>::into((
-                Account {
-                    nonce: 10,
-                    balance: U256::from(1000),
-                    bytecode_hash: Some(keccak256([0x60, 0x61]))
-                },
-                expected_storage_root
-            )),
-            trie_account
-        );
-
-        // Check that the AccountInfo converts to the same TrieAccount
-        assert_eq!(
-            Into::<TrieAccount>::into((
-                AccountInfo {
-                    nonce: 10,
-                    balance: U256::from(1000),
-                    code_hash: keccak256([0x60, 0x61]),
-                    ..Default::default()
-                },
-                expected_storage_root
-            )),
-            trie_account
-        );
     }
 
     #[test]
@@ -177,7 +90,7 @@ mod tests {
         assert_eq!(trie_account.nonce, 3);
         assert_eq!(trie_account.balance, U256::from(300));
         // Zero values in storage should result in EMPTY_ROOT_HASH
-        assert_eq!(trie_account.storage_root(), EMPTY_ROOT_HASH);
+        assert_eq!(trie_account.storage_root, EMPTY_ROOT_HASH);
         // No code provided, so code hash should be KECCAK_EMPTY
         assert_eq!(trie_account.code_hash, KECCAK_EMPTY);
     }

diff --git a/crates/trie/common/src/constants.rs b/crates/trie/common/src/constants.rs
@@ -5,9 +5,9 @@ pub const TRIE_ACCOUNT_RLP_MAX_SIZE: usize = 110;
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::TrieAccount;
     use alloy_primitives::{B256, U256};
     use alloy_rlp::Encodable;
+    use alloy_trie::TrieAccount;
 
     #[test]
     fn account_rlp_max_size() {

diff --git a/crates/trie/common/src/lib.rs b/crates/trie/common/src/lib.rs
@@ -16,8 +16,7 @@ mod constants;
 pub use constants::*;
 
 mod account;
-pub use account::TrieAccount;
-
+pub use account::*;
 mod key;
 pub use key::{KeccakKeyHasher, KeyHasher};
 

diff --git a/crates/trie/common/src/proofs.rs b/crates/trie/common/src/proofs.rs
@@ -1,6 +1,6 @@
 //! Merkle trie proofs.
 
-use crate::{Nibbles, TrieAccount};
+use crate::{account::AccountWithStorageRoot, Nibbles};
 use alloy_consensus::constants::KECCAK_EMPTY;
 use alloy_primitives::{
     keccak256,
@@ -11,7 +11,7 @@ use alloy_rlp::{encode_fixed_size, Decodable, EMPTY_STRING_CODE};
 use alloy_trie::{
     nodes::TrieNode,
     proof::{verify_proof, ProofNodes, ProofVerificationError},
-    TrieMask, EMPTY_ROOT_HASH,
+    TrieAccount, TrieMask, EMPTY_ROOT_HASH,
 };
 use itertools::Itertools;
 use reth_primitives_traits::Account;
@@ -255,7 +255,7 @@ impl AccountProof {
         let expected = if self.info.is_none() && self.storage_root == EMPTY_ROOT_HASH {
             None
         } else {
-            Some(alloy_rlp::encode(TrieAccount::from((
+            Some(alloy_rlp::encode(TrieAccount::from(AccountWithStorageRoot(
                 self.info.unwrap_or_default(),
                 self.storage_root,
             ))))

diff --git a/crates/trie/common/src/root.rs b/crates/trie/common/src/root.rs
@@ -1,9 +1,8 @@
 //! Common root computation functions.
 
-use crate::TrieAccount;
 use alloy_primitives::{keccak256, Address, B256, U256};
 use alloy_rlp::Encodable;
-use alloy_trie::HashBuilder;
+use alloy_trie::{HashBuilder, TrieAccount};
 use itertools::Itertools;
 use nybbles::Nibbles;
 

diff --git a/crates/trie/parallel/Cargo.toml b/crates/trie/parallel/Cargo.toml
@@ -24,6 +24,7 @@ reth-provider.workspace = true
 # alloy
 alloy-rlp.workspace = true
 alloy-primitives.workspace = true
+alloy-trie.workspace = true
 
 # tracing
 tracing.workspace = true

diff --git a/crates/trie/parallel/src/proof.rs b/crates/trie/parallel/src/proof.rs
@@ -1,9 +1,12 @@
+#[cfg(feature = "metrics")]
+use crate::metrics::ParallelStateRootMetrics;
 use crate::{root::ParallelStateRootError, stats::ParallelTrieTracker, StorageRootTargets};
 use alloy_primitives::{
     map::{HashMap, HashSet},
     B256,
 };
 use alloy_rlp::{BufMut, Encodable};
+use alloy_trie::TrieAccount;
 use itertools::Itertools;
 use reth_db::DatabaseError;
 use reth_execution_errors::StorageRootError;
@@ -18,16 +21,13 @@ use reth_trie::{
     proof::StorageProof,
     trie_cursor::{InMemoryTrieCursorFactory, TrieCursorFactory},
     walker::TrieWalker,
-    HashBuilder, MultiProof, Nibbles, TrieAccount, TrieInput, TRIE_ACCOUNT_RLP_MAX_SIZE,
+    HashBuilder, MultiProof, Nibbles, TrieInput, TRIE_ACCOUNT_RLP_MAX_SIZE,
 };
-use reth_trie_common::proof::ProofRetainer;
+use reth_trie_common::{proof::ProofRetainer, AccountWithStorageRoot};
 use reth_trie_db::{DatabaseHashedCursorFactory, DatabaseTrieCursorFactory};
 use std::sync::Arc;
 use tracing::{debug, error};
 
-#[cfg(feature = "metrics")]
-use crate::metrics::ParallelStateRootMetrics;
-
 /// TODO:
 #[derive(Debug)]
 pub struct ParallelProof<Factory> {
@@ -217,7 +217,8 @@ where
 
                     // Encode account
                     account_rlp.clear();
-                    let account = TrieAccount::from((account, storage_multiproof.root));
+                    let account =
+                        TrieAccount::from(AccountWithStorageRoot(account, storage_multiproof.root));
                     account.encode(&mut account_rlp as &mut dyn BufMut);
 
                     hash_builder.add_leaf(Nibbles::unpack(hashed_address), &account_rlp);

diff --git a/crates/trie/parallel/src/root.rs b/crates/trie/parallel/src/root.rs
@@ -3,6 +3,7 @@ use crate::metrics::ParallelStateRootMetrics;
 use crate::{stats::ParallelTrieTracker, storage_root_targets::StorageRootTargets};
 use alloy_primitives::B256;
 use alloy_rlp::{BufMut, Encodable};
+use alloy_trie::TrieAccount;
 use itertools::Itertools;
 use reth_db::DatabaseError;
 use reth_execution_errors::StorageRootError;
@@ -16,8 +17,9 @@ use reth_trie::{
     trie_cursor::{InMemoryTrieCursorFactory, TrieCursorFactory},
     updates::TrieUpdates,
     walker::TrieWalker,
-    HashBuilder, Nibbles, StorageRoot, TrieAccount, TrieInput, TRIE_ACCOUNT_RLP_MAX_SIZE,
+    HashBuilder, Nibbles, StorageRoot, TrieInput, TRIE_ACCOUNT_RLP_MAX_SIZE,
 };
+use reth_trie_common::AccountWithStorageRoot;
 use reth_trie_db::{DatabaseHashedCursorFactory, DatabaseTrieCursorFactory};
 use std::{collections::HashMap, sync::Arc};
 use thiserror::Error;
@@ -192,7 +194,7 @@ where
                     }
 
                     account_rlp.clear();
-                    let account = TrieAccount::from((account, storage_root));
+                    let account = TrieAccount::from(AccountWithStorageRoot(account, storage_root));
                     account.encode(&mut account_rlp as &mut dyn BufMut);
                     hash_builder.add_leaf(Nibbles::unpack(hashed_address), &account_rlp);
                 }

diff --git a/crates/trie/sparse/Cargo.toml b/crates/trie/sparse/Cargo.toml
@@ -22,6 +22,7 @@ reth-tracing.workspace = true
 # alloy
 alloy-primitives.workspace = true
 alloy-rlp.workspace = true
+alloy-trie.workspace = true
 
 # misc
 smallvec = { workspace = true, features = ["const_new"] }

diff --git a/crates/trie/sparse/src/state.rs b/crates/trie/sparse/src/state.rs
@@ -8,12 +8,14 @@ use alloy_primitives::{
     Bytes, B256,
 };
 use alloy_rlp::{Decodable, Encodable};
+use alloy_trie::TrieAccount;
 use reth_execution_errors::{SparseStateTrieError, SparseStateTrieResult, SparseTrieError};
 use reth_primitives_traits::Account;
 use reth_tracing::tracing::trace;
 use reth_trie_common::{
     updates::{StorageTrieUpdates, TrieUpdates},
-    MultiProof, Nibbles, TrieAccount, TrieNode, EMPTY_ROOT_HASH, TRIE_ACCOUNT_RLP_MAX_SIZE,
+    AccountWithStorageRoot, MultiProof, Nibbles, TrieNode, EMPTY_ROOT_HASH,
+    TRIE_ACCOUNT_RLP_MAX_SIZE,
 };
 use std::{fmt, iter::Peekable};
 
@@ -384,7 +386,8 @@ where
         } else {
             trace!(target: "trie::sparse", ?address, "Updating account");
             self.account_rlp_buf.clear();
-            TrieAccount::from((account, storage_root)).encode(&mut self.account_rlp_buf);
+            TrieAccount::from(AccountWithStorageRoot(account, storage_root))
+                .encode(&mut self.account_rlp_buf);
             self.update_account_leaf(nibbles, self.account_rlp_buf.clone())
         }
     }

diff --git a/crates/trie/trie/src/proof/mod.rs b/crates/trie/trie/src/proof/mod.rs
@@ -12,9 +12,10 @@ use alloy_primitives::{
     Address, B256,
 };
 use alloy_rlp::{BufMut, Encodable};
+use alloy_trie::TrieAccount;
 use reth_execution_errors::trie::StateProofError;
 use reth_trie_common::{
-    proof::ProofRetainer, AccountProof, MultiProof, StorageMultiProof, TrieAccount,
+    proof::ProofRetainer, AccountProof, AccountWithStorageRoot, MultiProof, StorageMultiProof,
 };
 
 mod blinded;
@@ -149,7 +150,8 @@ where
 
                     // Encode account
                     account_rlp.clear();
-                    let account = TrieAccount::from((account, storage_multiproof.root));
+                    let account =
+                        TrieAccount::from(AccountWithStorageRoot(account, storage_multiproof.root));
                     account.encode(&mut account_rlp as &mut dyn BufMut);
 
                     hash_builder.add_leaf(Nibbles::unpack(hashed_address), &account_rlp);