Skip to content

Key Traversal Abstraction #66

Closed
Closed
@frisitano

Description

@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:

/// 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>;
}

// 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
}

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 }
}

/// 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
}
}

/// 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

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions