feat(merkle tree): Provide Merkle proofs for tree entries and entry ranges

core/lib/merkle_tree/src/domain.rs

@@ -4,7 +4,7 @@ use rayon::{ThreadPool, ThreadPoolBuilder};
 
 use crate::{
     storage::{MerkleTreeColumnFamily, PatchSet, Patched, RocksDBWrapper},
-    types::{Key, LeafData, Root, TreeInstruction, TreeLogEntry, ValueHash, TREE_DEPTH},
+    types::{Key, Root, TreeInstruction, TreeLogEntry, ValueHash, TREE_DEPTH},
     BlockOutput, HashTree, MerkleTree,
 };
 use zksync_crypto::hasher::blake2::Blake2Hasher;
@@ -159,17 +159,6 @@ impl ZkSyncTree {
         });
     }
 
-    /// Reads leaf nodes with the specified keys from the tree storage. The nodes
-    /// are returned in a `Vec` in the same order as requested.
-    pub fn read_leaves(
-        &self,
-        l1_batch_number: L1BatchNumber,
-        leaf_keys: &[Key],
-    ) -> Vec<Option<LeafData>> {
-        let version = u64::from(l1_batch_number.0);
-        self.tree.read_leaves(version, leaf_keys)
-    }
-
     /// Processes an iterator of storage logs comprising a single L1 batch.
     pub fn process_l1_batch(&mut self, storage_logs: &[StorageLog]) -> TreeMetadata {
         match self.mode {

core/lib/merkle_tree/src/errors.rs

@@ -135,6 +135,35 @@ impl fmt::Display for DeserializeError {
 
 impl error::Error for DeserializeError {}
 
+/// Error accessing a specific tree version.
+#[derive(Debug)]
+pub struct NoVersionError {
+    pub(crate) missing_version: u64,
+    pub(crate) version_count: u64,
+}
+
+impl fmt::Display for NoVersionError {
+    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let &Self {
+            missing_version,
+            version_count,
+        } = self;
+        if missing_version >= version_count {
+            write!(
+                formatter,
+                "Version {missing_version} does not exist in Merkle tree; it has {version_count} versions"
+            )
+        } else {
+            write!(
+                formatter,
+                "Version {missing_version} was pruned from Merkle tree"
+            )
+        }
+    }
+}
+
+impl error::Error for NoVersionError {}
+
 #[cfg(test)]
 mod tests {
     use super::*;

core/lib/merkle_tree/src/getters.rs

@@ -1,37 +1,135 @@
 //! Getters for the Merkle tree.
 
 use crate::{
+    hasher::HasherWithStats,
     storage::{LoadAncestorsResult, SortedKeys, WorkingPatchSet},
-    types::{LeafData, Node},
-    Database, Key, MerkleTree,
+    types::{Nibbles, Node, TreeEntry, TreeEntryWithProof},
+    Database, Key, MerkleTree, NoVersionError, ValueHash,
 };
 
 impl<DB> MerkleTree<'_, DB>
 where
     DB: Database,
 {
-    /// Reads leaf nodes with the specified keys from the tree storage. The nodes
-    /// are returned in a `Vec` in the same order as requested.
-    pub fn read_leaves(&self, version: u64, leaf_keys: &[Key]) -> Vec<Option<LeafData>> {
-        let Some(root) = self.db.root(version) else {
-            return vec![None; leaf_keys.len()];
-        };
+    /// Reads entries with the specified keys from the tree. The entries are returned in the same order
+    /// as requested.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the tree `version` is missing.
+    pub fn entries(
+        &self,
+        version: u64,
+        leaf_keys: &[Key],
+    ) -> Result<Vec<TreeEntry>, NoVersionError> {
+        self.load_and_transform_entries(
+            version,
+            leaf_keys,
+            |patch_set, leaf_key, longest_prefix| {
+                let node = patch_set.get(longest_prefix);
+                match node {
+                    Some(Node::Leaf(leaf)) if &leaf.full_key == leaf_key => (*leaf).into(),
+                    _ => TreeEntry::empty(),
+                }
+            },
+        )
+    }
+
+    fn load_and_transform_entries<T>(
+        &self,
+        version: u64,
+        leaf_keys: &[Key],
+        mut transform: impl FnMut(&mut WorkingPatchSet, &Key, &Nibbles) -> T,
+    ) -> Result<Vec<T>, NoVersionError> {
+        let root = self.db.root(version).ok_or_else(|| {
+            let manifest = self.db.manifest().unwrap_or_default();
+            NoVersionError {
+                missing_version: version,
+                version_count: manifest.version_count,
+            }
+        })?;
         let sorted_keys = SortedKeys::new(leaf_keys.iter().copied());
         let mut patch_set = WorkingPatchSet::new(version, root);
         let LoadAncestorsResult {
             longest_prefixes, ..
         } = patch_set.load_ancestors(&sorted_keys, &self.db);
 
-        leaf_keys
+        Ok(leaf_keys
             .iter()
             .zip(&longest_prefixes)
-            .map(|(leaf_key, longest_prefix)| {
-                let node = patch_set.get(longest_prefix);
-                match node {
-                    Some(Node::Leaf(leaf)) if &leaf.full_key == leaf_key => Some((*leaf).into()),
-                    _ => None,
+            .map(|(leaf_key, longest_prefix)| transform(&mut patch_set, leaf_key, longest_prefix))
+            .collect())
+    }
+
+    /// Reads entries together with Merkle proofs with the specified keys from the tree. The entries are returned
+    /// in the same order as requested.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the tree `version` is missing.
+    pub fn entries_with_proofs(
+        &self,
+        version: u64,
+        leaf_keys: &[Key],
+    ) -> Result<Vec<TreeEntryWithProof>, NoVersionError> {
+        let mut hasher = HasherWithStats::from(self.hasher);
+        self.load_and_transform_entries(
+            version,
+            leaf_keys,
+            |patch_set, &leaf_key, longest_prefix| {
+                let (leaf, merkle_path) =
+                    patch_set.create_proof(&mut hasher, leaf_key, longest_prefix, 0);
+                let value_hash = leaf
+                    .as_ref()
+                    .map_or_else(ValueHash::zero, |leaf| leaf.value_hash);
+                TreeEntry {
+                    value_hash,
+                    leaf_index: leaf.map_or(0, |leaf| leaf.leaf_index),
                 }
-            })
-            .collect()
+                .with_merkle_path(merkle_path.into_inner())
+            },
+        )
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::PatchSet;
+
+    #[test]
+    fn entries_in_empty_tree() {
+        let mut tree = MerkleTree::new(PatchSet::default());
+        tree.extend(vec![]);
+        let missing_key = Key::from(123);
+
+        let entries = tree.entries(0, &[missing_key]).unwrap();
+        assert_eq!(entries.len(), 1);
+        assert!(entries[0].is_empty());
+
+        let entries = tree.entries_with_proofs(0, &[missing_key]).unwrap();
+        assert_eq!(entries.len(), 1);
+        assert!(entries[0].base.is_empty());
+        entries[0].verify(tree.hasher, missing_key, tree.hasher.empty_tree_hash());
+    }
+
+    #[test]
+    fn entries_in_single_node_tree() {
+        let mut tree = MerkleTree::new(PatchSet::default());
+        let key = Key::from(987_654);
+        let output = tree.extend(vec![(key, ValueHash::repeat_byte(1))]);
+        let missing_key = Key::from(123);
+
+        let entries = tree.entries(0, &[key, missing_key]).unwrap();
+        assert_eq!(entries.len(), 2);
+        assert_eq!(entries[0].value_hash, ValueHash::repeat_byte(1));
+        assert_eq!(entries[0].leaf_index, 1);
+
+        let entries = tree.entries_with_proofs(0, &[key, missing_key]).unwrap();
+        assert_eq!(entries.len(), 2);
+        assert!(!entries[0].base.is_empty());
+        entries[0].verify(tree.hasher, key, output.root_hash);
+        assert!(entries[1].base.is_empty());
+        entries[1].verify(tree.hasher, missing_key, output.root_hash);
     }
 }