Key Traversal Abstraction

[frisitano]

Describe the feature

In #36 we introduce zktrie which uses bits as a key traversal unit. This is in conflict with the key traversal unit of the native Ethereum trie which uses nibbles. As such, we introduce a BitsCompatability trait and associated feature gates to address this discrepancy. We should instead introduce a generics KeyTraversal trait which can be used to traverse a trie key. We should then remove the following trait and feature flags:

reth/crates/trie/trie/src/key.rs

Lines 15 to 32 in c960cb7

	/// 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`] (248) bits.
	fn unpack_and_truncate_bits<T: AsRef<[u8]>>(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<Self>;
	}

reth/crates/trie/trie/src/node_iter.rs

Lines 108 to 121 in c960cb7

	// If there's a hashed entry...
	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_and_truncate_bits(hashed_key);
	cmp
	}) {
	self.current_walker_key_checked = false;
	continue
	}

reth/crates/trie/trie/src/state.rs

Lines 119 to 150 in c960cb7

	pub fn construct_prefix_sets(&self) -> TriePrefixSetsMut {
	// Populate account prefix set.
	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_and_truncate_bits(hashed_address);
	#[cfg(not(feature = "scroll"))]
	let nibbles = Nibbles::unpack(hashed_address);
	account_prefix_set.insert(nibbles);
	if account.is_none() {
	destroyed_accounts.insert(*hashed_address);
	}
	}
	// Populate storage prefix sets.
	let mut storage_prefix_sets =
	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_and_truncate_bits(hashed_address);
	#[cfg(not(feature = "scroll"))]
	let nibbles = Nibbles::unpack(hashed_address);
	account_prefix_set.insert(nibbles);
	storage_prefix_sets.insert(*hashed_address, hashed_storage.construct_prefix_set());
	}
	TriePrefixSetsMut { account_prefix_set, storage_prefix_sets, destroyed_accounts }
	}

reth/crates/trie/trie/src/state.rs

Lines 262 to 278 in c960cb7

	/// Construct [`PrefixSetMut`] from hashed storage.
	pub fn construct_prefix_set(&self) -> PrefixSetMut {
	if self.wiped {
	PrefixSetMut::all()
	} 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_and_truncate_bits(hashed_slot);
	#[cfg(not(feature = "scroll"))]
	let nibbles = Nibbles::unpack(hashed_slot);
	prefix_set.insert(nibbles);
	}
	prefix_set
	}
	}

reth/crates/trie/trie/src/walker.rs

Lines 101 to 124 in c960cb7

	/// Returns the next unprocessed key in the trie.
	pub fn next_unprocessed_key(&self) -> Option<B256> {
	self.key()
	.and_then(|key| {
	if self.can_skip_current_node {
	// TODO(scroll): replace this with key abstraction.
	#[cfg(not(feature = "scroll"))]
	let key = key.increment().map(|inc| inc.pack());
	#[cfg(feature = "scroll")]
	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
	}
	})
	.map(|mut key| {
	key.resize(32, 0);
	B256::from_slice(key.as_slice())
	})
	}

Additional context

No response