M benches/store.rs → benches/store.rs

@@ -66,7 +66,8 @@ let leaves = &random_data[..size];
 
         let mtree_leaves: Vec<Word> = leaves.iter().map(|v| v.into()).collect();
         let mtree = MerkleTree::new(mtree_leaves.clone()).unwrap();
-        let store = MerkleStore::new().with_merkle_tree(mtree_leaves).unwrap();
+        let mut store = MerkleStore::new();
+        store.extend(mtree.inner_nodes());
         let depth = mtree.depth();
         let root = mtree.root();
         let size_u64 = size as u64;
@@ -108,9 +109,8 @@ let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH)
             .unwrap()
             .with_leaves(smt_leaves.clone())
             .unwrap();
-        let store = MerkleStore::new()
-            .with_sparse_merkle_tree(SimpleSmt::MAX_DEPTH, smt_leaves)
-            .unwrap();
+        let mut store = MerkleStore::new();
+        store.extend(smt.inner_nodes());
         let depth = smt.depth();
         let root = smt.root();
         let size_u64 = size as u64;
@@ -178,7 +178,8 @@ let leaves = &random_data[..size];
 
         let mtree_leaves: Vec<Word> = leaves.iter().map(|v| v.into()).collect();
         let mtree = MerkleTree::new(mtree_leaves.clone()).unwrap();
-        let store = MerkleStore::new().with_merkle_tree(mtree_leaves).unwrap();
+        let mut store = MerkleStore::new();
+        store.extend(mtree.inner_nodes());
         let root = mtree.root();
         let half_depth = mtree.depth() / 2;
         let half_size = 2_u64.pow(half_depth as u32);
@@ -221,9 +222,8 @@ let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH)
             .unwrap()
             .with_leaves(smt_leaves.clone())
             .unwrap();
-        let store = MerkleStore::new()
-            .with_sparse_merkle_tree(SimpleSmt::MAX_DEPTH, smt_leaves)
-            .unwrap();
+        let mut store = MerkleStore::new();
+        store.extend(smt.inner_nodes());
         let root = smt.root();
         let half_depth = smt.depth() / 2;
         let half_size = 2_u64.pow(half_depth as u32);
@@ -258,7 +258,8 @@ let leaves = &random_data[..size];
 
         let mtree_leaves: Vec<Word> = leaves.iter().map(|v| v.into()).collect();
         let mtree = MerkleTree::new(mtree_leaves.clone()).unwrap();
-        let store = MerkleStore::new().with_merkle_tree(mtree_leaves).unwrap();
+        let mut store = MerkleStore::new();
+        store.extend(mtree.inner_nodes());
         let depth = mtree.depth();
         let root = mtree.root();
         let size_u64 = size as u64;
@@ -300,9 +301,8 @@ let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH)
             .unwrap()
             .with_leaves(smt_leaves.clone())
             .unwrap();
-        let store = MerkleStore::new()
-            .with_sparse_merkle_tree(SimpleSmt::MAX_DEPTH, smt_leaves)
-            .unwrap();
+        let mut store = MerkleStore::new();
+        store.extend(smt.inner_nodes());
         let depth = smt.depth();
         let root = smt.root();
         let size_u64 = size as u64;
@@ -347,10 +347,13 @@ });
 
         // This could be done with `bench_with_input`, however to remove variables while comparing
         // with MerkleTree it is using `iter_batched`
-        group.bench_function(BenchmarkId::new("MerkleStore::with_merkle_tree", size), |b| {
+        group.bench_function(BenchmarkId::new("MerkleStore::extend::MerkleTree", size), |b| {
             b.iter_batched(
                 || leaves.iter().map(|v| v.into()).collect::<Vec<Word>>(),
-                |l| black_box(MerkleStore::new().with_merkle_tree(l)),
+                |l| {
+                    let mtree = MerkleTree::new(l).unwrap();
+                    black_box(MerkleStore::new().extend(mtree.inner_nodes()));
+                },
                 BatchSize::SmallInput,
             )
         });
@@ -369,7 +372,7 @@ BatchSize::SmallInput,
             )
         });
 
