Merge pull request #13 from nervosnetwork/release-v0.3.2

Release v0.3.2

.github/workflows/rust.yml

@@ -0,0 +1,13 @@
+name: Rust
+
+on: [push, pull_request]
+
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v1
+    - name: Run CI checks
+      run: make ci

.gitignore

@@ -1,2 +1,4 @@
 Cargo.lock
 target/
+.vscode/
+.idea

Cargo.toml

@@ -1,13 +1,18 @@
 [package]
 name = "ckb-merkle-mountain-range"
-version = "0.1.0"
+version = "0.3.2"
 authors = ["Nervos Core Dev <dev@nervos.org>"]
 edition = "2018"
 license = "MIT"
 description = "A generalized merkle mountain range implementation"
+repository = "https://github.com/nervosnetwork/merkle-mountain-range"
+
+[features]
+default = ["std"]
+std = []
 
 [dependencies]
-failure = "0.1.5"
+cfg-if = "0.1"
 
 [dev-dependencies]
 faster-hex = "0.3"

Makefile

@@ -1,6 +1,6 @@
 default: ci
 
-ci: fmt clippy test bench-test
+ci: fmt check-no-std clippy test bench-test
 
 test:
 	cargo test --all --all-features
@@ -13,3 +13,6 @@ clippy:
 
 fmt:
 	cargo fmt --all -- --check
+
+check-no-std:
+	cargo check --all --no-default-features

README.md

@@ -1,7 +1,16 @@
 # Merkle mountain range
