diff --git a/Cargo.lock b/Cargo.lock index efff0ef3d..0a61aed99 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6200,6 +6200,7 @@ dependencies = [ "serde", "serde_json", "sha2", + "smallvec", "strum", "sucds 0.6.0", "tempfile", @@ -7437,6 +7438,10 @@ name = "smallvec" version = "1.11.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4dccd0940a2dcdf68d092b8cbab7dc0ad8fa938bf95787e1b916b0e3d0e8e970" +dependencies = [ + "arbitrary", + "serde", +] [[package]] name = "smol_str" diff --git a/crates/primitives/Cargo.toml b/crates/primitives/Cargo.toml index 9487a1d46..eef82c7f0 100644 --- a/crates/primitives/Cargo.toml +++ b/crates/primitives/Cargo.toml @@ -31,8 +31,9 @@ c-kzg = { workspace = true, features = ["serde"], optional = true } tracing.workspace = true # misc -byteorder = "1" +smallvec = { version = "1.11", features = ["arbitrary", "serde", "union", "const_new"] } bytes.workspace = true +byteorder = "1" clap = { workspace = true, features = ["derive"], optional = true } derive_more = "0.99" itertools = "0.11" @@ -82,16 +83,21 @@ secp256k1.workspace = true [features] default = ["c-kzg"] arbitrary = [ - "revm-primitives/arbitrary", - "reth-rpc-types/arbitrary", - "reth-ethereum-forks/arbitrary", - "dep:arbitrary", - "dep:proptest", - "dep:proptest-derive", + "revm-primitives/arbitrary", + "reth-rpc-types/arbitrary", + "reth-ethereum-forks/arbitrary", + "dep:arbitrary", + "dep:proptest", + "dep:proptest-derive", ] c-kzg = ["dep:c-kzg", "revm/c-kzg", "revm-primitives/c-kzg"] clap = ["dep:clap"] -optimism = ["reth-codecs/optimism", "revm-primitives/optimism", "reth-ethereum-forks/optimism", "revm/optimism"] +optimism = [ + "reth-codecs/optimism", + "revm-primitives/optimism", + "reth-ethereum-forks/optimism", + "revm/optimism", +] test-utils = ["dep:plain_hasher", "dep:hash-db", "dep:ethers-core"] [[bench]] diff --git a/crates/primitives/benches/nibbles.rs b/crates/primitives/benches/nibbles.rs index 3d839fd7a..1db4f8378 100644 --- a/crates/primitives/benches/nibbles.rs +++ b/crates/primitives/benches/nibbles.rs @@ -1,13 +1,63 @@ use criterion::{criterion_group, criterion_main, Criterion}; +use proptest::{prelude::*, strategy::ValueTree}; use reth_primitives::trie::Nibbles; +use std::{hint::black_box, time::Duration}; /// Benchmarks the nibble unpacking. pub fn nibbles_benchmark(c: &mut Criterion) { let mut g = c.benchmark_group("nibbles"); - g.bench_function("unpack", |b| { - let raw = (1..=32).collect::>(); - b.iter(|| Nibbles::unpack(&raw)) + g.warm_up_time(Duration::from_secs(1)); + g.noise_threshold(0.02); + + g.bench_function("unpack/32", |b| { + let bytes = get_bytes(32); + b.iter(|| Nibbles::unpack(black_box(&bytes[..]))) }); + g.bench_function("unpack/256", |b| { + let bytes = get_bytes(256); + b.iter(|| Nibbles::unpack(black_box(&bytes[..]))) + }); + g.bench_function("unpack/2048", |b| { + let bytes = get_bytes(2048); + b.iter(|| Nibbles::unpack(black_box(&bytes[..]))) + }); + + g.bench_function("pack/32", |b| { + let nibbles = get_nibbles(32); + b.iter(|| black_box(&nibbles).pack()) + }); + g.bench_function("pack/256", |b| { + let nibbles = get_nibbles(256); + b.iter(|| black_box(&nibbles).pack()) + }); + g.bench_function("pack/2048", |b| { + let nibbles = get_nibbles(2048); + b.iter(|| black_box(&nibbles).pack()) + }); + + g.bench_function("encode_path_leaf/31", |b| { + let nibbles = get_nibbles(31); + b.iter(|| black_box(&nibbles).encode_path_leaf(false)) + }); + g.bench_function("encode_path_leaf/256", |b| { + let nibbles = get_nibbles(256); + b.iter(|| black_box(&nibbles).encode_path_leaf(false)) + }); + g.bench_function("encode_path_leaf/2048", |b| { + let nibbles = get_nibbles(2048); + b.iter(|| black_box(&nibbles).encode_path_leaf(false)) + }); +} + +fn get_nibbles(len: usize) -> Nibbles { + Nibbles::from_nibbles_unchecked(get_bytes(len)) +} + +fn get_bytes(len: usize) -> Vec { + proptest::collection::vec(proptest::arbitrary::any::(), len) + .new_tree(&mut Default::default()) + .unwrap() + .current() } criterion_group!(benches, nibbles_benchmark); diff --git a/crates/primitives/src/trie/hash_builder/mod.rs b/crates/primitives/src/trie/hash_builder/mod.rs index 89559819e..ed2c86de7 100644 --- a/crates/primitives/src/trie/hash_builder/mod.rs +++ b/crates/primitives/src/trie/hash_builder/mod.rs @@ -63,7 +63,7 @@ pub struct HashBuilder { impl From for HashBuilder { fn from(state: HashBuilderState) -> Self { Self { - key: Nibbles::new_unchecked(state.key), + key: Nibbles::from_nibbles_unchecked(state.key), stack: state.stack, value: state.value, groups: state.groups, @@ -564,7 +564,7 @@ mod tests { let (_, updates) = hb.split(); - let update = updates.get(&Nibbles::new_unchecked(hex!("01"))).unwrap(); + let update = updates.get(&Nibbles::from_nibbles_unchecked(hex!("01"))).unwrap(); assert_eq!(update.state_mask, TrieMask::new(0b1111)); // 1st nibble: 0, 1, 2, 3 assert_eq!(update.tree_mask, TrieMask::new(0)); assert_eq!(update.hash_mask, TrieMask::new(6)); // in the 1st nibble, the ones with 1 and 2 are branches with `hashes` @@ -634,7 +634,7 @@ mod tests { let mut hb2 = HashBuilder::default(); // Insert the branch with the `0x6` shared prefix. - hb2.add_branch(Nibbles::new_unchecked([0x6]), branch_node_hash, false); + hb2.add_branch(Nibbles::from_nibbles_unchecked([0x6]), branch_node_hash, false); let expected = trie_root(raw_input.clone()); assert_eq!(hb.root(), expected); diff --git a/crates/primitives/src/trie/nibbles.rs b/crates/primitives/src/trie/nibbles.rs index aa36f2955..7c6d92a2b 100644 --- a/crates/primitives/src/trie/nibbles.rs +++ b/crates/primitives/src/trie/nibbles.rs @@ -1,9 +1,15 @@ use crate::Bytes; -use alloy_rlp::RlpEncodableWrapper; +use bytes::Buf; use derive_more::{Deref, From, Index}; -use reth_codecs::{main_codec, Compact}; +use reth_codecs::Compact; use serde::{Deserialize, Serialize}; -use std::{borrow::Borrow, ops::RangeBounds}; +use smallvec::SmallVec; +use std::{ + borrow::Borrow, + fmt, + mem::MaybeUninit, + ops::{Bound, RangeBounds}, +}; /// The representation of nibbles of the merkle trie stored in the database. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, PartialOrd, Ord, Hash, Deref)] @@ -19,7 +25,7 @@ impl From for StoredNibblesSubKey { impl From> for StoredNibblesSubKey { #[inline] fn from(value: Vec) -> Self { - Self(Nibbles::new_unchecked(value)) + Self(Nibbles::from_nibbles_unchecked(value)) } } @@ -48,7 +54,7 @@ impl Compact for StoredNibblesSubKey { fn from_compact(buf: &[u8], _len: usize) -> (Self, &[u8]) { let len = buf[64] as usize; - (Self(Nibbles::new_unchecked(buf[..len].to_vec())), &buf[65..]) + (Self(Nibbles::from_nibbles_unchecked(&buf[..len])), &buf[65..]) } } @@ -58,44 +64,92 @@ impl Compact for StoredNibblesSubKey { /// the keys in a Merkle Patricia Trie (MPT). /// Using nibbles simplifies trie operations and enables consistent key representation in the MPT. /// -/// The internal representation is a shared heap-allocated vector ([`Bytes`]) that stores one nibble -/// per byte. This means that each byte has its upper 4 bits set to zero and the lower 4 bits -/// representing the nibble value. -#[main_codec] +/// The internal representation is a [`SmallVec`] that stores one nibble per byte. This means that +/// each byte has its upper 4 bits set to zero and the lower 4 bits representing the nibble value. #[derive( Clone, - Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash, - RlpEncodableWrapper, Index, From, Deref, + serde::Serialize, + serde::Deserialize, )] -pub struct Nibbles(Bytes); +#[cfg_attr(any(test, feature = "arbitrary"), derive(arbitrary::Arbitrary))] +pub struct Nibbles(SmallVec<[u8; 64]>); + +impl alloy_rlp::Encodable for Nibbles { + #[inline] + fn length(&self) -> usize { + alloy_rlp::Encodable::length(self.as_slice()) + } + + #[inline] + fn encode(&self, out: &mut dyn alloy_rlp::BufMut) { + alloy_rlp::Encodable::encode(self.as_slice(), out) + } +} + +#[cfg(any(test, feature = "arbitrary"))] +impl proptest::arbitrary::Arbitrary for Nibbles { + type Parameters = (); + type Strategy = proptest::strategy::Map< + proptest::collection::VecStrategy>, + fn(Vec) -> Self, + >; + + #[inline] + fn arbitrary_with((): ()) -> Self::Strategy { + use proptest::prelude::*; + proptest::collection::vec(0x0..=0xf, 0..64).prop_map(Self::from_nibbles_unchecked) + } +} + +impl Compact for Nibbles { + fn to_compact(self, buf: &mut B) -> usize + where + B: bytes::BufMut + AsMut<[u8]>, + { + buf.put_slice(self.as_slice()); + self.len() + } + + fn from_compact(mut buf: &[u8], len: usize) -> (Self, &[u8]) { + let nibbles = &buf[..len]; + buf.advance(len); + (Nibbles::from_nibbles_unchecked(nibbles), buf) + } +} + +impl fmt::Debug for Nibbles { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_tuple("Nibbles").field(&crate::hex::encode(self.as_slice())).finish() + } +} impl From> for Nibbles { #[inline] fn from(value: Vec) -> Self { - Self::new_unchecked(value) + Self(SmallVec::from_vec(value)) } } impl From for Vec { #[inline] fn from(value: Nibbles) -> Self { - value.0.into() + value.0.into_vec() } } impl From for Bytes { #[inline] fn from(value: Nibbles) -> Self { - value.into_bytes() + value.0.into_vec().into() } } @@ -134,23 +188,81 @@ impl Borrow<[u8]> for Nibbles { } } +impl Extend for Nibbles { + #[inline] + fn extend>(&mut self, iter: T) { + self.0.extend(iter) + } +} + impl Nibbles { + /// Creates a new empty [`Nibbles`] instance. + #[inline] + pub const fn new() -> Self { + Self(SmallVec::new_const()) + } + + /// Creates a new [`Nibbles`] instance with the given capacity. + #[inline] + pub fn with_capacity(capacity: usize) -> Self { + Self(SmallVec::with_capacity(capacity)) + } + /// Creates a new [`Nibbles`] instance from nibble bytes, without checking their validity. #[inline] - pub fn new_unchecked>(nibbles: T) -> Self { - Self(nibbles.into()) + pub fn from_nibbles_unchecked>(nibbles: T) -> Self { + Self(SmallVec::from_slice(nibbles.as_ref())) } /// Converts a byte slice into a [`Nibbles`] instance containing the nibbles (half-bytes or 4 /// bits) that make up the input byte data. + #[inline] pub fn unpack>(data: T) -> Self { let data = data.as_ref(); - let mut nibbles = Vec::with_capacity(data.len() * 2); - for &byte in data { - nibbles.push(byte >> 4); - nibbles.push(byte & 0x0f); + if data.len() <= 32 { + // SAFETY: checked length. + unsafe { Self::unpack_stack(data) } + } else { + Self::unpack_heap(data) + } + } + + /// Unpacks on the stack. + /// + /// # Safety + /// + /// `data.len()` must be less than or equal to 32. + unsafe fn unpack_stack(data: &[u8]) -> Self { + let mut nibbles = MaybeUninit::<[u8; 64]>::uninit(); + Self::unpack_to(data, nibbles.as_mut_ptr().cast()); + let unpacked_len = data.len() * 2; + Self(SmallVec::from_buf_and_len_unchecked(nibbles, unpacked_len)) + } + + /// Unpacks on the heap. + fn unpack_heap(data: &[u8]) -> Self { + // Collect into a vec directly to avoid the smallvec overhead since we know this is going on + // the heap. + let unpacked_len = data.len() * 2; + let mut nibbles = Vec::with_capacity(unpacked_len); + // SAFETY: enough capacity. + unsafe { Self::unpack_to(data, nibbles.as_mut_ptr()) }; + // SAFETY: within capacity and `unpack_to` initialized the memory. + unsafe { nibbles.set_len(unpacked_len) }; + Self(SmallVec::from_vec(nibbles)) + } + + /// Unpacks into the given pointer. + /// + /// # Safety + /// + /// `ptr` must be valid for at least `data.len() * 2` bytes. + #[inline] + unsafe fn unpack_to(data: &[u8], ptr: *mut u8) { + for (i, &byte) in data.iter().enumerate() { + ptr.add(i * 2).write(byte >> 4); + ptr.add(i * 2 + 1).write(byte & 0x0f); } - Self(nibbles.into()) } /// Packs the nibbles stored in the struct into a byte vector. @@ -159,15 +271,68 @@ impl Nibbles { /// 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 packed_len = (self.len() + 1) / 2; - let mut v = Vec::with_capacity(packed_len); - for i in 0..packed_len { - let hi = *unsafe { self.get_unchecked(i * 2) }; - let lo = self.get(i * 2 + 1).copied().unwrap_or(0); - v.push((hi << 4) | lo); + #[inline] + pub fn pack(&self) -> SmallVec<[u8; 32]> { + if self.len() <= 64 { + // SAFETY: checked length. + unsafe { self.pack_stack() } + } else { + self.pack_heap() } - v + } + + /// Packs on the stack. + /// + /// # Safety + /// + /// `self.len()` must be less than or equal to 32. + unsafe fn pack_stack(&self) -> SmallVec<[u8; 32]> { + let mut nibbles = MaybeUninit::<[u8; 32]>::uninit(); + self.pack_to(nibbles.as_mut_ptr().cast()); + let packed_len = (self.len() + 1) / 2; + SmallVec::from_buf_and_len_unchecked(nibbles, packed_len) + } + + /// Packs on the heap. + fn pack_heap(&self) -> SmallVec<[u8; 32]> { + // Collect into a vec directly to avoid the smallvec overhead since we know this is going on + // the heap. + let packed_len = (self.len() + 1) / 2; + let mut vec = Vec::with_capacity(packed_len); + // SAFETY: enough capacity. + unsafe { self.pack_to(vec.as_mut_ptr()) }; + // SAFETY: within capacity and `pack_to` initialized the memory. + unsafe { vec.set_len(packed_len) }; + SmallVec::from_vec(vec) + } + + /// Packs into the given pointer. + /// + /// # Safety + /// + /// `ptr` must be valid for at least `self.len() / 2 + IS_ODD` bytes. + #[inline] + unsafe fn pack_to(&self, ptr: *mut u8) { + for i in 0..self.len() / 2 { + ptr.add(i).write(self.get_byte_unchecked(i * 2)); + } + if self.len() % 2 != 0 { + let i = self.len() / 2; + ptr.add(i).write(self.last().unwrap_unchecked() << 4); + } + } + + /// Gets the byte at the given index by combining two consecutive nibbles. + /// + /// # Safety + /// + /// `i..i + 1` must be in range. + #[inline] + unsafe fn get_byte_unchecked(&self, i: usize) -> u8 { + debug_assert!(i + 1 < self.len(), "index {i}..{} out of bounds of {}", i + 1, self.len()); + let hi = *self.get_unchecked(i); + let lo = *self.get_unchecked(i + 1); + (hi << 4) | lo } /// Encodes a given path leaf as a compact array of bytes, where each byte represents two @@ -189,60 +354,77 @@ impl Nibbles { /// /// # Returns /// - /// A `Vec` containing the compact byte representation of the nibble sequence, including the + /// A vector containing the compact byte representation of the nibble sequence, including the /// header byte. /// - /// # Example + /// This vector's length is `self.len() / 2 + 1`. For stack-allocated nibbles, this is at most + /// 33 bytes, so 36 was chosen as the stack capacity to round up to the next usize-aligned + /// size. + /// + /// # Examples /// /// ``` /// # use reth_primitives::trie::Nibbles; /// /// // Extension node with an even path length: - /// let nibbles = Nibbles::new_unchecked(&[0x0A, 0x0B, 0x0C, 0x0D]); - /// assert_eq!(nibbles.encode_path_leaf(false), vec![0x00, 0xAB, 0xCD]); + /// let nibbles = Nibbles::from_nibbles_unchecked(&[0x0A, 0x0B, 0x0C, 0x0D]); + /// assert_eq!(nibbles.encode_path_leaf(false)[..], [0x00, 0xAB, 0xCD]); /// /// // Extension node with an odd path length: - /// let nibbles = Nibbles::new_unchecked(&[0x0A, 0x0B, 0x0C]); - /// assert_eq!(nibbles.encode_path_leaf(false), vec![0x1A, 0xBC]); + /// let nibbles = Nibbles::from_nibbles_unchecked(&[0x0A, 0x0B, 0x0C]); + /// assert_eq!(nibbles.encode_path_leaf(false)[..], [0x1A, 0xBC]); /// /// // Leaf node with an even path length: - /// let nibbles = Nibbles::new_unchecked(&[0x0A, 0x0B, 0x0C, 0x0D]); - /// assert_eq!(nibbles.encode_path_leaf(true), vec![0x20, 0xAB, 0xCD]); + /// let nibbles = Nibbles::from_nibbles_unchecked(&[0x0A, 0x0B, 0x0C, 0x0D]); + /// assert_eq!(nibbles.encode_path_leaf(true)[..], [0x20, 0xAB, 0xCD]); /// /// // Leaf node with an odd path length: - /// let nibbles = Nibbles::new_unchecked(&[0x0A, 0x0B, 0x0C]); - /// assert_eq!(nibbles.encode_path_leaf(true), vec![0x3A, 0xBC]); + /// let nibbles = Nibbles::from_nibbles_unchecked(&[0x0A, 0x0B, 0x0C]); + /// assert_eq!(nibbles.encode_path_leaf(true)[..], [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; + pub fn encode_path_leaf(&self, is_leaf: bool) -> SmallVec<[u8; 36]> { + let encoded_len = self.len() / 2 + 1; + let mut encoded = SmallVec::with_capacity(encoded_len); + // SAFETY: enough capacity. + unsafe { self.encode_path_leaf_to(is_leaf, encoded.as_mut_ptr()) }; + // SAFETY: within capacity and `encode_path_leaf_to` initialized the memory. + unsafe { encoded.set_len(encoded_len) }; + encoded + } - // Set the first byte of the encoded vector. - encoded[0] = match (is_leaf, odd_nibbles) { + /// # Safety + /// + /// `ptr` must be valid for at least `self.len() / 2 + 1` bytes. + #[inline] + unsafe fn encode_path_leaf_to(&self, is_leaf: bool, ptr: *mut u8) { + let odd_nibbles = self.len() % 2 != 0; + *ptr = self.encode_path_leaf_first_byte(is_leaf, odd_nibbles); + let mut nibble_idx = if odd_nibbles { 1 } else { 0 }; + for i in 0..self.len() / 2 { + ptr.add(i + 1).write(self.get_byte_unchecked(nibble_idx)); + nibble_idx += 2; + } + } + + #[inline] + fn encode_path_leaf_first_byte(&self, is_leaf: bool, odd_nibbles: bool) -> u8 { + 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.0.to_vec(); + let mut incremented = self.clone(); - for nibble in incremented.iter_mut().rev() { - debug_assert!(*nibble < 0x10); + for nibble in incremented.0.iter_mut().rev() { + debug_assert!(*nibble <= 0xf); if *nibble < 0xf { *nibble += 1; - return Some(Self::new_unchecked(incremented)) + return Some(incremented); } else { *nibble = 0; } @@ -272,7 +454,13 @@ impl Nibbles { #[inline] #[track_caller] pub fn at(&self, i: usize) -> usize { - self.0[i] as usize + self[i] as usize + } + + /// Returns the first nibble of the current nibble sequence. + #[inline] + pub fn first(&self) -> Option { + self.0.first().copied() } /// Returns the last nibble of the current nibble sequence. @@ -293,61 +481,58 @@ impl Nibbles { len } - /// Returns a reference to the underlying [`Bytes`]. - #[inline] - pub fn as_bytes(&self) -> &Bytes { - &self.0 - } - /// Returns the nibbles as a byte slice. #[inline] pub fn as_slice(&self) -> &[u8] { &self.0 } - /// Returns the underlying [`Bytes`]. - #[inline] - pub fn into_bytes(self) -> Bytes { - self.0 - } - /// Slice the current nibbles within the provided index range. + /// + /// # Panics + /// + /// Panics if the range is out of bounds. #[inline] + #[track_caller] pub fn slice(&self, range: impl RangeBounds) -> Self { - Self(self.0.slice(range)) + let start = match range.start_bound() { + Bound::Included(&n) => n, + Bound::Excluded(&n) => n.checked_add(1).expect("out of range"), + Bound::Unbounded => 0, + }; + let end = match range.end_bound() { + Bound::Included(&n) => n.checked_add(1).expect("out of range"), + Bound::Excluded(&n) => n, + Bound::Unbounded => self.len(), + }; + Self::from_nibbles_unchecked(&self[start..end]) } /// Join two nibbles together. #[inline] pub fn join(&self, b: &Self) -> Self { - let mut hex_data = Vec::with_capacity(self.len() + b.len()); - hex_data.extend_from_slice(self); - hex_data.extend_from_slice(b); - Self::new_unchecked(hex_data) + let mut nibbles = SmallVec::with_capacity(self.len() + b.len()); + nibbles.extend_from_slice(self); + nibbles.extend_from_slice(b); + Self(nibbles) } /// Pushes a nibble to the end of the current nibbles. - /// - /// **Note**: This method re-allocates on each call. #[inline] pub fn push(&mut self, nibble: u8) { - self.extend([nibble]); + self.0.push(nibble); } /// Extend the current nibbles with another nibbles. - /// - /// **Note**: This method re-allocates on each call. #[inline] - pub fn extend(&mut self, b: impl AsRef<[u8]>) { - let mut bytes = self.0.to_vec(); - bytes.extend_from_slice(b.as_ref()); - self.0 = bytes.into(); + pub fn extend_from_slice(&mut self, b: impl AsRef<[u8]>) { + self.0.extend_from_slice(b.as_ref()); } /// Truncates the current nibbles to the given length. #[inline] - pub fn truncate(&mut self, len: usize) { - self.0.truncate(len); + pub fn truncate(&mut self, new_len: usize) { + self.0.truncate(new_len); } /// Clears the current nibbles. @@ -365,40 +550,75 @@ mod tests { #[test] fn hashed_regression() { - let nibbles = Nibbles::new_unchecked(hex!("05010406040a040203030f010805020b050c04070003070e0909070f010b0a0805020301070c0a0902040b0f000f0006040a04050f020b090701000a0a040b")); + let nibbles = Nibbles::from_nibbles_unchecked(hex!("05010406040a040203030f010805020b050c04070003070e0909070f010b0a0805020301070c0a0902040b0f000f0006040a04050f020b090701000a0a040b")); let path = nibbles.encode_path_leaf(true); let expected = hex!("351464a4233f1852b5c47037e997f1ba852317ca924bf0f064a45f2b9710aa4b"); - assert_eq!(path, expected); + 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::new_unchecked(input); + let tests = [ + (&[][..], &[][..]), + (&[0xa], &[0xa0]), + (&[0xa, 0x0], &[0xa0]), + (&[0xa, 0xb], &[0xab]), + (&[0xa, 0xb, 0x2], &[0xab, 0x20]), + (&[0xa, 0xb, 0x2, 0x0], &[0xab, 0x20]), + (&[0xa, 0xb, 0x2, 0x7], &[0xab, 0x27]), + ]; + for (input, expected) in tests { + assert!(input.iter().all(|&x| x <= 0xf)); + let nibbles = Nibbles::from_nibbles_unchecked(input); let encoded = nibbles.pack(); - assert_eq!(encoded, expected); + assert_eq!(&encoded[..], expected); } } + #[test] + fn slice() { + const RAW: &[u8] = &hex!("05010406040a040203030f010805020b050c04070003070e0909070f010b0a0805020301070c0a0902040b0f000f0006040a04050f020b090701000a0a040b"); + + #[track_caller] + fn test_slice(range: impl RangeBounds, expected: &[u8]) { + let nibbles = Nibbles::from_nibbles_unchecked(RAW); + let sliced = nibbles.slice(range); + assert_eq!(sliced, Nibbles::from_nibbles_unchecked(expected)); + assert_eq!(sliced.as_slice(), expected); + } + + test_slice(0..0, &[]); + test_slice(0..1, &[0x05]); + test_slice(1..1, &[]); + test_slice(1..=1, &[0x01]); + test_slice(0..=1, &[0x05, 0x01]); + test_slice(0..2, &[0x05, 0x01]); + + test_slice(..0, &[]); + test_slice(..1, &[0x05]); + test_slice(..=1, &[0x05, 0x01]); + test_slice(..2, &[0x05, 0x01]); + + test_slice(.., RAW); + test_slice(..RAW.len(), RAW); + test_slice(0.., RAW); + test_slice(0..RAW.len(), RAW); + } + proptest! { #[test] fn pack_unpack_roundtrip(input in any::>()) { let nibbles = Nibbles::unpack(&input); + prop_assert!(nibbles.iter().all(|&nibble| nibble <= 0xf)); let packed = nibbles.pack(); - prop_assert_eq!(packed, input); + prop_assert_eq!(&packed[..], input); } #[test] fn encode_path_first_byte(input in any::>()) { prop_assume!(!input.is_empty()); let input = Nibbles::unpack(input); + prop_assert!(input.iter().all(|&nibble| nibble <= 0xf)); let input_is_odd = input.len() % 2 == 1; let compact_leaf = input.encode_path_leaf(true); @@ -407,7 +627,7 @@ mod tests { 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()); + assert_eq!(leaf_flag & 0x0f, input.first().unwrap()); } @@ -417,7 +637,7 @@ mod tests { 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()); + assert_eq!(extension_flag & 0x0f, input.first().unwrap()); } } } diff --git a/crates/primitives/src/trie/nodes/extension.rs b/crates/primitives/src/trie/nodes/extension.rs index 09fe529ec..35233fdff 100644 --- a/crates/primitives/src/trie/nodes/extension.rs +++ b/crates/primitives/src/trie/nodes/extension.rs @@ -1,5 +1,6 @@ use super::{super::Nibbles, rlp_node}; use alloy_rlp::{BufMut, Encodable}; +use smallvec::SmallVec; /// An intermediate node that exists solely to compress the trie's paths. It contains a path segment /// (a shared prefix of keys) and a single child pointer. Essentially, an extension node can be @@ -9,8 +10,8 @@ use alloy_rlp::{BufMut, Encodable}; /// with a single child into one node. This simplification reduces the space and computational /// complexity when performing operations on the trie. pub struct ExtensionNode<'a> { - /// A common prefix for keys. - pub prefix: Vec, + /// A common prefix for keys. See [`Nibbles::encode_path_leaf`] for more information. + pub prefix: SmallVec<[u8; 36]>, /// A pointer to the child. pub node: &'a [u8], } diff --git a/crates/primitives/src/trie/nodes/leaf.rs b/crates/primitives/src/trie/nodes/leaf.rs index 684bf913d..3de18de81 100644 --- a/crates/primitives/src/trie/nodes/leaf.rs +++ b/crates/primitives/src/trie/nodes/leaf.rs @@ -1,5 +1,6 @@ use super::{super::Nibbles, rlp_node}; use alloy_rlp::{BufMut, Encodable}; +use smallvec::SmallVec; /// A leaf node represents the endpoint or terminal node in the trie. In other words, a leaf node is /// where actual values are stored. @@ -10,9 +11,9 @@ use alloy_rlp::{BufMut, Encodable}; /// node means that the search has successfully found the value associated with that key. #[derive(Default)] pub struct LeafNode<'a> { - /// The key path. - pub key: Vec, - /// value: SmallVec<[u8; 36]> + /// The key path. See [`Nibbles::encode_path_leaf`] for more information. + pub key: SmallVec<[u8; 36]>, + /// The node value. pub value: &'a [u8], } @@ -59,14 +60,14 @@ mod tests { // From manual regression test #[test] fn encode_leaf_node_nibble() { - let nibble = Nibbles::new_unchecked(hex!("0604060f")); + let nibble = Nibbles::from_nibbles_unchecked(hex!("0604060f")); let encoded = nibble.encode_path_leaf(true); - assert_eq!(encoded, hex!("20646f")); + assert_eq!(encoded[..], hex!("20646f")); } #[test] fn rlp_leaf_node_roundtrip() { - let nibble = Nibbles::new_unchecked(hex!("0604060f")); + let nibble = Nibbles::from_nibbles_unchecked(hex!("0604060f")); let val = hex!("76657262"); let leaf = LeafNode::new(&nibble, &val); let rlp = leaf.rlp(&mut vec![]); diff --git a/crates/storage/codecs/src/lib.rs b/crates/storage/codecs/src/lib.rs index bfe096d12..2e9a8d3cb 100644 --- a/crates/storage/codecs/src/lib.rs +++ b/crates/storage/codecs/src/lib.rs @@ -70,7 +70,9 @@ macro_rules! impl_uint_compact { ($($name:tt),+) => { $( impl Compact for $name { - fn to_compact(self, buf: &mut B) -> usize where B: bytes::BufMut + AsMut<[u8]> { + fn to_compact(self, buf: &mut B) -> usize + where B: bytes::BufMut + AsMut<[u8]> + { let leading = self.leading_zeros() as usize / 8; buf.put_slice(&self.to_be_bytes()[leading..]); std::mem::size_of::<$name>() - leading @@ -262,44 +264,49 @@ impl Compact for Bytes { } } -/// Implements the [`Compact`] trait for fixed size hash types like [`B256`]. +impl Compact for [u8; N] { + fn to_compact(self, buf: &mut B) -> usize + where + B: bytes::BufMut + AsMut<[u8]>, + { + buf.put_slice(&self); + N + } + + fn from_compact(mut buf: &[u8], len: usize) -> (Self, &[u8]) { + if len == 0 { + return ([0; N], buf) + } + + let v = buf[..N].try_into().unwrap(); + buf.advance(N); + (v, buf) + } +} + +/// Implements the [`Compact`] trait for fixed size byte array types like [`B256`]. #[macro_export] -macro_rules! impl_hash_compact { +macro_rules! impl_compact_for_bytes { ($($name:tt),+) => { $( impl Compact for $name { - fn to_compact(self, buf: &mut B) -> usize where B: bytes::BufMut + AsMut<[u8]> { - buf.put_slice(self.as_slice()); - std::mem::size_of::<$name>() - } - - fn from_compact(mut buf: &[u8], len: usize) -> (Self,&[u8]) { - if len == 0 { - return ($name::default(), buf) - } - - let v = $name::from_slice( - buf.get(..std::mem::size_of::<$name>()).expect("size not matching"), - ); - buf.advance(std::mem::size_of::<$name>()); - (v, buf) - } - - fn specialized_to_compact(self, buf: &mut B) -> usize + fn to_compact(self, buf: &mut B) -> usize where - B: bytes::BufMut + AsMut<[u8]> { - self.to_compact(buf) + B: bytes::BufMut + AsMut<[u8]> + { + self.0.to_compact(buf) } - fn specialized_from_compact(buf: &[u8], len: usize) -> (Self, &[u8]) { - Self::from_compact(buf, len) + fn from_compact(buf: &[u8], len: usize) -> (Self, &[u8]) { + let (v, buf) = <[u8; std::mem::size_of::<$name>()]>::from_compact(buf, len); + (Self::from(v), buf) } } )+ }; } -impl_hash_compact!(Address, B256, B512, Bloom); +impl_compact_for_bytes!(Address, B256, B512, Bloom); impl Compact for bool { /// `bool` vars go directly to the `StructFlags` and are not written to the buffer. diff --git a/crates/trie/benches/prefix_set.rs b/crates/trie/benches/prefix_set.rs index c9e4dbbd6..2fd8ff28d 100644 --- a/crates/trie/benches/prefix_set.rs +++ b/crates/trie/benches/prefix_set.rs @@ -94,12 +94,12 @@ fn generate_test_data(size: usize) -> (Vec, Vec, Vec) { let mut preload = vec(vec(any::(), 32), size).new_tree(&mut runner).unwrap().current(); preload.dedup(); preload.sort(); - let preload = preload.into_iter().map(Nibbles::new_unchecked).collect::>(); + let preload = preload.into_iter().map(Nibbles::from_nibbles_unchecked).collect::>(); let mut input = vec(vec(any::(), 0..=32), size).new_tree(&mut runner).unwrap().current(); input.dedup(); input.sort(); - let input = input.into_iter().map(Nibbles::new_unchecked).collect::>(); + let input = input.into_iter().map(Nibbles::from_nibbles_unchecked).collect::>(); let expected = input .iter() diff --git a/crates/trie/src/prefix_set/mod.rs b/crates/trie/src/prefix_set/mod.rs index 0c262759a..9b80c703c 100644 --- a/crates/trie/src/prefix_set/mod.rs +++ b/crates/trie/src/prefix_set/mod.rs @@ -26,8 +26,8 @@ pub use loader::{LoadedPrefixSets, PrefixSetLoader}; /// use reth_trie::prefix_set::PrefixSetMut; /// /// let mut prefix_set = PrefixSetMut::default(); -/// prefix_set.insert(Nibbles::new_unchecked(&[0xa, 0xb])); -/// prefix_set.insert(Nibbles::new_unchecked(&[0xa, 0xb, 0xc])); +/// prefix_set.insert(Nibbles::from_nibbles_unchecked(&[0xa, 0xb])); +/// prefix_set.insert(Nibbles::from_nibbles_unchecked(&[0xa, 0xb, 0xc])); /// assert!(prefix_set.contains(&[0xa, 0xb])); /// assert!(prefix_set.contains(&[0xa, 0xb, 0xc])); /// ``` @@ -158,10 +158,10 @@ mod tests { #[test] fn test_contains_with_multiple_inserts_and_duplicates() { let mut prefix_set = PrefixSetMut::default(); - prefix_set.insert(Nibbles::new_unchecked(b"123")); - prefix_set.insert(Nibbles::new_unchecked(b"124")); - prefix_set.insert(Nibbles::new_unchecked(b"456")); - prefix_set.insert(Nibbles::new_unchecked(b"123")); // Duplicate + prefix_set.insert(Nibbles::from_nibbles_unchecked(b"123")); + prefix_set.insert(Nibbles::from_nibbles_unchecked(b"124")); + prefix_set.insert(Nibbles::from_nibbles_unchecked(b"456")); + prefix_set.insert(Nibbles::from_nibbles_unchecked(b"123")); // Duplicate assert!(prefix_set.contains(b"12")); assert!(prefix_set.contains(b"45")); diff --git a/crates/trie/src/trie_cursor/account_cursor.rs b/crates/trie/src/trie_cursor/account_cursor.rs index b98be8bd2..94e7be590 100644 --- a/crates/trie/src/trie_cursor/account_cursor.rs +++ b/crates/trie/src/trie_cursor/account_cursor.rs @@ -41,7 +41,6 @@ where #[cfg(test)] mod tests { - use super::*; use reth_db::{ cursor::{DbCursorRO, DbCursorRW}, diff --git a/crates/trie/src/trie_cursor/subnode.rs b/crates/trie/src/trie_cursor/subnode.rs index d49ccfe7f..27f332fed 100644 --- a/crates/trie/src/trie_cursor/subnode.rs +++ b/crates/trie/src/trie_cursor/subnode.rs @@ -39,7 +39,7 @@ impl From for CursorSubNode { Some(n) => n as i8, None => -1, }; - Self { key: Nibbles::new_unchecked(value.key), nibble, node: value.node } + Self { key: Nibbles::from_nibbles_unchecked(value.key), nibble, node: value.node } } } diff --git a/crates/trie/src/walker.rs b/crates/trie/src/walker.rs index 54e7fc7ec..67da36d74 100644 --- a/crates/trie/src/walker.rs +++ b/crates/trie/src/walker.rs @@ -131,7 +131,7 @@ impl TrieWalker { assert!(!node.state_mask.is_empty()); } - Ok(entry.map(|(k, v)| (Nibbles::new_unchecked(k), v))) + Ok(entry.map(|(k, v)| (Nibbles::from_nibbles_unchecked(k), v))) } /// Consumes the next node in the trie, updating the stack. @@ -313,7 +313,7 @@ mod tests { // We're traversing the path in lexigraphical order. for expected in expected { let got = walker.advance().unwrap(); - assert_eq!(got.unwrap(), Nibbles::new_unchecked(expected.clone())); + assert_eq!(got.unwrap(), Nibbles::from_nibbles_unchecked(expected.clone())); } // There should be 8 paths traversed in total from 3 branches. @@ -361,26 +361,26 @@ mod tests { // No changes let mut cursor = TrieWalker::new(&mut trie, Default::default()); - assert_eq!(cursor.key(), Some(Nibbles::new_unchecked([]))); // root + assert_eq!(cursor.key(), Some(Nibbles::from_nibbles_unchecked([]))); // root assert!(cursor.can_skip_current_node); // due to root_hash cursor.advance().unwrap(); // skips to the end of trie assert_eq!(cursor.key(), None); // We insert something that's not part of the existing trie/prefix. let mut changed = PrefixSetMut::default(); - changed.insert(Nibbles::new_unchecked([0xF, 0x1])); + changed.insert(Nibbles::from_nibbles_unchecked([0xF, 0x1])); let mut cursor = TrieWalker::new(&mut trie, changed.freeze()); // Root node - assert_eq!(cursor.key(), Some(Nibbles::new_unchecked([]))); + assert_eq!(cursor.key(), Some(Nibbles::from_nibbles_unchecked([]))); // Should not be able to skip state due to the changed values assert!(!cursor.can_skip_current_node); cursor.advance().unwrap(); - assert_eq!(cursor.key(), Some(Nibbles::new_unchecked([0x2]))); + assert_eq!(cursor.key(), Some(Nibbles::from_nibbles_unchecked([0x2]))); cursor.advance().unwrap(); - assert_eq!(cursor.key(), Some(Nibbles::new_unchecked([0x2, 0x1]))); + assert_eq!(cursor.key(), Some(Nibbles::from_nibbles_unchecked([0x2, 0x1]))); cursor.advance().unwrap(); - assert_eq!(cursor.key(), Some(Nibbles::new_unchecked([0x4]))); + assert_eq!(cursor.key(), Some(Nibbles::from_nibbles_unchecked([0x4]))); cursor.advance().unwrap(); assert_eq!(cursor.key(), None); // the end of trie