feat: Migrate trie account type and state root functions from alloy (#65

)

* feat: Migrate trie account type and state root functions from alloy

* Rename to TrieAccount / Use simple u64 serialize/deserialize

* fix cargo hack warnings

* Remove itertools and handle sort internally

* fix cfg_attr for TrieAccount

* Feature gate TrieAccount related code behind `ethereum`

* fix import

* chore: move into ethereum crate

---------

Co-authored-by: Matthias Seitz <[email protected]>

Cargo.toml

@@ -103,6 +103,7 @@ arbitrary = [
     "alloy-primitives/arbitrary",
     "nybbles/arbitrary",
 ]
+ethereum = []
 
 [[bench]]
 name = "bench"

src/account.rs

@@ -0,0 +1,60 @@
+use crate::{EMPTY_ROOT_HASH, KECCAK_EMPTY};
+use alloy_primitives::{keccak256, B256, U256};
+use alloy_rlp::{RlpDecodable, RlpEncodable};
+
+/// Represents an TrieAccount in the account trie.
+#[derive(Copy, Clone, Debug, PartialEq, Eq, RlpDecodable, RlpEncodable)]
+#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
+#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
+#[cfg_attr(feature = "serde", serde(rename_all = "camelCase"))]
+pub struct TrieAccount {
+    /// The account's nonce.
+    #[cfg_attr(feature = "serde", serde(with = "quantity"))]
+    pub nonce: u64,
+    /// The account's balance.
+    pub balance: U256,
+    /// The hash of the storage account data.
+    pub storage_root: B256,
+    /// The hash of the code of the account.
+    pub code_hash: B256,
+}
+
+impl Default for TrieAccount {
+    fn default() -> Self {
+        Self {
+            nonce: 0,
+            balance: U256::ZERO,
+            storage_root: EMPTY_ROOT_HASH,
+            code_hash: KECCAK_EMPTY,
+        }
+    }
+}
+
+impl TrieAccount {
+    /// Compute  hash as committed to in the MPT trie without memorizing.
+    pub fn trie_hash_slow(&self) -> B256 {
+        keccak256(alloy_rlp::encode(self))
+    }
+}
+
+#[cfg(feature = "serde")]
+mod quantity {
+    use alloy_primitives::U64;
+    use serde::{Deserialize, Deserializer, Serialize, Serializer};
+
+    /// Serializes a primitive number as a "quantity" hex string.
+    pub(crate) fn serialize<S>(value: &u64, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        U64::from(*value).serialize(serializer)
+    }
+
+    /// Deserializes a primitive number from a "quantity" hex string.
+    pub(crate) fn deserialize<'de, D>(deserializer: D) -> Result<u64, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        U64::deserialize(deserializer).map(|value| value.to())
+    }
+}

src/lib.rs

@@ -28,6 +28,11 @@ pub use hash_builder::HashBuilder;
 
 pub mod proof;
 
+#[cfg(feature = "ethereum")]
+mod account;
+#[cfg(feature = "ethereum")]
+pub use account::TrieAccount;
+
 mod mask;
 pub use mask::TrieMask;
 
@@ -40,20 +45,26 @@ pub use alloy_primitives::map::HashMap;
 #[doc(no_inline)]
 pub use nybbles::{self, Nibbles};
 
+use alloy_primitives::{b256, B256};
+
 /// Root hash of an empty trie.
-pub const EMPTY_ROOT_HASH: alloy_primitives::B256 =
-    alloy_primitives::b256!("56e81f171bcc55a6ff8345e692c0f86e5b48e01b996cadc001622fb5e363b421");
+pub const EMPTY_ROOT_HASH: B256 =
+    b256!("56e81f171bcc55a6ff8345e692c0f86e5b48e01b996cadc001622fb5e363b421");
+
+/// Keccak256 over empty array.
+pub const KECCAK_EMPTY: B256 =
+    b256!("c5d2460186f7233c927e7db2dcc703c0e500b653ca82273b7bfad8045d85a470");
 
 #[cfg(test)]
