Init merkle-mountain-range

The MMR code is copy from https://github.com/nervosnetwork/ckb/pull/1598

src/mmr.rs

@@ -0,0 +1,207 @@
+//! Merkle Mountain Range
+//!
+//! references:
+//! 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::helper::{get_peaks, parent_offset, pos_height_in_tree, sibling_offset};
+use crate::mmr_store::{MMRBatch, MMRStore};
+use crate::{Error, Merge, Result};
+use std::borrow::Cow;
+use std::fmt::Debug;
+use std::marker::PhantomData;
+
+pub struct MMR<T, M, S: MMRStore<T>> {
+    mmr_size: u64,
+    batch: MMRBatch<T, S>,
+    merge: PhantomData<M>,
+}
+
+impl<'a, T: Clone + PartialEq + Debug, M: Merge<Item = T>, S: MMRStore<T>> MMR<T, M, S> {
+    pub fn new(mmr_size: u64, store: S) -> Self {
+        MMR {
+            mmr_size,
+            batch: MMRBatch::new(store),
+            merge: PhantomData,
+        }
+    }
+
+    // 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
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.mmr_size == 0
+    }
+
+    // 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 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 {
+            pos += 1;
+            let left_pos = pos - parent_offset(height);
+            let right_pos = left_pos + sibling_offset(height);
+            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);
+            elems.push(parent_elem);
+            height += 1
+        }
+        // store hashes
+        self.batch.append(elem_pos, elems);
+        // update mmr_size
+        self.mmr_size = pos + 1;
+        Ok(elem_pos)
+    }
+
+    /// get_root
+    pub fn get_root(&self) -> Result<T> {
+        if self.mmr_size == 0 {
+            return Err(Error::GetRootOnEmpty);
+        } 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)
+    }
+
+    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));
+        }
+        Ok(rhs_peak_elems.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;
+            }
+            proof.push(
+                self.batch
+                    .get_elem(sib_pos)?
+                    .ok_or(Error::InconsistentStore)?,
+            );
+            pos = next_pos;
+            height += 1;
+        }
+        // now we get peak merkle proof
+        let peak_pos = pos;
+        // calculate bagging proof
+        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 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);
+        Ok(MerkleProof::new(self.mmr_size, proof))
+    }
+
+    pub fn commit(self) -> Result<()> {
+        self.batch.commit()
+    }
+}
+
+#[derive(Debug)]
+pub struct MerkleProof<T, M> {
+    mmr_size: u64,
+    proof: Vec<T>,
+    merge: PhantomData<M>,
+}
+
+impl<T: PartialEq + Debug, M: Merge<Item = T>> MerkleProof<T, M> {
+    pub fn new(mmr_size: u64, proof: Vec<T>) -> Self {
+        MerkleProof {
+            mmr_size,
+            proof,
+            merge: PhantomData,
+        }
+    }
+
+    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.contains(&pos) {
+            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
+        }
+
+        // bagging peaks
+        // bagging with left peaks if pos is last peak
+        let mut bagging_left = Some(&pos) == peaks.last();
+        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)
+            };
+        }
+        Ok(root == sum_elem)
+    }
+}