feat(sync-layer): adapt MiniMerkleTree to manage priority queue

checks-config/era.dic

@@ -969,3 +969,4 @@ preloaded
 e2e
 upcasting
 foundryup
+uncached

core/lib/mini_merkle_tree/benches/tree.rs

@@ -10,11 +10,7 @@ const TREE_SIZES: &[usize] = &[32, 64, 128, 256, 512, 1_024];
 fn compute_merkle_root(bencher: &mut Bencher<'_>, tree_size: usize) {
     let leaves = (0..tree_size).map(|i| [i as u8; 88]);
     let tree = MiniMerkleTree::new(leaves, None);
-    bencher.iter_batched(
-        || tree.clone(),
-        MiniMerkleTree::merkle_root,
-        BatchSize::SmallInput,
-    );
+    bencher.iter_batched(|| &tree, MiniMerkleTree::merkle_root, BatchSize::SmallInput);
 }
 
 fn compute_merkle_path(bencher: &mut Bencher<'_>, tree_size: usize) {

core/lib/mini_merkle_tree/src/lib.rs

@@ -5,7 +5,7 @@
 #![warn(clippy::all, clippy::pedantic)]
 #![allow(clippy::must_use_candidate, clippy::similar_names)]
 
-use std::iter;
+use std::{collections::VecDeque, iter};
 
 use once_cell::sync::Lazy;
 
@@ -19,16 +19,26 @@ use zksync_crypto::hasher::{keccak::KeccakHasher, Hasher};
 /// we unlikely to ever hit.
 const MAX_TREE_DEPTH: usize = 32;
 
-/// In-memory Merkle tree of bounded depth (no more than 10).
+/// In-memory Merkle tree of bounded depth (no more than 32).
 ///
 /// The tree is left-leaning, meaning that during its initialization, the size of a tree
 /// can be specified larger than the number of provided leaves. In this case, the remaining leaves
 /// will be considered to equal `[0_u8; LEAF_SIZE]`.
+///
+/// The tree has dynamic size, meaning that it can grow by a factor of 2 when the number of leaves
+/// exceeds the current tree size. It does not shrink.
+///
+/// The tree is optimized for the case when the queries are performed on the rightmost leaves
+/// and the leftmost leaves are cached. Caching enables the merkle roots and paths to be computed
+/// in `O(n)` time, where `n` is the number of uncached leaves (in contrast to the total number of
+/// leaves). Cache itself only takes up `O(depth)` space.
 #[derive(Debug, Clone)]
 pub struct MiniMerkleTree<const LEAF_SIZE: usize, H = KeccakHasher> {
     hasher: H,
-    hashes: Box<[H256]>,
+    hashes: VecDeque<H256>,
     binary_tree_size: usize,
+    head_index: usize,
+    left_cache: Vec<H256>,
 }
 
 impl<const LEAF_SIZE: usize> MiniMerkleTree<LEAF_SIZE>
@@ -62,12 +72,13 @@ where
     /// Panics if any of the following conditions applies:
     ///
     /// - `min_tree_size` (if supplied) is not a power of 2.
+    /// - The number of leaves is greater than `2^32`.
     pub fn with_hasher(
         hasher: H,
         leaves: impl Iterator<Item = [u8; LEAF_SIZE]>,
         min_tree_size: Option<usize>,
     ) -> Self {
-        let hashes: Box<[H256]> = leaves.map(|bytes| hasher.hash_bytes(&bytes)).collect();
+        let hashes: VecDeque<_> = leaves.map(|bytes| hasher.hash_bytes(&bytes)).collect();
         let mut binary_tree_size = hashes.len().next_power_of_two();
         if let Some(min_tree_size) = min_tree_size {
             assert!(
@@ -76,8 +87,9 @@ where
             );
             binary_tree_size = min_tree_size.max(binary_tree_size);
         }
+        let depth = tree_depth_by_size(binary_tree_size);
         assert!(
-            tree_depth_by_size(binary_tree_size) <= MAX_TREE_DEPTH,
+            depth <= MAX_TREE_DEPTH,
             "Tree contains more than {} items; this is not supported",
             1 << MAX_TREE_DEPTH
         );
@@ -86,67 +98,164 @@ where
             hasher,
             hashes,
             binary_tree_size,
+            head_index: 0,
+            left_cache: vec![],
         }
     }
 
