Test tree operation after recovery

core/lib/merkle_tree/tests/integration/common.rs

@@ -2,7 +2,10 @@
 
 use once_cell::sync::Lazy;
 
+use std::collections::HashMap;
+
 use zksync_crypto::hasher::{blake2::Blake2Hasher, Hasher};
+use zksync_merkle_tree::{HashTree, TreeInstruction};
 use zksync_types::{AccountTreeId, Address, StorageKey, H256, U256};
 
 pub fn generate_key_value_pairs(indexes: impl Iterator<Item = u64>) -> Vec<(U256, H256)> {
@@ -15,20 +18,27 @@ pub fn generate_key_value_pairs(indexes: impl Iterator<Item = u64>) -> Vec<(U256
     kvs.collect()
 }
 
-// The extended version of computations used in `InternalNode`.
-pub fn compute_tree_hash(kvs: &[(U256, H256)]) -> H256 {
-    assert!(!kvs.is_empty());
+pub fn compute_tree_hash(kvs: impl Iterator<Item = (U256, H256)>) -> H256 {
+    let kvs_with_indices = kvs
+        .enumerate()
+        .map(|(i, (key, value))| (key, value, i as u64 + 1));
+    compute_tree_hash_with_indices(kvs_with_indices)
+}
 
+// The extended version of computations used in `InternalNode`.
+fn compute_tree_hash_with_indices(kvs: impl Iterator<Item = (U256, H256, u64)>) -> H256 {
     let hasher = Blake2Hasher;
     let mut empty_tree_hash = hasher.hash_bytes(&[0_u8; 40]);
-    let level = kvs.iter().enumerate().map(|(i, (key, value))| {
-        let leaf_index = i as u64 + 1;
+    let level = kvs.map(|(key, value, leaf_index)| {
         let mut bytes = [0_u8; 40];
         bytes[..8].copy_from_slice(&leaf_index.to_be_bytes());
         bytes[8..].copy_from_slice(value.as_ref());
-        (*key, hasher.hash_bytes(&bytes))
+        (key, hasher.hash_bytes(&bytes))
     });
     let mut level: Vec<(U256, H256)> = level.collect();
+    if level.is_empty() {
+        return hasher.empty_subtree_hash(256);
+    }
     level.sort_unstable_by_key(|(key, _)| *key);
 
     for _ in 0..256 {
@@ -62,6 +72,47 @@ pub fn compute_tree_hash(kvs: &[(U256, H256)]) -> H256 {
 // Computing the expected hash takes some time in the debug mode, so we memoize it.
 pub static KVS_AND_HASH: Lazy<(Vec<(U256, H256)>, H256)> = Lazy::new(|| {
     let kvs = generate_key_value_pairs(0..100);
-    let expected_hash = compute_tree_hash(&kvs);
+    let expected_hash = compute_tree_hash(kvs.iter().copied());
     (kvs, expected_hash)
 });
+
+pub fn convert_to_writes(kvs: &[(U256, H256)]) -> Vec<(U256, TreeInstruction)> {
+    let kvs = kvs
+        .iter()
+        .map(|&(key, hash)| (key, TreeInstruction::Write(hash)));
+    kvs.collect()
+}
+
+/// Emulates leaf index assignment in a real Merkle tree.
+#[derive(Debug)]
+pub struct TreeMap(HashMap<U256, (H256, u64)>);
+
+impl TreeMap {
+    pub fn new(initial_entries: &[(U256, H256)]) -> Self {
+        let map = initial_entries
+            .iter()
+            .enumerate()
+            .map(|(i, (key, value))| (*key, (*value, i as u64 + 1)))
+            .collect();
+        Self(map)
+    }
+
+    pub fn extend(&mut self, kvs: &[(U256, H256)]) {
+        for &(key, new_value) in kvs {
+            if let Some((value, _)) = self.0.get_mut(&key) {
+                *value = new_value;
+            } else {
+                let leaf_index = self.0.len() as u64 + 1;
+                self.0.insert(key, (new_value, leaf_index));
+            }
+        }
+    }
+
+    pub fn root_hash(&self) -> H256 {
+        let entries = self
+            .0
+            .iter()
+            .map(|(key, (value, idx))| (*key, *value, *idx));
+        compute_tree_hash_with_indices(entries)
+    }
+}

core/lib/merkle_tree/tests/integration/merkle_tree.rs

@@ -11,14 +11,7 @@ use zksync_merkle_tree::{
 };
 use zksync_types::{AccountTreeId, Address, StorageKey, H256, U256};
 
-use crate::common::{compute_tree_hash, generate_key_value_pairs, KVS_AND_HASH};
-
-fn convert_to_writes(kvs: &[(U256, H256)]) -> Vec<(U256, TreeInstruction)> {
-    let kvs = kvs
-        .iter()
-        .map(|&(key, hash)| (key, TreeInstruction::Write(hash)));
-    kvs.collect()
-}
+use crate::common::{compute_tree_hash, convert_to_writes, generate_key_value_pairs, KVS_AND_HASH};
 
 #[test]
 fn compute_tree_hash_works_correctly() {
@@ -31,7 +24,7 @@ fn compute_tree_hash_works_correctly() {
     let address: Address = "4b3af74f66ab1f0da3f2e4ec7a3cb99baf1af7b2".parse().unwrap();
     let key = StorageKey::new(AccountTreeId::new(address), H256::zero());
     let key = key.hashed_key_u256();
-    let hash = compute_tree_hash(&[(key, H256([1; 32]))]);
+    let hash = compute_tree_hash([(key, H256([1; 32]))].into_iter());
     assert_eq!(hash, EXPECTED_HASH);
 }
 
@@ -42,7 +35,7 @@ fn root_hash_is_computed_correctly_on_empty_tree() {
 
         let mut tree = MerkleTree::new(PatchSet::default());
         let kvs = generate_key_value_pairs(0..kv_count);
-        let expected_hash = compute_tree_hash(&kvs);
+        let expected_hash = compute_tree_hash(kvs.iter().copied());
         let output = tree.extend(kvs);
         assert_eq!(output.root_hash, expected_hash);
     }
@@ -59,7 +52,7 @@ fn output_proofs_are_computed_correctly_on_empty_tree() {
 
         let mut tree = MerkleTree::new(PatchSet::default());
         let kvs = generate_key_value_pairs(0..kv_count);
-        let expected_hash = compute_tree_hash(&kvs);
+        let expected_hash = compute_tree_hash(kvs.iter().copied());
         let instructions = convert_to_writes(&kvs);
         let output = tree.extend_with_proofs(instructions.clone());
 
@@ -89,7 +82,7 @@ fn entry_proofs_are_computed_correctly_on_empty_tree() {
 
         let mut tree = MerkleTree::new(PatchSet::default());
         let kvs = generate_key_value_pairs(0..kv_count);
-        let expected_hash = compute_tree_hash(&kvs);
+        let expected_hash = compute_tree_hash(kvs.iter().copied());
         tree.extend(kvs.clone());
 
         let existing_keys: Vec<_> = kvs.iter().map(|(key, _)| *key).collect();
@@ -137,7 +130,7 @@ fn proofs_are_computed_correctly_for_mixed_instructions() {
     let mut instructions: Vec<_> = reads.collect();
     // Overwrite all keys in the tree.
     let writes: Vec<_> = kvs.iter().map(|(key, _)| (*key, H256::zero())).collect();
-    let expected_hash = compute_tree_hash(&writes);
+    let expected_hash = compute_tree_hash(writes.iter().copied());
     instructions.extend(convert_to_writes(&writes));
     instructions.shuffle(&mut rng);
 
@@ -322,7 +315,7 @@ fn test_root_hash_computing_with_key_updates(db: impl Database) {
 
     let mut kvs = generate_key_value_pairs(0..50);
     let mut tree = MerkleTree::new(db);
-    let expected_hash = compute_tree_hash(&kvs);
+    let expected_hash = compute_tree_hash(kvs.iter().copied());
     let output = tree.extend(kvs.clone());
     assert_eq!(output.root_hash, expected_hash);
 
@@ -337,7 +330,7 @@ fn test_root_hash_computing_with_key_updates(db: impl Database) {
     let changed_kvs: Vec<_> = changed_kvs.collect();
     let new_kvs = generate_key_value_pairs(50..75);
     kvs.extend_from_slice(&new_kvs);
-    let expected_hash = compute_tree_hash(&kvs);
+    let expected_hash = compute_tree_hash(kvs.iter().copied());
 
     // We can merge `changed_kvs` and `new_kvs` in any way that preserves `new_kvs` ordering.
     // We'll do multiple ways (which also will effectively test DB rollbacks).
@@ -452,7 +445,7 @@ fn test_root_hash_equals_to_previous_implementation(db: &mut impl Database) {
     let values = (0..100).map(H256::from_low_u64_be);
     let kvs: Vec<_> = keys.zip(values).collect();
 
-    let expected_hash = compute_tree_hash(&kvs);
+    let expected_hash = compute_tree_hash(kvs.iter().copied());
     assert_eq!(expected_hash, PREV_IMPL_HASH);
 
     let mut tree = MerkleTree::new(db);

core/lib/merkle_tree/tests/integration/recovery.rs

@@ -1,11 +1,14 @@
 //! Tests for tree recovery.
 
+use rand::{rngs::StdRng, seq::SliceRandom, SeedableRng};
+use zksync_crypto::hasher::blake2::Blake2Hasher;
+
 use zksync_merkle_tree::{
     recovery::{MerkleTreeRecovery, RecoveryEntry},
-    PatchSet, PruneDatabase,
+    Database, MerkleTree, PatchSet, PruneDatabase, ValueHash,
 };
 
-use crate::common::KVS_AND_HASH;
+use crate::common::{convert_to_writes, generate_key_value_pairs, TreeMap, KVS_AND_HASH};
 
 #[test]
 fn recovery_basics() {
@@ -64,8 +67,56 @@ fn test_recovery_in_chunks<DB: PruneDatabase>(mut create_db: impl FnMut() -> DB)
         assert_eq!(recovery.last_processed_key(), Some(greatest_key));
         assert_eq!(recovery.root_hash(), *expected_hash);
 
-        let tree = recovery.finalize();
+        let mut tree = recovery.finalize();
         tree.verify_consistency(recovered_version).unwrap();
+        // Check that new tree versions can be built and function as expected.
+        test_tree_after_recovery(&mut tree, recovered_version, *expected_hash);
+    }
+}
+
+fn test_tree_after_recovery<DB: Database>(
+    tree: &mut MerkleTree<DB>,
+    recovered_version: u64,
+    root_hash: ValueHash,
+) {
+    const RNG_SEED: u64 = 765;
+    const CHUNK_SIZE: usize = 18;
+
+    assert_eq!(tree.latest_version(), Some(recovered_version));
+    assert_eq!(tree.root_hash(recovered_version), Some(root_hash));
+    for ver in 0..recovered_version {
+        assert_eq!(tree.root_hash(ver), None);
+    }
+
+    // Check adding new and updating existing entries in the tree.
+    let mut rng = StdRng::seed_from_u64(RNG_SEED);
+    let mut kvs = generate_key_value_pairs(100..=150);
+    let mut modified_kvs = generate_key_value_pairs(50..=100);
+    for (_, value) in &mut modified_kvs {
+        *value = ValueHash::repeat_byte(1);
+    }
+    kvs.extend(modified_kvs);
+    kvs.shuffle(&mut rng);
+
+    let mut tree_map = TreeMap::new(&KVS_AND_HASH.0);
+    let mut prev_root_hash = root_hash;
+    for (i, chunk) in kvs.chunks(CHUNK_SIZE).enumerate() {
+        tree_map.extend(chunk);
+
+        let new_root_hash = if i % 2 == 0 {
+            let output = tree.extend(chunk.to_vec());
+            output.root_hash
+        } else {
+            let instructions = convert_to_writes(chunk);
+            let output = tree.extend_with_proofs(instructions.clone());
+            output.verify_proofs(&Blake2Hasher, prev_root_hash, &instructions);
+            output.root_hash().unwrap()
+        };
+
+        assert_eq!(new_root_hash, tree_map.root_hash());
+        tree.verify_consistency(recovered_version + i as u64)
+            .unwrap();
+        prev_root_hash = new_root_hash;
     }
 }