Merge pull request #16 from yangby-cryptape/pr/check-before-merge

feat: check nodes (or peaks) before merge them

benches/mmr_benchmark.rs

@@ -31,13 +31,13 @@ struct MergeNumberHash;
 
 impl Merge for MergeNumberHash {
     type Item = NumberHash;
-    fn merge(lhs: &Self::Item, rhs: &Self::Item) -> Self::Item {
+    fn merge(lhs: &Self::Item, rhs: &Self::Item) -> Result<Self::Item> {
         let mut hasher = new_blake2b();
         let mut hash = [0u8; 32];
         hasher.update(&lhs.0);
         hasher.update(&rhs.0);
         hasher.finalize(&mut hash);
-        NumberHash(hash.to_vec().into())
+        Ok(NumberHash(hash.to_vec().into()))
     }
 }
 

src/error.rs

@@ -9,6 +9,9 @@ pub enum Error {
     CorruptedProof,
     /// The leaves is an empty list, or beyond the mmr range
     GenProofForInvalidLeaves,
+
+    /// The two nodes couldn't merge into one.
+    MergeError(crate::string::String),
 }
 
 impl core::fmt::Display for Error {
@@ -20,6 +23,7 @@ impl core::fmt::Display for Error {
             StoreError(msg) => write!(f, "Store error {}", msg)?,
             CorruptedProof => write!(f, "Corrupted proof")?,
             GenProofForInvalidLeaves => write!(f, "Generate proof ofr invalid leaves")?,
+            MergeError(msg) => write!(f, "Merge error {}", msg)?,
         }
         Ok(())
     }

src/merge.rs

@@ -1,9 +1,11 @@
+use crate::Result;
+
 pub trait Merge {
     type Item;
 
-    fn merge(left: &Self::Item, right: &Self::Item) -> Self::Item;
+    fn merge(left: &Self::Item, right: &Self::Item) -> Result<Self::Item>;
 
-    fn merge_peaks(peak1: &Self::Item, peak2: &Self::Item) -> Self::Item {
+    fn merge_peaks(peak1: &Self::Item, peak2: &Self::Item) -> Result<Self::Item> {
         Self::merge(peak1, peak2)
     }
 }

src/mmr.rs

@@ -63,7 +63,7 @@ impl<'a, T: Clone + PartialEq + Debug, M: Merge<Item = T>, S: MMRStore<T>> MMR<T
             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);
+            let parent_elem = M::merge(&left_elem, &right_elem)?;
             elems.push(parent_elem);
             height += 1
         }
@@ -96,7 +96,7 @@ impl<'a, T: Clone + PartialEq + Debug, M: Merge<Item = T>, S: MMRStore<T>> MMR<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_peaks(&right_peak, &left_peak));
+            rhs_peaks.push(M::merge_peaks(&right_peak, &left_peak)?);
         }
         Ok(rhs_peaks.pop())
     }
@@ -320,7 +320,7 @@ fn calculate_peak_root<
             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));
@@ -394,7 +394,7 @@ fn bagging_peaks_hashes<'a, T: 'a + PartialEq + Debug + Clone, M: Merge<Item = T
     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_peaks(&right_peak, &left_peak));
+        peaks_hashes.push(M::merge_peaks(&right_peak, &left_peak)?);
     }
     peaks_hashes.pop().ok_or(Error::CorruptedProof)
 }

src/tests/mod.rs

@@ -3,7 +3,7 @@ mod test_helper;
 mod test_mmr;
 mod test_sequence;
 
-use crate::Merge;
+use crate::{Merge, Result};
 use blake2b_rs::{Blake2b, Blake2bBuilder};
 use bytes::Bytes;
 
@@ -27,12 +27,12 @@ struct MergeNumberHash;
 
 impl Merge for MergeNumberHash {
     type Item = NumberHash;
-    fn merge(lhs: &Self::Item, rhs: &Self::Item) -> Self::Item {
+    fn merge(lhs: &Self::Item, rhs: &Self::Item) -> Result<Self::Item> {
         let mut hasher = new_blake2b();
         let mut hash = [0u8; 32];
         hasher.update(&lhs.0);
         hasher.update(&rhs.0);
         hasher.finalize(&mut hash);
-        NumberHash(hash.to_vec().into())
+        Ok(NumberHash(hash.to_vec().into()))
     }
 }

src/tests/test_accumulate_headers.rs

@@ -72,17 +72,17 @@ struct MergeHashWithTD;
 
 impl Merge for MergeHashWithTD {
     type Item = HashWithTD;
-    fn merge(lhs: &Self::Item, rhs: &Self::Item) -> Self::Item {
+    fn merge(lhs: &Self::Item, rhs: &Self::Item) -> Result<Self::Item> {
         let mut hasher = new_blake2b();
         let mut hash = [0u8; 32];
         hasher.update(&lhs.serialize());
         hasher.update(&rhs.serialize());
         hasher.finalize(&mut hash);
         let td = lhs.td + rhs.td;
-        HashWithTD {
+        Ok(HashWithTD {
             hash: hash.to_vec().into(),
             td,
-        }
+        })
     }
 }
 

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, MMR};
+use crate::{util::MemStore, Merge, Result, MMR};
 
 #[derive(Eq, PartialEq, Clone, Default)]
 struct NumberRange {
@@ -42,13 +42,13 @@ impl NumberRange {
 
 impl Merge for MergeNumberRange {
     type Item = NumberRange;
-    fn merge(lhs: &Self::Item, rhs: &Self::Item) -> Self::Item {
-        Self::Item {
+    fn merge(lhs: &Self::Item, rhs: &Self::Item) -> Result<Self::Item> {
+        Ok(Self::Item {
             start: lhs.start,
             end: rhs.end,
-        }
+        })
     }
-    fn merge_peaks(lhs: &Self::Item, rhs: &Self::Item) -> Self::Item {
+    fn merge_peaks(lhs: &Self::Item, rhs: &Self::Item) -> Result<Self::Item> {
         Self::merge(rhs, lhs)
     }
 }