Accept only optimal ancestry proofs

Cargo.toml

@@ -1,10 +1,10 @@
 [package]
 name = "polkadot-ckb-merkle-mountain-range"
-version = "0.6.0"
+version = "0.8.0"
 authors = [
- "Nervos Core Dev <dev@nervos.org>",
- "Parity Technologies <admin@parity.io>",
- "Robert Hambrock <roberthambrock@gmail.com>"
+ "Nervos Core Dev <dev@nervos.org>",
+ "Parity Technologies <admin@parity.io>",
+ "Robert Hambrock <roberthambrock@gmail.com>"
 ]
 
 edition = "2018"
@@ -18,7 +18,7 @@ std = []
 
 [dependencies]
 cfg-if = "1.0"
-itertools = {version = "0.10.5", default-features = false, features = ["use_alloc"]}
+itertools = { version = "0.10.5", default-features = false, features = ["use_alloc"] }
 
 [dev-dependencies]
 faster-hex = "0.8.0"

src/ancestry_proof.rs

@@ -1,9 +1,10 @@
 use crate::collections::VecDeque;
 use crate::helper::{
- get_peak_map, get_peaks, is_descendant_pos, leaf_index_to_pos, parent_offset,
- pos_height_in_tree, sibling_offset,
+ get_peak_map, get_peaks, leaf_index_to_pos, parent_offset, pos_height_in_tree, sibling_offset,
 };
-use crate::mmr::{bagging_peaks_hashes, take_while_vec};
+pub use crate::mmr::bagging_peaks_hashes;
+use crate::mmr::take_while_vec;
+use crate::util::VeqDequeExt;
 use crate::vec::Vec;
 use crate::{Error, Merge, Result};
 use core::fmt::Debug;