+[![Crates.io](https://img.shields.io/crates/v/ckb-merkle-mountain-range.svg)](https://crates.io/crates/ckb-merkle-mountain-range)
 
 A generalized merkle mountain range implementation.
 
+**Notice** this library is not stable yet, API and proof format may changes. Make sure you know what you do before using this library.
+
+## Features
+
+* Leaves accumulation
+* Multi leaves merkle proof
+* Accumulate from last leaf's merkle proof
+
 ## Construct
 
 ``` txt
@@ -29,7 +38,7 @@ For example, we insert a leaf to the example MMR:
 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])`.
 5. insert new node to next position `21`.
-6. the node `20` have no left sibling, complete the insertion.
+6. the node `21` have no left sibling, complete the insertion.
 
 Example MMR after insertion of a new leaf:
 

benches/mmr_benchmark.rs

@@ -1,7 +1,7 @@
 #[macro_use]
 extern crate criterion;
 
-use criterion::Criterion;
+use criterion::{BenchmarkId, Criterion};
 
 use bytes::Bytes;
 use ckb_merkle_mountain_range::{util::MemStore, Error, MMRStore, Merge, Result, MMR};
@@ -53,19 +53,21 @@ fn prepare_mmr(count: u32) -> (u64, MemStore<NumberHash>, Vec<u64>) {
 }
 
 fn bench(c: &mut Criterion) {
-    c.bench_function_over_inputs(
-        "MMR insert",
-        |b, &&size| {
-            b.iter(|| prepare_mmr(size));
-        },
-        &[10_000, 100_000, 100_0000],
-    );
+    {
+        let mut group = c.benchmark_group("MMR insertion");
+        let inputs = [10_000, 100_000, 100_0000];
+        for input in inputs.iter() {
+            group.bench_with_input(BenchmarkId::new("times", input), &input, |b, &&size| {
+                b.iter(|| prepare_mmr(size));
+            });
+        }
+    }
 
     c.bench_function("MMR gen proof", |b| {
         let (mmr_size, store, positions) = prepare_mmr(100_0000);
         let mmr = MMR::<_, MergeNumberHash, _>::new(mmr_size, &store);
         let mut rng = thread_rng();
-        b.iter(|| mmr.gen_proof(*positions.choose(&mut rng).unwrap()));
+        b.iter(|| mmr.gen_proof(vec![*positions.choose(&mut rng).unwrap()]));
     });
 
     c.bench_function("MMR verify", |b| {
@@ -77,13 +79,15 @@ fn bench(c: &mut Criterion) {
             .map(|_| {
                 let pos = positions.choose(&mut rng).unwrap();
                 let elem = (&store).get_elem(*pos).unwrap().unwrap();
-                let proof = mmr.gen_proof(*pos).unwrap();
+                let proof = mmr.gen_proof(vec![*pos]).unwrap();
                 (pos, elem, proof)
             })
             .collect();
         b.iter(|| {
             let (pos, elem, proof) = proofs.choose(&mut rng).unwrap();
-            proof.verify(root.clone(), **pos, elem.clone()).unwrap();
+            proof
+                .verify(root.clone(), vec![(**pos, elem.clone())])
+                .unwrap();
         });
     });
 }

src/error.rs

@@ -1,12 +1,32 @@
-pub use failure::Fail;
-pub type Result<T> = ::std::result::Result<T, Error>;
+pub type Result<T> = core::result::Result<T, Error>;
 
-#[derive(Fail, Debug, PartialEq, Eq, Clone)]
+#[derive(Debug, PartialEq, Eq, Clone)]
 pub enum Error {
-    #[fail(display = "Get root on an empty MMR")]
     GetRootOnEmpty,
-    #[fail(display = "Inconsistent store")]
     InconsistentStore,
-    #[fail(display = "Store error {}", _0)]
-    StoreError(String),
+    StoreError(crate::string::String),
+    /// proof items is not enough to build a tree
+    CorruptedProof,
+    /// The leaves is an empty list, or beyond the mmr range
+    GenProofForInvalidLeaves,
+}
+
+impl core::fmt::Display for Error {
+    fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
+        use Error::*;
+        match self {
+            GetRootOnEmpty => write!(f, "Get root on an empty MMR")?,
+            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")?,
+        }
+        Ok(())
+    }
+}
+
+cfg_if::cfg_if! {
+    if #[cfg(feature = "std")] {
+        impl ::std::error::Error for Error {}
+    }
 }

src/helper.rs

@@ -1,40 +1,18 @@
+use crate::vec::Vec;
+
 pub fn leaf_index_to_pos(index: u64) -> u64 {
-    if index == 0 {
-        return 0;
-    }
-    // leaf_count
-    let mut leaves = index + 1;
-    let mut tree_node_count = 0;
-    let mut height = 0u32;
-    while leaves > 1 {
-        // get heighest peak height
-        height = (leaves as f64).log2() as u32;
-        // calculate leaves in peak
-        let peak_leaves = 1 << height;
-        // heighest positon
-        let sub_tree_node_count = get_peak_pos_by_height(height) + 1;
-        tree_node_count += sub_tree_node_count;
-        leaves -= peak_leaves;
-    }
-    // two leaves can construct a new peak, the only valid number of leaves is 0 or 1.
-    debug_assert!(leaves == 0 || leaves == 1, "remain leaves incorrect");
-    if leaves == 1 {
-        // add one pos for remain leaf
-        // equals to `tree_node_count - 1 + 1`
-        tree_node_count
-    } else {
-        let pos = tree_node_count - 1;
-        pos - u64::from(height)
-    }
+    // mmr_size - H - 1, H is the height(intervals) of last peak
+    leaf_index_to_mmr_size(index) - (index + 1).trailing_zeros() as u64 - 1
 }
 
-// TODO optimize
 pub fn leaf_index_to_mmr_size(index: u64) -> u64 {
-    let mut pos = leaf_index_to_pos(index);
-    while pos_height_in_tree(pos + 1) > pos_height_in_tree(pos) {
-        pos += 1
-    }
-    pos + 1
+    // leaf index start with 0
+    let leaves_count = index + 1;
+
+    // the peak count(k) is actually the count of 1 in leaves count's binary representation
+    let peak_count = leaves_count.count_ones() as u64;
+
+    2 * leaves_count - peak_count
 }
 
 pub fn pos_height_in_tree(mut pos: u64) -> u32 {

src/lib.rs

@@ -1,5 +1,7 @@
+#![cfg_attr(not(feature = "std"), no_std)]
+
 mod error;
-mod helper;
+pub mod helper;
 mod merge;
 mod mmr;
 mod mmr_store;
@@ -12,3 +14,18 @@ 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;
+
+cfg_if::cfg_if! {
+    if #[cfg(feature = "std")] {
+        use std::borrow;
+        use std::collections;
+        use std::vec;
+        use std::string;
+    } else {
+        extern crate alloc;
+        use alloc::borrow;
+        use alloc::collections;
+        use alloc::vec;
+        use alloc::string;
+    }
+}

src/mmr.rs

@@ -4,13 +4,17 @@
 //! https://github.com/mimblewimble/grin/blob/master/doc/mmr.md#structure
 //! https://github.com/mimblewimble/grin/blob/0ff6763ee64e5a14e70ddd4642b99789a1648a32/core/src/core/pmmr.rs#L606
 
+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::vec;
+use crate::vec::Vec;
 use crate::{Error, Merge, Result};
-use std::borrow::Cow;
-use std::fmt::Debug;
-use std::marker::PhantomData;
+use core::fmt::Debug;
+use core::marker::PhantomData;
 
+#[allow(clippy::upper_case_acronyms)]
 pub struct MMR<T, M, S: MMRStore<T>> {
     mmr_size: u64,
     batch: MMRBatch<T, S>,
@@ -77,67 +81,128 @@ impl<'a, T: Clone + PartialEq + Debug, M: Merge<Item = T>, S: MMRStore<T>> MMR<T
         } else if self.mmr_size == 1 {
             return self.batch.get_elem(0)?.ok_or(Error::InconsistentStore);
         }
-        let peaks = get_peaks(self.mmr_size);
-        self.bag_rhs_peaks(0, &peaks)?
-            .ok_or(Error::InconsistentStore)
+        let peaks: Vec<T> = get_peaks(self.mmr_size)
+            .into_iter()
+            .map(|peak_pos| {
+                self.batch
+                    .get_elem(peak_pos)
+                    .and_then(|elem| elem.ok_or(Error::InconsistentStore))
+            })
+            .collect::<Result<Vec<T>>>()?;
+        self.bag_rhs_peaks(peaks)?.ok_or(Error::InconsistentStore)
     }
 