-pub(crate) fn triehash_trie_root<I, K, V>(iter: I) -> alloy_primitives::B256
+pub(crate) fn triehash_trie_root<I, K, V>(iter: I) -> B256
 where
     I: IntoIterator<Item = (K, V)>,
     K: AsRef<[u8]> + Ord,
     V: AsRef<[u8]>,
 {
     struct Keccak256Hasher;
     impl hash_db::Hasher for Keccak256Hasher {
-        type Out = alloy_primitives::B256;
+        type Out = B256;
         type StdHasher = plain_hasher::PlainHasher;
 
         const LENGTH: usize = 32;

src/root.rs

@@ -1,10 +1,9 @@
+use crate::{HashBuilder, EMPTY_ROOT_HASH};
 use alloc::vec::Vec;
 use alloy_primitives::B256;
 use alloy_rlp::Encodable;
 use nybbles::Nibbles;
 
-use crate::{HashBuilder, EMPTY_ROOT_HASH};
-
 /// Adjust the index of an item for rlp encoding.
 pub const fn adjust_index_for_rlp(i: usize, len: usize) -> usize {
     if i > 0x7f {
@@ -47,3 +46,61 @@ where
 
     hb.root()
 }
+
+/// Ethereum specific trie root functions.
+#[cfg(feature = "ethereum")]
+pub use ethereum::*;
+#[cfg(feature = "ethereum")]
+mod ethereum {
+    use super::*;
+    use crate::TrieAccount;
+    use alloy_primitives::{keccak256, Address};
+
+    /// Hashes and sorts account keys, then proceeds to calculating the root hash of the state
+    /// represented as MPT.
+    /// See [`state_root_unsorted`] for more info.
+    pub fn state_root_ref_unhashed<'a, A: Into<TrieAccount> + Clone + 'a>(
+        state: impl IntoIterator<Item = (&'a Address, &'a A)>,
+    ) -> B256 {
+        state_root_unsorted(
+            state.into_iter().map(|(address, account)| (keccak256(address), account.clone())),
+        )
+    }
+
+    /// Hashes and sorts account keys, then proceeds to calculating the root hash of the state
+    /// represented as MPT.
+    /// See [`state_root_unsorted`] for more info.
+    pub fn state_root_unhashed<A: Into<TrieAccount>>(
+        state: impl IntoIterator<Item = (Address, A)>,
+    ) -> B256 {
+        state_root_unsorted(
+            state.into_iter().map(|(address, account)| (keccak256(address), account)),
+        )
+    }
+
+    /// Sorts the hashed account keys and calculates the root hash of the state represented as MPT.
+    /// See [`state_root`] for more info.
+    pub fn state_root_unsorted<A: Into<TrieAccount>>(
+        state: impl IntoIterator<Item = (B256, A)>,
+    ) -> B256 {
+        let mut vec = Vec::from_iter(state);
+        vec.sort_unstable_by_key(|(key, _)| *key);
+        state_root(vec)
+    }
+
+    /// Calculates the root hash of the state represented as MPT.
+    ///
+    /// Corresponds to [geth's `deriveHash`](https://github.com/ethereum/go-ethereum/blob/6c149fd4ad063f7c24d726a73bc0546badd1bc73/core/genesis.go#L119).
+    ///
+    /// # Panics
+    ///
+    /// If the items are not in sorted order.
+    pub fn state_root<A: Into<TrieAccount>>(state: impl IntoIterator<Item = (B256, A)>) -> B256 {
+        let mut hb = HashBuilder::default();
+        for (hashed_key, account) in state {
+            let account_rlp_buf = alloy_rlp::encode(account.into());
+            hb.add_leaf(Nibbles::unpack(hashed_key), &account_rlp_buf);
+        }
+        hb.root()
+    }
+}