feat: verify a old root and all incremental items after it

src/mmr.rs

@@ -5,8 +5,11 @@
 //! 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_peak_map, get_peaks, parent_offset, pos_height_in_tree, sibling_offset};
+use crate::collections::{HashMap, VecDeque};
+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;
@@ -291,6 +294,75 @@ impl<T: Clone + PartialEq, 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.
+    // TODO refactor: Use leaves count to instead of mmr size in MerkleProof.
+    //
+    // Reason:
+    // - Calculate mmr size with leaves count is easy; but the reverse is difficult and it may
+    //   cause errors.
+    // - The `prev_leaves_count` in this method could be removed, it could be calculated via
+    //   `current_leaves_count - incremental.len()`.
+    pub fn verify_incremental(
+        &self,
+        root: T,
+        prev_leaves_count: u64,
+        prev_root: T,
+        incremental: Vec<T>,
+    ) -> Result<bool> {
+        let (positions, leaves) = incremental.into_iter().enumerate().fold(
+            (Vec::new(), Vec::new()),
+            |(mut positions, mut leaves), (index, leaf)| {
+                let pos = leaf_index_to_pos(prev_leaves_count + index as u64);
+                positions.push(pos);
+                leaves.push((pos, leaf));
+                (positions, leaves)
+            },
+        );
+        // Test if previous root is correct.
+        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 Ok(false);
+            }
+            prev_peaks_positions
+        };
+        let proof_items_positions = calculate_positions_of_proof_items_via_peaks_hashes(
+            positions,
+            self.mmr_size,
+            self.proof.len(),
+        )?;
+        if proof_items_positions.len() != self.proof.len() {
+            return Err(Error::CorruptedProof);
+        }
+        let map = proof_items_positions
+            .into_iter()
+            .enumerate()
+            .map(|(index, pos)| (pos, index))
+            .collect::<HashMap<_, _>>();
+        let prev_peaks_res = prev_peaks_positions
+            .iter()
+            .map(|pos| map.get(pos).map(|index| self.proof[*index].clone()))
+            .collect::<Option<_>>();
+        if let Some(prev_peaks) = prev_peaks_res {
+            let calculated_prev_root = bagging_peaks_hashes::<T, M>(prev_peaks)?;
+            if calculated_prev_root != prev_root {
+                return Ok(false);
+            }
+        } else {
+            return Ok(false);
+        }
+        // Test if incremental leaves are correct.
+        self.verify(root, leaves)
+    }
 }
 
 fn calculate_peak_root<'a, T: 'a, M: Merge<Item = T>, I: Iterator<Item = &'a T>>(
@@ -356,6 +428,70 @@ fn calculate_peak_root<'a, T: 'a, M: Merge<Item = T>, I: Iterator<Item = &'a T>>
     Err(Error::CorruptedProof)
 }
 
+// Simplified version of `calculate_peak_root`.
+fn calculate_positions_of_proof_items_via_peak_root(
+    leaves: Vec<u64>,
+    peak_pos: u64,
+    proof_items_count: &mut usize,
+    return_data: &mut Vec<u64>,
+) -> Result<()> {
+    debug_assert!(!leaves.is_empty(), "can't be empty");
+
+    let mut queue: VecDeque<_> = leaves.into_iter().map(|pos| (pos, 0)).collect();
+
+    // calculate tree root from each items
+    while let Some((pos, height)) = queue.pop_front() {
+        if pos == peak_pos {
+            if queue.is_empty() {
+                return Ok(());
+            } else {
+                return Err(Error::CorruptedProof);
+            }
+        }
+        // calculate sibling
+        let next_height = pos_height_in_tree(pos + 1);
+        let parent_pos = {
+            let sibling_offset = sibling_offset(height);
+            if next_height > height {
+                // implies pos is right sibling
+                let sib_pos = pos - sibling_offset;
+                let parent_pos = pos + 1;
+                if Some(&sib_pos) == queue.front().map(|(pos, _)| pos) {
+                    let _ = queue.pop_front();
+                } else {
+                    if *proof_items_count == 0 {
+                        return Err(Error::CorruptedProof);
+                    }
+                    *proof_items_count -= 1;
+                    return_data.push(sib_pos);
+                }
+                parent_pos
+            } else {
+                // pos is left sibling
+                let sib_pos = pos + sibling_offset;
+                let parent_pos = pos + parent_offset(height);
+                if Some(&sib_pos) == queue.front().map(|(pos, _)| pos) {
+                    let _ = queue.pop_front();
+                } else {
+                    if *proof_items_count == 0 {
+                        return Err(Error::CorruptedProof);
+                    }
+                    *proof_items_count -= 1;
+                    return_data.push(sib_pos);
+                }
+                parent_pos
+            }
+        };
+
+        if parent_pos <= peak_pos {
+            queue.push_back((parent_pos, height + 1))
+        } else {
+            return Err(Error::CorruptedProof);
+        }
+    }
+    Err(Error::CorruptedProof)
+}
+
 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,
@@ -411,6 +547,64 @@ fn calculate_peaks_hashes<'a, T: 'a + Clone, M: Merge<Item = T>, I: Iterator<Ite
     Ok(peaks_hashes)
 }
 
+// Simplified version of `calculate_peaks_hashes`.
+fn calculate_positions_of_proof_items_via_peaks_hashes(
+    mut leaves: Vec<u64>,
+    mmr_size: u64,
+    mut proof_items_count: usize,
+) -> Result<Vec<u64>> {
+    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 {
+        return Ok(Default::default());
+    }
+    // ensure leaves are sorted and unique
+    leaves.sort_by_key(|pos| *pos);
+    leaves.dedup_by(|a, b| a == b);
+    let peaks = get_peaks(mmr_size);
+
+    let mut peaks_hashes: Vec<u64> = Vec::with_capacity(peaks.len() + 1);
+    for peak_pos in peaks {
+        let mut leaves: Vec<_> = take_while_vec(&mut leaves, |pos| *pos <= peak_pos);
+        if leaves.len() == 1 && leaves[0] == peak_pos {
+            // leaf is the peak
+            let _ = leaves.remove(0);
+        } else if leaves.is_empty() {
+            // if empty, means the next proof is a peak root or rhs bagged root
+            if proof_items_count > 0 {
+                proof_items_count -= 1;
+                peaks_hashes.push(peak_pos);
+            } 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_positions_of_proof_items_via_peak_root(
+                leaves,
+                peak_pos,
+                &mut proof_items_count,
+                &mut peaks_hashes,
+            )?;
+        };
+    }
+
+    // ensure nothing left in leaves
+    if !leaves.is_empty() {
+        return Err(Error::CorruptedProof);
+    }
+
+    // check rhs peaks
+    if proof_items_count > 0 {
+        // With lots of fuzz test, this condition is never entered.
+        unreachable!();
+    }
+    Ok(peaks_hashes)
+}
+
 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).

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_incremental.rs

@@ -0,0 +1,54 @@
+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 prev_leaves_count = u64::from(curr);
+        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_leaves_count, prev_root, leaves)
+            .unwrap();
+        assert!(
+            result,
+            "start: {}, steps: {}, turn: {}, curr: {}",
+            start, steps, turn, curr
+        );
+    }
+}