-        group.bench_function(BenchmarkId::new("MerkleStore::with_sparse_merkle_tree", size), |b| {
+        group.bench_function(BenchmarkId::new("MerkleStore::extend::SimpleSmt", size), |b| {
             b.iter_batched(
                 || {
                     leaves
@@ -378,7 +381,10 @@ .enumerate()
                         .map(|(c, v)| (c.try_into().unwrap(), v.into()))
                         .collect::<Vec<(u64, Word)>>()
                 },
-                |l| black_box(MerkleStore::new().with_sparse_merkle_tree(SimpleSmt::MAX_DEPTH, l)),
+                |l| {
+                    let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH).unwrap().with_leaves(l).unwrap();
+                    black_box(MerkleStore::new().extend(smt.inner_nodes()));
+                },
                 BatchSize::SmallInput,
             )
         });
@@ -397,7 +403,8 @@ let leaves = &random_data[..size];
 
         let mtree_leaves: Vec<Word> = leaves.iter().map(|v| v.into()).collect();
         let mut mtree = MerkleTree::new(mtree_leaves.clone()).unwrap();
-        let mut store = MerkleStore::new().with_merkle_tree(mtree_leaves).unwrap();
+        let mut store = MerkleStore::new();
+        store.extend(mtree.inner_nodes());
         let depth = mtree.depth();
         let root = mtree.root();
         let size_u64 = size as u64;
@@ -446,9 +453,8 @@ let mut smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH)
             .unwrap()
             .with_leaves(smt_leaves.clone())
             .unwrap();
-        let mut store = MerkleStore::new()
-            .with_sparse_merkle_tree(SimpleSmt::MAX_DEPTH, smt_leaves)
-            .unwrap();
+        let mut store = MerkleStore::new();
+        store.extend(smt.inner_nodes());
         let depth = smt.depth();
         let root = smt.root();
         let size_u64 = size as u64;

M src/merkle/store/mod.rs → src/merkle/store/mod.rs

@@ -1,7 +1,7 @@
-use super::mmr::{Mmr, MmrPeaks};
+use super::mmr::Mmr;
 use super::{
-    BTreeMap, EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, MerklePathSet, MerkleTree,
-    NodeIndex, RootPath, Rpo256, RpoDigest, SimpleSmt, ValuePath, Vec, Word,
+    BTreeMap, EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, MerklePathSet, NodeIndex,
+    RootPath, Rpo256, RpoDigest, ValuePath, Vec, Word,
 };
 use crate::utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable};
 
@@ -47,9 +47,13 @@ ///
 /// // the store is initialized with the SMT empty nodes
 /// assert_eq!(store.num_internal_nodes(), 255);
 ///
+/// let tree1 = MerkleTree::new(vec![A, B, C, D, E, F, G, H0]).unwrap();
+/// let tree2 = MerkleTree::new(vec![A, B, C, D, E, F, G, H1]).unwrap();
+///
 /// // populates the store with two merkle trees, common nodes are shared
-/// store.add_merkle_tree([A, B, C, D, E, F, G, H0]);
-/// store.add_merkle_tree([A, B, C, D, E, F, G, H1]);
+/// store
+///     .extend(tree1.inner_nodes())
+///     .extend(tree2.inner_nodes());
 ///
 /// // every leaf except the last are the same
 /// for i in 0..7 {
@@ -111,31 +115,6 @@
         MerkleStore { nodes }
     }
 
-    /// Appends the provided merkle tree represented by its `leaves` to the set.
-    pub fn with_merkle_tree<I>(mut self, leaves: I) -> Result<Self, MerkleError>
-    where
-        I: IntoIterator<Item = Word>,
-    {
-        self.add_merkle_tree(leaves)?;
-        Ok(self)
-    }
-
-    /// 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(depth, entries)?;
-        Ok(self)
-    }
-
     /// Appends the provided merkle path set.
     pub fn with_merkle_path(
         mut self,
@@ -153,15 +132,6 @@ where
         I: IntoIterator<Item = (u64, Word, MerklePath)>,
     {
         self.add_merkle_paths(paths)?;
-        Ok(self)
-    }
-
-    /// Appends the provided [Mmr] represented by its `leaves` to the set.
-    pub fn with_mmr<I>(mut self, leaves: I) -> Result<Self, MerkleError>
-    where
-        I: IntoIterator<Item = Word>,
-    {
-        self.add_mmr(leaves)?;
         Ok(self)
     }
 
@@ -324,54 +294,6 @@
         self
     }
 
