chore(primitives): extract trie HashBuilder implementation

crates/primitives/Cargo.toml

@@ -19,6 +19,7 @@ revm-primitives = { workspace = true, features = ["serde"] }
 # ethereum
 alloy-primitives = { workspace = true, features = ["rand", "rlp"] }
 alloy-rlp = { workspace = true, features = ["arrayvec"] }
+alloy-trie = { version = "0.1", features = ["serde"] }
 ethers-core = { workspace = true, default-features = false, optional = true }
 nybbles = { version = "0.1", features = ["serde", "rlp"] }
 
@@ -88,6 +89,7 @@ arbitrary = [
     "reth-rpc-types/arbitrary",
     "reth-ethereum-forks/arbitrary",
     "nybbles/arbitrary",
+    "alloy-trie/arbitrary",
     "dep:arbitrary",
     "dep:proptest",
     "dep:proptest-derive",

crates/primitives/src/trie/hash_builder/state.rs

@@ -1,5 +1,8 @@
-use super::{super::TrieMask, HashBuilderValue};
+use super::StoredHashBuilderValue;
+use crate::trie::{StoredTrieMask, TrieMask};
+use alloy_trie::{hash_builder::HashBuilderValue, HashBuilder};
 use bytes::Buf;
+use nybbles::Nibbles;
 use reth_codecs::{derive_arbitrary, Compact};
 use serde::{Deserialize, Serialize};
 
@@ -26,6 +29,37 @@ pub struct HashBuilderState {
     pub stored_in_database: bool,
 }
 
+impl From<HashBuilderState> for HashBuilder {
+    fn from(state: HashBuilderState) -> Self {
+        Self {
+            key: Nibbles::from_nibbles_unchecked(state.key),
+            stack: state.stack,
+            value: state.value,
+            groups: state.groups,
+            tree_masks: state.tree_masks,
+            hash_masks: state.hash_masks,
+            stored_in_database: state.stored_in_database,
+            updated_branch_nodes: None,
+            proof_retainer: None,
+            rlp_buf: Vec::with_capacity(32),
+        }
+    }
+}
+
+impl From<HashBuilder> for HashBuilderState {
+    fn from(state: HashBuilder) -> Self {
+        Self {
+            key: state.key.into(),
+            stack: state.stack,
+            value: state.value,
+            groups: state.groups,
+            tree_masks: state.tree_masks,
+            hash_masks: state.hash_masks,
+            stored_in_database: state.stored_in_database,
+        }
+    }
+}
+
 impl Compact for HashBuilderState {
     fn to_compact<B>(self, buf: &mut B) -> usize
     where
@@ -43,24 +77,24 @@ impl Compact for HashBuilderState {
             len += 2 + item.len();
         }
 
-        len += self.value.to_compact(buf);
+        len += StoredHashBuilderValue(self.value).to_compact(buf);
 
         buf.put_u16(self.groups.len() as u16);
         len += 2;
-        for item in self.groups.iter() {
-            len += item.to_compact(buf);
+        for item in &self.groups {
+            len += StoredTrieMask(*item).to_compact(buf);
         }
 
         buf.put_u16(self.tree_masks.len() as u16);
         len += 2;
-        for item in self.tree_masks.iter() {
-            len += item.to_compact(buf);
+        for item in &self.tree_masks {
+            len += StoredTrieMask(*item).to_compact(buf);
         }
 
         buf.put_u16(self.hash_masks.len() as u16);
         len += 2;
-        for item in self.hash_masks.iter() {
-            len += item.to_compact(buf);
+        for item in &self.hash_masks {
+            len += StoredTrieMask(*item).to_compact(buf);
         }
 
         buf.put_u8(self.stored_in_database as u8);
@@ -79,28 +113,28 @@ impl Compact for HashBuilderState {
             buf.advance(item_len);
         }
 
-        let (value, mut buf) = HashBuilderValue::from_compact(buf, 0);
+        let (StoredHashBuilderValue(value), mut buf) = StoredHashBuilderValue::from_compact(buf, 0);
 
         let groups_len = buf.get_u16() as usize;
         let mut groups = Vec::with_capacity(groups_len);
         for _ in 0..groups_len {
-            let (item, rest) = TrieMask::from_compact(buf, 0);
+            let (StoredTrieMask(item), rest) = StoredTrieMask::from_compact(buf, 0);
             groups.push(item);
             buf = rest;
         }
 
         let tree_masks_len = buf.get_u16() as usize;
         let mut tree_masks = Vec::with_capacity(tree_masks_len);
         for _ in 0..tree_masks_len {
-            let (item, rest) = TrieMask::from_compact(buf, 0);
+            let (StoredTrieMask(item), rest) = StoredTrieMask::from_compact(buf, 0);
             tree_masks.push(item);
             buf = rest;
         }
 
         let hash_masks_len = buf.get_u16() as usize;
         let mut hash_masks = Vec::with_capacity(hash_masks_len);
         for _ in 0..hash_masks_len {
-            let (item, rest) = TrieMask::from_compact(buf, 0);
+            let (StoredTrieMask(item), rest) = StoredTrieMask::from_compact(buf, 0);
             hash_masks.push(item);
             buf = rest;
         }

crates/primitives/src/trie/hash_builder/value.rs

@@ -1,78 +1,39 @@
-use crate::B256;
-use reth_codecs::{derive_arbitrary, Compact};
-use serde::{Deserialize, Serialize};
+use alloy_primitives::B256;
+use alloy_trie::hash_builder::HashBuilderValue;
+use bytes::Buf;
+use reth_codecs::Compact;
 
-/// The current value of the hash builder.
-#[derive_arbitrary(compact)]
-#[derive(Clone, PartialEq, Serialize, Deserialize)]
-pub enum HashBuilderValue {
-    /// Value of the leaf node.
-    Hash(B256),
-    /// Hash of adjacent nodes.
-    Bytes(Vec<u8>),
-}
+/// A wrapper around `HashBuilderValue` that implements `Compact`.
+pub(crate) struct StoredHashBuilderValue(pub(crate) HashBuilderValue);
 
-impl Compact for HashBuilderValue {
+impl Compact for StoredHashBuilderValue {
     fn to_compact<B>(self, buf: &mut B) -> usize
     where
         B: bytes::BufMut + AsMut<[u8]>,
     {
-        match self {
-            Self::Hash(hash) => {
+        match self.0 {
+            HashBuilderValue::Hash(hash) => {
                 buf.put_u8(0);
                 1 + hash.to_compact(buf)
             }
-            Self::Bytes(bytes) => {
+            HashBuilderValue::Bytes(bytes) => {
                 buf.put_u8(1);
                 1 + bytes.to_compact(buf)
             }
         }
     }
 
-    fn from_compact(buf: &[u8], _len: usize) -> (Self, &[u8]) {
-        match buf[0] {
+    fn from_compact(mut buf: &[u8], _len: usize) -> (Self, &[u8]) {
+        match buf.get_u8() {
             0 => {
-                let (hash, buf) = B256::from_compact(&buf[1..], 32);
-                (Self::Hash(hash), buf)
+                let (hash, buf) = B256::from_compact(buf, 32);
+                (Self(HashBuilderValue::Hash(hash)), buf)
             }
             1 => {
-                let (bytes, buf) = Vec::from_compact(&buf[1..], 0);
-                (Self::Bytes(bytes), buf)
+                let (bytes, buf) = Vec::from_compact(buf, 0);
+                (Self(HashBuilderValue::Bytes(bytes)), buf)
             }
             _ => unreachable!("Junk data in database: unknown HashBuilderValue variant"),
         }
     }
 }
-
-impl std::fmt::Debug for HashBuilderValue {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        match self {
-            Self::Bytes(bytes) => write!(f, "Bytes({:?})", crate::hex::encode(bytes)),
-            Self::Hash(hash) => write!(f, "Hash({:?})", hash),
-        }
-    }
-}
-
-impl From<Vec<u8>> for HashBuilderValue {
-    fn from(value: Vec<u8>) -> Self {
-        Self::Bytes(value)
-    }
-}
-
-impl From<&[u8]> for HashBuilderValue {
-    fn from(value: &[u8]) -> Self {
-        Self::Bytes(value.to_vec())
-    }
-}
-
-impl From<B256> for HashBuilderValue {
-    fn from(value: B256) -> Self {
-        Self::Hash(value)
-    }
-}
-
-impl Default for HashBuilderValue {
-    fn default() -> Self {
-        Self::Bytes(vec![])
-    }
-}

crates/primitives/src/trie/mask.rs

@@ -1,84 +1,20 @@
+use super::TrieMask;
 use bytes::Buf;
-use derive_more::{BitAnd, BitAndAssign, BitOr, BitOrAssign, Deref, From, Not};
-use reth_codecs::{derive_arbitrary, Compact};
-use serde::{Deserialize, Serialize};
+use reth_codecs::Compact;
 
-/// A struct representing a mask of 16 bits, used for Ethereum trie operations.
-///
-/// Masks in a trie are used to efficiently represent and manage information about the presence or
-/// absence of certain elements, such as child nodes, within a trie. Masks are usually implemented
-/// as bit vectors, where each bit represents the presence (1) or absence (0) of a corresponding
-/// element.
-#[derive(
-    Default,
-    Clone,
-    Copy,
-    PartialEq,
-    Eq,
-    Serialize,
-    Deserialize,
-    PartialOrd,
-    Ord,
-    Deref,
-    From,
-    BitAnd,
-    BitAndAssign,
-    BitOr,
-    BitOrAssign,
-    Not,
-)]
-#[derive_arbitrary(compact)]
-pub struct TrieMask(u16);
+pub(crate) struct StoredTrieMask(pub(crate) TrieMask);
 
-impl TrieMask {
-    /// Creates a new `TrieMask` from the given inner value.
-    #[inline]
-    pub fn new(inner: u16) -> Self {
-        Self(inner)
-    }
-
-    /// Creates a new `TrieMask` from the given nibble.
-    #[inline]
-    pub fn from_nibble(nibble: u8) -> Self {
-        Self(1u16 << nibble)
-    }
-
-    /// Returns `true` if the current `TrieMask` is a subset of `other`.
-    #[inline]
-    pub fn is_subset_of(self, other: Self) -> bool {
-        self & other == self
-    }
-
-    /// Returns `true` if a given bit is set in a mask.
-    #[inline]
-    pub fn is_bit_set(self, index: u8) -> bool {
-        self.0 & (1u16 << index) != 0
-    }
-
-    /// Returns `true` if the mask is empty.
-    #[inline]
-    pub fn is_empty(self) -> bool {
-        self.0 == 0
-    }
-}
-
-impl std::fmt::Debug for TrieMask {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "TrieMask({:016b})", self.0)
-    }
-}
-
-impl Compact for TrieMask {
+impl Compact for StoredTrieMask {
     fn to_compact<B>(self, buf: &mut B) -> usize
     where
         B: bytes::BufMut + AsMut<[u8]>,
     {
-        buf.put_u16(self.0);
+        buf.put_u16(self.0.get());
         2
     }
 
     fn from_compact(mut buf: &[u8], _len: usize) -> (Self, &[u8]) {
         let mask = buf.get_u16();
-        (Self(mask), buf)
+        (Self(TrieMask::new(mask)), buf)
     }
 }

crates/primitives/src/trie/mod.rs

@@ -1,27 +1,27 @@
 //! Collection of trie related types.
 
-/// Various branch nodes produced by the hash builder.
-pub mod nodes;
-pub use nodes::BranchNodeCompact;
-
 /// The implementation of hash builder.
 pub mod hash_builder;
-pub use hash_builder::HashBuilder;
-
-/// Merkle trie proofs.
-mod proofs;
-pub use proofs::{AccountProof, StorageProof};
 
 mod account;
+pub use account::TrieAccount;
+
 mod mask;
+pub(crate) use mask::StoredTrieMask;
+
 mod nibbles;
+pub use nibbles::{Nibbles, StoredNibbles, StoredNibblesSubKey};
+
+pub mod nodes;
+pub use nodes::StoredBranchNode;
+
+mod proofs;
+pub use proofs::{AccountProof, StorageProof};
+
 mod storage;
+pub use storage::StorageTrieEntry;
+
 mod subnode;
+pub use subnode::StoredSubNode;
 
-pub use self::{
-    account::TrieAccount,
-    mask::TrieMask,
-    nibbles::{Nibbles, StoredNibbles, StoredNibblesSubKey},
-    storage::StorageTrieEntry,
-    subnode::StoredSubNode,
-};
+pub use alloy_trie::{BranchNodeCompact, HashBuilder, TrieMask, EMPTY_ROOT_HASH};