-    fn bag_rhs_peaks(&self, skip_peak_pos: u64, peaks: &[u64]) -> Result<Option<T>> {
-        let mut rhs_peak_elems: Vec<T> = peaks
-            .iter()
-            .filter(|&&p| p > skip_peak_pos)
-            .map(|&p| self.batch.get_elem(p))
-            .collect::<Result<Option<_>>>()?
-            .ok_or(Error::InconsistentStore)?;
-        while rhs_peak_elems.len() > 1 {
-            let right_peak = rhs_peak_elems.pop().expect("pop");
-            let left_peak = rhs_peak_elems.pop().expect("pop");
-            rhs_peak_elems.push(M::merge(&right_peak, &left_peak));
+    fn bag_rhs_peaks(&self, mut rhs_peaks: Vec<T>) -> Result<Option<T>> {
+        while rhs_peaks.len() > 1 {
+            let right_peak = rhs_peaks.pop().expect("pop");
+            let left_peak = rhs_peaks.pop().expect("pop");
+            rhs_peaks.push(M::merge(&right_peak, &left_peak));
         }
-        Ok(rhs_peak_elems.pop())
+        Ok(rhs_peaks.pop())
     }
 
-    pub fn gen_proof(&self, mut pos: u64) -> Result<MerkleProof<T, M>> {
-        let mut proof: Vec<T> = Vec::new();
-        let mut height = 0;
-        while pos < self.mmr_size {
-            let pos_height = pos_height_in_tree(pos);
-            let next_height = pos_height_in_tree(pos + 1);
-            let (sib_pos, next_pos) = if next_height > pos_height {
-                // implies pos is right sibling
-                let sib_pos = pos - sibling_offset(height);
-                (sib_pos, pos + 1)
-            } else {
-                // pos is left sibling
-                let sib_pos = pos + sibling_offset(height);
-                (sib_pos, pos + parent_offset(height))
-            };
-            if sib_pos > self.mmr_size - 1 {
-                break;
-            }
+    /// generate merkle proof for a peak
+    /// the pos_list must be sorted, otherwise the behaviour is undefined
+    ///
+    /// 1. find a lower tree in peak that can generate a complete merkle proof for position
+    /// 2. find that tree by compare positions
+    /// 3. generate proof for each positions
+    fn gen_proof_for_peak(
+        &self,
+        proof: &mut Vec<T>,
+        pos_list: Vec<u64>,
+        peak_pos: u64,
+    ) -> Result<()> {
+        // do nothing if position itself is the peak
+        if pos_list.len() == 1 && pos_list == [peak_pos] {
+            return Ok(());
+        }
+        // take peak root from store if no positions need to be proof
+        if pos_list.is_empty() {
             proof.push(
                 self.batch
-                    .get_elem(sib_pos)?
+                    .get_elem(peak_pos)?
                     .ok_or(Error::InconsistentStore)?,
             );
-            pos = next_pos;
-            height += 1;
+            return Ok(());
         }
-        // now we get peak merkle proof
-        let peak_pos = pos;
-        // calculate bagging proof
+
+        let mut queue: VecDeque<_> = pos_list.into_iter().map(|pos| (pos, 0u32)).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;
+            }
+
+            // calculate sibling
+            let (sib_pos, parent_pos) = {
+                let next_height = pos_height_in_tree(pos + 1);
+                let sibling_offset = sibling_offset(height);
+                if next_height > height {
+                    // implies pos is right sibling
+                    (pos - sibling_offset, pos + 1)
+                } else {
+                    // pos is left sibling
+                    (pos + sibling_offset, pos + parent_offset(height))
+                }
+            };
+
+            if Some(&sib_pos) == queue.front().map(|(pos, _)| pos) {
+                // drop sibling
+                queue.pop_front();
+            } else {
+                proof.push(
+                    self.batch
+                        .get_elem(sib_pos)?
+                        .ok_or(Error::InconsistentStore)?,
+                );
+            }
+            if parent_pos < peak_pos {
+                // save pos to tree buf
+                queue.push_back((parent_pos, height + 1));
+            }
+        }
+        Ok(())
+    }
+
+    /// Generate merkle proof for positions
+    /// 1. sort positions
+    /// 2. push merkle proof to proof by peak from left to right
+    /// 3. push bagged right hand side root
+    pub fn gen_proof(&self, mut pos_list: Vec<u64>) -> Result<MerkleProof<T, M>> {
+        if pos_list.is_empty() {
+            return Err(Error::GenProofForInvalidLeaves);
+        }
+        if self.mmr_size == 1 && pos_list == [0] {
+            return Ok(MerkleProof::new(self.mmr_size, Vec::new()));
+        }
+        // ensure positions is sorted
+        pos_list.sort_unstable();
         let peaks = get_peaks(self.mmr_size);
-        if let Some(rhs_peak_hash) = self.bag_rhs_peaks(peak_pos, &peaks[..])? {
-            proof.push(rhs_peak_hash);
+        let mut proof: Vec<T> = Vec::new();
+        // generate merkle proof for each peaks
+        let mut bagging_track = 0;
+        for peak_pos in peaks {
+            let pos_list: Vec<_> = take_while_vec(&mut pos_list, |&pos| pos <= peak_pos);
+            if pos_list.is_empty() {
+                bagging_track += 1;
+            } else {
+                bagging_track = 0;
+            }
+            self.gen_proof_for_peak(&mut proof, pos_list, peak_pos)?;
         }
-        let lhs_peaks: Vec<_> = peaks
-            .iter()
-            .filter(|&&p| p < peak_pos)
-            .map(|&p| self.batch.get_elem(p))
-            .rev()
-            .collect::<Result<Option<_>>>()?
-            .ok_or(Error::InconsistentStore)?;
-        proof.extend(lhs_peaks);
+
+        // ensure no remain positions
+        if !pos_list.is_empty() {
+            return Err(Error::GenProofForInvalidLeaves);
+        }
+
+        if bagging_track > 1 {
+            let rhs_peaks = proof.split_off(proof.len() - bagging_track);
+            proof.push(self.bag_rhs_peaks(rhs_peaks)?.expect("bagging rhs peaks"));
+        }
+
         Ok(MerkleProof::new(self.mmr_size, proof))
     }
 
@@ -153,7 +218,7 @@ pub struct MerkleProof<T, M> {
     merge: PhantomData<M>,
 }
 
-impl<T: PartialEq + Debug, M: Merge<Item = T>> MerkleProof<T, M> {
+impl<T: PartialEq + Debug + Clone, M: Merge<Item = T>> MerkleProof<T, M> {
     pub fn new(mmr_size: u64, proof: Vec<T>) -> Self {
         MerkleProof {
             mmr_size,
@@ -162,46 +227,201 @@ impl<T: PartialEq + Debug, M: Merge<Item = T>> MerkleProof<T, M> {
         }
     }
 
-    pub fn verify(&self, root: T, mut pos: u64, elem: T) -> Result<bool> {
-        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 mmr_size(&self) -> u64 {
+        self.mmr_size
+    }
 
-        // 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)
-            };
+    pub fn proof_items(&self) -> &[T] {
+        &self.proof
+    }
+
+    pub fn calculate_root(&self, leaves: Vec<(u64, T)>) -> Result<T> {
+        calculate_root::<_, M, _>(leaves, self.mmr_size, self.proof.iter())
+    }
+
+    /// 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,
+        mut leaves: Vec<(u64, T)>,
+        new_pos: u64,
+        new_elem: T,
+        new_mmr_size: u64,
+    ) -> Result<T> {
+        let pos_height = pos_height_in_tree(new_pos);
+        let next_height = pos_height_in_tree(new_pos + 1);
+        if next_height > pos_height {
+            let mut peaks_hashes =
+                calculate_peaks_hashes::<_, M, _>(leaves, self.mmr_size, self.proof.iter())?;
+            let peaks_pos = get_peaks(new_mmr_size);
+            // reverse touched peaks
+            let mut i = 0;
+            while peaks_pos[i] < new_pos {
+                i += 1
+            }
+            peaks_hashes[i..].reverse();
+            calculate_root::<_, M, _>(vec![(new_pos, new_elem)], new_mmr_size, peaks_hashes.iter())
+        } else {
+            leaves.push((new_pos, new_elem));
+            calculate_root::<_, M, _>(leaves, new_mmr_size, self.proof.iter())
         }
-        Ok(root == sum_elem)
+    }
+
+    pub fn verify(&self, root: T, leaves: Vec<(u64, T)>) -> Result<bool> {
+        self.calculate_root(leaves)
+            .map(|calculated_root| calculated_root == root)
     }
 }
+
+fn calculate_peak_root<
+    'a,
+    T: 'a + PartialEq + Debug + Clone,
+    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))
+        .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);
+        }
+        // calculate sibling
+        let next_height = pos_height_in_tree(pos + 1);
+        let (sib_pos, parent_pos) = {
+            let sibling_offset = sibling_offset(height);
+            if next_height > height {
+                // implies pos is right sibling
+                (pos - sibling_offset, pos + 1)
+            } else {
+                // pos is left sibling
+                (pos + sibling_offset, pos + parent_offset(height))
+            }
+        };
+        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)
+        } 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);
+        }
+    }
+    Err(Error::CorruptedProof)
+}
+
+fn calculate_peaks_hashes<
+    'a,
+    T: 'a + PartialEq + Debug + 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>> {
+    // 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
+    leaves.sort_by_key(|(pos, _)| *pos);
+    let peaks = get_peaks(mmr_size);
+
+    let mut peaks_hashes: Vec<T> = Vec::with_capacity(peaks.len() + 1);
+    for peak_pos in peaks {
+        let mut leaves: Vec<_> = take_while_vec(&mut leaves, |(pos, _)| *pos <= peak_pos);
+        let peak_root = if leaves.len() == 1 && leaves[0].0 == peak_pos {
+            // leaf is the peak
+            leaves.remove(0).1
+        } else if leaves.is_empty() {
+            // if empty, means the next proof is a peak root or rhs bagged root
+            if let Some(peak_root) = proof_iter.next() {
+                peak_root.clone()
+            } else {
+                // means that either all right peaks are bagged, or proof is corrupted
+                // so we break loop and check no items left
+                break;
+            }
+        } else {
+            calculate_peak_root::<_, M, _>(leaves, peak_pos, &mut proof_iter)?
+        };
+        peaks_hashes.push(peak_root.clone());
+    }
+
+    // ensure nothing left in leaves
+    if !leaves.is_empty() {
+        return Err(Error::CorruptedProof);
+    }
+
+    // check rhs peaks
+    if let Some(rhs_peaks_hashes) = proof_iter.next() {
+        peaks_hashes.push(rhs_peaks_hashes.clone());
+    }
+    // ensure nothing left in proof_iter
+    if proof_iter.next().is_some() {
+        return Err(Error::CorruptedProof);
+    }
+    Ok(peaks_hashes)
+}
+
+fn bagging_peaks_hashes<'a, T: 'a + PartialEq + Debug + Clone, 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 {
+        let right_peak = peaks_hashes.pop().expect("pop");
+        let left_peak = peaks_hashes.pop().expect("pop");
+        peaks_hashes.push(M::merge(&right_peak, &left_peak));
+    }
+    peaks_hashes.pop().ok_or(Error::CorruptedProof)
+}
+
+/// merkle proof
+/// 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>,
+>(
+    leaves: Vec<(u64, T)>,
+    mmr_size: u64,
+    proof_iter: I,
+) -> Result<T> {
+    let peaks_hashes = calculate_peaks_hashes::<_, M, _>(leaves, mmr_size, proof_iter)?;
+    bagging_peaks_hashes::<_, M>(peaks_hashes)
+}
+
+fn take_while_vec<T, P: Fn(&T) -> bool>(v: &mut Vec<T>, p: P) -> Vec<T> {
+    for i in 0..v.len() {
+        if !p(&v[i]) {
+            return v.drain(..i).collect();
+        }
+    }
+    v.drain(..).collect()
+}

