feat(merkle tree): Remove enumeration index assignment from Merkle tr…

…ee (#551)

## What ❔

Since enumeration indices are now fully stored in Postgres, it makes
sense to not duplicate their assignment in the Merkle tree. Instead, the
tree could take enum indices as inputs.

## Why ❔

This allows simplifying tree logic and unify "normal" L1 batch
processing and tree recovery. (This unification is not a part of this
PR; it'll be implemented separately.)

## Checklist

- [x] PR title corresponds to the body of PR (we generate changelog
entries from PRs).
- [x] Tests for the changes have been added / updated.
- [x] Documentation comments have been added / updated.
- [x] Code has been formatted via `zk fmt` and `zk lint`.

core/lib/merkle_tree/examples/loadtest/main.rs

@@ -15,7 +15,8 @@ use std::{
 
 use zksync_crypto::hasher::blake2::Blake2Hasher;
 use zksync_merkle_tree::{
-    Database, HashTree, MerkleTree, MerkleTreePruner, PatchSet, RocksDBWrapper, TreeInstruction,
+    Database, HashTree, MerkleTree, MerkleTreePruner, PatchSet, RocksDBWrapper, TreeEntry,
+    TreeInstruction,
 };
 use zksync_storage::{RocksDB, RocksDBOptions};
 use zksync_types::{AccountTreeId, Address, StorageKey, H256, U256};
@@ -135,19 +136,22 @@ impl Cli {
             next_key_idx += new_keys.len() as u64;
 
             next_value_idx += (new_keys.len() + updated_indices.len()) as u64;
-            let values = (next_value_idx..).map(H256::from_low_u64_be);
             let updated_keys = Self::generate_keys(updated_indices.into_iter());
-            let kvs = new_keys.into_iter().chain(updated_keys).zip(values);
+            let kvs = new_keys
+                .into_iter()
+                .chain(updated_keys)
+                .zip(next_value_idx..);
+            let kvs = kvs.map(|(key, idx)| {
+                // The assigned leaf indices here are not always correct, but it's OK for load test purposes.
+                TreeEntry::new(key, idx, H256::from_low_u64_be(idx))
+            });
 
             tracing::info!("Processing block #{version}");
             let start = Instant::now();
             let root_hash = if self.proofs {
-                let reads = Self::generate_keys(read_indices.into_iter())
-                    .map(|key| (key, TreeInstruction::Read));
-                let instructions = kvs
-                    .map(|(key, hash)| (key, TreeInstruction::Write(hash)))
-                    .chain(reads)
-                    .collect();
+                let reads =
+                    Self::generate_keys(read_indices.into_iter()).map(TreeInstruction::Read);
+                let instructions = kvs.map(TreeInstruction::Write).chain(reads).collect();
                 let output = tree.extend_with_proofs(instructions);
                 output.root_hash().unwrap()
             } else {
@@ -160,7 +164,7 @@ impl Cli {
 
         tracing::info!("Verifying tree consistency...");
         let start = Instant::now();
-        tree.verify_consistency(self.commit_count - 1)
+        tree.verify_consistency(self.commit_count - 1, false)
             .expect("tree consistency check failed");
         let elapsed = start.elapsed();
         tracing::info!("Verified tree consistency in {elapsed:?}");

core/lib/merkle_tree/examples/recovery.rs

@@ -9,8 +9,8 @@ use std::time::Instant;
 
 use zksync_crypto::hasher::blake2::Blake2Hasher;
 use zksync_merkle_tree::{
-    recovery::{MerkleTreeRecovery, RecoveryEntry},
-    HashTree, Key, PatchSet, PruneDatabase, RocksDBWrapper, ValueHash,
+    recovery::MerkleTreeRecovery, HashTree, Key, PatchSet, PruneDatabase, RocksDBWrapper,
+    TreeEntry, ValueHash,
 };
 use zksync_storage::{RocksDB, RocksDBOptions};
 
@@ -94,15 +94,15 @@ impl Cli {
                 .map(|_| {
                     last_leaf_index += 1;
                     if self.random {
-                        RecoveryEntry {
+                        TreeEntry {
                             key: Key::from(rng.gen::<[u8; 32]>()),
                             value: ValueHash::zero(),
                             leaf_index: last_leaf_index,
                         }
                     } else {
                         last_key += key_step - Key::from(rng.gen::<u64>());
                         // ^ Increases the key by a random increment close to `key` step with some randomness.
-                        RecoveryEntry {
+                        TreeEntry {
                             key: last_key,
                             value: ValueHash::zero(),
                             leaf_index: last_leaf_index,
@@ -127,7 +127,7 @@ impl Cli {
             recovery_started_at.elapsed()
         );
         let started_at = Instant::now();
-        tree.verify_consistency(recovered_version).unwrap();
+        tree.verify_consistency(recovered_version, true).unwrap();
         tracing::info!("Verified consistency in {:?}", started_at.elapsed());
     }
 }

core/lib/merkle_tree/src/consistency.rs

@@ -69,10 +69,17 @@ pub enum ConsistencyError {
 impl<DB: Database, H: HashTree> MerkleTree<DB, H> {
     /// Verifies the internal tree consistency as stored in the database.
     ///
+    /// If `validate_indices` flag is set, it will be checked that indices for all tree leaves are unique
+    /// and are sequentially assigned starting from 1.
+    ///
     /// # Errors
     ///
     /// Returns an error (the first encountered one if there are multiple).
-    pub fn verify_consistency(&self, version: u64) -> Result<(), ConsistencyError> {
+    pub fn verify_consistency(
+        &self,
+        version: u64,
+        validate_indices: bool,
+    ) -> Result<(), ConsistencyError> {
         let manifest = self.db.try_manifest()?;
         let manifest = manifest.ok_or(ConsistencyError::MissingVersion(version))?;
         if version >= manifest.version_count {
@@ -91,16 +98,19 @@ impl<DB: Database, H: HashTree> MerkleTree<DB, H> {
         // We want to perform a depth-first walk of the tree in order to not keep
         // much in memory.
         let root_key = Nibbles::EMPTY.with_version(version);
-        let leaf_data = LeafConsistencyData::new(leaf_count);
-        self.validate_node(&root_node, root_key, &leaf_data)?;
-        leaf_data.validate_count()
+        let leaf_data = validate_indices.then(|| LeafConsistencyData::new(leaf_count));
+        self.validate_node(&root_node, root_key, leaf_data.as_ref())?;
+        if let Some(leaf_data) = leaf_data {
+            leaf_data.validate_count()?;
+        }
+        Ok(())
     }
 
     fn validate_node(
         &self,
         node: &Node,
         key: NodeKey,
-        leaf_data: &LeafConsistencyData,
+        leaf_data: Option<&LeafConsistencyData>,
     ) -> Result<ValueHash, ConsistencyError> {
         match node {
             Node::Leaf(leaf) => {
@@ -111,7 +121,9 @@ impl<DB: Database, H: HashTree> MerkleTree<DB, H> {
                         full_key: leaf.full_key,
                     });
                 }
-                leaf_data.insert_leaf(leaf)?;
+                if let Some(leaf_data) = leaf_data {
+                    leaf_data.insert_leaf(leaf)?;
+                }
             }
 
             Node::Internal(node) => {
@@ -261,7 +273,10 @@ mod tests {
     use std::num::NonZeroU64;
 
     use super::*;
-    use crate::{types::InternalNode, PatchSet};
+    use crate::{
+        types::{InternalNode, TreeEntry},
+        PatchSet,
+    };
     use zksync_types::{H256, U256};
 
     const FIRST_KEY: Key = U256([0, 0, 0, 0x_dead_beef_0000_0000]);
@@ -270,8 +285,8 @@ mod tests {
     fn prepare_database() -> PatchSet {
         let mut tree = MerkleTree::new(PatchSet::default());
         tree.extend(vec![
-            (FIRST_KEY, H256([1; 32])),
-            (SECOND_KEY, H256([2; 32])),
+            TreeEntry::new(FIRST_KEY, 1, H256([1; 32])),
+            TreeEntry::new(SECOND_KEY, 2, H256([2; 32])),
         ]);
         tree.db
     }
@@ -300,7 +315,7 @@ mod tests {
             .num_threads(1)
             .build()
             .expect("failed initializing `rayon` thread pool");
-        thread_pool.install(|| MerkleTree::new(db).verify_consistency(0))
+        thread_pool.install(|| MerkleTree::new(db).verify_consistency(0, true))
     }
 
     #[test]