Fix zktrie Key Traversal (#75)

Author: frisitano

crates/scroll/state-commitment/src/root.rs

@@ -233,10 +233,8 @@ where
 
                     let account = ScrollTrieAccount::from((account, storage_root));
                     let account_hash = PoseidonValueHasher::hash_account(account);
-                    hash_builder.add_leaf(
-                        Nibbles::unpack_and_truncate_bits(hashed_address),
-                        account_hash.as_slice(),
-                    );
+                    hash_builder
+                        .add_leaf(Nibbles::unpack_bits(hashed_address), account_hash.as_slice());
 
                     // Decide if we need to return intermediate progress.
                     let total_updates_len = updated_storage_nodes +
@@ -433,10 +431,7 @@ where
                 TrieElement::Leaf(hashed_slot, value) => {
                     let hashed_value = PoseidonValueHasher::hash_storage(value);
                     tracker.inc_leaf();
-                    hash_builder.add_leaf(
-                        Nibbles::unpack_and_truncate_bits(hashed_slot),
-                        hashed_value.as_ref(),
-                    );
+                    hash_builder.add_leaf(Nibbles::unpack_bits(hashed_slot), hashed_value.as_ref());
                 }
             }
         }

crates/scroll/trie/src/hash_builder.rs

@@ -468,7 +468,7 @@ mod test {
 
     #[test]
     fn test_convert_to_bit_representation() {
-        let nibbles = Nibbles::unpack_and_truncate_bits(vec![7, 8]);
+        let nibbles = Nibbles::unpack_bits(vec![7, 8]);
         let expected = [0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 1, 0, 0, 0];
         assert_eq!(nibbles.as_slice(), expected);
     }
@@ -478,8 +478,8 @@ mod test {
         // 64 byte nibble
         let hex = hex!("0102030405060708090a0b0c0d0e0f0102030405060708090a0b0c0d0e0f0102030405060708090a0b0c0d0e0f0102030405060708090a0b0c0d0e0f01020304");
         assert_eq!(hex.len(), 64);
-        let nibbles = Nibbles::unpack_and_truncate_bits(hex);
-        assert_eq!(nibbles.len(), 248);
+        let nibbles = Nibbles::unpack_bits(hex);
+        assert_eq!(nibbles.len(), 254);
     }
 
     #[test]

crates/trie/trie/src/key.rs

@@ -3,10 +3,10 @@ use smallvec::SmallVec;
 use crate::Nibbles;
 
 /// The maximum number of bits a key can contain.
-pub const MAX_BITS: usize = 248;
+pub const MAX_BITS: usize = 254;
 
 /// The maximum number of bytes a key can contain.
-const MAX_BYTES: usize = 31;
+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
@@ -17,8 +17,8 @@ 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;
+    /// We truncate the Nibbles such that we only have [`MAX_BITS`] (bn254 field size) bits.
+    fn unpack_bits<T: AsRef<[u8]>>(data: T) -> Self;
 
     /// Pack the bits into a byte representation.
     fn pack_bits(&self) -> SmallVec<[u8; 32]>;
@@ -32,16 +32,14 @@ pub trait BitsCompatibility: Sized {
 }
 
 impl BitsCompatibility for Nibbles {
-    fn unpack_and_truncate_bits<T: AsRef<[u8]>>(data: T) -> Self {
+    fn unpack_bits<T: AsRef<[u8]>>(data: T) -> Self {
         let data = data.as_ref();
-        let unpacked_len = core::cmp::min(data.len() * 8, MAX_BITS);
-        let mut bits = Vec::with_capacity(unpacked_len);
-
-        for byte in data.iter().take(MAX_BYTES) {
-            for i in (0..8).rev() {
-                bits.push(byte >> i & 1);
-            }
-        }
+        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)
     }

crates/trie/trie/src/node_iter.rs

@@ -113,7 +113,7 @@ where
                     #[cfg(not(feature = "scroll"))]
                     let cmp = key < &Nibbles::unpack(hashed_key);
                     #[cfg(feature = "scroll")]
-                    let cmp = key < &Nibbles::unpack_and_truncate_bits(hashed_key);
+                    let cmp = key < &Nibbles::unpack_bits(hashed_key);
                     cmp
                 }) {
                     self.current_walker_key_checked = false;

crates/trie/trie/src/state.rs

@@ -123,7 +123,7 @@ impl HashedPostState {
         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);
+            let nibbles = Nibbles::unpack_bits(hashed_address);
             #[cfg(not(feature = "scroll"))]
             let nibbles = Nibbles::unpack(hashed_address);
             account_prefix_set.insert(nibbles);
@@ -139,7 +139,7 @@ impl HashedPostState {
         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);
+            let nibbles = Nibbles::unpack_bits(hashed_address);
             #[cfg(not(feature = "scroll"))]
             let nibbles = Nibbles::unpack(hashed_address);
             account_prefix_set.insert(nibbles);
@@ -268,7 +268,7 @@ impl HashedStorage {
             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);
+                let nibbles = Nibbles::unpack_bits(hashed_slot);
                 #[cfg(not(feature = "scroll"))]
                 let nibbles = Nibbles::unpack(hashed_slot);
                 prefix_set.insert(nibbles);