src/mmr_store.rs

@@ -1,4 +1,4 @@
-use crate::Result;
+use crate::{vec::Vec, Result};
 
 #[derive(Default)]
 pub struct MMRBatch<Elem, Store: MMRStore<Elem>> {
@@ -45,7 +45,7 @@ impl<Elem: Clone, Store: MMRStore<Elem>> MMRBatch<Elem, Store> {
 
 impl<Elem, Store: MMRStore<Elem>> IntoIterator for MMRBatch<Elem, Store> {
     type Item = (u64, Vec<Elem>);
-    type IntoIter = ::std::vec::IntoIter<Self::Item>;
+    type IntoIter = crate::vec::IntoIter<Self::Item>;
 
     fn into_iter(self) -> Self::IntoIter {
         self.memory_batch.into_iter()

src/tests/test_accumulate_headers.rs

@@ -156,7 +156,7 @@ impl Prover {
             mmr.get_root()?.serialize(),
             self.headers[later_number as usize].0.chain_root
         );
-        mmr.gen_proof(pos)
+        mmr.gen_proof(vec![pos])
     }
 
     fn get_pos(&self, number: u64) -> u64 {
@@ -188,6 +188,6 @@ fn test_insert_header() {
     let pos = leaf_index_to_pos(h1);
     assert_eq!(pos, prover.get_pos(h1));
     assert_eq!(prove_elem, (&prover.store).get_elem(pos).unwrap().unwrap());
-    let result = proof.verify(root, pos, prove_elem).expect("verify");
+    let result = proof.verify(root, vec![(pos, prove_elem)]).expect("verify");
     assert!(result);
 }

src/tests/test_mmr.rs

@@ -1,9 +1,10 @@
 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::*;
+use rand::{seq::SliceRandom, thread_rng};
 
-fn test_mmr(count: u32, proof_elem: u32) {
+fn test_mmr(count: u32, proof_elem: Vec<u32>) {
     let store = MemStore::default();
     let mut mmr = MMR::<_, MergeNumberHash, _>::new(0, &store);
     let positions: Vec<u64> = (0u32..count)
@@ -11,19 +12,50 @@ fn test_mmr(count: u32, proof_elem: u32) {
         .collect();
     let root = mmr.get_root().expect("get root");
     let proof = mmr
-        .gen_proof(positions[proof_elem as usize])
+        .gen_proof(
+            proof_elem
+                .iter()
+                .map(|elem| positions[*elem as usize])
+                .collect(),
+        )
         .expect("gen proof");
     mmr.commit().expect("commit changes");
     let result = proof
         .verify(
             root,
-            positions[proof_elem as usize],
-            NumberHash::from(proof_elem),
+            proof_elem
+                .iter()
+                .map(|elem| (positions[*elem as usize], NumberHash::from(*elem)))
+                .collect(),
         )
         .unwrap();
     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(vec![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(
+            vec![(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();
@@ -48,38 +80,64 @@ fn test_empty_mmr_root() {
 
 #[test]
 fn test_mmr_3_peaks() {
-    test_mmr(11, 5);
+    test_mmr(11, vec![5]);
 }
 
 #[test]
 fn test_mmr_2_peaks() {
-    test_mmr(10, 5);
+    test_mmr(10, vec![5]);
 }
 
 #[test]
 fn test_mmr_1_peak() {
-    test_mmr(8, 5);
+    test_mmr(8, vec![5]);
 }
 
 #[test]
 fn test_mmr_first_elem_proof() {
-    test_mmr(11, 0);
+    test_mmr(11, vec![0]);
 }
 
 #[test]
 fn test_mmr_last_elem_proof() {
-    test_mmr(11, 10);
+    test_mmr(11, vec![10]);
 }
 
 #[test]
 fn test_mmr_1_elem() {
-    test_mmr(1, 0);
+    test_mmr(1, vec![0]);
 }
 
 #[test]
 fn test_mmr_2_elems() {
-    test_mmr(2, 0);
-    test_mmr(2, 1);
+    test_mmr(2, vec![0]);
+    test_mmr(2, vec![1]);
+}
+
+#[test]
+fn test_mmr_2_leaves_merkle_proof() {
+    test_mmr(11, vec![3, 7]);
+    test_mmr(11, vec![3, 4]);
+}
+
+#[test]
+fn test_mmr_2_sibling_leaves_merkle_proof() {
+    test_mmr(11, vec![4, 5]);
+    test_mmr(11, vec![5, 6]);
+    test_mmr(11, vec![6, 7]);
+}
+
+#[test]
+fn test_mmr_3_leaves_merkle_proof() {
+    test_mmr(11, vec![4, 5, 6]);
+    test_mmr(11, vec![3, 5, 7]);
+    test_mmr(11, vec![3, 4, 5]);
+    test_mmr(100, vec![3, 5, 13]);
+}
+
+#[test]
+fn test_gen_root_from_proof() {
+    test_gen_new_root_from_proof(11);
 }
 
 prop_compose! {
@@ -92,7 +150,17 @@ prop_compose! {
 
 proptest! {
     #[test]
-    fn test_random_mmr((count , elem) in count_elem(500)) {
-        test_mmr(count, elem);
+    fn test_random_mmr(count in 10u32..500u32) {
+        let mut leaves: Vec<u32> = (0..count).collect();
+        let mut rng = thread_rng();
+        leaves.shuffle(&mut rng);
+        let leaves_count = rng.gen_range(1, count - 1);
+        leaves.truncate(leaves_count as usize);
+        test_mmr(count, leaves);
+    }
+
+    #[test]
+    fn test_random_gen_root_with_new_leaf(count in 1u32..500u32) {
+        test_gen_new_root_from_proof(count);
     }
 }

src/util.rs

@@ -1,11 +1,11 @@
-use crate::{MMRStore, Merge, Result, MMR};
-use std::cell::RefCell;
-use std::collections::HashMap;
-use std::fmt::Debug;
-use std::marker::PhantomData;
+use crate::collections::BTreeMap;
+use crate::{vec::Vec, MMRStore, Merge, MerkleProof, Result, MMR};
+use core::cell::RefCell;
+use core::fmt::Debug;
+use core::marker::PhantomData;
 
 #[derive(Clone)]
-pub struct MemStore<T>(RefCell<HashMap<u64, T>>);
+pub struct MemStore<T>(RefCell<BTreeMap<u64, T>>);
 
 impl<T> Default for MemStore<T> {
     fn default() -> Self {
@@ -15,7 +15,7 @@ impl<T> Default for MemStore<T> {
 
 impl<T> MemStore<T> {
     fn new() -> Self {
-        MemStore(RefCell::new(HashMap::new()))
+        MemStore(RefCell::new(Default::default()))
     }
 }
 
@@ -70,4 +70,9 @@ impl<T: Clone + Debug + PartialEq, M: Merge<Item = T>> MemMMR<T, M> {
         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)
+    }
 }