Merge pull request #9 from vacp2p/set_range

feat: set_range for indexed batch_insertions

src/tree.rs

@@ -193,8 +193,9 @@ where
     }
 
     /// Batch insertion, updates the tree in parallel.
-    pub fn batch_insert(&mut self, leaves: &[H::Fr]) -> PmtreeResult<()> {
-        let end = self.next_index + leaves.len();
+    pub fn batch_insert(&mut self, start: Option<usize>, leaves: &[H::Fr]) -> PmtreeResult<()> {
+        let start = start.unwrap_or(self.next_index);
+        let end = start + leaves.len();
 
         if end > self.capacity() {
             return Err(PmtreeErrorKind::TreeError(TreeErrorKind::MerkleTreeIsFull));
@@ -205,14 +206,7 @@ where
         let root_key = Key(0, 0);
 
         subtree.insert(root_key, self.root);
-        self.fill_nodes(
-            root_key,
-            self.next_index,
-            end,
-            &mut subtree,
-            leaves,
-            self.next_index,
-        )?;
+        self.fill_nodes(root_key, start, end, &mut subtree, leaves, start)?;
 
         let subtree = Arc::new(RwLock::new(subtree));
 
@@ -231,17 +225,31 @@ where
                 .collect(),
         )?;
 
-        // Update root value and next_index in memory
-        self.root = root_val;
-        self.next_index = end;
-
         // Update next_index value in db
-        self.db
-            .put(NEXT_INDEX_KEY, self.next_index.to_be_bytes().to_vec())?;
+
+        if start + leaves.len() > self.next_index {
+            self.next_index = start + leaves.len();
+            self.db
+                .put(NEXT_INDEX_KEY, self.next_index.to_be_bytes().to_vec())?;
+        }
+
+        // Update root value in memory
+        self.root = root_val;
 
         Ok(())
     }
 
+    pub fn set_range<I: IntoIterator<Item = H::Fr>>(
+        &mut self,
+        start: usize,
+        leaves: I,
+    ) -> PmtreeResult<()> {
+        self.batch_insert(
+            Some(start),
+            leaves.into_iter().collect::<Vec<_>>().as_slice(),
+        )
+    }
+
     // Fills hashmap subtree
     fn fill_nodes(
         &self,

tests/memory_keccak.rs

@@ -122,7 +122,7 @@ fn batch_insertions() -> PmtreeResult<()> {
         hex!("0000000000000000000000000000000000000000000000000000000000000004"),
     ];
 
-    mt.batch_insert(&leaves)?;
+    mt.batch_insert(None, &leaves)?;
 
     assert_eq!(
         mt.root(),
@@ -131,3 +131,22 @@ fn batch_insertions() -> PmtreeResult<()> {
 
     Ok(())
 }
+
+#[test]
+fn set_range() -> PmtreeResult<()> {
+    let mut mt = MerkleTree::<MemoryDB, MyKeccak>::new(2, MemoryDBConfig)?;
+
+    let leaves = [
+        hex!("0000000000000000000000000000000000000000000000000000000000000001"),
+        hex!("0000000000000000000000000000000000000000000000000000000000000002"),
+    ];
+
+    mt.set_range(2, leaves)?;
+
+    assert_eq!(
+        mt.root(),
+        hex!("1e9f6c8d3fd5b7ae3a29792adb094c6d4cc6149d0c81c8c8e57cf06c161a92b8")
+    );
+
+    Ok(())
+}

tests/sled_keccak.rs

@@ -159,7 +159,7 @@ fn batch_insertions() -> PmtreeResult<()> {
         hex!("0000000000000000000000000000000000000000000000000000000000000004"),
     ];
 
-    mt.batch_insert(&leaves)?;
+    mt.batch_insert(None, &leaves)?;
 
     assert_eq!(
         mt.root(),
@@ -170,3 +170,29 @@ fn batch_insertions() -> PmtreeResult<()> {
 
     Ok(())
 }
+
+#[test]
+fn set_range() -> PmtreeResult<()> {
+    let mut mt = MerkleTree::<MySled, MyKeccak>::new(
+        2,
+        SledConfig {
+            path: String::from("abacabasab"),
+        },
+    )?;
+
+    let leaves = [
+        hex!("0000000000000000000000000000000000000000000000000000000000000001"),
+        hex!("0000000000000000000000000000000000000000000000000000000000000002"),
+    ];
+
+    mt.set_range(2, leaves)?;
+
+    assert_eq!(
+        mt.root(),
+        hex!("1e9f6c8d3fd5b7ae3a29792adb094c6d4cc6149d0c81c8c8e57cf06c161a92b8")
+    );
+
+    fs::remove_dir_all("abacabasab").expect("Error removing db");
+
+    Ok(())
+}