From 42a98a7194cacd12ff012a1a55444fa1e1b7aecb Mon Sep 17 00:00:00 2001 From: Roman Krasiuk Date: Wed, 12 Apr 2023 00:23:19 +0300 Subject: [PATCH] feat(trie): stored nibbles (#2182) --- Cargo.lock | 1 + crates/primitives/src/trie/mod.rs | 2 +- crates/primitives/src/trie/nibbles.rs | 322 ++---------------- .../storage/db/src/tables/codecs/compact.rs | 7 +- crates/storage/db/src/tables/models/mod.rs | 51 ++- crates/trie/Cargo.toml | 1 + crates/trie/src/hash_builder/mod.rs | 9 +- crates/trie/src/lib.rs | 6 + crates/trie/src/nibbles.rs | 296 ++++++++++++++++ crates/trie/src/nodes/extension.rs | 3 +- crates/trie/src/nodes/leaf.rs | 3 +- crates/trie/src/prefix_set.rs | 3 +- 12 files changed, 404 insertions(+), 300 deletions(-) create mode 100644 crates/trie/src/nibbles.rs diff --git a/Cargo.lock b/Cargo.lock index 8c74f6e67f..d00630ab0d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5317,6 +5317,7 @@ dependencies = [ name = "reth-trie" version = "0.1.0" dependencies = [ + "derive_more", "hex", "proptest", "reth-primitives", diff --git a/crates/primitives/src/trie/mod.rs b/crates/primitives/src/trie/mod.rs index 92fcb5814a..03fc48b770 100644 --- a/crates/primitives/src/trie/mod.rs +++ b/crates/primitives/src/trie/mod.rs @@ -1,7 +1,7 @@ //! Collection of trie related types. mod nibbles; -pub use nibbles::Nibbles; +pub use nibbles::{StoredNibbles, StoredNibblesSubKey}; mod branch_node; pub use branch_node::BranchNodeCompact; diff --git a/crates/primitives/src/trie/nibbles.rs b/crates/primitives/src/trie/nibbles.rs index 5292c313f8..72ee94fb99 100644 --- a/crates/primitives/src/trie/nibbles.rs +++ b/crates/primitives/src/trie/nibbles.rs @@ -1,296 +1,48 @@ -use derive_more::{Deref, DerefMut, From, Index}; -use reth_rlp::RlpEncodableWrapper; +use crate::Bytes; +use derive_more::Deref; +use reth_codecs::{main_codec, Compact}; +use serde::{Deserialize, Serialize}; -/// Structure representing a sequence of nibbles. -/// -/// A nibble is a 4-bit value, and this structure is used to store -/// the nibble sequence representing the keys in a Merkle Patricia Trie (MPT). -/// Using nibbles simplifies trie operations and enables consistent key -/// representation in the MPT. -/// -/// The `hex_data` field is a `Vec` that stores the nibbles, with each -/// `u8` value containing a single nibble. This means that each byte in -/// `hex_data` has its upper 4 bits set to zero and the lower 4 bits -/// representing the nibble value. -#[derive( - Default, - Clone, - Eq, - PartialEq, - RlpEncodableWrapper, - PartialOrd, - Ord, - Index, - From, - Deref, - DerefMut, -)] -pub struct Nibbles { - /// The inner representation of the nibble sequence. - pub hex_data: Vec, +/// The nibbles are the keys for the AccountsTrie and the subkeys for the StorageTrie. +#[main_codec] +#[derive(Debug, Default, Clone, PartialEq, Eq, PartialOrd, Ord)] +pub struct StoredNibbles { + /// The inner nibble bytes + pub inner: Bytes, } -impl From<&[u8]> for Nibbles { - fn from(slice: &[u8]) -> Self { - Nibbles::from_hex(slice.to_vec()) +impl From> for StoredNibbles { + fn from(inner: Vec) -> Self { + Self { inner: inner.into() } } } -impl From<&[u8; N]> for Nibbles { - fn from(arr: &[u8; N]) -> Self { - Nibbles::from_hex(arr.to_vec()) +/// The representation of nibbles of the merkle trie stored in the database. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, PartialOrd, Ord, Deref)] +pub struct StoredNibblesSubKey(StoredNibbles); + +impl From> for StoredNibblesSubKey { + fn from(inner: Vec) -> Self { + Self(StoredNibbles { inner: inner.into() }) } } -impl std::fmt::Debug for Nibbles { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("Nibbles").field("hex_data", &hex::encode(&self.hex_data)).finish() - } -} - -impl Nibbles { - /// Creates a new [Nibbles] instance from bytes. - pub fn from_hex(hex: Vec) -> Self { - Nibbles { hex_data: hex } - } - - /// Take a byte array (slice or vector) as input and convert it into a [Nibbles] struct - /// containing the nibbles (half-bytes or 4 bits) that make up the input byte data. - pub fn unpack>(data: T) -> Self { - Nibbles { hex_data: data.as_ref().iter().flat_map(|item| [item / 16, item % 16]).collect() } - } - - /// Packs the nibbles stored in the struct into a byte vector. - /// - /// This method combines each pair of consecutive nibbles into a single byte, - /// effectively reducing the size of the data by a factor of two. - /// If the number of nibbles is odd, the last nibble is shifted left by 4 bits and - /// added to the packed byte vector. - pub fn pack(&self) -> Vec { - let length = (self.len() + 1) / 2; - if length == 0 { - Vec::default() - } else { - self.iter() - .enumerate() - .filter_map(|(index, nibble)| { - if index % 2 == 0 { - let next_nibble = self.get(index + 1).unwrap_or(&0); - Some((*nibble << 4) + *next_nibble) - } else { - None - } - }) - .collect() - } - } - - /// Encodes a given path leaf as a compact array of bytes, where each byte represents two - /// "nibbles" (half-bytes or 4 bits) of the original hex data, along with additional information - /// about the leaf itself. - /// - /// The method takes the following input: - /// `is_leaf`: A boolean value indicating whether the current node is a leaf node or not. - /// - /// The first byte of the encoded vector is set based on the `is_leaf` flag and the parity of - /// the hex data length (even or odd number of nibbles). - /// - If the node is an extension with even length, the header byte is `0x00`. - /// - If the node is an extension with odd length, the header byte is `0x10 + `. - /// - If the node is a leaf with even length, the header byte is `0x20`. - /// - If the node is a leaf with odd length, the header byte is `0x30 + `. - /// - /// If there is an odd number of nibbles, store the first nibble in the lower 4 bits of the - /// first byte of encoded. - /// - /// # Returns - /// - /// A `Vec` containing the compact byte representation of the nibble sequence, including the - /// header byte. - /// - /// # Example - /// - /// ``` - /// # use reth_primitives::trie::Nibbles; - /// - /// // Extension node with an even path length: - /// let nibbles = Nibbles::from_hex(vec![0x0A, 0x0B, 0x0C, 0x0D]); - /// assert_eq!(nibbles.encode_path_leaf(false), vec![0x00, 0xAB, 0xCD]); - /// - /// // Extension node with an odd path length: - /// let nibbles = Nibbles::from_hex(vec![0x0A, 0x0B, 0x0C]); - /// assert_eq!(nibbles.encode_path_leaf(false), vec![0x1A, 0xBC]); - /// - /// // Leaf node with an even path length: - /// let nibbles = Nibbles::from_hex(vec![0x0A, 0x0B, 0x0C, 0x0D]); - /// assert_eq!(nibbles.encode_path_leaf(true), vec![0x20, 0xAB, 0xCD]); - /// - /// // Leaf node with an odd path length: - /// let nibbles = Nibbles::from_hex(vec![0x0A, 0x0B, 0x0C]); - /// assert_eq!(nibbles.encode_path_leaf(true), vec![0x3A, 0xBC]); - /// ``` - pub fn encode_path_leaf(&self, is_leaf: bool) -> Vec { - let mut encoded = vec![0u8; self.len() / 2 + 1]; - let odd_nibbles = self.len() % 2 != 0; - - // Set the first byte of the encoded vector. - encoded[0] = match (is_leaf, odd_nibbles) { - (true, true) => 0x30 | self[0], - (true, false) => 0x20, - (false, true) => 0x10 | self[0], - (false, false) => 0x00, - }; - - let mut nibble_idx = if odd_nibbles { 1 } else { 0 }; - for byte in encoded.iter_mut().skip(1) { - *byte = (self[nibble_idx] << 4) + self[nibble_idx + 1]; - nibble_idx += 2; - } - - encoded - } - - /// Increments the nibble sequence by one. - pub fn increment(&self) -> Option { - let mut incremented = self.hex_data.clone(); - - for nibble in incremented.iter_mut().rev() { - assert!(*nibble < 0x10); - if *nibble < 0xf { - *nibble += 1; - return Some(Nibbles::from(incremented)) - } else { - *nibble = 0; - } - } - - None - } - - /// The last element of the hex vector is used to determine whether the nibble sequence - /// represents a leaf or an extension node. If the last element is 0x10 (16), then it's a leaf. - pub fn is_leaf(&self) -> bool { - self.hex_data[self.hex_data.len() - 1] == 16 - } - - /// Returns `true` if the current nibble sequence starts with the given prefix. - pub fn has_prefix(&self, other: &Self) -> bool { - self.starts_with(other) - } - - /// Returns the nibble at the given index. - pub fn at(&self, i: usize) -> usize { - self.hex_data[i] as usize - } - - /// Returns the last nibble of the current nibble sequence. - pub fn last(&self) -> Option { - self.hex_data.last().copied() - } - - /// Returns the length of the common prefix between the current nibble sequence and the given. - pub fn common_prefix_length(&self, other: &Nibbles) -> usize { - let len = std::cmp::min(self.len(), other.len()); - for i in 0..len { - if self[i] != other[i] { - return i - } - } - len - } - - /// Slice the current nibbles from the given start index to the end. - pub fn slice_from(&self, index: usize) -> Nibbles { - self.slice(index, self.hex_data.len()) - } - - /// Slice the current nibbles within the provided index range. - pub fn slice(&self, start: usize, end: usize) -> Nibbles { - Nibbles::from_hex(self.hex_data[start..end].to_vec()) - } - - /// Join two nibbles together. - pub fn join(&self, b: &Nibbles) -> Nibbles { - let mut hex_data = Vec::with_capacity(self.len() + b.len()); - hex_data.extend_from_slice(self); - hex_data.extend_from_slice(b); - Nibbles::from_hex(hex_data) - } - - /// Extend the current nibbles with another nibbles. - pub fn extend(&mut self, b: &Nibbles) { - self.hex_data.extend_from_slice(b); - } - - /// Truncate the current nibbles to the given length. - pub fn truncate(&mut self, len: usize) { - self.hex_data.truncate(len) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use proptest::prelude::*; - - #[test] - fn hashed_regression() { - let nibbles = hex::decode("05010406040a040203030f010805020b050c04070003070e0909070f010b0a0805020301070c0a0902040b0f000f0006040a04050f020b090701000a0a040b").unwrap(); - let nibbles = Nibbles::from(nibbles); - let path = nibbles.encode_path_leaf(true); - let expected = - hex::decode("351464a4233f1852b5c47037e997f1ba852317ca924bf0f064a45f2b9710aa4b") - .unwrap(); - assert_eq!(path, expected); - } - - #[test] - fn pack_nibbles() { - for (input, expected) in [ - (vec![], vec![]), - (vec![0xa], vec![0xa0]), - (vec![0xa, 0xb], vec![0xab]), - (vec![0xa, 0xb, 0x2], vec![0xab, 0x20]), - (vec![0xa, 0xb, 0x2, 0x0], vec![0xab, 0x20]), - (vec![0xa, 0xb, 0x2, 0x7], vec![0xab, 0x27]), - ] { - let nibbles = Nibbles::from(input); - let encoded = nibbles.pack(); - assert_eq!(encoded, expected); - } - } - - proptest! { - #[test] - fn pack_unpack_roundtrip(input in any::>()) { - let nibbles = Nibbles::unpack(&input); - let packed = nibbles.pack(); - prop_assert_eq!(packed, input); - } - - #[test] - fn encode_path_first_byte(input in any::>()) { - prop_assume!(!input.is_empty()); - let input = Nibbles::unpack(input); - let input_is_odd = input.len() % 2 == 1; - - let compact_leaf = input.encode_path_leaf(true); - let leaf_flag = compact_leaf[0]; - // Check flag - assert_ne!(leaf_flag & 0x20, 0); - assert_eq!(input_is_odd, (leaf_flag & 0x10) != 0); - if input_is_odd { - assert_eq!(leaf_flag & 0x0f, *input.first().unwrap()); - } - - - let compact_extension = input.encode_path_leaf(false); - let extension_flag = compact_extension[0]; - // Check first byte - assert_eq!(extension_flag & 0x20, 0); - assert_eq!(input_is_odd, (extension_flag & 0x10) != 0); - if input_is_odd { - assert_eq!(extension_flag & 0x0f, *input.first().unwrap()); - } - } +impl Compact for StoredNibblesSubKey { + fn to_compact(self, buf: &mut impl bytes::BufMut) -> usize { + assert!(self.inner.len() <= 64); + let mut padded = vec![0; 64]; + padded[..self.inner.len()].copy_from_slice(&self.inner[..]); + buf.put_slice(&padded); + buf.put_u8(self.inner.len() as u8); + 64 + 1 + } + + fn from_compact(buf: &[u8], _len: usize) -> (Self, &[u8]) + where + Self: Sized, + { + let len = buf[64] as usize; + let inner = Vec::from(&buf[..len]).into(); + (Self(StoredNibbles { inner }), &buf[65..]) } } diff --git a/crates/storage/db/src/tables/codecs/compact.rs b/crates/storage/db/src/tables/codecs/compact.rs index 25e5d59399..196a3ec4b3 100644 --- a/crates/storage/db/src/tables/codecs/compact.rs +++ b/crates/storage/db/src/tables/codecs/compact.rs @@ -4,7 +4,10 @@ use crate::{ Error, }; use reth_codecs::{main_codec, Compact}; -use reth_primitives::*; +use reth_primitives::{ + trie::{StoredNibbles, StoredNibblesSubKey}, + *, +}; /// Implements compression for Compact type. macro_rules! impl_compression_for_compact { @@ -40,6 +43,8 @@ impl_compression_for_compact!( Receipt, TxType, StorageEntry, + StoredNibbles, + StoredNibblesSubKey, StorageTrieEntry, StoredBlockBodyIndices, StoredBlockOmmers, diff --git a/crates/storage/db/src/tables/models/mod.rs b/crates/storage/db/src/tables/models/mod.rs index 0329d0f202..d161725920 100644 --- a/crates/storage/db/src/tables/models/mod.rs +++ b/crates/storage/db/src/tables/models/mod.rs @@ -1,4 +1,13 @@ //! Implements data structures specific to the database +use crate::{ + table::{Decode, Encode}, + Error, +}; +use reth_codecs::Compact; +use reth_primitives::{ + trie::{StoredNibbles, StoredNibblesSubKey}, + Address, H256, +}; pub mod accounts; pub mod blocks; @@ -10,12 +19,6 @@ pub use accounts::*; pub use blocks::*; pub use sharded_key::ShardedKey; -use crate::{ - table::{Decode, Encode}, - Error, -}; -use reth_primitives::{Address, H256}; - /// Macro that implements [`Encode`] and [`Decode`] for uint types. macro_rules! impl_uints { ($($name:tt),+) => { @@ -95,3 +98,39 @@ impl Decode for String { String::from_utf8(value.as_ref().to_vec()).map_err(|_| Error::DecodeError) } } + +impl Encode for StoredNibbles { + type Encoded = Vec; + + // Delegate to the Compact implementation + fn encode(self) -> Self::Encoded { + let mut buf = Vec::with_capacity(self.inner.len()); + self.to_compact(&mut buf); + buf + } +} + +impl Decode for StoredNibbles { + fn decode>(value: B) -> Result { + let buf = value.as_ref(); + Ok(Self::from_compact(buf, buf.len()).0) + } +} + +impl Encode for StoredNibblesSubKey { + type Encoded = Vec; + + // Delegate to the Compact implementation + fn encode(self) -> Self::Encoded { + let mut buf = Vec::with_capacity(65); + self.to_compact(&mut buf); + buf + } +} + +impl Decode for StoredNibblesSubKey { + fn decode>(value: B) -> Result { + let buf = value.as_ref(); + Ok(Self::from_compact(buf, buf.len()).0) + } +} diff --git a/crates/trie/Cargo.toml b/crates/trie/Cargo.toml index ceefac27dd..9167b06219 100644 --- a/crates/trie/Cargo.toml +++ b/crates/trie/Cargo.toml @@ -22,6 +22,7 @@ tracing = "0.1" # misc hex = "0.4" +derive_more = "0.99" [dev-dependencies] # reth diff --git a/crates/trie/src/hash_builder/mod.rs b/crates/trie/src/hash_builder/mod.rs index bf44da6c4d..5657a7ce43 100644 --- a/crates/trie/src/hash_builder/mod.rs +++ b/crates/trie/src/hash_builder/mod.rs @@ -1,8 +1,11 @@ -use crate::nodes::{rlp_hash, BranchNode, ExtensionNode, LeafNode}; +use crate::{ + nodes::{rlp_hash, BranchNode, ExtensionNode, LeafNode}, + Nibbles, +}; use reth_primitives::{ keccak256, proofs::EMPTY_ROOT, - trie::{BranchNodeCompact, Nibbles, TrieMask}, + trie::{BranchNodeCompact, TrieMask}, H256, }; use std::fmt::Debug; @@ -358,7 +361,7 @@ impl HashBuilder { mod tests { use super::*; use proptest::prelude::*; - use reth_primitives::{hex_literal::hex, proofs::KeccakHasher, trie::Nibbles, H256, U256}; + use reth_primitives::{hex_literal::hex, proofs::KeccakHasher, H256, U256}; use std::collections::{BTreeMap, HashMap}; use tokio::sync::mpsc::unbounded_channel; use tokio_stream::{wrappers::UnboundedReceiverStream, StreamExt}; diff --git a/crates/trie/src/lib.rs b/crates/trie/src/lib.rs index c67f22e87d..853030eb09 100644 --- a/crates/trie/src/lib.rs +++ b/crates/trie/src/lib.rs @@ -9,6 +9,12 @@ //! authenticated radix trie that is used to store key-value bindings. //! +mod nibbles; +pub use nibbles::Nibbles; + +/// The Ethereum account as represented in the trie. +pub mod account; + /// Various branch nodes producde by the hash builder. pub mod nodes; diff --git a/crates/trie/src/nibbles.rs b/crates/trie/src/nibbles.rs new file mode 100644 index 0000000000..6896b9b5fc --- /dev/null +++ b/crates/trie/src/nibbles.rs @@ -0,0 +1,296 @@ +use derive_more::{Deref, DerefMut, From, Index}; +use reth_rlp::RlpEncodableWrapper; + +/// Structure representing a sequence of nibbles. +/// +/// A nibble is a 4-bit value, and this structure is used to store +/// the nibble sequence representing the keys in a Merkle Patricia Trie (MPT). +/// Using nibbles simplifies trie operations and enables consistent key +/// representation in the MPT. +/// +/// The `hex_data` field is a `Vec` that stores the nibbles, with each +/// `u8` value containing a single nibble. This means that each byte in +/// `hex_data` has its upper 4 bits set to zero and the lower 4 bits +/// representing the nibble value. +#[derive( + Default, + Clone, + Eq, + PartialEq, + RlpEncodableWrapper, + PartialOrd, + Ord, + Index, + From, + Deref, + DerefMut, +)] +pub struct Nibbles { + /// The inner representation of the nibble sequence. + pub hex_data: Vec, +} + +impl From<&[u8]> for Nibbles { + fn from(slice: &[u8]) -> Self { + Nibbles::from_hex(slice.to_vec()) + } +} + +impl From<&[u8; N]> for Nibbles { + fn from(arr: &[u8; N]) -> Self { + Nibbles::from_hex(arr.to_vec()) + } +} + +impl std::fmt::Debug for Nibbles { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Nibbles").field("hex_data", &hex::encode(&self.hex_data)).finish() + } +} + +impl Nibbles { + /// Creates a new [Nibbles] instance from bytes. + pub fn from_hex(hex: Vec) -> Self { + Nibbles { hex_data: hex } + } + + /// Take a byte array (slice or vector) as input and convert it into a [Nibbles] struct + /// containing the nibbles (half-bytes or 4 bits) that make up the input byte data. + pub fn unpack>(data: T) -> Self { + Nibbles { hex_data: data.as_ref().iter().flat_map(|item| [item / 16, item % 16]).collect() } + } + + /// Packs the nibbles stored in the struct into a byte vector. + /// + /// This method combines each pair of consecutive nibbles into a single byte, + /// effectively reducing the size of the data by a factor of two. + /// If the number of nibbles is odd, the last nibble is shifted left by 4 bits and + /// added to the packed byte vector. + pub fn pack(&self) -> Vec { + let length = (self.len() + 1) / 2; + if length == 0 { + Vec::new() + } else { + self.iter() + .enumerate() + .filter_map(|(index, nibble)| { + if index % 2 == 0 { + let next_nibble = self.get(index + 1).unwrap_or(&0); + Some((*nibble << 4) + *next_nibble) + } else { + None + } + }) + .collect() + } + } + + /// Encodes a given path leaf as a compact array of bytes, where each byte represents two + /// "nibbles" (half-bytes or 4 bits) of the original hex data, along with additional information + /// about the leaf itself. + /// + /// The method takes the following input: + /// `is_leaf`: A boolean value indicating whether the current node is a leaf node or not. + /// + /// The first byte of the encoded vector is set based on the `is_leaf` flag and the parity of + /// the hex data length (even or odd number of nibbles). + /// - If the node is an extension with even length, the header byte is `0x00`. + /// - If the node is an extension with odd length, the header byte is `0x10 + `. + /// - If the node is a leaf with even length, the header byte is `0x20`. + /// - If the node is a leaf with odd length, the header byte is `0x30 + `. + /// + /// If there is an odd number of nibbles, store the first nibble in the lower 4 bits of the + /// first byte of encoded. + /// + /// # Returns + /// + /// A `Vec` containing the compact byte representation of the nibble sequence, including the + /// header byte. + /// + /// # Example + /// + /// ``` + /// # use reth_trie::Nibbles; + /// + /// // Extension node with an even path length: + /// let nibbles = Nibbles::from_hex(vec![0x0A, 0x0B, 0x0C, 0x0D]); + /// assert_eq!(nibbles.encode_path_leaf(false), vec![0x00, 0xAB, 0xCD]); + /// + /// // Extension node with an odd path length: + /// let nibbles = Nibbles::from_hex(vec![0x0A, 0x0B, 0x0C]); + /// assert_eq!(nibbles.encode_path_leaf(false), vec![0x1A, 0xBC]); + /// + /// // Leaf node with an even path length: + /// let nibbles = Nibbles::from_hex(vec![0x0A, 0x0B, 0x0C, 0x0D]); + /// assert_eq!(nibbles.encode_path_leaf(true), vec![0x20, 0xAB, 0xCD]); + /// + /// // Leaf node with an odd path length: + /// let nibbles = Nibbles::from_hex(vec![0x0A, 0x0B, 0x0C]); + /// assert_eq!(nibbles.encode_path_leaf(true), vec![0x3A, 0xBC]); + /// ``` + pub fn encode_path_leaf(&self, is_leaf: bool) -> Vec { + let mut encoded = vec![0u8; self.len() / 2 + 1]; + let odd_nibbles = self.len() % 2 != 0; + + // Set the first byte of the encoded vector. + encoded[0] = match (is_leaf, odd_nibbles) { + (true, true) => 0x30 | self[0], + (true, false) => 0x20, + (false, true) => 0x10 | self[0], + (false, false) => 0x00, + }; + + let mut nibble_idx = if odd_nibbles { 1 } else { 0 }; + for byte in encoded.iter_mut().skip(1) { + *byte = (self[nibble_idx] << 4) + self[nibble_idx + 1]; + nibble_idx += 2; + } + + encoded + } + + /// Increments the nibble sequence by one. + pub fn increment(&self) -> Option { + let mut incremented = self.hex_data.clone(); + + for nibble in incremented.iter_mut().rev() { + assert!(*nibble < 0x10); + if *nibble < 0xf { + *nibble += 1; + return Some(Nibbles::from(incremented)) + } else { + *nibble = 0; + } + } + + None + } + + /// The last element of the hex vector is used to determine whether the nibble sequence + /// represents a leaf or an extension node. If the last element is 0x10 (16), then it's a leaf. + pub fn is_leaf(&self) -> bool { + self.hex_data[self.hex_data.len() - 1] == 16 + } + + /// Returns `true` if the current nibble sequence starts with the given prefix. + pub fn has_prefix(&self, other: &Self) -> bool { + self.starts_with(other) + } + + /// Returns the nibble at the given index. + pub fn at(&self, i: usize) -> usize { + self.hex_data[i] as usize + } + + /// Returns the last nibble of the current nibble sequence. + pub fn last(&self) -> Option { + self.hex_data.last().copied() + } + + /// Returns the length of the common prefix between the current nibble sequence and the given. + pub fn common_prefix_length(&self, other: &Nibbles) -> usize { + let len = std::cmp::min(self.len(), other.len()); + for i in 0..len { + if self[i] != other[i] { + return i + } + } + len + } + + /// Slice the current nibbles from the given start index to the end. + pub fn slice_from(&self, index: usize) -> Nibbles { + self.slice(index, self.hex_data.len()) + } + + /// Slice the current nibbles within the provided index range. + pub fn slice(&self, start: usize, end: usize) -> Nibbles { + Nibbles::from_hex(self.hex_data[start..end].to_vec()) + } + + /// Join two nibbles together. + pub fn join(&self, b: &Nibbles) -> Nibbles { + let mut hex_data = Vec::with_capacity(self.len() + b.len()); + hex_data.extend_from_slice(self); + hex_data.extend_from_slice(b); + Nibbles::from_hex(hex_data) + } + + /// Extend the current nibbles with another nibbles. + pub fn extend(&mut self, b: &Nibbles) { + self.hex_data.extend_from_slice(b); + } + + /// Truncate the current nibbles to the given length. + pub fn truncate(&mut self, len: usize) { + self.hex_data.truncate(len) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use proptest::prelude::*; + + #[test] + fn hashed_regression() { + let nibbles = hex::decode("05010406040a040203030f010805020b050c04070003070e0909070f010b0a0805020301070c0a0902040b0f000f0006040a04050f020b090701000a0a040b").unwrap(); + let nibbles = Nibbles::from(nibbles); + let path = nibbles.encode_path_leaf(true); + let expected = + hex::decode("351464a4233f1852b5c47037e997f1ba852317ca924bf0f064a45f2b9710aa4b") + .unwrap(); + assert_eq!(path, expected); + } + + #[test] + fn pack_nibbles() { + for (input, expected) in [ + (vec![], vec![]), + (vec![0xa], vec![0xa0]), + (vec![0xa, 0xb], vec![0xab]), + (vec![0xa, 0xb, 0x2], vec![0xab, 0x20]), + (vec![0xa, 0xb, 0x2, 0x0], vec![0xab, 0x20]), + (vec![0xa, 0xb, 0x2, 0x7], vec![0xab, 0x27]), + ] { + let nibbles = Nibbles::from(input); + let encoded = nibbles.pack(); + assert_eq!(encoded, expected); + } + } + + proptest! { + #[test] + fn pack_unpack_roundtrip(input in any::>()) { + let nibbles = Nibbles::unpack(&input); + let packed = nibbles.pack(); + prop_assert_eq!(packed, input); + } + + #[test] + fn encode_path_first_byte(input in any::>()) { + prop_assume!(!input.is_empty()); + let input = Nibbles::unpack(input); + let input_is_odd = input.len() % 2 == 1; + + let compact_leaf = input.encode_path_leaf(true); + let leaf_flag = compact_leaf[0]; + // Check flag + assert_ne!(leaf_flag & 0x20, 0); + assert_eq!(input_is_odd, (leaf_flag & 0x10) != 0); + if input_is_odd { + assert_eq!(leaf_flag & 0x0f, *input.first().unwrap()); + } + + + let compact_extension = input.encode_path_leaf(false); + let extension_flag = compact_extension[0]; + // Check first byte + assert_eq!(extension_flag & 0x20, 0); + assert_eq!(input_is_odd, (extension_flag & 0x10) != 0); + if input_is_odd { + assert_eq!(extension_flag & 0x0f, *input.first().unwrap()); + } + } + } +} diff --git a/crates/trie/src/nodes/extension.rs b/crates/trie/src/nodes/extension.rs index 024dc48ed6..a998d4d8a7 100644 --- a/crates/trie/src/nodes/extension.rs +++ b/crates/trie/src/nodes/extension.rs @@ -1,5 +1,6 @@ use super::rlp_node; -use reth_primitives::{bytes::BytesMut, trie::Nibbles}; +use crate::Nibbles; +use reth_primitives::bytes::BytesMut; use reth_rlp::{BufMut, Encodable}; /// An intermediate node that exists solely to compress the trie's paths. It contains a path segment diff --git a/crates/trie/src/nodes/leaf.rs b/crates/trie/src/nodes/leaf.rs index 7bfb174047..516be22e64 100644 --- a/crates/trie/src/nodes/leaf.rs +++ b/crates/trie/src/nodes/leaf.rs @@ -1,5 +1,6 @@ use super::rlp_node; -use reth_primitives::{bytes::BytesMut, trie::Nibbles}; +use crate::Nibbles; +use reth_primitives::bytes::BytesMut; use reth_rlp::{BufMut, Encodable}; /// A leaf node represents the endpoint or terminal node in the trie. In other words, a leaf node is diff --git a/crates/trie/src/prefix_set.rs b/crates/trie/src/prefix_set.rs index 749796db1f..26210882df 100644 --- a/crates/trie/src/prefix_set.rs +++ b/crates/trie/src/prefix_set.rs @@ -1,7 +1,6 @@ +use crate::Nibbles; use std::collections::BTreeSet; -use reth_primitives::trie::Nibbles; - /// A container for efficiently storing and checking for the presence of key prefixes. /// /// This data structure stores a set of `Nibbles` and provides methods to insert