@@ -31,15 +32,12 @@ impl<T: PartialEq + Debug + Clone, M: Merge<Item = T>> AncestryProof<T, M> {
  if current_leaves_count <= self.prev_peaks.len() as u64 {
  return Err(Error::CorruptedProof);
  }
- // Test if previous root is correct.
- let prev_peaks_positions = {
- let prev_peaks_positions = get_peaks(self.prev_mmr_size);
- if prev_peaks_positions.len() != self.prev_peaks.len() {
- return Err(Error::CorruptedProof);
- }
- prev_peaks_positions
- };
 
+ // Test if previous root is correct.
+ let prev_peaks_positions = get_peaks(self.prev_mmr_size);
+ if prev_peaks_positions.len() != self.prev_peaks.len() {
+ return Err(Error::CorruptedProof);
+ }
  let calculated_prev_root = bagging_peaks_hashes::<T, M>(self.prev_peaks.clone())?;
  if calculated_prev_root != prev_root {
  return Ok(false);
@@ -74,8 +72,8 @@ impl<T: Clone + PartialEq, M: Merge<Item = T>> NodeMerkleProof<T, M> {
  &self.proof
  }
 
- pub fn calculate_root(&self, leaves: Vec<(u64, T)>) -> Result<T> {
- calculate_root::<_, M, _>(leaves, self.mmr_size, self.proof.iter())
+ pub fn calculate_root(&self, nodes: Vec<(u64, T)>) -> Result<T> {
+ calculate_root::<_, M, _>(nodes, self.mmr_size, self.proof.iter())
  }
 
  /// from merkle proof of leaf n to calculate merkle root of n + 1 leaves.
@@ -136,109 +134,71 @@ impl<T: Clone + PartialEq, M: Merge<Item = T>> NodeMerkleProof<T, M> {
  }
 }
 
+fn get_queue_elem_key<T>(item: &(u64, T, u8)) -> (u8, u64) {
+ let (pos, _item, height) = item;
+ (*height, *pos)
+}
+
 fn calculate_peak_root<
  'a,
  T: 'a + PartialEq,
  M: Merge<Item = T>,
- // I: Iterator<Item = &'a T>
+ I: Iterator<Item = &'a (u64, T)>,
 >(
  nodes: Vec<(u64, T)>,
  peak_pos: u64,
- // proof_iter: &mut I,
+ proof_iter: &mut I,
 ) -> Result<T> {
- debug_assert!(!nodes.is_empty(), "can't be empty");
- // (position, hash, height)
-
- let mut queue: VecDeque<_> = nodes
+ let mut queue: VecDeque<_> = VecDeque::new();
+ for value in nodes
  .into_iter()
  .map(|(pos, item)| (pos, item, pos_height_in_tree(pos)))
-  .collect();
-
- let mut sibs_processed_from_back = Vec::new();
+ {
+ queue.insert_sorted_by(value, get_queue_elem_key);
+ }
 
- // calculate tree root from each items
  while let Some((pos, item, height)) = queue.pop_front() {
  if pos == peak_pos {
  if queue.is_empty() {
  // return root once queue is consumed
  return Ok(item);
- }
- if queue
- .iter()
- .any(|entry| entry.0 == peak_pos && entry.1 != item)
- {
+ } else {
  return Err(Error::CorruptedProof);
  }
- if queue
- .iter()
- .all(|entry| entry.0 == peak_pos && &entry.1 == &item && entry.2 == height)
- {
- // return root if remaining queue consists only of duplicate root entries
- return Ok(item);
- }
- // if queue not empty, push peak back to the end
- queue.push_back((pos, item, height));
- continue;
  }
  // calculate sibling
  let next_height = pos_height_in_tree(pos + 1);
  let (parent_pos, parent_item) = {
  let sibling_offset = sibling_offset(height);
  if next_height > height {
  // implies pos is right sibling
- let (sib_pos, parent_pos) = (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 if Some(&sib_pos) == queue.back().map(|(pos, _, _)| pos) {
- let sibling_item = queue.pop_back().map(|(_, item, _)| item).unwrap();
- M::merge(&sibling_item, &item)?
- }
- // handle special if next queue item is descendant of sibling
- else if let Some(&(front_pos, ..)) = queue.front() {
- if height > 0 && is_descendant_pos(sib_pos, front_pos) {
- queue.push_back((pos, item, height));
- continue;
- } else {
- return Err(Error::CorruptedProof);
- }
  } else {
- return Err(Error::CorruptedProof);
+ let sibling_item = &proof_iter.next().ok_or(Error::CorruptedProof)?.1;
+ M::merge(sibling_item, &item)?
  };
  (parent_pos, parent_item)
  } else {
  // pos is left sibling
- let (sib_pos, parent_pos) = (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 if Some(&sib_pos) == queue.back().map(|(pos, _, _)| pos) {
- let sibling_item = queue.pop_back().map(|(_, item, _)| item).unwrap();
- let parent = M::merge(&item, &sibling_item)?;
- sibs_processed_from_back.push((sib_pos, sibling_item, height));
- parent
- } else if let Some(&(front_pos, ..)) = queue.front() {
- if height > 0 && is_descendant_pos(sib_pos, front_pos) {
- queue.push_back((pos, item, height));
- continue;
- } else {
- return Err(Error::CorruptedProof);
- }
  } else {
- return Err(Error::CorruptedProof);
+ let sibling_item = &proof_iter.next().ok_or(Error::CorruptedProof)?.1;
+ M::merge(&item, sibling_item)?
  };
  (parent_pos, parent_item)
  }
  };
 
  if parent_pos <= peak_pos {
- let parent = (parent_pos, parent_item, height + 1);
- if peak_pos == parent_pos
- || queue.front() != Some(&parent)
- && !sibs_processed_from_back.iter().any(|item| item == &parent)
- {
- queue.push_front(parent)
- };
+ queue.insert_sorted_by((parent_pos, parent_item, height + 1), get_queue_elem_key);
  } else {
  return Err(Error::CorruptedProof);
  }
@@ -254,7 +214,7 @@ fn calculate_peaks_hashes<
 >(
  nodes: Vec<(u64, T)>,
  mmr_size: u64,
- proof_iter: I,
+ mut proof_iter: I,
 ) -> Result<Vec<T>> {
  // special handle the only 1 leaf MMR
  if mmr_size == 1 && nodes.len() == 1 && nodes[0].0 == 0 {
@@ -264,7 +224,6 @@ fn calculate_peaks_hashes<
  // ensure nodes are sorted and unique
  let mut nodes: Vec<_> = nodes
  .into_iter()
- .chain(proof_iter.cloned())
  .sorted_by_key(|(pos, _)| *pos)
  .dedup_by(|a, b| a.0 == b.0)
  .collect();
@@ -283,7 +242,7 @@ fn calculate_peaks_hashes<
  // so we break loop and check no items left
  break;
  } else {
- calculate_peak_root::<_, M>(nodes, peak_pos)?
+ calculate_peak_root::<_, M, _>(nodes, peak_pos, &mut proof_iter)?
  };
  peaks_hashes.push(peak_root.clone());
  }
@@ -294,13 +253,13 @@ fn calculate_peaks_hashes<
  }
 
  // check rhs peaks
- // if let Some((_, rhs_peaks_hashes)) = proof_iter.next() {
- //  peaks_hashes.push(rhs_peaks_hashes.clone());
- // }
+ 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);
- // }
+ if proof_iter.next().is_some() {
+ return Err(Error::CorruptedProof);
+ }
  Ok(peaks_hashes)
 }
 

src/helper.rs

@@ -75,16 +75,6 @@ pub fn get_peak_map(mmr_size: u64) -> u64 {
  peak_map
 }
 
-/// Returns whether `descendant_contender` is a descendant of `ancestor_contender` in a tree of the MMR.
-pub fn is_descendant_pos(ancestor_contender: u64, descendant_contender: u64) -> bool {
- // NOTE: "ancestry" here refers to the hierarchy within an MMR tree, not temporal hierarchy.
- // the descendant needs to have been added to the mmr prior to the ancestor
- descendant_contender <= ancestor_contender
- // the descendant needs to be within the cone of positions descendant from the ancestor
- && descendant_contender
- >= (ancestor_contender + 1 - sibling_offset(pos_height_in_tree(ancestor_contender)))
-}
-
 /// 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

src/mmr.rs

@@ -692,7 +692,7 @@ fn calculate_peaks_hashes<'a, T: 'a + Clone, M: Merge<Item = T>, I: Iterator<Ite
  Ok(peaks_hashes)
 }
 
-pub(crate) fn bagging_peaks_hashes<T, M: Merge<Item = T>>(mut peaks_hashes: Vec<T>) -> Result<T> {
+pub 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 {

src/util.rs

@@ -35,17 +35,33 @@ impl<T> MMRStoreWriteOps<T> for &MemStore<T> {
 
 pub type MemMMR<'a, T, M> = MMR<T, M, &'a MemStore<T>>;
 
-pub trait VeqDequeExt<T: Ord> {
- fn insert_sorted(&mut self, value: T);
+pub trait VeqDequeExt<T> {
+ fn insert_sorted(&mut self, value: T)
+ where
+ T: Ord;
+
+ fn insert_sorted_by<K: Ord, F: FnMut(&T) -> K>(&mut self, value: T, f: F);
 }
 
-impl<T: Ord> VeqDequeExt<T> for VecDeque<T> {
- fn insert_sorted(&mut self, value: T) {
+impl<T> VeqDequeExt<T> for VecDeque<T> {
+ fn insert_sorted(&mut self, value: T)
+ where
+ T: Ord,
+ {
  match self.binary_search(&value) {
  Ok(_pos) => {
  // element already in vector @ `pos`
  }
  Err(pos) => self.insert(pos, value),
  }
  }
+
+ fn insert_sorted_by<K: Ord, F: FnMut(&T) -> K>(&mut self, value: T, mut f: F) {
+ match self.binary_search_by(|x| f(x).cmp(&f(&value))) {
+ Ok(_pos) => {
+ // element already in vector @ `pos`
+ }
+ Err(pos) => self.insert(pos, value),
+ }
+ }
 }