-    /// Adds all the nodes of a Merkle tree represented by `leaves`.
-    ///
-    /// This will instantiate a Merkle tree using `leaves` and include all the nodes into the
-    /// store.
-    ///
-    /// # Errors
-    ///
-    /// This method may return the following errors:
-    /// - `DepthTooSmall` if leaves is empty or contains only 1 element
-    /// - `NumLeavesNotPowerOfTwo` if the number of leaves is not a power-of-two
-    pub fn add_merkle_tree<I>(&mut self, leaves: I) -> Result<Word, MerkleError>
-    where
-        I: IntoIterator<Item = Word>,
-    {
-        let leaves: Vec<_> = leaves.into_iter().collect();
-        if leaves.len() < 2 {
-            return Err(MerkleError::DepthTooSmall(leaves.len() as u8));
-        }
-
-        let tree = MerkleTree::new(leaves)?;
-        self.extend(tree.inner_nodes());
-
-        Ok(tree.root())
-    }
-
-    /// Adds a Sparse Merkle tree defined by the specified `entries` to the store, and returns the
-    /// root of the added tree.
-    ///
-    /// The entries are expected to contain tuples of `(index, node)` describing nodes in the tree
-    /// at `depth`.
-    ///
-    /// # Errors
-    /// 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(depth)?.with_leaves(entries)?;
-        self.extend(smt.inner_nodes());
-
-        Ok(smt.root())
-    }
-
     /// Adds all the nodes of a Merkle path represented by `path`, opening to `node`. Returns the
     /// new root.
     ///
@@ -431,17 +353,6 @@ for (index, path) in path_set.to_paths() {
             self.add_merkle_path(index, path.value, path.path)?;
         }
         Ok(root)
-    }
-
-    /// Appends the provided [Mmr] into the store.
-    pub fn add_mmr<I>(&mut self, leaves: I) -> Result<MmrPeaks, MerkleError>
-    where
-        I: IntoIterator<Item = Word>,
-    {
-        let mmr = Mmr::from(leaves);
-        self.extend(mmr.inner_nodes());
-
-        Ok(mmr.accumulator())
     }
 
     /// Sets a node to `value`.

M src/merkle/store/tests.rs → src/merkle/store/tests.rs

@@ -1,7 +1,7 @@
 use super::*;
 use crate::{
     hash::rpo::Rpo256,
-    merkle::{int_to_node, MerklePathSet},
+    merkle::{int_to_node, MerklePathSet, MerkleTree, SimpleSmt},
     Felt, Word, WORD_SIZE, ZERO,
 };
 
@@ -15,7 +15,8 @@
 #[test]
 fn test_root_not_in_store() -> Result<(), MerkleError> {
     let mtree = MerkleTree::new(LEAVES4.to_vec())?;
-    let store = MerkleStore::default().with_merkle_tree(LEAVES4)?;
+    let mut store = MerkleStore::new();
+    store.extend(mtree.inner_nodes());
     assert_eq!(
         store.get_node(LEAVES4[0], NodeIndex::make(mtree.depth(), 0)),
         Err(MerkleError::RootNotInStore(LEAVES4[0])),
@@ -32,10 +33,9 @@ }
 
 #[test]
 fn test_merkle_tree() -> Result<(), MerkleError> {
+    let mtree = MerkleTree::new(LEAVES4.to_vec())?;
     let mut store = MerkleStore::default();
-
-    let mtree = MerkleTree::new(LEAVES4.to_vec())?;
-    store.add_merkle_tree(LEAVES4.to_vec())?;
+    store.extend(mtree.inner_nodes());
 
     // STORE LEAVES ARE CORRECT ==============================================================
     // checks the leaves in the store corresponds to the expected values
@@ -176,24 +176,22 @@ }
 
 #[test]
 fn test_get_invalid_node() {
-    let mut store = MerkleStore::default();
     let mtree = MerkleTree::new(LEAVES4.to_vec()).expect("creating a merkle tree must work");
-    store
-        .add_merkle_tree(LEAVES4.to_vec())
-        .expect("adding a merkle tree to the store must work");
+    let mut store = MerkleStore::default();
+    store.extend(mtree.inner_nodes());
     let _ = store.get_node(mtree.root(), NodeIndex::make(mtree.depth(), 3));
 }
 
 #[test]
 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(48, keys2.into_iter().zip(leaves2.into_iter()))?;
     let smt = SimpleSmt::new(1)
         .unwrap()
         .with_leaves(keys2.into_iter().zip(leaves2.into_iter()))
         .unwrap();
+    let mut store = MerkleStore::default();
+    store.extend(smt.inner_nodes());
 
     let idx = NodeIndex::make(1, 0);
     assert_eq!(smt.get_node(idx).unwrap(), leaves2[0]);
@@ -208,15 +206,14 @@ }
 
 #[test]
 fn test_sparse_merkle_tree() -> Result<(), MerkleError> {
-    let mut store = MerkleStore::default();
-    store
-        .add_sparse_merkle_tree(SimpleSmt::MAX_DEPTH, KEYS4.into_iter().zip(LEAVES4.into_iter()))?;
-
     let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH)
         .unwrap()
         .with_leaves(KEYS4.into_iter().zip(LEAVES4.into_iter()))
         .unwrap();
 
