From c5c7b38d5e0c782adc1e2fc377f197320a98465f Mon Sep 17 00:00:00 2001 From: jjy Date: Wed, 22 Jan 2020 15:38:05 +0800 Subject: [PATCH] Add calculate_root_with_new_leaf on MerkleProof --- src/mmr.rs | 138 ++++++++++++++++++++++++++++++------------ src/tests/test_mmr.rs | 32 +++++++++- 2 files changed, 130 insertions(+), 40 deletions(-) diff --git a/src/mmr.rs b/src/mmr.rs index d1446ee..8889283 100644 --- a/src/mmr.rs +++ b/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> MerkleProof { &self.proof } - pub fn calculate_root(&self, mut pos: u64, elem: T) -> Result { - 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 { + 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 { + 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 { @@ -219,3 +217,65 @@ impl> MerkleProof { .map(|calculated_root| calculated_root == root) } } + +fn calculate_peak_root< + 'a, + T: 'a + PartialEq + Debug, + M: Merge, + I: Iterator, +>( + 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, I: Iterator>( + pos: u64, + elem: T, + mmr_size: u64, + mut proof_iter: I, +) -> Result { + 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) +} diff --git a/src/tests/test_mmr.rs b/src/tests/test_mmr.rs index d9242b2..4f33b04 100644 --- a/src/tests/test_mmr.rs +++ b/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 = (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); + } }