diff --git a/crates/primitives/src/trie/branch_node.rs b/crates/primitives/src/trie/branch_node.rs new file mode 100644 index 0000000000..40e29642da --- /dev/null +++ b/crates/primitives/src/trie/branch_node.rs @@ -0,0 +1,146 @@ +use super::TrieMask; +use crate::H256; +use reth_codecs::Compact; +use serde::{Deserialize, Serialize}; + +/// A struct representing a branch node in an Ethereum trie. +/// +/// A branch node can have up to 16 children, each corresponding to one of the possible nibble +/// values (0 to 15) in the trie's path. +/// +/// The masks in a BranchNode are used to efficiently represent and manage information about the +/// presence and types of its children. They are bitmasks, where each bit corresponds to a nibble +/// (half-byte, or 4 bits) value from 0 to 15. +#[derive(Debug, Default, Clone, PartialEq, Eq, Serialize, Deserialize, PartialOrd, Ord)] +pub struct BranchNodeCompact { + /// The bitmask indicating the presence of children at the respective nibble positions in the + /// trie. If the bit at position i (counting from the right) is set (1), it indicates that a + /// child exists for the nibble value i. If the bit is unset (0), it means there is no child + /// for that nibble value. + pub state_mask: TrieMask, + /// The bitmask representing the internal (unhashed) children at the + /// respective nibble positions in the trie. If the bit at position `i` (counting from the + /// right) is set (1) and also present in the state_mask, it indicates that the + /// corresponding child at the nibble value `i` is an internal child. If the bit is unset + /// (0), it means the child is not an internal child. + pub tree_mask: TrieMask, + /// The bitmask representing the hashed children at the respective nibble + /// positions in the trie. If the bit at position `i` (counting from the right) is set (1) and + /// also present in the state_mask, it indicates that the corresponding child at the nibble + /// value `i` is a hashed child. If the bit is unset (0), it means the child is not a + /// hashed child. + pub hash_mask: TrieMask, + /// Collection of hashes associated with the children of the branch node. + /// Each child hash is calculated by hashing two consecutive sub-branch roots. + pub hashes: Vec, + /// An optional root hash of the subtree rooted at this branch node. + pub root_hash: Option, +} + +impl BranchNodeCompact { + /// Creates a new [BranchNodeCompact] from the given parameters. + pub fn new( + state_mask: TrieMask, + tree_mask: TrieMask, + hash_mask: TrieMask, + hashes: Vec, + root_hash: Option, + ) -> Self { + assert!(tree_mask.is_subset_of(&state_mask)); + assert!(hash_mask.is_subset_of(&state_mask)); + assert_eq!(hash_mask.count_ones() as usize, hashes.len()); + Self { state_mask, tree_mask, hash_mask, hashes, root_hash } + } + + /// Returns the hash associated with the given nibble. + pub fn hash_for_nibble(&self, nibble: i8) -> H256 { + let mask = *TrieMask::from_nibble(nibble) - 1; + let index = (*self.hash_mask & mask).count_ones(); + self.hashes[index as usize] + } +} + +impl Compact for BranchNodeCompact { + fn to_compact(self, buf: &mut impl bytes::BufMut) -> usize { + let BranchNodeCompact { state_mask, tree_mask, hash_mask, root_hash, hashes } = self; + + let mut buf_size = 0; + + buf_size += state_mask.to_compact(buf); + buf_size += tree_mask.to_compact(buf); + buf_size += hash_mask.to_compact(buf); + + if let Some(root_hash) = root_hash { + buf_size += H256::len_bytes(); + buf.put_slice(root_hash.as_bytes()); + } + + for hash in &hashes { + buf_size += H256::len_bytes(); + buf.put_slice(hash.as_bytes()); + } + + buf_size + } + + fn from_compact(buf: &[u8], len: usize) -> (Self, &[u8]) + where + Self: Sized, + { + let hash_len = H256::len_bytes(); + + // Assert the buffer is long enough to contain the masks and the hashes. + assert_eq!(buf.len() % hash_len, 6); + + // Consume the masks. + let (state_mask, buf) = TrieMask::from_compact(buf, len); + let (tree_mask, buf) = TrieMask::from_compact(buf, len); + let (hash_mask, buf) = TrieMask::from_compact(buf, len); + + let mut buf = buf; + let mut num_hashes = buf.len() / hash_len; + let mut root_hash = None; + + // Check if the root hash is present + if hash_mask.count_ones() as usize + 1 == num_hashes { + let (hash, remaining) = H256::from_compact(buf, hash_len); + root_hash = Some(hash); + buf = remaining; + num_hashes -= 1; + } + + // Consume all remaining hashes. + let mut hashes = Vec::::with_capacity(num_hashes); + for _ in 0..num_hashes { + let (hash, remaining) = H256::from_compact(buf, hash_len); + hashes.push(hash); + buf = remaining; + } + + (Self::new(state_mask, tree_mask, hash_mask, hashes, root_hash), buf) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use hex_literal::hex; + + #[test] + fn node_encoding() { + let n = BranchNodeCompact::new( + 0xf607.into(), + 0x0005.into(), + 0x4004.into(), + vec![ + hex!("90d53cd810cc5d4243766cd4451e7b9d14b736a1148b26b3baac7617f617d321").into(), + hex!("cc35c964dda53ba6c0b87798073a9628dbc9cd26b5cce88eb69655a9c609caf1").into(), + ], + Some(hex!("aaaabbbb0006767767776fffffeee44444000005567645600000000eeddddddd").into()), + ); + + let mut out = Vec::new(); + let compact_len = BranchNodeCompact::to_compact(n.clone(), &mut out); + assert_eq!(BranchNodeCompact::from_compact(&out, compact_len).0, n); + } +} diff --git a/crates/primitives/src/trie/mask.rs b/crates/primitives/src/trie/mask.rs new file mode 100644 index 0000000000..79f1eaa85b --- /dev/null +++ b/crates/primitives/src/trie/mask.rs @@ -0,0 +1,47 @@ +use derive_more::{BitAnd, Deref, From}; +use reth_codecs::Compact; +use serde::{Deserialize, Serialize}; + +/// A struct representing a mask of 16 bits, used for Ethereum trie operations. +#[derive( + Debug, + Default, + Clone, + Copy, + PartialEq, + Eq, + Serialize, + Deserialize, + PartialOrd, + Ord, + Deref, + From, + BitAnd, +)] +pub struct TrieMask(u16); + +impl TrieMask { + /// Creates a new `TrieMask` from the given nibble. + pub fn from_nibble(nibble: i8) -> Self { + Self(1u16 << nibble) + } + + /// Returns `true` if the current `TrieMask` is a subset of `other`. + pub fn is_subset_of(&self, other: &Self) -> bool { + *self & *other == *self + } +} + +impl Compact for TrieMask { + fn to_compact(self, buf: &mut impl bytes::BufMut) -> usize { + buf.put_slice(self.to_be_bytes().as_slice()); + 2 + } + + fn from_compact(buf: &[u8], _len: usize) -> (Self, &[u8]) + where + Self: Sized, + { + (Self(u16::from_be_bytes(buf[..2].try_into().unwrap())), &buf[2..]) + } +} diff --git a/crates/primitives/src/trie/mod.rs b/crates/primitives/src/trie/mod.rs index 7bead65f9d..92fcb5814a 100644 --- a/crates/primitives/src/trie/mod.rs +++ b/crates/primitives/src/trie/mod.rs @@ -2,3 +2,9 @@ mod nibbles; pub use nibbles::Nibbles; + +mod branch_node; +pub use branch_node::BranchNodeCompact; + +mod mask; +pub use mask::TrieMask;