Add calculate_root_with_new_leaf on MerkleProof

src/mmr.rs

@@ -7,6 +7,7 @@
 use crate::borrow::Cow;
 use crate::helper::{get_peaks, parent_offset, pos_height_in_tree, sibling_offset};
 use crate::mmr_store::{MMRBatch, MMRStore};
+use crate::vec;
 use crate::vec::Vec;
 use crate::{Error, Merge, Result};
 use core::fmt::Debug;
@@ -171,47 +172,44 @@ impl<T: PartialEq + Debug, M: Merge<Item = T>> MerkleProof<T, M> {
         &self.proof
     }
 
-    pub fn calculate_root(&self, mut pos: u64, elem: T) -> Result<T> {
-        let peaks = get_peaks(self.mmr_size);
-        let mut sum_elem = elem;
-        let mut height = 0;
-        let mut proof_iter = self.proof.iter();
-        // calculate peak's merkle root
-        // start bagging peaks if pos reach a peak pos
-        while peaks.binary_search(&pos).is_err() {
-            let proof = match proof_iter.next() {
-                Some(proof) => proof,
-                None => break,
-            };
-            // verify merkle path
-            let pos_height = pos_height_in_tree(pos);
-            let next_height = pos_height_in_tree(pos + 1);
-            sum_elem = if next_height > pos_height {
-                // to next pos
-                pos += 1;
-                M::merge(proof, &sum_elem)
-            } else {
-                pos += parent_offset(height);
-                M::merge(&sum_elem, proof)
-            };
-            height += 1
-        }
+    pub fn calculate_root(&self, pos: u64, elem: T) -> Result<T> {
+        calculate_root::<_, M, _>(pos, elem, self.mmr_size, self.proof.iter())
+    }
 
-        // bagging peaks
-        // bagging with left peaks if pos is last peak(last pos)
-        let mut bagging_left = pos == self.mmr_size - 1;
-        for proof in &mut proof_iter {
-            sum_elem = if bagging_left {
-                M::merge(&sum_elem, &proof)
-            } else {
-                // we are not in the last peak, so bag with right peaks first
-                // notice the right peaks is already bagging into one hash in proof,
-                // so after this merge, the remain proofs are always left peaks.
-                bagging_left = true;
-                M::merge(&proof, &sum_elem)
-            };
+    /// from merkle proof of leaf n to calculate merkle root of n + 1 leaves.
+    /// by observe the MMR construction graph we know it is possible.
+    /// https://github.com/jjyr/merkle-mountain-range#construct
+    /// this is kinda tricky, but it works, and useful
+    pub fn calculate_root_with_new_leaf(
+        &self,
+        pos: u64,
+        elem: T,
+        new_pos: u64,
+        new_elem: T,
+        new_mmr_size: u64,
+    ) -> Result<T> {
+        if self.mmr_size == 0 {
+            return Ok(elem);
+        }
+        let pos_height = pos_height_in_tree(new_pos);
+        let next_height = pos_height_in_tree(new_pos + 1);
+
+        if next_height > pos_height {
+            // new elem on right branch
+            let new_proof = vec![elem];
+            let new_proof_iter = new_proof.iter().chain(self.proof.iter());
+            calculate_root::<_, M, _>(new_pos, new_elem, new_mmr_size, new_proof_iter)
+        } else {
+            // new elem on left branch
+            debug_assert_eq!(self.mmr_size + 1, new_mmr_size);
+            let peaks = get_peaks(self.mmr_size);
+            let mut proof_iter = self.proof.iter();
+            let (root_elem, _) =
+                calculate_peak_root::<_, M, _>(pos, &peaks, elem, &mut proof_iter)?;
+            let new_proof = vec![root_elem];
+            let new_proof_iter = new_proof.iter().chain(proof_iter);
+            calculate_root::<_, M, _>(new_pos, new_elem, new_mmr_size, new_proof_iter)
         }
-        Ok(sum_elem)
     }
 
     pub fn verify(&self, root: T, pos: u64, elem: T) -> Result<bool> {
@@ -219,3 +217,65 @@ impl<T: PartialEq + Debug, M: Merge<Item = T>> MerkleProof<T, M> {
             .map(|calculated_root| calculated_root == root)
     }
 }
+
+fn calculate_peak_root<
+    'a,
+    T: 'a + PartialEq + Debug,
+    M: Merge<Item = T>,
+    I: Iterator<Item = &'a T>,
+>(
+    mut pos: u64,
+    peaks: &[u64],
+    elem: T,
+    proof_iter: &mut I,
+) -> Result<(T, u64)> {
+    let mut root_elem = elem;
+    let mut height = 0;
+    // calculate peak's merkle root
+    // start bagging peaks if pos reach a peak pos
+    while peaks.binary_search(&pos).is_err() {
+        let proof = match proof_iter.next() {
+            Some(proof) => proof,
+            None => break,
+        };
+        // verify merkle path
+        let pos_height = pos_height_in_tree(pos);
+        let next_height = pos_height_in_tree(pos + 1);
+        root_elem = if next_height > pos_height {
+            // to next pos
+            pos += 1;
+            M::merge(proof, &root_elem)
+        } else {
+            pos += parent_offset(height);
+            M::merge(&root_elem, proof)
+        };
+        height += 1
+    }
+    Ok((root_elem, pos))
+}
+
+fn calculate_root<'a, T: 'a + PartialEq + Debug, M: Merge<Item = T>, I: Iterator<Item = &'a T>>(
+    pos: u64,
+    elem: T,
+    mmr_size: u64,
+    mut proof_iter: I,
+) -> Result<T> {
+    let peaks = get_peaks(mmr_size);
+    let (mut root_elem, pos) = calculate_peak_root::<_, M, _>(pos, &peaks, elem, &mut proof_iter)?;
+
+    // bagging peaks
+    // bagging with left peaks if pos is last peak(last pos)
+    let mut bagging_left = pos == mmr_size - 1;
+    for proof in &mut proof_iter {
+        root_elem = if bagging_left {
+            M::merge(&root_elem, &proof)
+        } else {
+            // we are not in the last peak, so bag with right peaks first
+            // notice the right peaks is already bagging into one hash in proof,
+            // so after this merge, the remain proofs are always left peaks.
+            bagging_left = true;
+            M::merge(&proof, &root_elem)
+        };
+    }
+    Ok(root_elem)
+}

src/tests/test_mmr.rs

@@ -1,5 +1,5 @@
 use super::{MergeNumberHash, NumberHash};
-use crate::{util::MemStore, Error, MMR};
+use crate::{leaf_index_to_mmr_size, util::MemStore, Error, MMR};
 use faster_hex::hex_string;
 use proptest::prelude::*;
 
@@ -24,6 +24,31 @@ fn test_mmr(count: u32, proof_elem: u32) {
     assert!(result);
 }
 
+fn test_gen_new_root_from_proof(count: u32) {
+    let store = MemStore::default();
+    let mut mmr = MMR::<_, MergeNumberHash, _>::new(0, &store);
+    let positions: Vec<u64> = (0u32..count)
+        .map(|i| mmr.push(NumberHash::from(i)).unwrap())
+        .collect();
+    let elem = count - 1;
+    let pos = positions[elem as usize];
+    let proof = mmr.gen_proof(pos).expect("gen proof");
+    let new_elem = count;
+    let new_pos = mmr.push(NumberHash::from(new_elem)).unwrap();
+    let root = mmr.get_root().expect("get root");
+    mmr.commit().expect("commit changes");
+    let calculated_root = proof
+        .calculate_root_with_new_leaf(
+            pos,
+            NumberHash::from(elem),
+            new_pos,
+            NumberHash::from(new_elem),
+            leaf_index_to_mmr_size(new_elem.into()),
+        )
+        .unwrap();
+    assert_eq!(calculated_root, root);
+}
+
 #[test]
 fn test_mmr_root() {
     let store = MemStore::default();
@@ -95,4 +120,9 @@ proptest! {
     fn test_random_mmr((count , elem) in count_elem(500)) {
         test_mmr(count, elem);
     }
+
+    #[test]
+    fn test_random_gen_root_with_new_leaf(count in 1u32..500u32) {
+        test_gen_new_root_from_proof(count);
+    }
 }