Merge pull request #115 from 0xPolygonMiden/vlopes11-store-smt-depth

feat: Add `depth` as store SMT argument

README.md

@@ -13,7 +13,7 @@ For performance benchmarks of these hash functions and their comparison to other
 [Merkle module](./src/merkle/) provides a set of data structures related to Merkle trees. All these data structures are implemented using the RPO hash function described above. The data structures are:
 
 * `MerkleTree`: a regular fully-balanced binary Merkle tree. The depth of this tree can be at most 64.
-* `SimpleSmt`: a Sparse Merkle Tree, mapping 63-bit keys to 4-element leaf values.
+* `SimpleSmt`: a Sparse Merkle Tree, mapping 64-bit keys to 4-element leaf values.
 * `MerklePathSet`: a collection of Merkle authentication paths all resolving to the same root. The length of the paths can be at most 64.
 * `MerkleStore`: a collection of Merkle trees of different heights designed to efficiently store trees with common subtrees.
 * `Mmr`: a Merkle mountain range structure designed to function as an append-only log.

benches/store.rs

@@ -27,7 +27,7 @@ fn random_index(range: u64) -> u64 {
 fn get_empty_leaf_simplesmt(c: &mut Criterion) {
     let mut group = c.benchmark_group("get_empty_leaf_simplesmt");
 
-    let depth = 63u8;
+    let depth = SimpleSmt::MAX_DEPTH;
     let size = 2u64.pow(depth as u32);
 
     // both SMT and the store are pre-populated with empty hashes, accessing these values is what is
@@ -103,12 +103,12 @@ fn get_leaf_simplesmt(c: &mut Criterion) {
             .enumerate()
             .map(|(c, v)| (c.try_into().unwrap(), v.into()))
             .collect::<Vec<(u64, Word)>>();
-        let smt = SimpleSmt::new(63)
+        let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH)
             .unwrap()
             .with_leaves(smt_leaves.clone())
             .unwrap();
         let store = MerkleStore::new()
-            .with_sparse_merkle_tree(smt_leaves)
+            .with_sparse_merkle_tree(SimpleSmt::MAX_DEPTH, smt_leaves)
             .unwrap();
         let depth = smt.depth();
         let root = smt.root();
@@ -136,7 +136,7 @@ fn get_leaf_simplesmt(c: &mut Criterion) {
 fn get_node_of_empty_simplesmt(c: &mut Criterion) {
     let mut group = c.benchmark_group("get_node_of_empty_simplesmt");
 
-    let depth = 63u8;
+    let depth = SimpleSmt::MAX_DEPTH;
     let size = 2u64.pow(depth as u32);
 
     // both SMT and the store are pre-populated with the empty hashes, accessing the internal nodes
@@ -216,12 +216,12 @@ fn get_node_simplesmt(c: &mut Criterion) {
             .enumerate()
             .map(|(c, v)| (c.try_into().unwrap(), v.into()))
             .collect::<Vec<(u64, Word)>>();
-        let smt = SimpleSmt::new(63)
+        let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH)
             .unwrap()
             .with_leaves(smt_leaves.clone())
             .unwrap();
         let store = MerkleStore::new()
-            .with_sparse_merkle_tree(smt_leaves)
+            .with_sparse_merkle_tree(SimpleSmt::MAX_DEPTH, smt_leaves)
             .unwrap();
         let root = smt.root();
         let size_u64 = size as u64;
@@ -295,12 +295,12 @@ fn get_leaf_path_simplesmt(c: &mut Criterion) {
             .enumerate()
             .map(|(c, v)| (c.try_into().unwrap(), v.into()))
             .collect::<Vec<(u64, Word)>>();
-        let smt = SimpleSmt::new(63)
+        let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH)
             .unwrap()
             .with_leaves(smt_leaves.clone())
             .unwrap();
         let store = MerkleStore::new()
-            .with_sparse_merkle_tree(smt_leaves)
+            .with_sparse_merkle_tree(SimpleSmt::MAX_DEPTH, smt_leaves)
             .unwrap();
         let depth = smt.depth();
         let root = smt.root();
@@ -366,7 +366,7 @@ fn new(c: &mut Criterion) {
                         .map(|(c, v)| (c.try_into().unwrap(), v.into()))
                         .collect::<Vec<(u64, Word)>>()
                 },
-                |l| black_box(SimpleSmt::new(63).unwrap().with_leaves(l)),
+                |l| black_box(SimpleSmt::new(SimpleSmt::MAX_DEPTH).unwrap().with_leaves(l)),
                 BatchSize::SmallInput,
             )
         });
@@ -382,7 +382,11 @@ fn new(c: &mut Criterion) {
                             .map(|(c, v)| (c.try_into().unwrap(), v.into()))
                             .collect::<Vec<(u64, Word)>>()
                     },
-                    |l| black_box(MerkleStore::new().with_sparse_merkle_tree(l)),
+                    |l| {
+                        black_box(
+                            MerkleStore::new().with_sparse_merkle_tree(SimpleSmt::MAX_DEPTH, l),
+                        )
+                    },
                     BatchSize::SmallInput,
                 )
             },
@@ -450,12 +454,12 @@ fn update_leaf_simplesmt(c: &mut Criterion) {
             .enumerate()
             .map(|(c, v)| (c.try_into().unwrap(), v.into()))
             .collect::<Vec<(u64, Word)>>();
-        let mut smt = SimpleSmt::new(63)
+        let mut smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH)
             .unwrap()
             .with_leaves(smt_leaves.clone())
             .unwrap();
         let mut store = MerkleStore::new()
-            .with_sparse_merkle_tree(smt_leaves)
+            .with_sparse_merkle_tree(SimpleSmt::MAX_DEPTH, smt_leaves)
             .unwrap();
         let depth = smt.depth();
         let root = smt.root();