+    let mut store = MerkleStore::default();
+    store.extend(smt.inner_nodes());
+
     // STORE LEAVES ARE CORRECT ==============================================================
     // checks the leaves in the store corresponds to the expected values
     assert_eq!(
@@ -472,7 +469,9 @@
     // For this example, the depth of the Merkle tree is 1, as we have only two leaves. Here we
     // attempt to fetch a node on the maximum depth, and it should fail because the root shouldn't
     // exist for the set.
-    let store = MerkleStore::default().with_merkle_tree([a, b]).unwrap();
+    let mtree = MerkleTree::new(vec![a, b]).unwrap();
+    let mut store = MerkleStore::new();
+    store.extend(mtree.inner_nodes());
     let index = NodeIndex::root();
     let err = store.get_node(root, index).err().unwrap();
     assert_eq!(err, MerkleError::RootNotInStore(root));
@@ -499,7 +498,9 @@ let n = Rpo256::merge(&[k.into(), l.into()]);
 
     let root = Rpo256::merge(&[m.into(), n.into()]);
 
-    let store = MerkleStore::default().with_merkle_tree([a, b, c, d, e, f, g, h]).unwrap();
+    let mtree = MerkleTree::new(vec![a, b, c, d, e, f, g, h]).unwrap();
+    let mut store = MerkleStore::new();
+    store.extend(mtree.inner_nodes());
     let path = store.get_path(root.into(), NodeIndex::make(3, 1)).unwrap().path;
 
     let expected = MerklePath::new([a.into(), j.into(), n.into()].to_vec());
@@ -509,7 +510,8 @@
 #[test]
 fn test_set_node() -> Result<(), MerkleError> {
     let mtree = MerkleTree::new(LEAVES4.to_vec())?;
-    let mut store = MerkleStore::default().with_merkle_tree(LEAVES4)?;
+    let mut store = MerkleStore::new();
+    store.extend(mtree.inner_nodes());
     let value = int_to_node(42);
     let index = NodeIndex::make(mtree.depth(), 0);
     let new_root = store.set_node(mtree.root(), index, value)?.root;
@@ -520,8 +522,9 @@ }
 
 #[test]
 fn test_constructors() -> Result<(), MerkleError> {
-    let store = MerkleStore::new().with_merkle_tree(LEAVES4)?;
     let mtree = MerkleTree::new(LEAVES4.to_vec())?;
+    let mut store = MerkleStore::new();
+    store.extend(mtree.inner_nodes());
 
     let depth = mtree.depth();
     let leaves = 2u64.pow(depth.into());
@@ -532,12 +535,12 @@ assert_eq!(mtree.get_path(index)?, value_path.path);
     }
 
     let depth = 32;
-    let store = MerkleStore::default()
-        .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();
+    let mut store = MerkleStore::new();
+    store.extend(smt.inner_nodes());
     let depth = smt.depth();
 
     for key in KEYS4 {
@@ -726,7 +729,9 @@
 #[cfg(std)]
 #[test]
 fn test_serialization() -> Result<(), Box<dyn Error>> {
-    let original = MerkleStore::new().with_merkle_tree(LEAVES4)?;
+    let mtree = MerkleTree::new(LEAVES4.to_vec())?;
+    let mut store = MerkleStore::new();
+    store.extend(mtree.inner_nodes());
     let decoded = MerkleStore::read_from_bytes(&original.to_bytes())?;
     assert_eq!(original, decoded);
     Ok(())