feat(zkos): Implement ZK OS tree manager (#3730)

## What ❔

Implements a tree manager component responsible for updating the ZK OS
Merkle tree based on Postgres data and providing REST API for it (to be
used for witness input generation).

## Why ❔

Necessary to keep the ZK OS Merkle tree synced.

## Is this a breaking change?

- [ ] Yes
- [x] No

## Operational changes

No operational changes; the new component is not tied to the node
binaries.

## 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 `zkstack dev fmt` and `zkstack dev
lint`.

Author: slowli

core/Cargo.lock

@@ -40,6 +40,7 @@ members = [
   "node/external_proof_integration_api",
   "node/logs_bloom_backfill",
   "node/da_clients",
+  "node/zk_os_tree_manager",
   # Libraries
   "lib/db_connection",
   "lib/zksync_core_leftovers",
@@ -303,6 +304,7 @@ zksync_web3_decl = { version = "27.0.0-non-semver-compat", path = "lib/web3_decl
 zksync_crypto_primitives = { version = "27.0.0-non-semver-compat", path = "lib/crypto_primitives" }
 zksync_external_price_api = { version = "27.0.0-non-semver-compat", path = "lib/external_price_api" }
 zksync_task_management = { version = "27.0.0-non-semver-compat", path = "lib/task_management" }
+zk_os_merkle_tree = { version = "27.0.0-non-semver-compat", path = "lib/zk_os_merkle_tree" }
 
 # Framework and components
 zksync_node_framework = { version = "27.0.0-non-semver-compat", path = "node/node_framework" }

core/Cargo.toml

@@ -27,14 +27,26 @@ impl DeployedContract {
     }
 }
 
-/// Holder for l1 batches data, used in eth sender metrics
+/// Holder for l1 batches data.
+#[derive(Debug, Clone, Copy, PartialEq)]
 pub struct L1BatchStatistics {
     pub number: L1BatchNumber,
     pub timestamp: u64,
     pub l2_tx_count: u32,
     pub l1_tx_count: u32,
 }
 
+impl From<L1BatchHeader> for L1BatchStatistics {
+    fn from(header: L1BatchHeader) -> Self {
+        Self {
+            number: header.number,
+            timestamp: header.timestamp,
+            l1_tx_count: header.l1_tx_count.into(),
+            l2_tx_count: header.l2_tx_count.into(),
+        }
+    }
+}
+
 /// Holder for the block metadata that is not available from transactions themselves.
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
 pub struct L1BatchHeader {

core/lib/types/src/block.rs

@@ -9,8 +9,7 @@ use zksync_basic_types::H256;
 use crate::{types::Leaf, BatchOutput, HashTree, TreeEntry};
 
 /// Operation on a Merkle tree entry used in [`BatchTreeProof`].
-#[derive(Debug, Clone, Copy)]
-#[cfg_attr(test, derive(PartialEq))]
+#[derive(Debug, Clone, Copy, PartialEq)]
 pub enum TreeOperation {
     /// Operation hitting an existing entry (i.e., an update or read).
     Hit { index: u64 },
@@ -29,6 +28,13 @@ pub struct IntermediateHash {
     pub location: (u8, u64),
 }
 
+#[cfg(not(test))]
+impl From<H256> for IntermediateHash {
+    fn from(value: H256) -> Self {
+        Self { value }
+    }
+}
+
 /// Partial view of the Merkle tree returned from [`BatchTreeProof::verify()`].
 #[derive(Debug)]
 pub struct MerkleTreeView {

core/lib/zk_os_merkle_tree/src/hasher/proofs.rs

@@ -8,21 +8,22 @@ pub use zksync_crypto_primitives::hasher::blake2::Blake2Hasher;
 
 pub use self::{
     errors::DeserializeError,
-    hasher::HashTree,
+    hasher::{BatchTreeProof, HashTree, TreeOperation},
+    reader::MerkleTreeReader,
     storage::{Database, MerkleTreeColumnFamily, PatchSet, Patched, RocksDBWrapper},
     types::{BatchOutput, TreeEntry},
 };
 use crate::{
-    hasher::BatchTreeProof,
     metrics::{BatchProofStage, LoadStage, MerkleTreeInfo, METRICS},
-    storage::{TreeUpdate, WorkingPatchSet},
-    types::MAX_TREE_DEPTH,
+    storage::{AsEntry, TreeUpdate, WorkingPatchSet},
+    types::{Leaf, MAX_TREE_DEPTH},
 };
 
 mod consistency;
 mod errors;
 mod hasher;
 mod metrics;
+mod reader;
 mod storage;
 #[cfg(test)]
 mod tests;
@@ -33,7 +34,7 @@ mod types;
 /// these types will remain stable.
 #[doc(hidden)]
 pub mod unstable {
-    pub use crate::types::{KeyLookup, Manifest, Node, NodeKey, Root};
+    pub use crate::types::{KeyLookup, Leaf, Manifest, Node, NodeKey, RawNode, Root};
 }
 
 /// Marker trait for tree parameters.
@@ -101,6 +102,23 @@ impl<DB: Database> MerkleTree<DB> {
     }
 }
 
+impl<DB: Database, P: TreeParams> MerkleTree<DB, P>
+where
+    P::Hasher: Default,
+{
+    /// Returns the hash of the empty tree.
+    pub fn empty_tree_hash() -> H256 {
+        let hasher = P::Hasher::default();
+        let min_guard_hash = hasher.hash_leaf(&Leaf::MIN_GUARD);
+        let max_guard_hash = hasher.hash_leaf(&Leaf::MAX_GUARD);
+        let mut hash = hasher.hash_branch(&min_guard_hash, &max_guard_hash);
+        for depth in 1..P::TREE_DEPTH {
+            hash = hasher.hash_branch(&hash, &hasher.empty_subtree_hash(depth));
+        }
+        hash
+    }
+}
+
 impl<DB: Database, P: TreeParams> MerkleTree<DB, P> {
     /// Loads a tree with the specified hasher.
     ///
@@ -125,6 +143,11 @@ impl<DB: Database, P: TreeParams> MerkleTree<DB, P> {
         Ok(Self { db, hasher })
     }
 
+    /// Returns a reference to the database.
+    pub fn db(&self) -> &DB {
+        &self.db
+    }
+
     /// Returns the root hash of a tree at the specified `version`, or `None` if the version
     /// was not written yet.
     pub fn root_hash(&self, version: u64) -> anyhow::Result<Option<H256>> {
@@ -134,6 +157,15 @@ impl<DB: Database, P: TreeParams> MerkleTree<DB, P> {
         Ok(Some(root.hash::<P>(&self.hasher)))
     }
 
+    /// Returns the root hash and leaf count at the specified version.
+    pub fn root_info(&self, version: u64) -> anyhow::Result<Option<(H256, u64)>> {
+        let Some(root) = self.db.try_root(version)? else {
+            return Ok(None);
+        };
+        let root_hash = root.hash::<P>(&self.hasher);
+        Ok(Some((root_hash, root.leaf_count)))
+    }
+
     /// Returns the latest version of the tree present in the database, or `None` if
     /// no versions are present yet.
     pub fn latest_version(&self) -> anyhow::Result<Option<u64>> {
@@ -160,7 +192,7 @@ impl<DB: Database, P: TreeParams> MerkleTree<DB, P> {
     /// - Returns an error if the version doesn't exist.
     /// - Proxies database errors.
     pub fn prove(&self, version: u64, keys: &[H256]) -> anyhow::Result<BatchTreeProof> {
-        let (patch, mut update) = self.create_patch(version, &[], keys)?;
+        let (patch, mut update) = self.create_patch::<TreeEntry>(version, &[], keys)?;
         Ok(patch.create_batch_proof(&self.hasher, vec![], update.take_read_operations()))
     }
 
@@ -180,10 +212,22 @@ impl<DB: Database, P: TreeParams> MerkleTree<DB, P> {
         Ok(output)
     }
 
+    /// Same as [`Self::extend()`], but with a cross-check for the leaf indices of `entries`.
+    ///
+    /// `reference_indices` must have the same length as `entries`. Note that because indices are *only* used for cross-checking,
+    /// entries (more specifically, new insertions) must still be correctly ordered!
+    pub fn extend_with_reference(
+        &mut self,
+        entries: &[(u64, TreeEntry)],
+    ) -> anyhow::Result<BatchOutput> {
+        let (output, _) = self.extend_inner(entries, None)?;
+        Ok(output)
+    }
+
     #[tracing::instrument(level = "debug", name = "extend", skip_all, fields(latest_version))]
     fn extend_inner(
         &mut self,
-        entries: &[TreeEntry],
+        entries: &[impl AsEntry],
         read_keys: Option<&[H256]>,
     ) -> anyhow::Result<(BatchOutput, Option<BatchTreeProof>)> {
         let latest_version = self
@@ -291,3 +335,10 @@ impl<DB: Database, P: TreeParams> MerkleTree<DB, P> {
         Ok(())
     }
 }
+
+impl<DB: Database, P: TreeParams> MerkleTree<Patched<DB>, P> {
+    /// Flushes changes to the underlying storage.
+    pub fn flush(&mut self) -> anyhow::Result<()> {
+        self.db.flush()
+    }
+}

core/lib/zk_os_merkle_tree/src/lib.rs

@@ -0,0 +1,94 @@
+//! Merkle tree reader.
+
+use std::fmt;
+
+use zksync_basic_types::H256;
+
+use crate::{
+    hasher::BatchTreeProof,
+    types::{NodeKey, RawNode},
+    Database, DefaultTreeParams, MerkleTree, RocksDBWrapper, TreeParams,
+};
+
+pub struct MerkleTreeReader<DB, P: TreeParams = DefaultTreeParams>(MerkleTree<DB, P>);
+
+impl<DB: fmt::Debug, P: TreeParams> fmt::Debug for MerkleTreeReader<DB, P>
+where
+    P::Hasher: fmt::Debug,
+{
+    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
+        fmt::Debug::fmt(&self.0, formatter)
+    }
+}
+
+impl<DB: Database + Clone, P: TreeParams> Clone for MerkleTreeReader<DB, P>
+where
+    P::Hasher: Clone,
+{
+    fn clone(&self) -> Self {
+        Self(MerkleTree {
+            db: self.0.db.clone(),
+            hasher: self.0.hasher.clone(),
+        })
+    }
+}
+
+impl<DB: Database> MerkleTreeReader<DB> {
+    /// Creates a tree reader based on the provided database.
+    ///
+    /// # Errors
+    ///
+    /// Errors if sanity checks fail.
+    pub fn new(db: DB) -> anyhow::Result<Self> {
+        MerkleTree::new(db).map(Self)
+    }
+}
+
+impl<DB: Database, P: TreeParams> MerkleTreeReader<DB, P> {
+    /// Returns a reference to the database.
+    pub fn db(&self) -> &DB {
+        &self.0.db
+    }
+
+    /// Converts this reader to the underlying DB.
+    pub fn into_db(self) -> DB {
+        self.0.db
+    }
+
+    /// Returns the latest version of the tree present in the database, or `None` if
+    /// no versions are present yet.
+    pub fn latest_version(&self) -> anyhow::Result<Option<u64>> {
+        self.0.latest_version()
+    }
+
+    /// Returns the root hash and leaf count at the specified version.
+    pub fn root_info(&self, version: u64) -> anyhow::Result<Option<(H256, u64)>> {
+        self.0.root_info(version)
+    }
+
+    /// Creates a batch proof for `keys` at the specified tree version.
+    ///
+    /// # Errors
+    ///
+    /// - Returns an error if the version doesn't exist.
+    /// - Proxies database errors.
+    pub fn prove(&self, version: u64, keys: &[H256]) -> anyhow::Result<BatchTreeProof> {
+        self.0.prove(version, keys)
+    }
+}
+
+impl<P: TreeParams> MerkleTreeReader<RocksDBWrapper, P> {
+    /// Returns raw nodes for the specified `keys`.
+    pub fn raw_nodes(&self, node_keys: &[NodeKey]) -> anyhow::Result<Vec<Option<RawNode>>> {
+        let raw_nodes = self.0.db.raw_nodes(node_keys).into_iter();
+        let raw_nodes = node_keys.iter().zip(raw_nodes).map(|(key, slice)| {
+            let slice = slice?;
+            Some(if key.nibble_count == 0 {
+                RawNode::deserialize_root(&slice)
+            } else {
+                RawNode::deserialize(&slice)
+            })
+        });
+        Ok(raw_nodes.collect())
+    }
+}

core/lib/zk_os_merkle_tree/src/reader.rs

@@ -10,7 +10,7 @@ pub(crate) use self::patch::{TreeUpdate, WorkingPatchSet};
 pub use self::rocksdb::{MerkleTreeColumnFamily, RocksDBWrapper};
 use crate::{
     errors::{DeserializeContext, DeserializeError, DeserializeErrorKind},
-    types::{InternalNode, KeyLookup, Leaf, Manifest, Node, NodeKey, Root},
+    types::{InternalNode, KeyLookup, Leaf, Manifest, Node, NodeKey, Root, TreeEntry},
 };
 
 mod patch;
@@ -19,6 +19,38 @@ mod serialization;
 #[cfg(test)]
 mod tests;
 
+/// [`TreeEntry`] generalization allowing to check leaf index assignment when loading data from the tree.
+pub(crate) trait AsEntry {
+    fn as_entry(&self) -> &TreeEntry;
+
+    fn check_index(&self, index: u64) -> anyhow::Result<()>;
+}
+
+impl AsEntry for TreeEntry {
+    fn as_entry(&self) -> &TreeEntry {
+        self
+    }
+
+    fn check_index(&self, _index: u64) -> anyhow::Result<()> {
+        Ok(())
+    }
+}
+
+impl AsEntry for (u64, TreeEntry) {
+    fn as_entry(&self) -> &TreeEntry {
+        &self.1
+    }
+
+    fn check_index(&self, index: u64) -> anyhow::Result<()> {
+        let (ref_index, entry) = self;
+        anyhow::ensure!(
+            index == *ref_index,
+            "Unexpected index for {entry:?}: reference is {ref_index}, but tree implies {index}",
+        );
+        Ok(())
+    }
+}
+
 /// Generic database functionality. Its main implementation is [`RocksDB`].
 pub trait Database: Send + Sync {
     fn indices(&self, version: u64, keys: &[H256]) -> Result<Vec<KeyLookup>, DeserializeError>;
@@ -303,6 +335,11 @@ impl<DB: Database> Patched<DB> {
         self.patch = None;
     }
 
+    /// Returns a reference to the underlying database. It is unsound to modify the database using this reference.
+    pub fn inner(&self) -> &DB {
+        &self.inner
+    }
+
     /// Returns the wrapped database.
     ///
     /// # Panics