src/merkle/simple_smt/mod.rs

@@ -8,7 +8,7 @@ mod tests;
 // SPARSE MERKLE TREE
 // ================================================================================================
 
-/// A sparse Merkle tree with 63-bit keys and 4-element leaf values, without compaction.
+/// A sparse Merkle tree with 64-bit keys and 4-element leaf values, without compaction.
 /// Manipulation and retrieval of leaves and internal nodes is provided by its internal `Store`.
 /// The root of the tree is recomputed on each new leaf update.
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -26,7 +26,7 @@ impl SimpleSmt {
     pub const MIN_DEPTH: u8 = 1;
 
     /// Maximum supported depth.
-    pub const MAX_DEPTH: u8 = 63;
+    pub const MAX_DEPTH: u8 = 64;
 
     // CONSTRUCTORS
     // --------------------------------------------------------------------------------------------
@@ -57,7 +57,7 @@ impl SimpleSmt {
     {
         // check if the leaves count will fit the depth setup
         let mut entries = entries.into_iter();
-        let max = 1 << self.depth;
+        let max = 1 << self.depth.min(63);
         if entries.len() > max {
             return Err(MerkleError::InvalidEntriesCount(max, entries.len()));
         }

src/merkle/store/mod.rs

@@ -115,13 +115,19 @@ impl MerkleStore {
         Ok(self)
     }
 
-    /// Appends the provided sparse merkle tree represented by its `entries` to the set.
-    pub fn with_sparse_merkle_tree<R, I>(mut self, entries: R) -> Result<Self, MerkleError>
+    /// Appends the provided Sparse Merkle tree represented by its `entries` to the set.
+    ///
+    /// For more information, check [MerkleStore::add_sparse_merkle_tree].
+    pub fn with_sparse_merkle_tree<R, I>(
+        mut self,
+        depth: u8,
+        entries: R,
+    ) -> Result<Self, MerkleError>
     where
         R: IntoIterator<IntoIter = I>,
         I: Iterator<Item = (u64, Word)> + ExactSizeIterator,
     {
-        self.add_sparse_merkle_tree(entries)?;
+        self.add_sparse_merkle_tree(depth, entries)?;
         Ok(self)
     }
 
@@ -272,20 +278,24 @@ impl MerkleStore {
         Ok(tree.nodes[1])
     }
 
-    /// Adds all the nodes of a Sparse Merkle tree represented by `entries`.
+    /// Adds a Sparse Merkle tree defined by the specified `entries` to the store, and returns the
+    /// root of the added tree.
     ///
-    /// This will instantiate a Sparse Merkle tree using `entries` and include all the nodes into
-    /// the store.
+    /// The entries are expected to contain tuples of `(index, node)` describing nodes in the tree
+    /// at `depth`.
     ///
     /// # Errors
-    ///
-    /// This will return `InvalidEntriesCount` if the length of `entries` is not `63`.
-    pub fn add_sparse_merkle_tree<R, I>(&mut self, entries: R) -> Result<Word, MerkleError>
+    /// Returns an error if the provided `depth` is greater than [SimpleSmt::MAX_DEPTH].
+    pub fn add_sparse_merkle_tree<R, I>(
+        &mut self,
+        depth: u8,
+        entries: R,
+    ) -> Result<Word, MerkleError>
     where
         R: IntoIterator<IntoIter = I>,
         I: Iterator<Item = (u64, Word)> + ExactSizeIterator,
     {
-        let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH)?.with_leaves(entries)?;
+        let smt = SimpleSmt::new(depth)?.with_leaves(entries)?;
         for branch in smt.store.branches.values() {
             let parent = Rpo256::merge(&[branch.left, branch.right]);
             self.nodes.insert(

src/merkle/store/tests.rs

@@ -205,7 +205,7 @@ fn test_add_sparse_merkle_tree_one_level() -> Result<(), MerkleError> {
     let mut store = MerkleStore::default();
     let keys2: [u64; 2] = [0, 1];
     let leaves2: [Word; 2] = [int_to_node(1), int_to_node(2)];
-    store.add_sparse_merkle_tree(keys2.into_iter().zip(leaves2.into_iter()))?;
+    store.add_sparse_merkle_tree(48, keys2.into_iter().zip(leaves2.into_iter()))?;
     let smt = SimpleSmt::new(1)
         .unwrap()
         .with_leaves(keys2.into_iter().zip(leaves2.into_iter()))
@@ -231,7 +231,10 @@ fn test_add_sparse_merkle_tree_one_level() -> Result<(), MerkleError> {
 #[test]
 fn test_sparse_merkle_tree() -> Result<(), MerkleError> {
     let mut store = MerkleStore::default();
-    store.add_sparse_merkle_tree(KEYS4.into_iter().zip(LEAVES4.into_iter()))?;
+    store.add_sparse_merkle_tree(
+        SimpleSmt::MAX_DEPTH,
+        KEYS4.into_iter().zip(LEAVES4.into_iter()),
+    )?;
 
     let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH)
         .unwrap()
@@ -584,9 +587,10 @@ fn test_constructors() -> Result<(), MerkleError> {
         assert_eq!(mtree.get_path(index)?, value_path.path);
     }
 
+    let depth = 32;
     let store = MerkleStore::default()
-        .with_sparse_merkle_tree(KEYS4.into_iter().zip(LEAVES4.into_iter()))?;
-    let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH)
+        .with_sparse_merkle_tree(depth, KEYS4.into_iter().zip(LEAVES4.into_iter()))?;
+    let smt = SimpleSmt::new(depth)
         .unwrap()
         .with_leaves(KEYS4.into_iter().zip(LEAVES4.into_iter()))
         .unwrap();