diff --git a/src/merkle/mmr/accumulator.rs b/src/merkle/mmr/accumulator.rs new file mode 100644 index 0000000..f27355f --- /dev/null +++ b/src/merkle/mmr/accumulator.rs @@ -0,0 +1,44 @@ +use super::{super::Vec, MmrProof, Rpo256, Word}; + +#[derive(Debug, Clone, PartialEq)] +pub struct MmrPeaks { + /// The number of leaves is used to differentiate accumulators that have the same number of + /// peaks. This happens because the number of peaks goes up-and-down as the structure is used + /// causing existing trees to be merged and new ones to be created. As an example, every time + /// the MMR has a power-of-two number of leaves there is a single peak. + /// + /// Every tree in the MMR forest has a distinct power-of-two size, this means only the right + /// most tree can have an odd number of elements (1). Additionally this means that the bits in + /// `num_leaves` conveniently encode the size of each individual tree. + /// + /// Examples: + /// + /// Example 1: With 5 leaves, the binary 0b101. The number of set bits is equal the number + /// of peaks, in this case there are 2 peaks. The 0-indexed least-significant position of + /// the bit determines the number of elements of a tree, so the rightmost tree has 2**0 + /// elements and the left most has 2**2. + /// + /// Example 2: With 12 leaves, the binary is 0b1100, this case also has 2 peaks, the + /// leftmost tree has 2**3=8 elements, and the right most has 2**2=4 elements. + pub num_leaves: usize, + + /// All the peaks of every tree in the MMR forest. The peaks are always ordered by number of + /// leaves, starting from the peak with most children, to the one with least. + /// + /// Invariant: The length of `peaks` must be equal to the number of true bits in `num_leaves`. + pub peaks: Vec, +} + +impl MmrPeaks { + /// Hashes the peaks sequentially, compacting it to a single digest + pub fn hash_peaks(&self) -> Word { + Rpo256::hash_elements(&self.peaks.as_slice().concat()).into() + } + + pub fn verify(&self, value: Word, opening: MmrProof) -> bool { + let root = &self.peaks[opening.peak_index()]; + opening + .merkle_path + .verify(opening.relative_pos() as u64, value, root) + } +} diff --git a/src/merkle/mmr/bit.rs b/src/merkle/mmr/bit.rs new file mode 100644 index 0000000..85376ff --- /dev/null +++ b/src/merkle/mmr/bit.rs @@ -0,0 +1,46 @@ +/// Iterate over the bits of a `usize` and yields the bit positions for the true bits. +pub struct TrueBitPositionIterator { + value: usize, +} + +impl TrueBitPositionIterator { + pub fn new(value: usize) -> TrueBitPositionIterator { + TrueBitPositionIterator { value } + } +} + +impl Iterator for TrueBitPositionIterator { + type Item = u32; + + fn next(&mut self) -> Option<::Item> { + // trailing_zeros is computed with the intrinsic cttz. [Rust 1.67.0] x86 uses the `bsf` + // instruction. AArch64 uses the `rbit clz` instructions. + let zeros = self.value.trailing_zeros(); + + if zeros == usize::BITS { + None + } else { + let bit_position = zeros; + let mask = 1 << bit_position; + self.value ^= mask; + Some(bit_position) + } + } +} + +impl DoubleEndedIterator for TrueBitPositionIterator { + fn next_back(&mut self) -> Option<::Item> { + // trailing_zeros is computed with the intrinsic ctlz. [Rust 1.67.0] x86 uses the `bsr` + // instruction. AArch64 uses the `clz` instruction. + let zeros = self.value.leading_zeros(); + + if zeros == usize::BITS { + None + } else { + let bit_position = usize::BITS - zeros - 1; + let mask = 1 << bit_position; + self.value ^= mask; + Some(bit_position) + } + } +} diff --git a/src/merkle/mmr/full.rs b/src/merkle/mmr/full.rs new file mode 100644 index 0000000..92c350c --- /dev/null +++ b/src/merkle/mmr/full.rs @@ -0,0 +1,299 @@ +//! A fully materialized Merkle mountain range (MMR). +//! +//! A MMR is a forest structure, i.e. it is an ordered set of disjoint rooted trees. The trees are +//! ordered by size, from the most to least number of leaves. Every tree is a perfect binary tree, +//! meaning a tree has all its leaves at the same depth, and every inner node has a branch-factor +//! of 2 with both children set. +//! +//! Additionally the structure only supports adding leaves to the right-most tree, the one with the +//! least number of leaves. The structure preserves the invariant that each tree has different +//! depths, i.e. as part of adding adding a new element to the forest the trees with same depth are +//! merged, creating a new tree with depth d+1, this process is continued until the property is +//! restabilished. +use super::bit::TrueBitPositionIterator; +use super::{super::Vec, MmrPeaks, MmrProof, Rpo256, Word}; +use crate::merkle::MerklePath; +use core::fmt::{Display, Formatter}; + +#[cfg(feature = "std")] +use std::error::Error; + +// MMR +// =============================================================================================== + +/// A fully materialized Merkle Mountain Range, with every tree in the forest and all their +/// elements. +/// +/// Since this is a full representation of the MMR, elements are never removed and the MMR will +/// grow roughly `O(2n)` in number of leaf elements. +pub struct Mmr { + /// Refer to the `forest` method documentation for details of the semantics of this value. + pub(super) forest: usize, + + /// Contains every element of the forest. + /// + /// The trees are in postorder sequential representation. This representation allows for all + /// the elements of every tree in the forest to be stored in the same sequential buffer. It + /// also means new elements can be added to the forest, and merging of trees is very cheap with + /// no need to copy elements. + pub(super) nodes: Vec, +} + +#[derive(Debug, PartialEq, Eq, Copy, Clone)] +pub enum MmrError { + InvalidPosition(usize), +} + +impl Display for MmrError { + fn fmt(&self, fmt: &mut Formatter<'_>) -> Result<(), core::fmt::Error> { + match self { + MmrError::InvalidPosition(pos) => write!(fmt, "Mmr does not contain position {pos}"), + } + } +} + +#[cfg(feature = "std")] +impl Error for MmrError {} + +impl Default for Mmr { + fn default() -> Self { + Self::new() + } +} + +impl Mmr { + // CONSTRUCTORS + // ============================================================================================ + + /// Constructor for an empty `Mmr`. + pub fn new() -> Mmr { + Mmr { + forest: 0, + nodes: Vec::new(), + } + } + + // ACCESSORS + // ============================================================================================ + + /// Returns the MMR forest representation. + /// + /// The forest value has the following interpretations: + /// - its value is the number of elements in the forest + /// - bit count corresponds to the number of trees in the forest + /// - each true bit position determines the depth of a tree in the forest + pub const fn forest(&self) -> usize { + self.forest + } + + // FUNCTIONALITY + // ============================================================================================ + + /// Given a leaf position, returns the Merkle path to its corresponding peak. If the position + /// is greater-or-equal than the tree size an error is returned. + /// + /// Note: The leaf position is the 0-indexed number corresponding to the order the leaves were + /// added, this corresponds to the MMR size _prior_ to adding the element. So the 1st element + /// has position 0, the second position 1, and so on. + pub fn open(&self, pos: usize) -> Result { + // find the target tree responsible for the MMR position + let tree_bit = + leaf_to_corresponding_tree(pos, self.forest).ok_or(MmrError::InvalidPosition(pos))?; + let forest_target = 1usize << tree_bit; + + // isolate the trees before the target + let forest_before = self.forest & high_bitmask(tree_bit + 1); + let index_offset = nodes_in_forest(forest_before); + + // find the root + let index = nodes_in_forest(forest_target) - 1; + + // update the value position from global to the target tree + let relative_pos = pos - forest_before; + + // collect the path and the final index of the target value + let (_, path) = + self.collect_merkle_path_and_value(tree_bit, relative_pos, index_offset, index); + + Ok(MmrProof { + forest: self.forest, + position: pos, + merkle_path: MerklePath::new(path), + }) + } + + /// Returns the leaf value at position `pos`. + /// + /// Note: The leaf position is the 0-indexed number corresponding to the order the leaves were + /// added, this corresponds to the MMR size _prior_ to adding the element. So the 1st element + /// has position 0, the second position 1, and so on. + pub fn get(&self, pos: usize) -> Result { + // find the target tree responsible for the MMR position + let tree_bit = + leaf_to_corresponding_tree(pos, self.forest).ok_or(MmrError::InvalidPosition(pos))?; + let forest_target = 1usize << tree_bit; + + // isolate the trees before the target + let forest_before = self.forest & high_bitmask(tree_bit + 1); + let index_offset = nodes_in_forest(forest_before); + + // find the root + let index = nodes_in_forest(forest_target) - 1; + + // update the value position from global to the target tree + let relative_pos = pos - forest_before; + + // collect the path and the final index of the target value + let (value, _) = + self.collect_merkle_path_and_value(tree_bit, relative_pos, index_offset, index); + + Ok(value) + } + + /// Adds a new element to the MMR. + pub fn add(&mut self, el: Word) { + // Note: every node is also a tree of size 1, adding an element to the forest creates a new + // rooted-tree of size 1. This may temporarily break the invariant that every tree in the + // forest has different sizes, the loop below will eagerly merge trees of same size and + // restore the invariant. + self.nodes.push(el); + + let mut left_offset = self.nodes.len().saturating_sub(2); + let mut right = el; + let mut left_tree = 1; + while self.forest & left_tree != 0 { + right = *Rpo256::merge(&[self.nodes[left_offset].into(), right.into()]); + self.nodes.push(right); + + left_offset = left_offset.saturating_sub(nodes_in_forest(left_tree)); + left_tree <<= 1; + } + + self.forest += 1; + } + + /// Returns an accumulator representing the current state of the MMMR. + pub fn accumulator(&self) -> MmrPeaks { + let peaks: Vec = TrueBitPositionIterator::new(self.forest) + .rev() + .map(|bit| nodes_in_forest(1 << bit)) + .scan(0, |offset, el| { + *offset += el; + Some(*offset) + }) + .map(|offset| self.nodes[offset - 1]) + .collect(); + + MmrPeaks { + num_leaves: self.forest, + peaks, + } + } + + // UTILITIES + // ============================================================================================ + + /// Internal function used to collect the Merkle path of a value. + fn collect_merkle_path_and_value( + &self, + tree_bit: u32, + relative_pos: usize, + index_offset: usize, + mut index: usize, + ) -> (Word, Vec) { + // collect the Merkle path + let mut tree_depth = tree_bit as usize; + let mut path = Vec::with_capacity(tree_depth + 1); + while tree_depth > 0 { + let bit = relative_pos & tree_depth; + let right_offset = index - 1; + let left_offset = right_offset - nodes_in_forest(tree_depth); + + // Elements to the right have a higher position because they were + // added later. Therefore when the bit is true the node's path is + // to the right, and its sibling to the left. + let sibling = if bit != 0 { + index = right_offset; + self.nodes[index_offset + left_offset] + } else { + index = left_offset; + self.nodes[index_offset + right_offset] + }; + + tree_depth >>= 1; + path.push(sibling); + } + + // the rest of the codebase has the elements going from leaf to root, adjust it here for + // easy of use/consistency sake + path.reverse(); + + let value = self.nodes[index_offset + index]; + (value, path) + } +} + +impl From for Mmr +where + T: IntoIterator, +{ + fn from(values: T) -> Self { + let mut mmr = Mmr::new(); + for v in values { + mmr.add(v) + } + mmr + } +} + +// UTILITIES +// =============================================================================================== + +/// Given a 0-indexed leaf position and the current forest, return the tree number responsible for +/// the position. +/// +/// Note: +/// The result is a tree position `p`, it has the following interpretations. $p+1$ is the depth of +/// the tree, which corresponds to the size of a Merkle proof for that tree. $2^p$ is equal to the +/// number of leaves in this particular tree. and $2^(p+1)-1$ corresponds to size of the tree. +pub(crate) const fn leaf_to_corresponding_tree(pos: usize, forest: usize) -> Option { + if pos >= forest { + None + } else { + // - each bit in the forest is a unique tree and the bit position its power-of-two size + // - each tree owns a consecutive range of positions equal to its size from left-to-right + // - this means the first tree owns from `0` up to the `2^k_0` first positions, where `k_0` + // is the highest true bit position, the second tree from `2^k_0 + 1` up to `2^k_1` where + // `k_1` is the second higest bit, so on. + // - this means the highest bits work as a category marker, and the position is owned by + // the first tree which doesn't share a high bit with the position + let before = forest & pos; + let after = forest ^ before; + let tree = after.ilog2(); + + Some(tree) + } +} + +/// Return a bitmask for the bits including and above the given position. +pub(crate) const fn high_bitmask(bit: u32) -> usize { + if bit > usize::BITS - 1 { + 0 + } else { + usize::MAX << bit + } +} + +/// Return the total number of nodes of a given forest +/// +/// Panics: +/// +/// This will panic if the forest has size greater than `usize::MAX / 2` +pub(crate) const fn nodes_in_forest(forest: usize) -> usize { + // - the size of a perfect binary tree is $2^{k+1}-1$ or $2*2^k-1$ + // - the forest represents the sum of $2^k$ so a single multiplication is necessary + // - the number of `-1` is the same as the number of trees, which is the same as the number + // bits set + let tree_count = forest.count_ones() as usize; + forest * 2 - tree_count +} diff --git a/src/merkle/mmr/mod.rs b/src/merkle/mmr/mod.rs new file mode 100644 index 0000000..d8903ca --- /dev/null +++ b/src/merkle/mmr/mod.rs @@ -0,0 +1,15 @@ +mod accumulator; +mod bit; +mod full; +mod proof; + +#[cfg(test)] +mod tests; + +use super::{Rpo256, Word}; + +// REEXPORTS +// ================================================================================================ +pub use accumulator::MmrPeaks; +pub use full::Mmr; +pub use proof::MmrProof; diff --git a/src/merkle/mmr/proof.rs b/src/merkle/mmr/proof.rs new file mode 100644 index 0000000..0904b83 --- /dev/null +++ b/src/merkle/mmr/proof.rs @@ -0,0 +1,33 @@ +/// The representation of a single Merkle path. +use super::super::MerklePath; +use super::full::{high_bitmask, leaf_to_corresponding_tree}; + +#[derive(Debug, Clone, PartialEq)] +pub struct MmrProof { + /// The state of the MMR when the MmrProof was created. + pub forest: usize, + + /// The position of the leaf value on this MmrProof. + pub position: usize, + + /// The Merkle opening, starting from the value's sibling up to and excluding the root of the + /// responsible tree. + pub merkle_path: MerklePath, +} + +impl MmrProof { + /// Converts the leaf global position into a local position that can be used to verify the + /// merkle_path. + pub fn relative_pos(&self) -> usize { + let tree_bit = leaf_to_corresponding_tree(self.position, self.forest) + .expect("position must be part of the forest"); + let forest_before = self.forest & high_bitmask(tree_bit + 1); + self.position - forest_before + } + + pub fn peak_index(&self) -> usize { + let root = leaf_to_corresponding_tree(self.position, self.forest) + .expect("position must be part of the forest"); + (self.forest.count_ones() - root - 1) as usize + } +} diff --git a/src/merkle/mmr/tests.rs b/src/merkle/mmr/tests.rs new file mode 100644 index 0000000..577c6c4 --- /dev/null +++ b/src/merkle/mmr/tests.rs @@ -0,0 +1,440 @@ +use super::bit::TrueBitPositionIterator; +use super::full::{high_bitmask, leaf_to_corresponding_tree, nodes_in_forest}; +use super::{super::Vec, Mmr, Rpo256, Word}; +use crate::merkle::{int_to_node, MerklePath}; + +#[test] +fn test_position_equal_or_higher_than_leafs_is_never_contained() { + let empty_forest = 0; + for pos in 1..1024 { + // pos is index, 0 based + // tree is a length counter, 1 based + // so a valid pos is always smaller, not equal, to tree + assert_eq!(leaf_to_corresponding_tree(pos, pos), None); + assert_eq!(leaf_to_corresponding_tree(pos, pos - 1), None); + // and empty forest has no trees, so no position is valid + assert_eq!(leaf_to_corresponding_tree(pos, empty_forest), None); + } +} + +#[test] +fn test_position_zero_is_always_contained_by_the_highest_tree() { + for leaves in 1..1024usize { + let tree = leaves.ilog2(); + assert_eq!(leaf_to_corresponding_tree(0, leaves), Some(tree)); + } +} + +#[test] +fn test_leaf_to_corresponding_tree() { + assert_eq!(leaf_to_corresponding_tree(0, 0b0001), Some(0)); + assert_eq!(leaf_to_corresponding_tree(0, 0b0010), Some(1)); + assert_eq!(leaf_to_corresponding_tree(0, 0b0011), Some(1)); + assert_eq!(leaf_to_corresponding_tree(0, 0b1011), Some(3)); + + // position one is always owned by the left-most tree + assert_eq!(leaf_to_corresponding_tree(1, 0b0010), Some(1)); + assert_eq!(leaf_to_corresponding_tree(1, 0b0011), Some(1)); + assert_eq!(leaf_to_corresponding_tree(1, 0b1011), Some(3)); + + // position two starts as its own root, and then it is merged with the left-most tree + assert_eq!(leaf_to_corresponding_tree(2, 0b0011), Some(0)); + assert_eq!(leaf_to_corresponding_tree(2, 0b0100), Some(2)); + assert_eq!(leaf_to_corresponding_tree(2, 0b1011), Some(3)); + + // position tree is merged on the left-most tree + assert_eq!(leaf_to_corresponding_tree(3, 0b0011), None); + assert_eq!(leaf_to_corresponding_tree(3, 0b0100), Some(2)); + assert_eq!(leaf_to_corresponding_tree(3, 0b1011), Some(3)); + + assert_eq!(leaf_to_corresponding_tree(4, 0b0101), Some(0)); + assert_eq!(leaf_to_corresponding_tree(4, 0b0110), Some(1)); + assert_eq!(leaf_to_corresponding_tree(4, 0b0111), Some(1)); + assert_eq!(leaf_to_corresponding_tree(4, 0b1000), Some(3)); + + assert_eq!(leaf_to_corresponding_tree(12, 0b01101), Some(0)); + assert_eq!(leaf_to_corresponding_tree(12, 0b01110), Some(1)); + assert_eq!(leaf_to_corresponding_tree(12, 0b01111), Some(1)); + assert_eq!(leaf_to_corresponding_tree(12, 0b10000), Some(4)); +} + +#[test] +fn test_high_bitmask() { + assert_eq!(high_bitmask(0), usize::MAX); + assert_eq!(high_bitmask(1), usize::MAX << 1); + assert_eq!(high_bitmask(usize::BITS - 2), 0b11usize.rotate_right(2)); + assert_eq!(high_bitmask(usize::BITS - 1), 0b1usize.rotate_right(1)); + assert_eq!(high_bitmask(usize::BITS), 0, "overflow should be handled"); +} + +#[test] +fn test_nodes_in_forest() { + assert_eq!(nodes_in_forest(0b0000), 0); + assert_eq!(nodes_in_forest(0b0001), 1); + assert_eq!(nodes_in_forest(0b0010), 3); + assert_eq!(nodes_in_forest(0b0011), 4); + assert_eq!(nodes_in_forest(0b0100), 7); + assert_eq!(nodes_in_forest(0b0101), 8); + assert_eq!(nodes_in_forest(0b0110), 10); + assert_eq!(nodes_in_forest(0b0111), 11); + assert_eq!(nodes_in_forest(0b1000), 15); + assert_eq!(nodes_in_forest(0b1001), 16); + assert_eq!(nodes_in_forest(0b1010), 18); + assert_eq!(nodes_in_forest(0b1011), 19); +} + +#[test] +fn test_nodes_in_forest_single_bit() { + assert_eq!(nodes_in_forest(2usize.pow(0)), 2usize.pow(1) - 1); + assert_eq!(nodes_in_forest(2usize.pow(1)), 2usize.pow(2) - 1); + assert_eq!(nodes_in_forest(2usize.pow(2)), 2usize.pow(3) - 1); + assert_eq!(nodes_in_forest(2usize.pow(3)), 2usize.pow(4) - 1); + + for bit in 0..(usize::BITS - 1) { + let size = 2usize.pow(bit + 1) - 1; + assert_eq!(nodes_in_forest(1usize << bit), size); + } +} + +const LEAVES: [Word; 7] = [ + int_to_node(0), + int_to_node(1), + int_to_node(2), + int_to_node(3), + int_to_node(4), + int_to_node(5), + int_to_node(6), +]; + +#[test] +fn test_mmr_simple() { + let mut postorder = Vec::new(); + postorder.push(LEAVES[0]); + postorder.push(LEAVES[1]); + postorder.push(*Rpo256::hash_elements(&[LEAVES[0], LEAVES[1]].concat())); + postorder.push(LEAVES[2]); + postorder.push(LEAVES[3]); + postorder.push(*Rpo256::hash_elements(&[LEAVES[2], LEAVES[3]].concat())); + postorder.push(*Rpo256::hash_elements( + &[postorder[2], postorder[5]].concat(), + )); + postorder.push(LEAVES[4]); + postorder.push(LEAVES[5]); + postorder.push(*Rpo256::hash_elements(&[LEAVES[4], LEAVES[5]].concat())); + postorder.push(LEAVES[6]); + + let mut mmr = Mmr::new(); + assert_eq!(mmr.forest(), 0); + assert_eq!(mmr.nodes.len(), 0); + + mmr.add(LEAVES[0]); + assert_eq!(mmr.forest(), 1); + assert_eq!(mmr.nodes.len(), 1); + assert_eq!(mmr.nodes.as_slice(), &postorder[0..mmr.nodes.len()]); + + let acc = mmr.accumulator(); + assert_eq!(acc.num_leaves, 1); + assert_eq!(acc.peaks, &[postorder[0]]); + + mmr.add(LEAVES[1]); + assert_eq!(mmr.forest(), 2); + assert_eq!(mmr.nodes.len(), 3); + assert_eq!(mmr.nodes.as_slice(), &postorder[0..mmr.nodes.len()]); + + let acc = mmr.accumulator(); + assert_eq!(acc.num_leaves, 2); + assert_eq!(acc.peaks, &[postorder[2]]); + + mmr.add(LEAVES[2]); + assert_eq!(mmr.forest(), 3); + assert_eq!(mmr.nodes.len(), 4); + assert_eq!(mmr.nodes.as_slice(), &postorder[0..mmr.nodes.len()]); + + let acc = mmr.accumulator(); + assert_eq!(acc.num_leaves, 3); + assert_eq!(acc.peaks, &[postorder[2], postorder[3]]); + + mmr.add(LEAVES[3]); + assert_eq!(mmr.forest(), 4); + assert_eq!(mmr.nodes.len(), 7); + assert_eq!(mmr.nodes.as_slice(), &postorder[0..mmr.nodes.len()]); + + let acc = mmr.accumulator(); + assert_eq!(acc.num_leaves, 4); + assert_eq!(acc.peaks, &[postorder[6]]); + + mmr.add(LEAVES[4]); + assert_eq!(mmr.forest(), 5); + assert_eq!(mmr.nodes.len(), 8); + assert_eq!(mmr.nodes.as_slice(), &postorder[0..mmr.nodes.len()]); + + let acc = mmr.accumulator(); + assert_eq!(acc.num_leaves, 5); + assert_eq!(acc.peaks, &[postorder[6], postorder[7]]); + + mmr.add(LEAVES[5]); + assert_eq!(mmr.forest(), 6); + assert_eq!(mmr.nodes.len(), 10); + assert_eq!(mmr.nodes.as_slice(), &postorder[0..mmr.nodes.len()]); + + let acc = mmr.accumulator(); + assert_eq!(acc.num_leaves, 6); + assert_eq!(acc.peaks, &[postorder[6], postorder[9]]); + + mmr.add(LEAVES[6]); + assert_eq!(mmr.forest(), 7); + assert_eq!(mmr.nodes.len(), 11); + assert_eq!(mmr.nodes.as_slice(), &postorder[0..mmr.nodes.len()]); + + let acc = mmr.accumulator(); + assert_eq!(acc.num_leaves, 7); + assert_eq!(acc.peaks, &[postorder[6], postorder[9], postorder[10]]); +} + +#[test] +fn test_mmr_open() { + let mmr: Mmr = LEAVES.into(); + let h01: Word = Rpo256::hash_elements(&LEAVES[0..2].concat()).into(); + let h23: Word = Rpo256::hash_elements(&LEAVES[2..4].concat()).into(); + + // node at pos 7 is the root + assert!( + mmr.open(7).is_err(), + "Element 7 is not in the tree, result should be None" + ); + + // node at pos 6 is the root + let empty: MerklePath = MerklePath::new(vec![]); + let opening = mmr + .open(6) + .expect("Element 6 is contained in the tree, expected an opening result."); + assert_eq!(opening.merkle_path, empty); + assert_eq!(opening.forest, mmr.forest); + assert_eq!(opening.position, 6); + assert!( + mmr.accumulator().verify(LEAVES[6], opening), + "MmrProof should be valid for the current accumulator." + ); + + // nodes 4,5 are detph 1 + let root_to_path = MerklePath::new(vec![LEAVES[4]]); + let opening = mmr + .open(5) + .expect("Element 5 is contained in the tree, expected an opening result."); + assert_eq!(opening.merkle_path, root_to_path); + assert_eq!(opening.forest, mmr.forest); + assert_eq!(opening.position, 5); + assert!( + mmr.accumulator().verify(LEAVES[5], opening), + "MmrProof should be valid for the current accumulator." + ); + + let root_to_path = MerklePath::new(vec![LEAVES[5]]); + let opening = mmr + .open(4) + .expect("Element 4 is contained in the tree, expected an opening result."); + assert_eq!(opening.merkle_path, root_to_path); + assert_eq!(opening.forest, mmr.forest); + assert_eq!(opening.position, 4); + assert!( + mmr.accumulator().verify(LEAVES[4], opening), + "MmrProof should be valid for the current accumulator." + ); + + // nodes 0,1,2,3 are detph 2 + let root_to_path = MerklePath::new(vec![LEAVES[2], h01]); + let opening = mmr + .open(3) + .expect("Element 3 is contained in the tree, expected an opening result."); + assert_eq!(opening.merkle_path, root_to_path); + assert_eq!(opening.forest, mmr.forest); + assert_eq!(opening.position, 3); + assert!( + mmr.accumulator().verify(LEAVES[3], opening), + "MmrProof should be valid for the current accumulator." + ); + + let root_to_path = MerklePath::new(vec![LEAVES[3], h01]); + let opening = mmr + .open(2) + .expect("Element 2 is contained in the tree, expected an opening result."); + assert_eq!(opening.merkle_path, root_to_path); + assert_eq!(opening.forest, mmr.forest); + assert_eq!(opening.position, 2); + assert!( + mmr.accumulator().verify(LEAVES[2], opening), + "MmrProof should be valid for the current accumulator." + ); + + let root_to_path = MerklePath::new(vec![LEAVES[0], h23]); + let opening = mmr + .open(1) + .expect("Element 1 is contained in the tree, expected an opening result."); + assert_eq!(opening.merkle_path, root_to_path); + assert_eq!(opening.forest, mmr.forest); + assert_eq!(opening.position, 1); + assert!( + mmr.accumulator().verify(LEAVES[1], opening), + "MmrProof should be valid for the current accumulator." + ); + + let root_to_path = MerklePath::new(vec![LEAVES[1], h23]); + let opening = mmr + .open(0) + .expect("Element 0 is contained in the tree, expected an opening result."); + assert_eq!(opening.merkle_path, root_to_path); + assert_eq!(opening.forest, mmr.forest); + assert_eq!(opening.position, 0); + assert!( + mmr.accumulator().verify(LEAVES[0], opening), + "MmrProof should be valid for the current accumulator." + ); +} + +#[test] +fn test_mmr_get() { + let mmr: Mmr = LEAVES.into(); + assert_eq!( + mmr.get(0).unwrap(), + LEAVES[0], + "value at pos 0 must correspond" + ); + assert_eq!( + mmr.get(1).unwrap(), + LEAVES[1], + "value at pos 1 must correspond" + ); + assert_eq!( + mmr.get(2).unwrap(), + LEAVES[2], + "value at pos 2 must correspond" + ); + assert_eq!( + mmr.get(3).unwrap(), + LEAVES[3], + "value at pos 3 must correspond" + ); + assert_eq!( + mmr.get(4).unwrap(), + LEAVES[4], + "value at pos 4 must correspond" + ); + assert_eq!( + mmr.get(5).unwrap(), + LEAVES[5], + "value at pos 5 must correspond" + ); + assert_eq!( + mmr.get(6).unwrap(), + LEAVES[6], + "value at pos 6 must correspond" + ); + assert!(mmr.get(7).is_err()); +} + +#[test] +fn test_mmr_invariants() { + let mut mmr = Mmr::new(); + for v in 1..=1028 { + mmr.add(int_to_node(v)); + let accumulator = mmr.accumulator(); + assert_eq!( + v as usize, + mmr.forest(), + "MMR leaf count must increase by one on every add" + ); + assert_eq!( + v as usize, accumulator.num_leaves, + "MMR and its accumulator must match leaves count" + ); + assert_eq!( + accumulator.num_leaves.count_ones() as usize, + accumulator.peaks.len(), + "bits on leaves must match the number of peaks" + ); + + let expected_nodes: usize = TrueBitPositionIterator::new(mmr.forest()) + .map(|bit_pos| nodes_in_forest(1 << bit_pos)) + .sum(); + + assert_eq!( + expected_nodes, + mmr.nodes.len(), + "the sum of every tree size must be equal to the number of nodes in the MMR (forest: {:b})", + mmr.forest(), + ); + } +} + +#[test] +fn test_bit_position_iterator() { + assert_eq!(TrueBitPositionIterator::new(0).count(), 0); + assert_eq!(TrueBitPositionIterator::new(0).rev().count(), 0); + + assert_eq!( + TrueBitPositionIterator::new(1).collect::>(), + vec![0] + ); + assert_eq!( + TrueBitPositionIterator::new(1).rev().collect::>(), + vec![0], + ); + + assert_eq!( + TrueBitPositionIterator::new(2).collect::>(), + vec![1] + ); + assert_eq!( + TrueBitPositionIterator::new(2).rev().collect::>(), + vec![1], + ); + + assert_eq!( + TrueBitPositionIterator::new(3).collect::>(), + vec![0, 1], + ); + assert_eq!( + TrueBitPositionIterator::new(3).rev().collect::>(), + vec![1, 0], + ); + + assert_eq!( + TrueBitPositionIterator::new(0b11010101).collect::>(), + vec![0, 2, 4, 6, 7], + ); + assert_eq!( + TrueBitPositionIterator::new(0b11010101) + .rev() + .collect::>(), + vec![7, 6, 4, 2, 0], + ); +} + +mod property_tests { + use super::leaf_to_corresponding_tree; + use proptest::prelude::*; + + proptest! { + #[test] + fn test_last_position_is_always_contained_in_the_last_tree(leaves in any::().prop_filter("cant have an empty tree", |v| *v != 0)) { + let last_pos = leaves - 1; + let lowest_bit = leaves.trailing_zeros(); + + assert_eq!( + leaf_to_corresponding_tree(last_pos, leaves), + Some(lowest_bit), + ); + } + } + + proptest! { + #[test] + fn test_contained_tree_is_always_power_of_two((leaves, pos) in any::().prop_flat_map(|v| (Just(v), 0..v))) { + let tree = leaf_to_corresponding_tree(pos, leaves).expect("pos is smaller than leaves, there should always be a corresponding tree"); + let mask = 1usize << tree; + + assert!(tree < usize::BITS, "the result must be a bit in usize"); + assert!(mask & leaves != 0, "the result should be a tree in leaves"); + } + } +} diff --git a/src/merkle/mod.rs b/src/merkle/mod.rs index 0b82752..6d53a71 100644 --- a/src/merkle/mod.rs +++ b/src/merkle/mod.rs @@ -5,6 +5,8 @@ use super::{ }; use core::fmt; +// REEXPORTS +// ================================================================================================ mod index; pub use index::NodeIndex; @@ -20,6 +22,9 @@ pub use path_set::MerklePathSet; mod simple_smt; pub use simple_smt::SimpleSmt; +mod mmr; +pub use mmr::{Mmr, MmrPeaks}; + // ERRORS // ================================================================================================ diff --git a/src/merkle/simple_smt/tests.rs b/src/merkle/simple_smt/tests.rs index 2096fd1..e449014 100644 --- a/src/merkle/simple_smt/tests.rs +++ b/src/merkle/simple_smt/tests.rs @@ -1,8 +1,7 @@ use super::{ - super::{MerkleTree, RpoDigest, SimpleSmt}, + super::{int_to_node, MerkleTree, RpoDigest, SimpleSmt}, NodeIndex, Rpo256, Vec, Word, }; -use crate::{Felt, FieldElement}; use core::iter; use proptest::prelude::*; use rand_utils::prng_array; @@ -275,7 +274,3 @@ fn compute_internal_nodes() -> (Word, Word, Word) { (root.into(), node2.into(), node3.into()) } - -const fn int_to_node(value: u64) -> Word { - [Felt::new(value), Felt::ZERO, Felt::ZERO, Felt::ZERO] -}