Merge pull request #45 from eval-exec/exec/bump-0.7.0

bump version to 0.7.0

Cargo.toml

@@ -1,8 +1,9 @@
 [package]
 name = "ckb-merkle-mountain-range"
-version = "0.5.0"
+version = "0.7.0"
 authors = ["Nervos Core Dev <dev@nervos.org>"]
-edition = "2018"
+edition = "2024"
+rust-version = "1.85.0"
 license = "MIT"
 description = "A generalized merkle mountain range implementation"
 repository = "https://github.com/nervosnetwork/merkle-mountain-range"
@@ -15,13 +16,13 @@ std = []
 cfg-if = "1.0"
 
 [dev-dependencies]
-faster-hex = "0.6.1"
-criterion = "0.3"
+faster-hex = "0.8.0"
+criterion = "0.5.1"
 rand = "0.8.5"
-proptest = "1.0.0"
-lazy_static = "1.3.0"
-bytes = "1.2.0"
-blake2b-rs = "0.1.4"
+proptest = "1.2.0"
+lazy_static = "1.4.0"
+bytes = "1.4.0"
+blake2b-rs = "0.2.0"
 
 [[bench]]
 name = "mmr_benchmark"

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 nervosnetwork
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -24,7 +24,8 @@ A generalized merkle mountain range implementation.
 ```
 
 In MMR, we use the insertion order to reference leaves and nodes.
-we inserting a new leaf to MMR by the following:
+
+We insert a new leaf to MMR by the following:
 
 1. insert leaf or node to next position.
 2. if the current position has a left sibling, we merge the left and right nodes to produce a new parent node, then go back to step 1 to insert the node.
@@ -34,9 +35,9 @@ For example, we insert a leaf to the example MMR:
 1. insert leaf to next position: `19`.
 2. now check the left sibling `18` and calculate parent node: `merge(mmr[18], mmr[19])`.
 3. insert parent node to position `20`.
-4. the node `20` also has a left sibling `17`, calculate parent node: `merge(mmr[17], mmr[20])`.
+4. since the node `20` also has a left sibling `17`, calculate parent node: `merge(mmr[17], mmr[20])`.
 5. insert new node to next position `21`.
-6. the node `21` have no left sibling, complete the insertion.
+6. since the node `21` have no left sibling, complete the insertion.
 
 Example MMR after insertion of a new leaf:
 
@@ -54,17 +55,17 @@ Example MMR after insertion of a new leaf:
 
 An MMR is constructed by one or more sub merkle trees (or mountains). Each sub merkle tree's root is a peak in MMR, we calculate the MMR root by bagging these peaks from right to left.
 
-For example, we have a MMR with 3 peaks: `14, 17, 18`, we bagging thses peaks from right to left to get the root: `merge(merge(mmr[18], mmr[17]), mmr[14])`.
+For example, in the 11 leaf MMR we have 3 peaks: `14, 17, 18`, we bag these peaks from right to left to get the root: `merge(mmr[14], merge(mmr[17], mmr[18]))`.
 
 ## Merkle proof
 
-The merkle proof is an array of hashes constructed by the follows parts:
+The merkle proof is an array of hashes constructed with the following parts:
 
 1. A merkle proof from the leaf's sibling to the peak that contains the leaf.
-2. A hash that bagging all right-hand side peaks, skip this part if no right-hand peaks.
-3. Hashes of all left-hand peaks, the from right to left, skip this part if no left-hand peaks.
+2. A hash that bags all right-hand side peaks, skip this part if no right-hand peaks.
+3. Hashes of all left-hand peaks from right to left, skip this part if no left-hand peaks.
 
-We can reconstruct the merkle root from the proofs. Pre calculate the peaks positions from the size of MMR may help us do the bagging.
+We can reconstruct the merkle root from the proofs. Pre-calculating the peak positions from the size of MMR may help us do the bagging.
 
 ## References
 

benches/helper_benchmark.rs

@@ -4,7 +4,7 @@ use criterion::Criterion;
 
 use ckb_merkle_mountain_range::{leaf_index_to_mmr_size, leaf_index_to_pos};
 
-use rand::{thread_rng, Rng};
+use rand::{Rng, thread_rng};
 
 fn bench(c: &mut Criterion) {
     c.bench_function("left_index_to_pos", |b| {

benches/mmr_benchmark.rs

@@ -4,7 +4,7 @@ extern crate criterion;
 use criterion::{BenchmarkId, Criterion};
 
 use bytes::Bytes;
-use ckb_merkle_mountain_range::{util::MemStore, Error, MMRStore, Merge, Result, MMR};
+use ckb_merkle_mountain_range::{Error, MMR, MMRStoreReadOps, Merge, Result, util::MemStore};
 use rand::{seq::SliceRandom, thread_rng};
 use std::convert::TryFrom;
 

src/error.rs

@@ -7,6 +7,8 @@ pub enum Error {
     StoreError(crate::string::String),
     /// proof items is not enough to build a tree
     CorruptedProof,
+    /// tried to verify proof of a non-leaf
+    NodeProofsNotSupported,
     /// The leaves is an empty list, or beyond the mmr range
     GenProofForInvalidLeaves,
 
@@ -22,7 +24,8 @@ impl core::fmt::Display for Error {
             InconsistentStore => write!(f, "Inconsistent store")?,
             StoreError(msg) => write!(f, "Store error {}", msg)?,
             CorruptedProof => write!(f, "Corrupted proof")?,
-            GenProofForInvalidLeaves => write!(f, "Generate proof ofr invalid leaves")?,
+            NodeProofsNotSupported => write!(f, "Tried to verify membership of a non-leaf")?,
+            GenProofForInvalidLeaves => write!(f, "Generate proof for invalid leaves")?,
             MergeError(msg) => write!(f, "Merge error {}", msg)?,
         }
         Ok(())

src/helper.rs

@@ -1,3 +1,4 @@
+use crate::vec;
 use crate::vec::Vec;
 
 pub fn leaf_index_to_pos(index: u64) -> u64 {
@@ -15,75 +16,97 @@ pub fn leaf_index_to_mmr_size(index: u64) -> u64 {
     2 * leaves_count - peak_count
 }
 
-pub fn pos_height_in_tree(mut pos: u64) -> u32 {
-    pos += 1;
-    fn all_ones(num: u64) -> bool {
-        num != 0 && num.count_zeros() == num.leading_zeros()
-    }
-    fn jump_left(pos: u64) -> u64 {
-        let bit_length = 64 - pos.leading_zeros();
-        let most_significant_bits = 1 << (bit_length - 1);
-        pos - (most_significant_bits - 1)
+pub fn pos_height_in_tree(mut pos: u64) -> u8 {
+    if pos == 0 {
+        return 0;
     }
 
-    while !all_ones(pos) {
-        pos = jump_left(pos)
+    let mut peak_size = u64::MAX >> pos.leading_zeros();
+    while peak_size > 0 {
+        if pos >= peak_size {
+            pos -= peak_size;
+        }
+        peak_size >>= 1;
     }
-
-    64 - pos.leading_zeros() - 1
+    pos as u8
 }
 
-pub fn parent_offset(height: u32) -> u64 {
+pub fn parent_offset(height: u8) -> u64 {
     2 << height
 }
 
-pub fn sibling_offset(height: u32) -> u64 {
+pub fn sibling_offset(height: u8) -> u64 {
     (2 << height) - 1
 }
 
-pub fn get_peaks(mmr_size: u64) -> Vec<u64> {
-    let mut pos_s = Vec::new();
-    let (mut height, mut pos) = left_peak_height_pos(mmr_size);
-    pos_s.push(pos);
-    while height > 0 {
-        let peak = match get_right_peak(height, pos, mmr_size) {
-            Some(peak) => peak,
-            None => break,
-        };
-        height = peak.0;
-        pos = peak.1;
-        pos_s.push(pos);
+/// Returns the height of the peaks in the mmr, presented by a bitmap.
+/// for example, for a mmr with 11 leaves, the mmr_size is 19, it will return 0b1011.
+/// 0b1011 indicates that the left peaks are at height 0, 1 and 3.
+///           14
+///        /       \
+///      6          13
+///    /   \       /   \
+///   2     5     9     12     17
+///  / \   /  \  / \   /  \   /  \
+/// 0   1 3   4 7   8 10  11 15  16 18
+///
+/// please note that when the mmr_size is invalid, it will return the bitmap of the last valid mmr.
+/// in the below example, the mmr_size is 6, but it's not a valid mmr, it will return 0b11.
+///   2     5
+///  / \   /  \
+/// 0   1 3   4
+pub fn get_peak_map(mmr_size: u64) -> u64 {
+    if mmr_size == 0 {
+        return 0;
     }
-    pos_s
-}
 
-fn get_right_peak(mut height: u32, mut pos: u64, mmr_size: u64) -> Option<(u32, u64)> {
-    // move to right sibling pos
-    pos += sibling_offset(height);
-    // loop until we find a pos in mmr
-    while pos > mmr_size - 1 {
-        if height == 0 {
-            return None;
+    let mut pos = mmr_size;
+    let mut peak_size = u64::MAX >> pos.leading_zeros();
+    let mut peak_map = 0;
+    while peak_size > 0 {
+        peak_map <<= 1;
+        if pos >= peak_size {
+            pos -= peak_size;
+            peak_map |= 1;
         }
-        // move to left child
-        pos -= parent_offset(height - 1);
-        height -= 1;
+        peak_size >>= 1;
     }
-    Some((height, pos))
+
+    peak_map
 }
 
-fn get_peak_pos_by_height(height: u32) -> u64 {
-    (1 << (height + 1)) - 2
-}
-
-fn left_peak_height_pos(mmr_size: u64) -> (u32, u64) {
-    let mut height = 1;
-    let mut prev_pos = 0;
-    let mut pos = get_peak_pos_by_height(height);
-    while pos < mmr_size {
-        height += 1;
-        prev_pos = pos;
-        pos = get_peak_pos_by_height(height);
+/// Returns the pos of the peaks in the mmr.
+/// for example, for a mmr with 11 leaves, the mmr_size is 19, it will return [14, 17, 18].
+///           14
+///        /       \
+///      6          13
+///    /   \       /   \
+///   2     5     9     12     17
+///  / \   /  \  / \   /  \   /  \
+/// 0   1 3   4 7   8 10  11 15  16 18
+///
+/// please note that when the mmr_size is invalid, it will return the peaks of the last valid mmr.
+/// in the below example, the mmr_size is 6, but it's not a valid mmr, it will return [2, 3].
+///   2     5
+///  / \   /  \
+/// 0   1 3   4
+pub fn get_peaks(mmr_size: u64) -> Vec<u64> {
+    if mmr_size == 0 {
+        return vec![];
     }
-    (height - 1, prev_pos)
+
+    let leading_zeros = mmr_size.leading_zeros();
+    let mut pos = mmr_size;
+    let mut peak_size = u64::MAX >> leading_zeros;
+    let mut peaks = Vec::with_capacity(64 - leading_zeros as usize);
+    let mut peaks_sum = 0;
+    while peak_size > 0 {
+        if pos >= peak_size {
+            pos -= peak_size;
+            peaks.push(peaks_sum + peak_size - 1);
+            peaks_sum += peak_size;
+        }
+        peak_size >>= 1;
+    }
+    peaks
 }

src/lib.rs

@@ -12,8 +12,8 @@ pub mod util;
 pub use error::{Error, Result};
 pub use helper::{leaf_index_to_mmr_size, leaf_index_to_pos};
 pub use merge::Merge;
-pub use mmr::{MerkleProof, MMR};
-pub use mmr_store::MMRStore;
+pub use mmr::{MMR, MerkleProof};
+pub use mmr_store::{MMRStoreReadOps, MMRStoreWriteOps};
 
 cfg_if::cfg_if! {
     if #[cfg(feature = "std")] {

src/mmr.rs

@@ -6,13 +6,17 @@
 
 use crate::borrow::Cow;
 use crate::collections::VecDeque;
-use crate::helper::{get_peaks, parent_offset, pos_height_in_tree, sibling_offset};
-use crate::mmr_store::{MMRBatch, MMRStore};
+use crate::helper::{
+    get_peak_map, get_peaks, leaf_index_to_mmr_size, leaf_index_to_pos, parent_offset,
+    pos_height_in_tree, sibling_offset,
+};
+use crate::mmr_store::{MMRBatch, MMRStoreReadOps, MMRStoreWriteOps};
 use crate::vec;
 use crate::vec::Vec;
 use crate::{Error, Merge, Result};
 use core::fmt::Debug;
 use core::marker::PhantomData;
+use core::mem;
 
 #[allow(clippy::upper_case_acronyms)]
 pub struct MMR<T, M, S> {
@@ -21,7 +25,7 @@ pub struct MMR<T, M, S> {
     merge: PhantomData<M>,
 }
 
-impl<'a, T: Clone + PartialEq + Debug, M: Merge<Item = T>, S: MMRStore<T>> MMR<T, M, S> {
+impl<T, M, S> MMR<T, M, S> {
     pub fn new(mmr_size: u64, store: S) -> Self {
         MMR {
             mmr_size,
@@ -30,16 +34,6 @@ impl<'a, T: Clone + PartialEq + Debug, M: Merge<Item = T>, S: MMRStore<T>> MMR<T
         }
     }
 
-    // find internal MMR elem, the pos must exists, otherwise a error will return
-    fn find_elem<'b>(&self, pos: u64, hashes: &'b [T]) -> Result<Cow<'b, T>> {
-        let pos_offset = pos.checked_sub(self.mmr_size);
-        if let Some(elem) = pos_offset.and_then(|i| hashes.get(i as usize)) {
-            return Ok(Cow::Borrowed(elem));
-        }
-        let elem = self.batch.get_elem(pos)?.ok_or(Error::InconsistentStore)?;
-        Ok(Cow::Owned(elem))
-    }
-
     pub fn mmr_size(&self) -> u64 {
         self.mmr_size
     }
@@ -48,24 +42,41 @@ impl<'a, T: Clone + PartialEq + Debug, M: Merge<Item = T>, S: MMRStore<T>> MMR<T
         self.mmr_size == 0
     }
 
+    pub fn batch(&self) -> &MMRBatch<T, S> {
+        &self.batch
+    }
+
+    pub fn store(&self) -> &S {
+        self.batch.store()
+    }
+}
+
+impl<T: Clone + PartialEq, M: Merge<Item = T>, S: MMRStoreReadOps<T>> MMR<T, M, S> {
+    // find internal MMR elem, the pos must exists, otherwise a error will return
+    fn find_elem<'b>(&self, pos: u64, hashes: &'b [T]) -> Result<Cow<'b, T>> {
+        let pos_offset = pos.checked_sub(self.mmr_size);
+        if let Some(elem) = pos_offset.and_then(|i| hashes.get(i as usize)) {
+            return Ok(Cow::Borrowed(elem));
+        }
+        let elem = self.batch.get_elem(pos)?.ok_or(Error::InconsistentStore)?;
+        Ok(Cow::Owned(elem))
+    }
+
     // push a element and return position
     pub fn push(&mut self, elem: T) -> Result<u64> {
-        let mut elems: Vec<T> = Vec::new();
-        // position of new elem
+        let mut elems = vec![elem];
         let elem_pos = self.mmr_size;
-        elems.push(elem);
-        let mut height = 0u32;
-        let mut pos = elem_pos;
-        // continue to merge tree node if next pos heigher than current
-        while pos_height_in_tree(pos + 1) > height {
+        let peak_map = get_peak_map(self.mmr_size);
+        let mut pos = self.mmr_size;
+        let mut peak = 1;
+        while (peak_map & peak) != 0 {
+            peak <<= 1;
             pos += 1;
-            let left_pos = pos - parent_offset(height);
-            let right_pos = left_pos + sibling_offset(height);
+            let left_pos = pos - peak;
             let left_elem = self.find_elem(left_pos, &elems)?;
-            let right_elem = self.find_elem(right_pos, &elems)?;
-            let parent_elem = M::merge(&left_elem, &right_elem)?;
+            let right_elem = elems.last().expect("checked");
+            let parent_elem = M::merge(&left_elem, right_elem)?;
             elems.push(parent_elem);
-            height += 1
         }
         // store hashes
         self.batch.append(elem_pos, elems);
@@ -127,12 +138,17 @@ impl<'a, T: Clone + PartialEq + Debug, M: Merge<Item = T>, S: MMRStore<T>> MMR<T
             return Ok(());
         }
 
-        let mut queue: VecDeque<_> = pos_list.into_iter().map(|pos| (pos, 0u32)).collect();
+        let mut queue: VecDeque<_> = pos_list.into_iter().map(|pos| (pos, 0)).collect();
+
         // Generate sub-tree merkle proof for positions
         while let Some((pos, height)) = queue.pop_front() {
             debug_assert!(pos <= peak_pos);
             if pos == peak_pos {
-                break;
+                if queue.is_empty() {
+                    break;
+                } else {
+                    return Err(Error::NodeProofsNotSupported);
+                }
             }
 
             // calculate sibling
@@ -177,8 +193,12 @@ impl<'a, T: Clone + PartialEq + Debug, M: Merge<Item = T>, S: MMRStore<T>> MMR<T
         if self.mmr_size == 1 && pos_list == [0] {
             return Ok(MerkleProof::new(self.mmr_size, Vec::new()));
         }
-        // ensure positions is sorted
+        if pos_list.iter().any(|pos| pos_height_in_tree(*pos) > 0) {
+            return Err(Error::NodeProofsNotSupported);
+        }
+        // ensure positions are sorted and unique
         pos_list.sort_unstable();
+        pos_list.dedup();
         let peaks = get_peaks(self.mmr_size);
         let mut proof: Vec<T> = Vec::new();
         // generate merkle proof for each peaks
@@ -205,8 +225,10 @@ impl<'a, T: Clone + PartialEq + Debug, M: Merge<Item = T>, S: MMRStore<T>> MMR<T
 
         Ok(MerkleProof::new(self.mmr_size, proof))
     }
+}
 
-    pub fn commit(self) -> Result<()> {
+impl<T, M, S: MMRStoreWriteOps<T>> MMR<T, M, S> {
+    pub fn commit(&mut self) -> Result<()> {
         self.batch.commit()
     }
 }
@@ -218,7 +240,7 @@ pub struct MerkleProof<T, M> {
     merge: PhantomData<M>,
 }
 
-impl<T: PartialEq + Debug + Clone, M: Merge<Item = T>> MerkleProof<T, M> {
+impl<T: Clone + PartialEq, M: Merge<Item = T>> MerkleProof<T, M> {
     pub fn new(mmr_size: u64, proof: Vec<T>) -> Self {
         MerkleProof {
             mmr_size,
@@ -273,80 +295,141 @@ impl<T: PartialEq + Debug + Clone, M: Merge<Item = T>> MerkleProof<T, M> {
         self.calculate_root(leaves)
             .map(|calculated_root| calculated_root == root)
     }
+
+    /// Verifies a old root and all incremental leaves.
+    ///
+    /// If this method returns `true`, it means the following assertion are true:
+    /// - The old root could be generated in the history of the current MMR.
+    /// - All incremental leaves are on the current MMR.
+    /// - The MMR, which could generate the old root, appends all incremental leaves, becomes the
+    ///   current MMR.
+    pub fn verify_incremental(&self, root: T, prev_root: T, incremental: Vec<T>) -> Result<bool> {
+        let current_leaves_count = get_peak_map(self.mmr_size);
+        if current_leaves_count <= incremental.len() as u64 {
+            return Err(Error::CorruptedProof);
+        }
+        // Test if previous root is correct.
+        let prev_leaves_count = current_leaves_count - incremental.len() as u64;
+        let prev_peaks_positions = {
+            let prev_index = prev_leaves_count - 1;
+            let prev_mmr_size = leaf_index_to_mmr_size(prev_index);
+            let prev_peaks_positions = get_peaks(prev_mmr_size);
+            if prev_peaks_positions.len() != self.proof.len() {
+                return Err(Error::CorruptedProof);
+            }
+            prev_peaks_positions
+        };
+        let current_peaks_positions = get_peaks(self.mmr_size);
+
+        let mut reverse_index = prev_peaks_positions.len() - 1;
+        for (i, position) in prev_peaks_positions.iter().enumerate() {
+            if *position < current_peaks_positions[i] {
+                reverse_index = i;
+                break;
+            }
+        }
+        let mut prev_peaks: Vec<_> = self.proof_items().to_vec();
+        let mut reverse_peaks = prev_peaks.split_off(reverse_index);
+        reverse_peaks.reverse();
+        prev_peaks.extend(reverse_peaks);
+
+        let calculated_prev_root = bagging_peaks_hashes::<T, M>(prev_peaks)?;
+        if calculated_prev_root != prev_root {
+            return Ok(false);
+        }
+
+        // Test if incremental leaves are correct.
+        let leaves = incremental
+            .into_iter()
+            .enumerate()
+            .map(|(index, leaf)| {
+                let pos = leaf_index_to_pos(prev_leaves_count + index as u64);
+                (pos, leaf)
+            })
+            .collect();
+        self.verify(root, leaves)
+    }
 }
 
-fn calculate_peak_root<
-    'a,
-    T: 'a + PartialEq + Debug + Clone,
-    M: Merge<Item = T>,
-    I: Iterator<Item = &'a T>,
->(
+fn calculate_peak_root<'a, T: 'a, M: Merge<Item = T>, I: Iterator<Item = &'a T>>(
     leaves: Vec<(u64, T)>,
     peak_pos: u64,
     proof_iter: &mut I,
 ) -> Result<T> {
     debug_assert!(!leaves.is_empty(), "can't be empty");
     // (position, hash, height)
+
     let mut queue: VecDeque<_> = leaves
         .into_iter()
-        .map(|(pos, item)| (pos, item, 0u32))
+        .map(|(pos, item)| (pos, item, 0))
         .collect();
 
     // calculate tree root from each items
     while let Some((pos, item, height)) = queue.pop_front() {
         if pos == peak_pos {
-            // return root
-            return Ok(item);
+            if queue.is_empty() {
+                // return root once queue is consumed
+                return Ok(item);
+            } else {
+                return Err(Error::CorruptedProof);
+            }
         }
         // calculate sibling
         let next_height = pos_height_in_tree(pos + 1);
-        let (sib_pos, parent_pos) = {
+        let (parent_pos, parent_item) = {
             let sibling_offset = sibling_offset(height);
             if next_height > height {
                 // implies pos is right sibling
-                (pos - sibling_offset, pos + 1)
+                let sib_pos = pos - sibling_offset;
+                let parent_pos = pos + 1;
+                let parent_item = if Some(&sib_pos) == queue.front().map(|(pos, _, _)| pos) {
+                    let sibling_item = queue.pop_front().map(|(_, item, _)| item).unwrap();
+                    M::merge(&sibling_item, &item)?
+                } else {
+                    let sibling_item = proof_iter.next().ok_or(Error::CorruptedProof)?;
+                    M::merge(sibling_item, &item)?
+                };
+                (parent_pos, parent_item)
             } else {
                 // pos is left sibling
-                (pos + sibling_offset, pos + parent_offset(height))
+                let sib_pos = pos + sibling_offset;
+                let parent_pos = pos + parent_offset(height);
+                let parent_item = if Some(&sib_pos) == queue.front().map(|(pos, _, _)| pos) {
+                    let sibling_item = queue.pop_front().map(|(_, item, _)| item).unwrap();
+                    M::merge(&item, &sibling_item)?
+                } else {
+                    let sibling_item = proof_iter.next().ok_or(Error::CorruptedProof)?;
+                    M::merge(&item, sibling_item)?
+                };
+                (parent_pos, parent_item)
             }
         };
-        let sibling_item = if Some(&sib_pos) == queue.front().map(|(pos, _, _)| pos) {
-            queue.pop_front().map(|(_, item, _)| item).unwrap()
-        } else {
-            proof_iter.next().ok_or(Error::CorruptedProof)?.clone()
-        };
 
-        let parent_item = if next_height > height {
-            M::merge(&sibling_item, &item)
+        if parent_pos <= peak_pos {
+            queue.push_back((parent_pos, parent_item, height + 1))
         } else {
-            M::merge(&item, &sibling_item)
-        }?;
-
-        if parent_pos < peak_pos {
-            queue.push_back((parent_pos, parent_item, height + 1));
-        } else {
-            return Ok(parent_item);
+            return Err(Error::CorruptedProof);
         }
     }
     Err(Error::CorruptedProof)
 }
 
-fn calculate_peaks_hashes<
-    'a,
-    T: 'a + PartialEq + Debug + Clone,
-    M: Merge<Item = T>,
-    I: Iterator<Item = &'a T>,
->(
+fn calculate_peaks_hashes<'a, T: 'a + Clone, M: Merge<Item = T>, I: Iterator<Item = &'a T>>(
     mut leaves: Vec<(u64, T)>,
     mmr_size: u64,
     mut proof_iter: I,
 ) -> Result<Vec<T>> {
+    if leaves.iter().any(|(pos, _)| pos_height_in_tree(*pos) > 0) {
+        return Err(Error::NodeProofsNotSupported);
+    }
+
     // special handle the only 1 leaf MMR
     if mmr_size == 1 && leaves.len() == 1 && leaves[0].0 == 0 {
         return Ok(leaves.into_iter().map(|(_pos, item)| item).collect());
     }
-    // sort items by position
+    // ensure leaves are sorted and unique
     leaves.sort_by_key(|(pos, _)| *pos);
+    leaves.dedup_by(|a, b| a.0 == b.0);
     let peaks = get_peaks(mmr_size);
 
     let mut peaks_hashes: Vec<T> = Vec::with_capacity(peaks.len() + 1);
@@ -386,9 +469,7 @@ fn calculate_peaks_hashes<
     Ok(peaks_hashes)
 }
 
-fn bagging_peaks_hashes<'a, T: 'a + PartialEq + Debug + Clone, M: Merge<Item = T>>(
-    mut peaks_hashes: Vec<T>,
-) -> Result<T> {
+fn bagging_peaks_hashes<T, M: Merge<Item = T>>(mut peaks_hashes: Vec<T>) -> Result<T> {
     // bagging peaks
     // bagging from right to left via hash(right, left).
     while peaks_hashes.len() > 1 {
@@ -403,12 +484,7 @@ fn bagging_peaks_hashes<'a, T: 'a + PartialEq + Debug + Clone, M: Merge<Item = T
 /// 1. sort items by position
 /// 2. calculate root of each peak
 /// 3. bagging peaks
-fn calculate_root<
-    'a,
-    T: 'a + PartialEq + Debug + Clone,
-    M: Merge<Item = T>,
-    I: Iterator<Item = &'a T>,
->(
+fn calculate_root<'a, T: 'a + Clone, M: Merge<Item = T>, I: Iterator<Item = &'a T>>(
     leaves: Vec<(u64, T)>,
     mmr_size: u64,
     proof_iter: I,
@@ -423,5 +499,5 @@ fn take_while_vec<T, P: Fn(&T) -> bool>(v: &mut Vec<T>, p: P) -> Vec<T> {
             return v.drain(..i).collect();
         }
     }
-    v.drain(..).collect()
+    mem::take(v)
 }

src/mmr_store.rs

@@ -1,4 +1,4 @@
-use crate::{vec::Vec, Result};
+use crate::{Result, vec::Vec};
 
 #[derive(Default)]
 pub struct MMRBatch<Elem, Store> {
@@ -6,7 +6,7 @@ pub struct MMRBatch<Elem, Store> {
     store: Store,
 }
 
-impl<Elem: Clone, Store: MMRStore<Elem>> MMRBatch<Elem, Store> {
+impl<Elem, Store> MMRBatch<Elem, Store> {
     pub fn new(store: Store) -> Self {
         MMRBatch {
             memory_batch: Vec::new(),
@@ -18,6 +18,12 @@ impl<Elem: Clone, Store: MMRStore<Elem>> MMRBatch<Elem, Store> {
         self.memory_batch.push((pos, elems));
     }
 
+    pub fn store(&self) -> &Store {
+        &self.store
+    }
+}
+
+impl<Elem: Clone, Store: MMRStoreReadOps<Elem>> MMRBatch<Elem, Store> {
     pub fn get_elem(&self, pos: u64) -> Result<Option<Elem>> {
         for (start_pos, elems) in self.memory_batch.iter().rev() {
             if pos < *start_pos {
@@ -30,20 +36,18 @@ impl<Elem: Clone, Store: MMRStore<Elem>> MMRBatch<Elem, Store> {
         }
         self.store.get_elem(pos)
     }
+}
 
-    pub fn commit(self) -> Result<()> {
-        let Self {
-            mut store,
-            memory_batch,
-        } = self;
-        for (pos, elems) in memory_batch {
-            store.append(pos, elems)?;
+impl<Elem, Store: MMRStoreWriteOps<Elem>> MMRBatch<Elem, Store> {
+    pub fn commit(&mut self) -> Result<()> {
+        for (pos, elems) in self.memory_batch.drain(..) {
+            self.store.append(pos, elems)?;
         }
         Ok(())
     }
 }
 
-impl<Elem, Store: MMRStore<Elem>> IntoIterator for MMRBatch<Elem, Store> {
+impl<Elem, Store> IntoIterator for MMRBatch<Elem, Store> {
     type Item = (u64, Vec<Elem>);
     type IntoIter = crate::vec::IntoIter<Self::Item>;
 
@@ -52,7 +56,10 @@ impl<Elem, Store: MMRStore<Elem>> IntoIterator for MMRBatch<Elem, Store> {
     }
 }
 
-pub trait MMRStore<Elem> {
+pub trait MMRStoreReadOps<Elem> {
     fn get_elem(&self, pos: u64) -> Result<Option<Elem>>;
+}
+
+pub trait MMRStoreWriteOps<Elem> {
     fn append(&mut self, pos: u64, elems: Vec<Elem>) -> Result<()>;
 }

src/tests/mod.rs

@@ -1,5 +1,6 @@
 mod test_accumulate_headers;
 mod test_helper;
+mod test_incremental;
 mod test_mmr;
 mod test_sequence;
 

src/tests/test_accumulate_headers.rs

@@ -1,5 +1,5 @@
 use super::new_blake2b;
-use crate::{leaf_index_to_pos, util::MemStore, MMRStore, Merge, MerkleProof, Result, MMR};
+use crate::{MMR, MMRStoreReadOps, Merge, MerkleProof, Result, leaf_index_to_pos, util::MemStore};
 use bytes::{Bytes, BytesMut};
 use std::fmt;
 
@@ -106,7 +106,7 @@ impl Prover {
         let mut mmr = MMR::<_, MergeHashWithTD, _>::new(self.positions.len() as u64, &self.store);
         // get previous element
         let mut previous = if let Some(pos) = self.positions.last() {
-            MMRStore::<_>::get_elem(&&self.store, *pos)?.expect("exists")
+            mmr.store().get_elem(*pos)?.expect("exists")
         } else {
             let genesis = Header::default();
 

src/tests/test_helper.rs

@@ -1,9 +1,9 @@
 use super::{MergeNumberHash, NumberHash};
 use crate::{
-    helper::{get_peaks, pos_height_in_tree},
+    MMR,
+    helper::{get_peak_map, get_peaks, pos_height_in_tree},
     leaf_index_to_mmr_size, leaf_index_to_pos,
     util::MemStore,
-    MMR,
 };
 use lazy_static::lazy_static;
 use proptest::prelude::*;
@@ -55,16 +55,41 @@ fn test_pos_height_in_tree() {
     assert_eq!(pos_height_in_tree(7), 0);
 }
 
+#[test]
+fn test_get_peak_map() {
+    assert_eq!(get_peak_map(0), 0b0);
+    assert_eq!(get_peak_map(1), 0b1);
+    assert_eq!(get_peak_map(3), 0b10);
+    assert_eq!(get_peak_map(4), 0b11);
+    // 5 and 6 are not valid mmr_size, it will return the bitmap of the last valid mmr (size 4)
+    assert_eq!(get_peak_map(5), 0b11);
+    assert_eq!(get_peak_map(6), 0b11);
+    assert_eq!(get_peak_map(7), 0b100);
+    assert_eq!(get_peak_map(8), 0b101);
+    // 9 is not valid mmr_size, it will return the bitmap of the last valid mmr (size 8)
+    assert_eq!(get_peak_map(9), 0b101);
+    assert_eq!(get_peak_map(15), 0b1000);
+    assert_eq!(get_peak_map(16), 0b1001);
+    assert_eq!(get_peak_map(18), 0b1010);
+    assert_eq!(get_peak_map(19), 0b1011);
+}
+
 #[test]
 fn test_get_peaks() {
-    assert_eq!(get_peaks(0), vec![0]);
+    assert_eq!(get_peaks(0), vec![]);
     assert_eq!(get_peaks(1), vec![0]);
-    assert_eq!(get_peaks(2), vec![0]);
     assert_eq!(get_peaks(3), vec![2]);
     assert_eq!(get_peaks(4), vec![2, 3]);
+    // 5 and 6 are not valid mmr_size, it will return the peaks of the last valid mmr (size 4)
     assert_eq!(get_peaks(5), vec![2, 3]);
-    assert_eq!(get_peaks(6), vec![2, 5]);
+    assert_eq!(get_peaks(6), vec![2, 3]);
     assert_eq!(get_peaks(7), vec![6]);
+    assert_eq!(get_peaks(8), vec![6, 7]);
+    // 9 is not valid mmr_size, it will return the peaks of the last valid mmr (size 8)
+    assert_eq!(get_peaks(9), vec![6, 7]);
+    assert_eq!(get_peaks(15), vec![14]);
+    assert_eq!(get_peaks(16), vec![14, 15]);
+    assert_eq!(get_peaks(18), vec![14, 17]);
     assert_eq!(get_peaks(19), vec![14, 17, 18]);
 }
 

src/tests/test_incremental.rs

@@ -0,0 +1,51 @@
+use proptest::proptest;
+
+use super::{MergeNumberHash, NumberHash};
+use crate::util::{MemMMR, MemStore};
+
+proptest! {
+    #[test]
+    fn test_incremental(start in 1u32..500, steps in 1usize..50, turns in 10usize..20) {
+        test_incremental_with_params(start, steps, turns);
+    }
+}
+
+fn test_incremental_with_params(start: u32, steps: usize, turns: usize) {
+    let store = MemStore::default();
+    let mut mmr = MemMMR::<_, MergeNumberHash>::new(0, &store);
+
+    let mut curr = 0;
+
+    let _positions: Vec<u64> = (0u32..start)
+        .map(|_| {
+            let pos = mmr.push(NumberHash::from(curr)).unwrap();
+            curr += 1;
+            pos
+        })
+        .collect();
+    mmr.commit().expect("commit changes");
+
+    for turn in 0..turns {
+        let prev_root = mmr.get_root().expect("get root");
+        let (positions, leaves) = (0..steps).fold(
+            (Vec::new(), Vec::new()),
+            |(mut positions, mut leaves), _| {
+                let leaf = NumberHash::from(curr);
+                let pos = mmr.push(leaf.clone()).unwrap();
+                curr += 1;
+                positions.push(pos);
+                leaves.push(leaf);
+                (positions, leaves)
+            },
+        );
+        mmr.commit().expect("commit changes");
+        let proof = mmr.gen_proof(positions).expect("gen proof");
+        let root = mmr.get_root().expect("get root");
+        let result = proof.verify_incremental(root, prev_root, leaves).unwrap();
+        assert!(
+            result,
+            "start: {}, steps: {}, turn: {}, curr: {}",
+            start, steps, turn, curr
+        );
+    }
+}

src/tests/test_mmr.rs

@@ -1,12 +1,17 @@
 use super::{MergeNumberHash, NumberHash};
-use crate::{leaf_index_to_mmr_size, util::MemStore, Error, MMR};
+use crate::{
+    Error,
+    helper::pos_height_in_tree,
+    leaf_index_to_mmr_size,
+    util::{MemMMR, MemStore},
+};
 use faster_hex::hex_string;
 use proptest::prelude::*;
 use rand::{seq::SliceRandom, thread_rng};
 
 fn test_mmr(count: u32, proof_elem: Vec<u32>) {
     let store = MemStore::default();
-    let mut mmr = MMR::<_, MergeNumberHash, _>::new(0, &store);
+    let mut mmr = MemMMR::<_, MergeNumberHash>::new(0, &store);
     let positions: Vec<u64> = (0u32..count)
         .map(|i| mmr.push(NumberHash::from(i)).unwrap())
         .collect();
@@ -34,7 +39,7 @@ fn test_mmr(count: u32, proof_elem: Vec<u32>) {
 
 fn test_gen_new_root_from_proof(count: u32) {
     let store = MemStore::default();
-    let mut mmr = MMR::<_, MergeNumberHash, _>::new(0, &store);
+    let mut mmr = MemMMR::<_, MergeNumberHash>::new(0, &store);
     let positions: Vec<u64> = (0u32..count)
         .map(|i| mmr.push(NumberHash::from(i)).unwrap())
         .collect();
@@ -59,7 +64,7 @@ fn test_gen_new_root_from_proof(count: u32) {
 #[test]
 fn test_mmr_root() {
     let store = MemStore::default();
-    let mut mmr = MMR::<_, MergeNumberHash, _>::new(0, &store);
+    let mut mmr = MemMMR::<_, MergeNumberHash>::new(0, &store);
     (0u32..11).for_each(|i| {
         mmr.push(NumberHash::from(i)).unwrap();
     });
@@ -74,7 +79,7 @@ fn test_mmr_root() {
 #[test]
 fn test_empty_mmr_root() {
     let store = MemStore::<NumberHash>::default();
-    let mmr = MMR::<_, MergeNumberHash, _>::new(0, &store);
+    let mmr = MemMMR::<_, MergeNumberHash>::new(0, &store);
     assert_eq!(Err(Error::GetRootOnEmpty), mmr.get_root());
 }
 
@@ -140,6 +145,126 @@ fn test_gen_root_from_proof() {
     test_gen_new_root_from_proof(11);
 }
 
+#[test]
+fn test_gen_proof_with_duplicate_leaves() {
+    test_mmr(10, vec![5, 5]);
+}
+
+fn test_invalid_proof_verification(
+    leaf_count: u32,
+    positions_to_verify: Vec<u64>,
+    // positions of entries that should be tampered
+    tampered_positions: Vec<usize>,
+    // optionally handroll proof from these positions
+    handrolled_proof_positions: Option<Vec<u64>>,
+) {
+    use crate::{Merge, MerkleProof};
+    use std::fmt::{Debug, Formatter};
+
+    // Simple item struct to allow debugging the contents of MMR nodes/peaks
+    #[derive(Clone, PartialEq)]
+    enum MyItem {
+        Number(u32),
+        Merged(Box<MyItem>, Box<MyItem>),
+    }
+
+    impl Debug for MyItem {
+        fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+            match self {
+                MyItem::Number(x) => f.write_fmt(format_args!("{}", x)),
+                MyItem::Merged(a, b) => f.write_fmt(format_args!("Merged({:#?}, {:#?})", a, b)),
+            }
+        }
+    }
+
+    #[derive(Debug)]
+    struct MyMerge;
+
+    impl Merge for MyMerge {
+        type Item = MyItem;
+        fn merge(lhs: &Self::Item, rhs: &Self::Item) -> Result<Self::Item, crate::Error> {
+            Ok(MyItem::Merged(Box::new(lhs.clone()), Box::new(rhs.clone())))
+        }
+    }
+
+    let store = MemStore::default();
+    let mut mmr = MemMMR::<_, MyMerge>::new(0, &store);
+    let mut positions: Vec<u64> = Vec::new();
+    for i in 0u32..leaf_count {
+        let pos = mmr.push(MyItem::Number(i)).unwrap();
+        positions.push(pos);
+    }
+    let root = mmr.get_root().unwrap();
+
+    let entries_to_verify: Vec<(u64, MyItem)> = positions_to_verify
+        .iter()
+        .map(|pos| (*pos, mmr.batch().get_elem(*pos).unwrap().unwrap()))
+        .collect();
+
+    let mut tampered_entries_to_verify = entries_to_verify.clone();
+    tampered_positions.iter().for_each(|proof_pos| {
+        tampered_entries_to_verify[*proof_pos] = (
+            tampered_entries_to_verify[*proof_pos].0,
+            MyItem::Number(31337),
+        )
+    });
+
+    let handrolled_proof: Option<MerkleProof<MyItem, MyMerge>> =
+        handrolled_proof_positions.map(|handrolled_proof_positions| {
+            MerkleProof::new(
+                mmr.mmr_size(),
+                handrolled_proof_positions
+                    .iter()
+                    .map(|pos| mmr.batch().get_elem(*pos).unwrap().unwrap())
+                    .collect(),
+            )
+        });
+
+    // verification should fail whenever trying to prove membership of a non-member
+    if let Some(handrolled_proof) = handrolled_proof {
+        let handrolled_proof_result =
+            handrolled_proof.verify(root.clone(), tampered_entries_to_verify.clone());
+        assert!(handrolled_proof_result.is_err() || !handrolled_proof_result.unwrap());
+    }
+
+    match mmr.gen_proof(positions_to_verify.clone()) {
+        Ok(proof) => {
+            assert!(proof.verify(root.clone(), entries_to_verify).unwrap());
+            assert!(!proof.verify(root, tampered_entries_to_verify).unwrap());
+        }
+        Err(Error::NodeProofsNotSupported) => {
+            // if couldn't generate proof, then it contained a non-leaf
+            assert!(
+                positions_to_verify
+                    .iter()
+                    .any(|pos| pos_height_in_tree(*pos) > 0)
+            );
+        }
+        Err(e) => panic!("Unexpected error: {}", e),
+    }
+}
+
+#[test]
+fn test_generic_proofs() {
+    test_invalid_proof_verification(7, vec![5], vec![0], Some(vec![2, 9, 10]));
+    test_invalid_proof_verification(7, vec![1, 2], vec![0], Some(vec![5, 9, 10]));
+    test_invalid_proof_verification(7, vec![1, 5], vec![0], Some(vec![0, 9, 10]));
+    test_invalid_proof_verification(7, vec![1, 6], vec![0], Some(vec![0, 5, 9, 10]));
+    test_invalid_proof_verification(7, vec![5, 6], vec![0], Some(vec![2, 9, 10]));
+    test_invalid_proof_verification(7, vec![1, 5, 6], vec![0], Some(vec![0, 9, 10]));
+    test_invalid_proof_verification(7, vec![1, 5, 7], vec![0], Some(vec![0, 8, 10]));
+    test_invalid_proof_verification(7, vec![5, 6, 7], vec![0], Some(vec![2, 8, 10]));
+    test_invalid_proof_verification(7, vec![5, 6, 7, 8, 9, 10], vec![0], Some(vec![2]));
+    test_invalid_proof_verification(7, vec![1, 5, 7, 8, 9, 10], vec![0], Some(vec![0]));
+    test_invalid_proof_verification(7, vec![0, 1, 5, 7, 8, 9, 10], vec![0], Some(vec![]));
+    test_invalid_proof_verification(7, vec![0, 1, 5, 6, 7, 8, 9, 10], vec![0], Some(vec![]));
+    test_invalid_proof_verification(7, vec![0, 1, 2, 5, 6, 7, 8, 9, 10], vec![0], Some(vec![]));
+    test_invalid_proof_verification(7, vec![0, 1, 2, 3, 7, 8, 9, 10], vec![0], Some(vec![4]));
+    test_invalid_proof_verification(7, vec![0, 2, 3, 7, 8, 9, 10], vec![0], Some(vec![1, 4]));
+    test_invalid_proof_verification(7, vec![0, 3, 7, 8, 9, 10], vec![0], Some(vec![1, 4]));
+    test_invalid_proof_verification(7, vec![0, 2, 3, 7, 8, 9, 10], vec![0], Some(vec![1, 4]));
+}
+
 prop_compose! {
     fn count_elem(count: u32)
                 (elem in 0..count)

src/tests/test_sequence.rs

@@ -3,7 +3,7 @@ use std::fmt;
 use proptest::proptest;
 use rand::{prelude::*, thread_rng};
 
-use crate::{util::MemStore, Merge, Result, MMR};
+use crate::{MMR, Merge, Result, util::MemStore};
 
 #[derive(Eq, PartialEq, Clone, Default)]
 struct NumberRange {

src/util.rs

@@ -1,8 +1,6 @@
 use crate::collections::BTreeMap;
-use crate::{vec::Vec, MMRStore, Merge, MerkleProof, Result, MMR};
+use crate::{MMR, MMRStoreReadOps, MMRStoreWriteOps, Result, vec::Vec};
 use core::cell::RefCell;
-use core::fmt::Debug;
-use core::marker::PhantomData;
 
 #[derive(Clone)]
 pub struct MemStore<T>(RefCell<BTreeMap<u64, T>>);
@@ -19,11 +17,13 @@ impl<T> MemStore<T> {
     }
 }
 
-impl<T: Clone> MMRStore<T> for &MemStore<T> {
+impl<T: Clone> MMRStoreReadOps<T> for &MemStore<T> {
     fn get_elem(&self, pos: u64) -> Result<Option<T>> {
         Ok(self.0.borrow().get(&pos).cloned())
     }
+}
 
+impl<T> MMRStoreWriteOps<T> for &MemStore<T> {
     fn append(&mut self, pos: u64, elems: Vec<T>) -> Result<()> {
         let mut store = self.0.borrow_mut();
         for (i, elem) in elems.into_iter().enumerate() {
@@ -33,46 +33,4 @@ impl<T: Clone> MMRStore<T> for &MemStore<T> {
     }
 }
 
-pub struct MemMMR<T, M> {
-    store: MemStore<T>,
-    mmr_size: u64,
-    merge: PhantomData<M>,
-}
-
-impl<T: Clone + Debug + PartialEq, M: Merge<Item = T>> Default for MemMMR<T, M> {
-    fn default() -> Self {
-        Self::new(0, Default::default())
-    }
-}
-
-impl<T: Clone + Debug + PartialEq, M: Merge<Item = T>> MemMMR<T, M> {
-    pub fn new(mmr_size: u64, store: MemStore<T>) -> Self {
-        MemMMR {
-            mmr_size,
-            store,
-            merge: PhantomData,
-        }
-    }
-
-    pub fn store(&self) -> &MemStore<T> {
-        &self.store
-    }
-
-    pub fn get_root(&self) -> Result<T> {
-        let mmr = MMR::<T, M, &MemStore<T>>::new(self.mmr_size, &self.store);
-        mmr.get_root()
-    }
-
-    pub fn push(&mut self, elem: T) -> Result<u64> {
-        let mut mmr = MMR::<T, M, &MemStore<T>>::new(self.mmr_size, &self.store);
-        let pos = mmr.push(elem)?;
-        self.mmr_size = mmr.mmr_size();
-        mmr.commit()?;
-        Ok(pos)
-    }
-
-    pub fn gen_proof(&self, pos_list: Vec<u64>) -> Result<MerkleProof<T, M>> {
-        let mmr = MMR::<T, M, &MemStore<T>>::new(self.mmr_size, &self.store);
-        mmr.gen_proof(pos_list)
-    }
-}
+pub type MemMMR<'a, T, M> = MMR<T, M, &'a MemStore<T>>;