From 09bf1c84ed4e97c206a5462f4987fb0e3e12749c Mon Sep 17 00:00:00 2001 From: frisitano Date: Thu, 12 Dec 2024 19:42:34 +0000 Subject: [PATCH 1/4] add additional test coverage for state root --- Cargo.lock | 1 + crates/primitives-traits/src/account.rs | 14 ++- crates/scroll/state-commitment/Cargo.toml | 1 + crates/scroll/state-commitment/src/lib.rs | 5 +- .../src/{root.rs => root/mod.rs} | 3 + .../scroll/state-commitment/src/root/utils.rs | 66 ++++++++++++++ crates/scroll/state-commitment/src/test.rs | 86 +++++++++++++++++-- 7 files changed, 168 insertions(+), 8 deletions(-) rename crates/scroll/state-commitment/src/{root.rs => root/mod.rs} (99%) create mode 100644 crates/scroll/state-commitment/src/root/utils.rs diff --git a/Cargo.lock b/Cargo.lock index ef166e9a5d3c..eaa45ddd1574 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9423,6 +9423,7 @@ dependencies = [ "alloy-consensus", "alloy-primitives", "alloy-rlp", + "itertools 0.13.0", "metrics", "poseidon-bn254 0.1.0 (git+https://github.com/scroll-tech/poseidon-bn254?rev=526a64a)", "proptest", diff --git a/crates/primitives-traits/src/account.rs b/crates/primitives-traits/src/account.rs index a881dcbdbc31..d1914c3f8697 100644 --- a/crates/primitives-traits/src/account.rs +++ b/crates/primitives-traits/src/account.rs @@ -25,7 +25,7 @@ pub mod compact_ids { /// An Ethereum account. #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] -#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)] +#[derive(Clone, Copy, Debug, PartialEq, Eq)] #[cfg_attr(any(test, feature = "arbitrary"), derive(arbitrary::Arbitrary))] #[cfg_attr(any(test, feature = "reth-codec"), derive(reth_codecs::Compact))] #[cfg_attr(any(test, feature = "reth-codec"), reth_codecs::add_arbitrary_tests(compact))] @@ -47,6 +47,18 @@ pub struct Account { pub account_extension: Option, } +impl Default for Account { + fn default() -> Self { + Self { + nonce: 0, + balance: U256::ZERO, + bytecode_hash: None, + #[cfg(feature = "scroll")] + account_extension: Some(Default::default()), + } + } +} + impl Account { /// Whether the account has bytecode. pub const fn has_bytecode(&self) -> bool { diff --git a/crates/scroll/state-commitment/Cargo.toml b/crates/scroll/state-commitment/Cargo.toml index cb90692ccaaa..987842193fb3 100644 --- a/crates/scroll/state-commitment/Cargo.toml +++ b/crates/scroll/state-commitment/Cargo.toml @@ -38,6 +38,7 @@ zktrie = { git = "https://github.com/scroll-tech/zktrie.git", rev = "309160464c1 # misc tracing.workspace = true +itertools.workspace = true [dev-dependencies] reth-db-api.workspace = true diff --git a/crates/scroll/state-commitment/src/lib.rs b/crates/scroll/state-commitment/src/lib.rs index a6995479e175..207de1965a3b 100644 --- a/crates/scroll/state-commitment/src/lib.rs +++ b/crates/scroll/state-commitment/src/lib.rs @@ -7,7 +7,7 @@ mod commitment; pub use commitment::BinaryMerklePatriciaTrie; mod root; -pub use root::{StateRoot, StorageRoot}; +pub use root::*; mod key; pub use key::PoseidonKeyHasher; @@ -19,5 +19,6 @@ pub use value::PoseidonValueHasher; #[cfg(feature = "test-utils")] pub mod test_utils; -#[cfg(all(test, feature = "scroll"))] +// #[cfg(all(test, feature = "scroll"))] +#[cfg(test)] mod test; diff --git a/crates/scroll/state-commitment/src/root.rs b/crates/scroll/state-commitment/src/root/mod.rs similarity index 99% rename from crates/scroll/state-commitment/src/root.rs rename to crates/scroll/state-commitment/src/root/mod.rs index 25969d33299a..1b350afd19cd 100644 --- a/crates/scroll/state-commitment/src/root.rs +++ b/crates/scroll/state-commitment/src/root/mod.rs @@ -21,6 +21,9 @@ use tracing::{debug, trace}; #[cfg(feature = "metrics")] use reth_trie::metrics::{StateRootMetrics, TrieRootMetrics, TrieType}; +mod utils; +pub use utils::*; + // TODO(scroll): Instead of introducing this new type we should make StateRoot generic over // the [`HashBuilder`] and key traversal types diff --git a/crates/scroll/state-commitment/src/root/utils.rs b/crates/scroll/state-commitment/src/root/utils.rs new file mode 100644 index 000000000000..9d482460ef89 --- /dev/null +++ b/crates/scroll/state-commitment/src/root/utils.rs @@ -0,0 +1,66 @@ +//! Utilities for computing state root and storage root from a plain provided state. + +use crate::{PoseidonKeyHasher, PoseidonValueHasher, ScrollTrieAccount}; +use alloy_primitives::{Address, B256, U256}; +use itertools::Itertools; +use reth_scroll_trie::HashBuilder; +use reth_trie::{key::BitsCompatibility, KeyHasher, Nibbles}; + +/// Hashes and sorts account keys, then proceeds to calculating the root hash of the state +/// represented as MPT. +pub fn state_root_ref_unhashed<'a, A: Into + Clone + 'a>( + state: impl IntoIterator, +) -> B256 { + state_root_unsorted( + state + .into_iter() + .map(|(address, account)| (PoseidonKeyHasher::hash_key(address), account.clone())), + ) +} + +/// Sorts the hashed account keys and calculates the root hash of the state represented as MPT. +pub fn state_root_unsorted>( + state: impl IntoIterator, +) -> B256 { + state_root(state.into_iter().sorted_unstable_by_key(|(key, _)| *key)) +} + +/// Calculates the root hash of the state represented as BMPT. +/// +/// # Panics +/// +/// If the items are not in sorted order. +pub fn state_root>(state: impl IntoIterator) -> B256 { + let mut hb = HashBuilder::default(); + for (hashed_key, account) in state { + let account_hash = PoseidonValueHasher::hash_account(account.into()); + hb.add_leaf(Nibbles::unpack_bits(hashed_key), account_hash.as_slice()); + } + hb.root() +} + +/// Hashes storage keys, sorts them and them calculates the root hash of the storage trie. +pub fn storage_root_unhashed(storage: impl IntoIterator) -> B256 { + storage_root_unsorted( + storage.into_iter().map(|(slot, value)| (PoseidonKeyHasher::hash_key(slot), value)), + ) +} + +/// Sorts and calculates the root hash of account storage trie. +pub fn storage_root_unsorted(storage: impl IntoIterator) -> B256 { + storage_root(storage.into_iter().sorted_unstable_by_key(|(key, _)| *key)) +} + +/// Calculates the root hash of account storage trie. +/// +/// # Panics +/// +/// If the items are not in sorted order. +pub fn storage_root(storage: impl IntoIterator) -> B256 { + let mut hb = HashBuilder::default(); + for (hashed_slot, value) in storage { + let hashed_value = PoseidonValueHasher::hash_storage(value); + hb.add_leaf(Nibbles::unpack_bits(hashed_slot), hashed_value.as_ref()); + } + hb.root() +} diff --git a/crates/scroll/state-commitment/src/test.rs b/crates/scroll/state-commitment/src/test.rs index 0423d6ed3975..264106a5ec54 100644 --- a/crates/scroll/state-commitment/src/test.rs +++ b/crates/scroll/state-commitment/src/test.rs @@ -9,23 +9,26 @@ use reth_db::{ transaction::DbTxMut, }; use reth_primitives::{Account, StorageEntry}; -use reth_provider::test_utils::create_test_provider_factory; +use reth_provider::{test_utils::create_test_provider_factory, StorageTrieWriter, TrieWriter}; use reth_scroll_state_commitment::{ test_utils::{b256_clear_last_byte, b256_reverse_bits, u256_clear_msb}, PoseidonKeyHasher, StateRoot, StorageRoot, }; -use reth_scroll_state_commitment::test_utils::b256_clear_first_byte; +use reth_scroll_state_commitment::state_root_unsorted; use reth_trie::{ - trie_cursor::InMemoryTrieCursorFactory, updates::TrieUpdates, HashedPostState, HashedStorage, - KeyHasher, + key::BitsCompatibility, + prefix_set::{PrefixSetMut, TriePrefixSets}, + trie_cursor::InMemoryTrieCursorFactory, + updates::TrieUpdates, + HashedPostState, HashedStorage, KeyHasher, Nibbles, }; use reth_trie_db::{DatabaseStateRoot, DatabaseStorageRoot, DatabaseTrieCursorFactory}; use std::collections::BTreeMap; proptest! { #![proptest_config(ProptestConfig { - cases: 6, ..ProptestConfig::default() + cases: 12, ..ProptestConfig::default() })] #[test] @@ -169,6 +172,79 @@ proptest! { assert_eq!(expected_root, storage_root); } } + + #[test] + fn fuzz_state_root_incremental_database(account_changes: [BTreeMap; 5]) { + let factory = create_test_provider_factory(); + let tx = factory.provider_rw().unwrap(); + let mut hashed_account_cursor = tx.tx_ref().cursor_write::().unwrap(); + + let mut state = BTreeMap::default(); + for accounts in account_changes { + let should_generate_changeset = !state.is_empty(); + let mut changes = PrefixSetMut::default(); + let accounts = accounts.into_iter().map(|(hashed_address, balance)| { + let balance = u256_clear_msb(balance); + let hashed_address = b256_clear_last_byte(hashed_address); + (hashed_address, balance) + }).collect::>(); + + for (hashed_address, balance) in accounts.clone() { + hashed_account_cursor.upsert(hashed_address, Account { balance, ..Default::default() }).unwrap(); + if should_generate_changeset { + changes.insert(Nibbles::unpack_bits(hashed_address)); + } + } + + let (state_root, trie_updates) = StateRoot::from_tx(tx.tx_ref()) + .with_prefix_sets(TriePrefixSets { account_prefix_set: changes.freeze(), ..Default::default() }) + .root_with_updates() + .unwrap(); + + state.append(&mut accounts.clone()); + let expected_root = state_root_unsorted( + state.iter().map(|(&key, &balance)| (key, (Account { balance, ..Default::default() }, B256::default()))) + ); + assert_eq!(expected_root, state_root); + tx.write_trie_updates(&trie_updates).unwrap(); + } + } + + #[test] + fn incremental_storage_root_database(address: Address, storage_changes: [BTreeMap; 5]) { + let hashed_address = PoseidonKeyHasher::hash_key(address); + let factory = create_test_provider_factory(); + let tx = factory.provider_rw().unwrap(); + let mut hashed_storage_cursor = tx.tx_ref().cursor_write::().unwrap(); + + let mut storage = BTreeMap::default(); + for change_set in storage_changes { + let should_generate_changeset = !storage.is_empty(); + let mut changes = PrefixSetMut::default(); + let change_set = change_set.into_iter().map(|(slot, value)| { + let slot = b256_clear_last_byte(slot); + (slot, value) + }).collect::>(); + + for (hashed_slot, value) in change_set.clone() { + hashed_storage_cursor.upsert(hashed_address, StorageEntry { key: hashed_slot, value }).unwrap(); + if should_generate_changeset { + changes.insert(Nibbles::unpack_bits(hashed_slot)); + } + } + + let (storage_root, _, trie_updates) = StorageRoot::from_tx_hashed(tx.tx_ref(), hashed_address) + .with_prefix_set(changes.freeze()) + .root_with_updates() + .unwrap(); + + storage.append(&mut change_set.clone()); + let expected_root = reth_scroll_state_commitment::storage_root_unsorted(storage.clone()); + assert_eq!(expected_root, storage_root); + + tx.write_individual_storage_trie_updates(hashed_address, &trie_updates).unwrap(); + } + } } #[test] From 49465a9e93125c269caa1ac0c5f31ff76253989a Mon Sep 17 00:00:00 2001 From: frisitano Date: Fri, 13 Dec 2024 12:14:07 +0000 Subject: [PATCH 2/4] fix: unpack and pack nibbles --- Cargo.lock | 2 +- crates/scroll/state-commitment/src/lib.rs | 3 +- .../scroll/state-commitment/src/root/mod.rs | 5 +- .../scroll/state-commitment/src/root/utils.rs | 2 +- crates/scroll/state-commitment/src/test.rs | 8 +- crates/scroll/trie/src/hash_builder.rs | 2 +- crates/scroll/trie/src/leaf.rs | 2 +- .../src/providers/database/provider.rs | 15 +-- crates/trie/common/Cargo.toml | 5 +- crates/trie/common/src/key.rs | 111 ++++++++++++++++++ crates/trie/common/src/lib.rs | 2 +- crates/trie/common/src/updates.rs | 16 +-- crates/trie/db/src/prefix_set.rs | 8 +- crates/trie/trie/Cargo.toml | 4 - crates/trie/trie/src/key.rs | 87 -------------- crates/trie/trie/src/lib.rs | 3 - crates/trie/trie/src/node_iter.rs | 12 +- crates/trie/trie/src/state.rs | 31 +---- crates/trie/trie/src/walker.rs | 9 +- 19 files changed, 157 insertions(+), 170 deletions(-) delete mode 100644 crates/trie/trie/src/key.rs diff --git a/Cargo.lock b/Cargo.lock index eaa45ddd1574..85315f80e065 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9768,7 +9768,6 @@ dependencies = [ "reth-storage-errors", "reth-trie-common", "serde_json", - "smallvec", "tracing", "triehash", ] @@ -9800,6 +9799,7 @@ dependencies = [ "serde", "serde_json", "serde_with", + "smallvec", ] [[package]] diff --git a/crates/scroll/state-commitment/src/lib.rs b/crates/scroll/state-commitment/src/lib.rs index 207de1965a3b..c9b7c300b548 100644 --- a/crates/scroll/state-commitment/src/lib.rs +++ b/crates/scroll/state-commitment/src/lib.rs @@ -19,6 +19,5 @@ pub use value::PoseidonValueHasher; #[cfg(feature = "test-utils")] pub mod test_utils; -// #[cfg(all(test, feature = "scroll"))] -#[cfg(test)] +#[cfg(all(test, feature = "scroll"))] mod test; diff --git a/crates/scroll/state-commitment/src/root/mod.rs b/crates/scroll/state-commitment/src/root/mod.rs index 1b350afd19cd..5336955521a5 100644 --- a/crates/scroll/state-commitment/src/root/mod.rs +++ b/crates/scroll/state-commitment/src/root/mod.rs @@ -6,15 +6,14 @@ use reth_scroll_primitives::poseidon::EMPTY_ROOT_HASH; use reth_scroll_trie::HashBuilder; use reth_trie::{ hashed_cursor::{HashedCursorFactory, HashedPostStateCursorFactory, HashedStorageCursor}, - key::BitsCompatibility, node_iter::{TrieElement, TrieNodeIter}, prefix_set::{PrefixSet, TriePrefixSets}, stats::TrieTracker, trie_cursor::{InMemoryTrieCursorFactory, TrieCursorFactory}, updates::{StorageTrieUpdates, TrieUpdates}, walker::TrieWalker, - HashedPostState, HashedPostStateSorted, HashedStorage, IntermediateStateRootState, KeyHasher, - Nibbles, StateRootProgress, TrieInput, + BitsCompatibility, HashedPostState, HashedPostStateSorted, HashedStorage, + IntermediateStateRootState, KeyHasher, Nibbles, StateRootProgress, TrieInput, }; use tracing::{debug, trace}; diff --git a/crates/scroll/state-commitment/src/root/utils.rs b/crates/scroll/state-commitment/src/root/utils.rs index 9d482460ef89..e51bbef53b7f 100644 --- a/crates/scroll/state-commitment/src/root/utils.rs +++ b/crates/scroll/state-commitment/src/root/utils.rs @@ -4,7 +4,7 @@ use crate::{PoseidonKeyHasher, PoseidonValueHasher, ScrollTrieAccount}; use alloy_primitives::{Address, B256, U256}; use itertools::Itertools; use reth_scroll_trie::HashBuilder; -use reth_trie::{key::BitsCompatibility, KeyHasher, Nibbles}; +use reth_trie::{BitsCompatibility, KeyHasher, Nibbles}; /// Hashes and sorts account keys, then proceeds to calculating the root hash of the state /// represented as MPT. diff --git a/crates/scroll/state-commitment/src/test.rs b/crates/scroll/state-commitment/src/test.rs index 264106a5ec54..98335247a963 100644 --- a/crates/scroll/state-commitment/src/test.rs +++ b/crates/scroll/state-commitment/src/test.rs @@ -10,18 +10,14 @@ use reth_db::{ }; use reth_primitives::{Account, StorageEntry}; use reth_provider::{test_utils::create_test_provider_factory, StorageTrieWriter, TrieWriter}; -use reth_scroll_state_commitment::{ - test_utils::{b256_clear_last_byte, b256_reverse_bits, u256_clear_msb}, - PoseidonKeyHasher, StateRoot, StorageRoot, -}; +use reth_scroll_state_commitment::{test_utils::*, PoseidonKeyHasher, StateRoot, StorageRoot}; use reth_scroll_state_commitment::state_root_unsorted; use reth_trie::{ - key::BitsCompatibility, prefix_set::{PrefixSetMut, TriePrefixSets}, trie_cursor::InMemoryTrieCursorFactory, updates::TrieUpdates, - HashedPostState, HashedStorage, KeyHasher, Nibbles, + BitsCompatibility, HashedPostState, HashedStorage, KeyHasher, Nibbles, }; use reth_trie_db::{DatabaseStateRoot, DatabaseStorageRoot, DatabaseTrieCursorFactory}; use std::collections::BTreeMap; diff --git a/crates/scroll/trie/src/hash_builder.rs b/crates/scroll/trie/src/hash_builder.rs index ed60e8cda552..8eacab8d1544 100644 --- a/crates/scroll/trie/src/hash_builder.rs +++ b/crates/scroll/trie/src/hash_builder.rs @@ -464,7 +464,7 @@ mod test { use alloc::collections::BTreeMap; use hex_literal::hex; use reth_scroll_primitives::poseidon::{hash_with_domain, Fr, PrimeField}; - use reth_trie::key::BitsCompatibility; + use reth_trie::BitsCompatibility; #[test] fn test_convert_to_bit_representation() { diff --git a/crates/scroll/trie/src/leaf.rs b/crates/scroll/trie/src/leaf.rs index 17be70943834..55a3df25a3b2 100644 --- a/crates/scroll/trie/src/leaf.rs +++ b/crates/scroll/trie/src/leaf.rs @@ -1,7 +1,7 @@ use super::LEAF_NODE_DOMAIN; use alloy_primitives::B256; use reth_scroll_primitives::poseidon::{hash_with_domain, Fr, PrimeField}; -use reth_trie::{key::BitsCompatibility, LeafNodeRef}; +use reth_trie::{BitsCompatibility, LeafNodeRef}; /// A trait used to hash the leaf node. pub(crate) trait HashLeaf { diff --git a/crates/storage/provider/src/providers/database/provider.rs b/crates/storage/provider/src/providers/database/provider.rs index a81242eb0223..a6237e26b439 100644 --- a/crates/storage/provider/src/providers/database/provider.rs +++ b/crates/storage/provider/src/providers/database/provider.rs @@ -66,8 +66,9 @@ use reth_storage_api::{ use reth_storage_errors::provider::{ProviderResult, RootMismatch}; use reth_trie::{ prefix_set::{PrefixSet, PrefixSetMut, TriePrefixSets}, + unpack_nibbles, updates::{StorageTrieUpdates, TrieUpdates}, - HashedPostStateSorted, KeyHasher, Nibbles, StoredNibbles, + HashedPostStateSorted, KeyHasher, StoredNibbles, }; use reth_trie_db::{DatabaseStorageTrieCursor, StateCommitment}; use revm::{ @@ -278,7 +279,7 @@ impl DatabaseProvider DatabaseProvider::default(); let storage_entries = self.unwind_storage_hashing(changed_storages.iter().copied())?; for (hashed_address, hashed_slots) in storage_entries { - account_prefix_set.insert(Nibbles::unpack(hashed_address)); + account_prefix_set.insert(unpack_nibbles(hashed_address)); let mut storage_prefix_set = PrefixSetMut::with_capacity(hashed_slots.len()); for slot in hashed_slots { - storage_prefix_set.insert(Nibbles::unpack(slot)); + storage_prefix_set.insert(unpack_nibbles(slot)); } storage_prefix_sets.insert(hashed_address, storage_prefix_set.freeze()); } @@ -2551,12 +2552,12 @@ impl HashingWriter for DatabaseProvi let storages = self.plain_state_storages(lists)?; let storage_entries = self.insert_storage_for_hashing(storages)?; for (hashed_address, hashed_slots) in storage_entries { - account_prefix_set.insert(Nibbles::unpack(hashed_address)); + account_prefix_set.insert(unpack_nibbles(hashed_address)); for slot in hashed_slots { storage_prefix_sets .entry(hashed_address) .or_default() - .insert(Nibbles::unpack(slot)); + .insert(unpack_nibbles(slot)); } } } @@ -2568,7 +2569,7 @@ impl HashingWriter for DatabaseProvi let accounts = self.basic_accounts(lists)?; let hashed_addresses = self.insert_account_for_hashing(accounts)?; for (hashed_address, account) in hashed_addresses { - account_prefix_set.insert(Nibbles::unpack(hashed_address)); + account_prefix_set.insert(unpack_nibbles(hashed_address)); if account.is_none() { destroyed_accounts.insert(hashed_address); } diff --git a/crates/trie/common/Cargo.toml b/crates/trie/common/Cargo.toml index bb80a50f0e2f..80a0efc8d8c8 100644 --- a/crates/trie/common/Cargo.toml +++ b/crates/trie/common/Cargo.toml @@ -30,10 +30,12 @@ bytes.workspace = true derive_more.workspace = true itertools.workspace = true nybbles = { workspace = true, features = ["rlp"] } +smallvec = { version = "1.0", default-features = false, features = [ + "const_new", +] } # `serde` feature serde = { workspace = true, optional = true } - serde_with = { workspace = true, optional = true } # `test-utils` feature @@ -62,6 +64,7 @@ serde = [ "dep:serde", "bytes/serde", "nybbles/serde", + "smallvec/serde", "alloy-primitives/serde", "alloy-consensus/serde", "alloy-trie/serde", diff --git a/crates/trie/common/src/key.rs b/crates/trie/common/src/key.rs index 71f8019bff54..f775c78ab60b 100644 --- a/crates/trie/common/src/key.rs +++ b/crates/trie/common/src/key.rs @@ -1,4 +1,6 @@ +use crate::Nibbles; use alloy_primitives::{keccak256, B256}; +use smallvec::SmallVec; /// Trait for hashing keys in state. pub trait KeyHasher: Default + Clone + Send + Sync + 'static { @@ -15,3 +17,112 @@ impl KeyHasher for KeccakKeyHasher { keccak256(bytes) } } + +/// The maximum number of bits a key can contain. +const MAX_BITS: usize = 254; + +/// The maximum number of bytes a key can contain. +const MAX_BYTES: usize = 32; + +// TODO(scroll): Refactor this into a trait that is more generic and can be used by any +// implementation that requires converting between nibbles and bits. Better yet we should use a +// trait that allows for defining the key type via a GAT as opposed to using Nibbles. + +/// A trait for converting a `Nibbles` representation to a bit representation. +pub trait BitsCompatibility: Sized { + /// Unpacks the bits from the provided bytes such that there is a byte for each bit in the + /// input. The representation is big-endian with respect to the input. + /// + /// We truncate the Nibbles such that we only have [`MAX_BITS`] (bn254 field size) bits. + fn unpack_bits>(data: T) -> Self; + + /// Pack the bits into a byte representation. + fn pack_bits(&self) -> SmallVec<[u8; 32]>; + + /// Encodes a leaf key represented as [`Nibbles`] into it's canonical little-endian + /// representation. + fn encode_leaf_key(&self) -> [u8; 32]; + + /// Increment the key to the next key. + fn increment_bit(&self) -> Option; +} + +impl BitsCompatibility for Nibbles { + fn unpack_bits>(data: T) -> Self { + let data = data.as_ref(); + let bits = data + .iter() + .take(MAX_BYTES) + .flat_map(|&byte| (0..8).rev().map(move |i| (byte >> i) & 1)) + .take(MAX_BITS) + .collect(); + + Self::from_vec_unchecked(bits) + } + + fn pack_bits(&self) -> SmallVec<[u8; 32]> { + let mut result = SmallVec::with_capacity((self.len() + 7) / 8); + + for bits in self.as_slice().chunks(8) { + let mut byte = 0; + for (bit_index, bit) in bits.iter().enumerate() { + byte |= *bit << (7 - bit_index); + } + result.push(byte); + } + + result + } + + fn encode_leaf_key(&self) -> [u8; 32] { + // This is strange we are now representing the leaf key using big endian?? + let mut result = [0u8; 32]; + for (byte_index, bytes) in self.as_slice().chunks(8).enumerate() { + for (bit_index, byte) in bytes.iter().enumerate() { + result[byte_index] |= byte << bit_index; + } + } + + result + } + + fn increment_bit(&self) -> Option { + let mut incremented = self.clone(); + + for nibble in incremented.as_mut_vec_unchecked().iter_mut().rev() { + if *nibble < 1 { + *nibble += 1; + return Some(incremented); + } + + *nibble = 0; + } + + None + } +} + +/// Helper method to unpack into [`Nibbles`] from a byte slice. +/// +/// For the `scroll` feature, this method will unpack the bits from the provided bytes such that +/// there is a byte for each bit in the input. The representation is big-endian with respect to the +/// input. When the `scroll` feature is not enabled, this method will unpack the bytes into nibbles. +pub fn unpack_nibbles>(data: T) -> Nibbles { + #[cfg(feature = "scroll")] + let nibbles = Nibbles::unpack_bits(data); + #[cfg(not(feature = "scroll"))] + let nibbles = Nibbles::unpack(data); + nibbles +} + +/// Helper method to pack into a byte slice from [`Nibbles`]. +/// +/// For the `scroll` feature, this method will pack the bits into a byte representation. When the +/// `scroll` feature is not enabled, this method will pack the nibbles into bytes. +pub fn pack_nibbles(nibbles: &Nibbles) -> SmallVec<[u8; 32]> { + #[cfg(feature = "scroll")] + let packed = nibbles.pack_bits(); + #[cfg(not(feature = "scroll"))] + let packed = nibbles.pack(); + packed +} diff --git a/crates/trie/common/src/lib.rs b/crates/trie/common/src/lib.rs index 6647de67811c..8fd52ef4d42f 100644 --- a/crates/trie/common/src/lib.rs +++ b/crates/trie/common/src/lib.rs @@ -19,7 +19,7 @@ mod account; pub use account::TrieAccount; mod key; -pub use key::{KeccakKeyHasher, KeyHasher}; +pub use key::{pack_nibbles, unpack_nibbles, BitsCompatibility, KeccakKeyHasher, KeyHasher}; mod nibbles; pub use nibbles::{Nibbles, StoredNibbles, StoredNibblesSubKey}; diff --git a/crates/trie/common/src/updates.rs b/crates/trie/common/src/updates.rs index 6f80eb16553e..1dc38f22656e 100644 --- a/crates/trie/common/src/updates.rs +++ b/crates/trie/common/src/updates.rs @@ -229,7 +229,7 @@ impl StorageTrieUpdates { /// This also sorts the set before serializing. #[cfg(any(test, feature = "serde"))] mod serde_nibbles_set { - use crate::Nibbles; + use crate::{pack_nibbles, unpack_nibbles, Nibbles}; use alloy_primitives::map::HashSet; use serde::{de::Error, Deserialize, Deserializer, Serialize, Serializer}; @@ -237,8 +237,10 @@ mod serde_nibbles_set { where S: Serializer, { - let mut storage_nodes = - map.iter().map(|elem| alloy_primitives::hex::encode(elem.pack())).collect::>(); + let mut storage_nodes = map + .iter() + .map(|elem| alloy_primitives::hex::encode(pack_nibbles(elem))) + .collect::>(); storage_nodes.sort_unstable(); storage_nodes.serialize(serializer) } @@ -250,7 +252,7 @@ mod serde_nibbles_set { Vec::::deserialize(deserializer)? .into_iter() .map(|node| { - Ok(Nibbles::unpack( + Ok(unpack_nibbles( alloy_primitives::hex::decode(node) .map_err(|err| D::Error::custom(err.to_string()))?, )) @@ -265,7 +267,7 @@ mod serde_nibbles_set { /// This also sorts the map's keys before encoding and serializing. #[cfg(any(test, feature = "serde"))] mod serde_nibbles_map { - use crate::Nibbles; + use crate::{pack_nibbles, unpack_nibbles, Nibbles}; use alloy_primitives::{hex, map::HashMap}; use serde::{ de::{Error, MapAccess, Visitor}, @@ -287,7 +289,7 @@ mod serde_nibbles_map { storage_nodes.sort_unstable_by_key(|node| node.0); for (k, v) in storage_nodes { // pack, then hex encode the Nibbles - let packed = alloy_primitives::hex::encode(k.pack()); + let packed = alloy_primitives::hex::encode(pack_nibbles(k)); map_serializer.serialize_entry(&packed, &v)?; } map_serializer.end() @@ -325,7 +327,7 @@ mod serde_nibbles_map { let decoded_key = hex::decode(&key).map_err(|err| Error::custom(err.to_string()))?; - let nibbles = Nibbles::unpack(&decoded_key); + let nibbles = unpack_nibbles(&decoded_key); result.insert(nibbles, value); } diff --git a/crates/trie/db/src/prefix_set.rs b/crates/trie/db/src/prefix_set.rs index 95ff6d91f374..b4a55ddb6851 100644 --- a/crates/trie/db/src/prefix_set.rs +++ b/crates/trie/db/src/prefix_set.rs @@ -10,7 +10,7 @@ use reth_db_api::{ use reth_primitives::StorageEntry; use reth_trie::{ prefix_set::{PrefixSetMut, TriePrefixSets}, - KeyHasher, Nibbles, + unpack_nibbles, KeyHasher, }; use std::{ collections::{HashMap, HashSet}, @@ -51,7 +51,7 @@ impl PrefixSetLoader<'_, TX, KH> { for account_entry in account_changeset_cursor.walk_range(range.clone())? { let (_, AccountBeforeTx { address, .. }) = account_entry?; let hashed_address = KH::hash_key(address); - account_prefix_set.insert(Nibbles::unpack(hashed_address)); + account_prefix_set.insert(unpack_nibbles(hashed_address)); if account_hashed_state_cursor.seek_exact(hashed_address)?.is_none() { destroyed_accounts.insert(hashed_address); @@ -65,11 +65,11 @@ impl PrefixSetLoader<'_, TX, KH> { for storage_entry in storage_cursor.walk_range(storage_range)? { let (BlockNumberAddress((_, address)), StorageEntry { key, .. }) = storage_entry?; let hashed_address = KH::hash_key(address); - account_prefix_set.insert(Nibbles::unpack(hashed_address)); + account_prefix_set.insert(unpack_nibbles(hashed_address)); storage_prefix_sets .entry(hashed_address) .or_default() - .insert(Nibbles::unpack(KH::hash_key(key))); + .insert(unpack_nibbles(KH::hash_key(key))); } Ok(TriePrefixSets { diff --git a/crates/trie/trie/Cargo.toml b/crates/trie/trie/Cargo.toml index 4306fb2f262b..901941e744b6 100644 --- a/crates/trie/trie/Cargo.toml +++ b/crates/trie/trie/Cargo.toml @@ -36,9 +36,6 @@ tracing.workspace = true rayon.workspace = true auto_impl.workspace = true itertools.workspace = true -smallvec = { version = "1.0", default-features = false, features = [ - "const_new", -] } # `metrics` feature reth-metrics = { workspace = true, optional = true } @@ -70,7 +67,6 @@ serde = [ "revm/serde", "reth-trie-common/serde", "reth-primitives-traits/serde", - "smallvec/serde" ] test-utils = [ "triehash", diff --git a/crates/trie/trie/src/key.rs b/crates/trie/trie/src/key.rs deleted file mode 100644 index 17809995ee47..000000000000 --- a/crates/trie/trie/src/key.rs +++ /dev/null @@ -1,87 +0,0 @@ -use smallvec::SmallVec; - -use crate::Nibbles; - -/// The maximum number of bits a key can contain. -pub const MAX_BITS: usize = 254; - -/// The maximum number of bytes a key can contain. -const MAX_BYTES: usize = 32; - -// TODO(scroll): Refactor this into a trait that is more generic and can be used by any -// implementation that requires converting between nibbles and bits. Better yet we should use a -// trait that allows for defining the key type via a GAT as opposed to using Nibbles. - -/// A trait for converting a `Nibbles` representation to a bit representation. -pub trait BitsCompatibility: Sized { - /// Unpacks the bits from the provided bytes such that there is a byte for each bit in the - /// input. The representation is big-endian with respect to the input. - /// - /// We truncate the Nibbles such that we only have [`MAX_BITS`] (bn254 field size) bits. - fn unpack_bits>(data: T) -> Self; - - /// Pack the bits into a byte representation. - fn pack_bits(&self) -> SmallVec<[u8; 32]>; - - /// Encodes a leaf key represented as [`Nibbles`] into it's canonical little-endian - /// representation. - fn encode_leaf_key(&self) -> [u8; 32]; - - /// Increment the key to the next key. - fn increment_bit(&self) -> Option; -} - -impl BitsCompatibility for Nibbles { - fn unpack_bits>(data: T) -> Self { - let data = data.as_ref(); - let bits = data - .iter() - .take(MAX_BYTES) - .flat_map(|&byte| (0..8).rev().map(move |i| (byte >> i) & 1)) - .take(MAX_BITS) - .collect(); - - Self::from_vec_unchecked(bits) - } - - fn pack_bits(&self) -> SmallVec<[u8; 32]> { - let mut result = SmallVec::with_capacity((self.len() + 7) / 8); - - for bits in self.as_slice().chunks(8) { - let mut byte = 0; - for (bit_index, bit) in bits.iter().enumerate() { - byte |= *bit << (7 - bit_index); - } - result.push(byte); - } - - result - } - - fn encode_leaf_key(&self) -> [u8; 32] { - // This is strange we are now representing the leaf key using big endian?? - let mut result = [0u8; 32]; - for (byte_index, bytes) in self.as_slice().chunks(8).enumerate() { - for (bit_index, byte) in bytes.iter().enumerate() { - result[byte_index] |= byte << bit_index; - } - } - - result - } - - fn increment_bit(&self) -> Option { - let mut incremented = self.clone(); - - for nibble in incremented.as_mut_vec_unchecked().iter_mut().rev() { - if *nibble < 1 { - *nibble += 1; - return Some(incremented); - } - - *nibble = 0; - } - - None - } -} diff --git a/crates/trie/trie/src/lib.rs b/crates/trie/trie/src/lib.rs index c823ffc77de0..1e7eeb9b52b8 100644 --- a/crates/trie/trie/src/lib.rs +++ b/crates/trie/trie/src/lib.rs @@ -13,9 +13,6 @@ )] #![cfg_attr(docsrs, feature(doc_cfg, doc_auto_cfg))] -/// A module for working with trie keys. -pub mod key; - /// The implementation of forward-only in-memory cursor. pub mod forward_cursor; diff --git a/crates/trie/trie/src/node_iter.rs b/crates/trie/trie/src/node_iter.rs index 7da75ba98323..c38c3cc3155b 100644 --- a/crates/trie/trie/src/node_iter.rs +++ b/crates/trie/trie/src/node_iter.rs @@ -1,9 +1,7 @@ use crate::{hashed_cursor::HashedCursor, trie_cursor::TrieCursor, walker::TrieWalker, Nibbles}; use alloy_primitives::B256; use reth_storage_errors::db::DatabaseError; - -#[cfg(feature = "scroll")] -use crate::key::BitsCompatibility; +use reth_trie_common::unpack_nibbles; /// Represents a branch node in the trie. #[derive(Debug)] @@ -109,13 +107,7 @@ where if let Some((hashed_key, value)) = self.current_hashed_entry.take() { // If the walker's key is less than the unpacked hashed key, // reset the checked status and continue - if self.walker.key().is_some_and(|key| { - #[cfg(not(feature = "scroll"))] - let cmp = key < &Nibbles::unpack(hashed_key); - #[cfg(feature = "scroll")] - let cmp = key < &Nibbles::unpack_bits(hashed_key); - cmp - }) { + if self.walker.key().is_some_and(|key| key < &unpack_nibbles(hashed_key)) { self.current_walker_key_checked = false; continue } diff --git a/crates/trie/trie/src/state.rs b/crates/trie/trie/src/state.rs index 99b58889289e..e3a9c457c821 100644 --- a/crates/trie/trie/src/state.rs +++ b/crates/trie/trie/src/state.rs @@ -1,7 +1,4 @@ -use crate::{ - prefix_set::{PrefixSetMut, TriePrefixSetsMut}, - Nibbles, -}; +use crate::prefix_set::{PrefixSetMut, TriePrefixSetsMut}; use alloy_primitives::{ map::{hash_map, HashMap, HashSet}, Address, B256, U256, @@ -9,13 +6,10 @@ use alloy_primitives::{ use itertools::Itertools; use rayon::prelude::{IntoParallelIterator, ParallelIterator}; use reth_primitives::Account; -use reth_trie_common::KeyHasher; +use reth_trie_common::{unpack_nibbles, KeyHasher}; use revm::db::{states::CacheAccount, AccountStatus, BundleAccount}; use std::borrow::Cow; -#[cfg(feature = "scroll")] -use crate::key::BitsCompatibility; - /// Representation of in-memory hashed state. #[derive(PartialEq, Eq, Clone, Default, Debug)] pub struct HashedPostState { @@ -121,12 +115,8 @@ impl HashedPostState { let mut account_prefix_set = PrefixSetMut::with_capacity(self.accounts.len()); let mut destroyed_accounts = HashSet::default(); for (hashed_address, account) in &self.accounts { - // TODO(scroll): replace with key abstraction - #[cfg(feature = "scroll")] - let nibbles = Nibbles::unpack_bits(hashed_address); - #[cfg(not(feature = "scroll"))] - let nibbles = Nibbles::unpack(hashed_address); - account_prefix_set.insert(nibbles); + // TODO(scroll): replace this with abstraction. + account_prefix_set.insert(unpack_nibbles(hashed_address)); if account.is_none() { destroyed_accounts.insert(*hashed_address); @@ -138,11 +128,7 @@ impl HashedPostState { HashMap::with_capacity_and_hasher(self.storages.len(), Default::default()); for (hashed_address, hashed_storage) in &self.storages { // TODO(scroll): replace this with abstraction. - #[cfg(feature = "scroll")] - let nibbles = Nibbles::unpack_bits(hashed_address); - #[cfg(not(feature = "scroll"))] - let nibbles = Nibbles::unpack(hashed_address); - account_prefix_set.insert(nibbles); + account_prefix_set.insert(unpack_nibbles(hashed_address)); storage_prefix_sets.insert(*hashed_address, hashed_storage.construct_prefix_set()); } @@ -266,12 +252,7 @@ impl HashedStorage { } else { let mut prefix_set = PrefixSetMut::with_capacity(self.storage.len()); for hashed_slot in self.storage.keys() { - // TODO(scroll): replace this with key abstraction. - #[cfg(feature = "scroll")] - let nibbles = Nibbles::unpack_bits(hashed_slot); - #[cfg(not(feature = "scroll"))] - let nibbles = Nibbles::unpack(hashed_slot); - prefix_set.insert(nibbles); + prefix_set.insert(unpack_nibbles(hashed_slot)); } prefix_set } diff --git a/crates/trie/trie/src/walker.rs b/crates/trie/trie/src/walker.rs index db76049279fd..654d20f5541e 100644 --- a/crates/trie/trie/src/walker.rs +++ b/crates/trie/trie/src/walker.rs @@ -5,12 +5,13 @@ use crate::{ }; use alloy_primitives::{map::HashSet, B256}; use reth_storage_errors::db::DatabaseError; +use reth_trie_common::pack_nibbles; #[cfg(feature = "metrics")] use crate::metrics::WalkerMetrics; #[cfg(feature = "scroll")] -use crate::key::BitsCompatibility; +use crate::BitsCompatibility; /// `TrieWalker` is a structure that enables traversal of a Merkle trie. /// It allows moving through the trie in a depth-first manner, skipping certain branches @@ -110,11 +111,7 @@ impl TrieWalker { let key = key.increment_bit().map(|inc| inc.pack_bits()); key } else { - #[cfg(feature = "scroll")] - let key = Some(key.pack_bits()); - #[cfg(not(feature = "scroll"))] - let key = Some(key.pack()); - key + Some(pack_nibbles(key)) } }) .map(|mut key| { From 7258d7a3ec558bddbe21be915270a3da199bb4da Mon Sep 17 00:00:00 2001 From: frisitano Date: Fri, 13 Dec 2024 12:19:37 +0000 Subject: [PATCH 3/4] lint: lint propagation --- crates/trie/common/Cargo.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/crates/trie/common/Cargo.toml b/crates/trie/common/Cargo.toml index 80a0efc8d8c8..35bb4e78be7a 100644 --- a/crates/trie/common/Cargo.toml +++ b/crates/trie/common/Cargo.toml @@ -94,6 +94,7 @@ arbitrary = [ "alloy-consensus/arbitrary", "alloy-primitives/arbitrary", "nybbles/arbitrary", + "smallvec/arbitrary", "revm-primitives/arbitrary", "reth-codecs/arbitrary", "reth-scroll-primitives?/arbitrary" From 19c2d590ae92a1f2ade9e32b287890517d5c53e7 Mon Sep 17 00:00:00 2001 From: frisitano Date: Fri, 13 Dec 2024 12:21:42 +0000 Subject: [PATCH 4/4] lint: private item in docs --- crates/trie/common/src/key.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/trie/common/src/key.rs b/crates/trie/common/src/key.rs index f775c78ab60b..c439c7eaf94e 100644 --- a/crates/trie/common/src/key.rs +++ b/crates/trie/common/src/key.rs @@ -33,7 +33,7 @@ pub trait BitsCompatibility: Sized { /// Unpacks the bits from the provided bytes such that there is a byte for each bit in the /// input. The representation is big-endian with respect to the input. /// - /// We truncate the Nibbles such that we only have [`MAX_BITS`] (bn254 field size) bits. + /// We truncate the Nibbles such that we only have 254 (bn254 field size) bits. fn unpack_bits>(data: T) -> Self; /// Pack the bits into a byte representation.