+    /// Returns `true` if the tree is empty.
+    pub fn is_empty(&self) -> bool {
+        self.head_index == 0 && self.hashes.is_empty()
+    }
+
     /// Returns the root hash of this tree.
-    /// # Panics
-    /// Will panic if the constant below is invalid.
-    pub fn merkle_root(self) -> H256 {
+    pub fn merkle_root(&self) -> H256 {
         if self.hashes.is_empty() {
             let depth = tree_depth_by_size(self.binary_tree_size);
-            self.hasher.empty_subtree_hash(depth)
+            if self.head_index == 0 {
+                self.hasher.empty_subtree_hash(depth)
+            } else {
+                self.left_cache[depth]
+            }
         } else {
-            self.compute_merkle_root_and_path(0, None)
+            self.compute_merkle_root_and_path(0, None, None)
         }
     }
 
     /// Returns the root hash and the Merkle proof for a leaf with the specified 0-based `index`.
-    pub fn merkle_root_and_path(self, index: usize) -> (H256, Vec<H256>) {
-        let mut merkle_path = vec![];
-        let root_hash = self.compute_merkle_root_and_path(index, Some(&mut merkle_path));
-        (root_hash, merkle_path)
+    /// `index` is relative to the leftmost uncached leaf.
+    pub fn merkle_root_and_path(&self, index: usize) -> (H256, Vec<H256>) {
+        let (mut left_path, mut right_path) = (vec![], vec![]);
+        let root_hash =
+            self.compute_merkle_root_and_path(index, Some((&mut left_path, &mut right_path)), None);
+        (root_hash, right_path)
+    }
+
+    /// Returns the root hash and the Merkle proofs for an interval of leafs.
+    /// The interval is [0, `length`), where `0` is the leftmost uncached leaf.
+    pub fn merkle_root_and_paths_for_interval(
+        &self,
+        length: usize,
+    ) -> (H256, Vec<H256>, Vec<H256>) {
+        let (mut left_path, mut right_path) = (vec![], vec![]);
+        let root_hash = self.compute_merkle_root_and_path(
+            length - 1,
+            Some((&mut left_path, &mut right_path)),
+            None,
+        );
+        (root_hash, left_path, right_path)
+    }
+
+    /// Adds a new leaf to the tree (replaces leftmost empty leaf).
+    /// If the tree is full, its size is doubled.
+    /// Note: empty leaves != zero leaves.
+    pub fn push(&mut self, leaf: [u8; LEAF_SIZE]) {
+        let leaf_hash = self.hasher.hash_bytes(&leaf);
+        self.hashes.push_back(leaf_hash);
+        if self.head_index + self.hashes.len() > self.binary_tree_size {
+            self.binary_tree_size *= 2;
+        }
+    }
+
+    /// Caches the rightmost `count` leaves.
+    /// Does not affect the root hash, but makes it impossible to get the paths to the cached leaves.
+    /// # Panics
+    /// Panics if `count` is greater than the number of non-cached leaves in the tree.
+    pub fn cache(&mut self, count: usize) {
+        assert!(self.hashes.len() >= count, "not enough leaves to cache");
+        let mut new_cache = vec![];
+        self.compute_merkle_root_and_path(count - 1, None, Some(&mut new_cache));
+        self.hashes.drain(0..count);
+        self.head_index += count;
+        self.left_cache = new_cache;
     }
 
     fn compute_merkle_root_and_path(
-        self,
-        mut index: usize,
-        mut merkle_path: Option<&mut Vec<H256>>,
+        &self,
+        mut right_index: usize,
+        mut merkle_paths: Option<(&mut Vec<H256>, &mut Vec<H256>)>,
+        mut new_cache: Option<&mut Vec<H256>>,
     ) -> H256 {
-        assert!(index < self.hashes.len(), "invalid tree leaf index");
+        assert!(right_index < self.hashes.len(), "invalid tree leaf index");
 
         let depth = tree_depth_by_size(self.binary_tree_size);
-        if let Some(merkle_path) = merkle_path.as_deref_mut() {
-            merkle_path.reserve(depth);
+        if let Some((left_path, right_path)) = &mut merkle_paths {
+            left_path.reserve(depth);
+            right_path.reserve(depth);
+        }
+        if let Some(new_cache) = new_cache.as_deref_mut() {
+            new_cache.reserve(depth + 1);
         }
 
-        let mut hashes = self.hashes;
+        let mut hashes = self.hashes.clone();
         let mut level_len = hashes.len();
+        let mut head_index = self.head_index;
+
         for level in 0..depth {
             let empty_hash_at_level = self.hasher.empty_subtree_hash(level);
 
-            if let Some(merkle_path) = merkle_path.as_deref_mut() {
-                let adjacent_idx = index ^ 1;
-                let adjacent_hash = if adjacent_idx < level_len {
-                    hashes[adjacent_idx]
-                } else {
+            let sibling_hash = |index: usize| {
+                if index == 0 && head_index % 2 == 1 {
+                    self.left_cache[level]
+                } else if index == level_len - 1 && (head_index + index) % 2 == 0 {
                     empty_hash_at_level
-                };
-                merkle_path.push(adjacent_hash);
+                } else {
+                    // `index` is relative to `head_index`
+                    let sibling = ((head_index + index) ^ 1) - head_index;
+                    hashes[sibling]
+                }
+            };
+
+            if let Some((left_path, right_path)) = &mut merkle_paths {
+                left_path.push(sibling_hash(0));
+                right_path.push(sibling_hash(right_index));
             }
 
-            for i in 0..(level_len / 2) {
-                hashes[i] = self.hasher.compress(&hashes[2 * i], &hashes[2 * i + 1]);
+            if let Some(new_cache) = new_cache.as_deref_mut() {
+                // We cache the rightmost left child on the current level
+                // within the given interval.
+                let cache = if (head_index + right_index) % 2 == 0 {
+                    hashes[right_index]
+                } else if right_index == 0 {
+                    self.left_cache[level]
+                } else {
+                    hashes[right_index - 1]
+                };
+                new_cache.push(cache);
             }
-            if level_len % 2 == 1 {
-                hashes[level_len / 2] = self
-                    .hasher
-                    .compress(&hashes[level_len - 1], &empty_hash_at_level);
+
+            let parity = head_index % 2;
+            // If our queue starts from the right child or ends with the left child,
+            // we need to round the length up.
+            let next_level_len = level_len / 2 + (parity | (level_len % 2));
+
+            for i in 0..next_level_len {
+                let lhs = if i == 0 && parity == 1 {
+                    // If the leftmost element is a right child, we need to use cache.
+                    self.left_cache[level]
+                } else {
+                    hashes[2 * i - parity]
+                };
+                let rhs = if i == next_level_len - 1 && (level_len - parity) % 2 == 1 {
+                    // If the rightmost element is a left child, we need to use the empty hash.
+                    empty_hash_at_level
+                } else {
+                    hashes[2 * i + 1 - parity]
+                };
+                hashes[i] = self.hasher.compress(&lhs, &rhs);
             }
 
-            index /= 2;
-            level_len = level_len / 2 + level_len % 2;
+            right_index = (right_index + parity) / 2;
+            level_len = next_level_len;
+            head_index /= 2;
+        }
+
+        if let Some(new_cache) = new_cache {
+            // It is important to cache the root as well, in case
+            // we just cached all elements and will grow on the next push.
+            new_cache.push(hashes[0]);
         }
+
         hashes[0]
     }
 }