From 40612859810a70e45a5667fbcaef8227192c406d Mon Sep 17 00:00:00 2001 From: "Augusto F. Hack" Date: Tue, 17 Oct 2023 17:38:18 +0200 Subject: [PATCH] mmr: added partial mmr --- src/merkle/mmr/accumulator.rs | 33 ++- src/merkle/mmr/delta.rs | 16 ++ src/merkle/mmr/error.rs | 35 ++++ src/merkle/mmr/full.rs | 179 +++++++++------- src/merkle/mmr/inorder.rs | 136 ++++++++++++ src/merkle/mmr/mod.rs | 38 ++++ src/merkle/mmr/partial.rs | 377 ++++++++++++++++++++++++++++++++++ src/merkle/mmr/proof.rs | 2 +- src/merkle/mmr/tests.rs | 328 +++++++++++++++++++++++++---- src/merkle/mod.rs | 2 +- src/merkle/path.rs | 14 ++ 11 files changed, 1043 insertions(+), 117 deletions(-) create mode 100644 src/merkle/mmr/delta.rs create mode 100644 src/merkle/mmr/error.rs create mode 100644 src/merkle/mmr/inorder.rs create mode 100644 src/merkle/mmr/partial.rs diff --git a/src/merkle/mmr/accumulator.rs b/src/merkle/mmr/accumulator.rs index a610fe72..82fec2d1 100644 --- a/src/merkle/mmr/accumulator.rs +++ b/src/merkle/mmr/accumulator.rs @@ -1,6 +1,6 @@ use super::{ super::{RpoDigest, Vec, ZERO}, - Felt, MmrProof, Rpo256, Word, + Felt, MmrError, MmrProof, Rpo256, Word, }; #[derive(Debug, Clone, PartialEq)] @@ -9,9 +9,9 @@ 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. + /// 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 + /// 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 (e.g. `1`). Additionally this means that the bits in /// `num_leaves` conveniently encode the size of each individual tree. /// @@ -23,16 +23,37 @@ pub struct MmrPeaks { /// elements and the left most has `2**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, + num_leaves: usize, - /// All the peaks of every tree in the MMR forest. The peaks are always ordered by number of + /// 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, + peaks: Vec, } impl MmrPeaks { + pub fn new(num_leaves: usize, peaks: Vec) -> Result { + if num_leaves.count_ones() as usize != peaks.len() { + return Err(MmrError::InvalidPeaks); + } + + Ok(Self { num_leaves, peaks }) + } + + // ACCESSORS + // -------------------------------------------------------------------------------------------- + + /// Returns a count of the [Mmr]'s leaves. + pub fn num_leaves(&self) -> usize { + self.num_leaves + } + + /// Returns the current peaks of the [Mmr]. + pub fn peaks(&self) -> &Vec { + &self.peaks + } + /// Hashes the peaks. /// /// The procedure will: diff --git a/src/merkle/mmr/delta.rs b/src/merkle/mmr/delta.rs new file mode 100644 index 00000000..4bd9961e --- /dev/null +++ b/src/merkle/mmr/delta.rs @@ -0,0 +1,16 @@ +use super::super::{RpoDigest, Vec}; + +/// Container for the update data of a [PartialMmr] +#[derive(Debug)] +pub struct MmrDelta { + /// The new version of the [Mmr] + pub forest: usize, + + /// Update data. + /// + /// The data is packed as follows: + /// 1. All the elements needed to perform authentication path updates. These are the right + /// siblings required to perform tree merges on the [PartialMmr]. + /// 2. The new peaks. + pub data: Vec, +} diff --git a/src/merkle/mmr/error.rs b/src/merkle/mmr/error.rs new file mode 100644 index 00000000..8c090501 --- /dev/null +++ b/src/merkle/mmr/error.rs @@ -0,0 +1,35 @@ +use crate::merkle::MerkleError; +use core::fmt::{Display, Formatter}; + +#[cfg(feature = "std")] +use std::error::Error; + +#[derive(Debug, PartialEq, Eq, Clone)] +pub enum MmrError { + InvalidPosition(usize), + InvalidPeaks, + InvalidPeak, + InvalidUpdate, + UnknownPeak, + MerkleError(MerkleError), +} + +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}"), + MmrError::InvalidPeaks => write!(fmt, "Invalid peaks count"), + MmrError::InvalidPeak => { + write!(fmt, "Peak values does not match merkle path computed root") + } + MmrError::InvalidUpdate => write!(fmt, "Invalid mmr update"), + MmrError::UnknownPeak => { + write!(fmt, "Peak not in Mmr") + } + MmrError::MerkleError(err) => write!(fmt, "{}", err), + } + } +} + +#[cfg(feature = "std")] +impl Error for MmrError {} diff --git a/src/merkle/mmr/full.rs b/src/merkle/mmr/full.rs index af7c212e..89b6027e 100644 --- a/src/merkle/mmr/full.rs +++ b/src/merkle/mmr/full.rs @@ -13,12 +13,8 @@ use super::{ super::{InnerNodeInfo, MerklePath, RpoDigest, Vec}, bit::TrueBitPositionIterator, - MmrPeaks, MmrProof, Rpo256, + leaf_to_corresponding_tree, MmrDelta, MmrError, MmrPeaks, MmrProof, Rpo256, }; -use core::fmt::{Display, Formatter}; - -#[cfg(feature = "std")] -use std::error::Error; // MMR // =============================================================================================== @@ -43,22 +39,6 @@ pub struct Mmr { 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() @@ -100,21 +80,16 @@ impl Mmr { // 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); + let (_, path) = self.collect_merkle_path_and_value(tree_bit, relative_pos, index_offset); Ok(MmrProof { forest: self.forest, @@ -132,21 +107,16 @@ impl Mmr { // 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); + let (value, _) = self.collect_merkle_path_and_value(tree_bit, relative_pos, index_offset); Ok(value) } @@ -185,7 +155,82 @@ impl Mmr { .map(|offset| self.nodes[offset - 1]) .collect(); - MmrPeaks { num_leaves: self.forest, peaks } + // Safety: the invariant is maintained by the [Mmr] + MmrPeaks::new(self.forest, peaks).unwrap() + } + + /// Compute the required update to `original_forest`. + /// + /// The result is a packed sequence of the authentication elements required to update the trees + /// that have been merged together, followed by the new peaks of the [Mmr]. + pub fn updates(&self, original_forest: usize) -> Result { + if original_forest > self.forest { + return Err(MmrError::InvalidPeaks); + } + + if original_forest == self.forest { + return Ok(MmrDelta { forest: self.forest, data: Vec::new() }); + } + + let mut result = Vec::new(); + + // Find the largest tree in this [Mmr] which is new to `original_forest`. + let candidate_trees = self.forest ^ original_forest; + let mut new_high = 1 << candidate_trees.ilog2(); + + // Collect authentication nodes used for to tree merges + // ---------------------------------------------------------------------------------------- + + // Find the trees from `original_forest` that have been merged into `new_high`. + let mut merges = original_forest & (new_high - 1); + + // Find the peaks that are common to `original_forest` and this [Mmr] + let common_trees = original_forest ^ merges; + + if merges != 0 { + // Skip the smallest trees unknown to `original_forest`. + let mut target = 1 << merges.trailing_zeros(); + + // Collect siblings required to computed the merged tree's peak + while target < new_high { + // Computes the offset to the smallest know peak + // - common_trees: peaks unchanged in the current update, target comes after these. + // - merges: peaks that have not been merged so far, target comes after these. + // - target: tree from which to load the sibling. On the first iteration this is a + // value known by the partial mmr, on subsequent iterations this value is to be + // computed from the known peaks and provided authentication nodes. + let known = nodes_in_forest(common_trees | merges | target); + let sibling = nodes_in_forest(target); + result.push(self.nodes[known + sibling - 1]); + + // Update the target and account for tree merges + target <<= 1; + while merges & target != 0 { + target <<= 1; + } + // Remove the merges done so far + merges ^= merges & (target - 1); + } + } else { + // The new high tree may not be the result of any merges, if it is smaller than all the + // trees of `original_forest`. + new_high = 0; + } + + // Collect the new [Mmr] peaks + // ---------------------------------------------------------------------------------------- + + let mut new_peaks = self.forest ^ common_trees ^ new_high; + let old_peaks = self.forest ^ new_peaks; + let mut offset = nodes_in_forest(old_peaks); + while new_peaks != 0 { + let target = 1 << new_peaks.ilog2(); + offset += nodes_in_forest(target); + result.push(self.nodes[offset - 1]); + new_peaks ^= target; + } + + Ok(MmrDelta { forest: self.forest, data: result }) } /// An iterator over inner nodes in the MMR. The order of iteration is unspecified. @@ -202,36 +247,52 @@ impl Mmr { // ============================================================================================ /// Internal function used to collect the Merkle path of a value. + /// + /// The arguments are relative to the target tree. To compute the opening of the second leaf + /// for a tree with depth 2 in the forest `0b110`: + /// + /// - `tree_bit`: Depth of the target tree, e.g. 2 for the smallest tree. + /// - `relative_pos`: 0-indexed leaf position in the target tree, e.g. 1 for the second leaf. + /// - `index_offset`: Node count prior to the target tree, e.g. 7 for the tree of depth 3. fn collect_merkle_path_and_value( &self, tree_bit: u32, relative_pos: usize, index_offset: usize, - mut index: usize, ) -> (RpoDigest, 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; + // see documentation of `leaf_to_corresponding_tree` for details + let tree_depth = (tree_bit + 1) as usize; + let mut path = Vec::with_capacity(tree_depth); + + // The tree walk below goes from the root to the leaf, compute the root index to start + let mut forest_target = 1usize << tree_bit; + let mut index = nodes_in_forest(forest_target) - 1; + + // Loop until the leaf is reached + while forest_target > 1 { + // Update the depth of the tree to correspond to a subtree + forest_target >>= 1; + + // compute the indeces of the right and left subtrees based on the post-order let right_offset = index - 1; - let left_offset = right_offset - nodes_in_forest(tree_depth); + let left_offset = right_offset - nodes_in_forest(forest_target); - // 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 { + let left_or_right = relative_pos & forest_target; + let sibling = if left_or_right != 0 { + // going down the right subtree, the right child becomes the new root index = right_offset; + // and the left child is the authentication self.nodes[index_offset + left_offset] } else { index = left_offset; self.nodes[index_offset + right_offset] }; - tree_depth >>= 1; path.push(sibling); } + debug_assert!(path.len() == tree_depth - 1); + // the rest of the codebase has the elements going from leaf to root, adjust it here for // easy of use/consistency sake path.reverse(); @@ -335,32 +396,6 @@ impl<'a> Iterator for MmrNodes<'a> { // 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 { diff --git a/src/merkle/mmr/inorder.rs b/src/merkle/mmr/inorder.rs new file mode 100644 index 00000000..2f607124 --- /dev/null +++ b/src/merkle/mmr/inorder.rs @@ -0,0 +1,136 @@ +//! Index for nodes of a binary tree based on an in-order tree walk. +//! +//! In-order walks have the parent node index split its left and right subtrees. All the left +//! children have indexes lower than the parent, meanwhile all the right subtree higher indexes. +//! This property makes it is easy to compute changes to the index by adding or subtracting the +//! leaves count. +use core::num::NonZeroUsize; + +/// Index of nodes in a perfectly balanced binary tree based on an in-order tree walk. +#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)] +pub struct InOrderIndex { + idx: usize, +} + +impl InOrderIndex { + /// Constructor for a new [InOrderIndex]. + pub fn new(idx: NonZeroUsize) -> InOrderIndex { + InOrderIndex { idx: idx.get() } + } + + /// Constructs an index from a leaf position. + /// + /// Panics: + /// + /// If `leaf` is higher than or equal to `usize::MAX / 2`. + pub fn from_leaf_pos(leaf: usize) -> InOrderIndex { + // Convert the position from 0-indexed to 1-indexed, since the bit manipulation in this + // implementation only works 1-indexed counting. + let pos = leaf + 1; + InOrderIndex { idx: pos * 2 - 1 } + } + + /// True if the index is pointing at a leaf. + /// + /// Every odd number represents a leaf. + pub fn is_leaf(&self) -> bool { + self.idx & 1 == 1 + } + + /// Returns the level of the index. + /// + /// Starts at level zero for leaves and increases by one for each parent. + pub fn level(&self) -> u32 { + self.idx.trailing_zeros() + } + + /// Returns the index of the left child. + /// + /// Panics: + /// + /// If the index corresponds to a leaf. + pub fn left_child(&self) -> InOrderIndex { + // The left child is itself a parent, with an index that splits its left/right subtrees. To + // go from the parent index to its left child, it is only necessary to subtract the count + // of elements on the child's right subtree + 1. + let els = 1 << (self.level() - 1); + InOrderIndex { idx: self.idx - els } + } + + /// Returns the index of the right child. + /// + /// Panics: + /// + /// If the index corresponds to a leaf. + pub fn right_child(&self) -> InOrderIndex { + // To compute the index of the parent of the right subtree it is sufficient to add the size + // of its left subtree + 1. + let els = 1 << (self.level() - 1); + InOrderIndex { idx: self.idx + els } + } + + /// Returns the index of the parent node. + pub fn parent(&self) -> InOrderIndex { + // If the current index corresponds to a node in a left tree, to go up a level it is + // required to add the number of nodes of the right sibling, analogously if the node is a + // right child, going up requires subtracting the number of nodes in its left subtree. + // + // Both of the above operations can be performed by bitwise manipulation. Below the mask + // sets the number of trailing zeros to be equal the new level of the index, and the bit + // marks the parent. + let target = self.level() + 1; + let bit = 1 << target; + let mask = bit - 1; + let idx = self.idx ^ (self.idx & mask); + InOrderIndex { idx: idx | bit } + } + + /// Returns the index of the sibling node. + pub fn sibling(&self) -> InOrderIndex { + let parent = self.parent(); + if *self > parent { + parent.left_child() + } else { + parent.right_child() + } + } +} + +#[cfg(test)] +mod test { + use super::InOrderIndex; + use proptest::prelude::*; + + proptest! { + #[test] + fn proptest_inorder_index_random(count in 1..1000usize) { + let left_pos = count * 2; + let right_pos = count * 2 + 1; + + let left = InOrderIndex::from_leaf_pos(left_pos); + let right = InOrderIndex::from_leaf_pos(right_pos); + + assert!(left.is_leaf()); + assert!(right.is_leaf()); + assert_eq!(left.parent(), right.parent()); + assert_eq!(left.parent().right_child(), right); + assert_eq!(left, right.parent().left_child()); + assert_eq!(left.sibling(), right); + assert_eq!(left, right.sibling()); + } + } + + #[test] + fn test_inorder_index_basic() { + let left = InOrderIndex::from_leaf_pos(0); + let right = InOrderIndex::from_leaf_pos(1); + + assert!(left.is_leaf()); + assert!(right.is_leaf()); + assert_eq!(left.parent(), right.parent()); + assert_eq!(left.parent().right_child(), right); + assert_eq!(left, right.parent().left_child()); + assert_eq!(left.sibling(), right); + assert_eq!(left, right.sibling()); + } +} diff --git a/src/merkle/mmr/mod.rs b/src/merkle/mmr/mod.rs index 118bb120..0301311e 100644 --- a/src/merkle/mmr/mod.rs +++ b/src/merkle/mmr/mod.rs @@ -1,6 +1,10 @@ mod accumulator; mod bit; +mod delta; +mod error; mod full; +mod inorder; +mod partial; mod proof; #[cfg(test)] @@ -11,5 +15,39 @@ use super::{Felt, Rpo256, Word}; // REEXPORTS // ================================================================================================ pub use accumulator::MmrPeaks; +pub use delta::MmrDelta; +pub use error::MmrError; pub use full::Mmr; +pub use inorder::InOrderIndex; +pub use partial::PartialMmr; pub use proof::MmrProof; + +// 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. Because the root element is not part of the proof, $p$ is the length of the +/// authentication path. $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) + } +} diff --git a/src/merkle/mmr/partial.rs b/src/merkle/mmr/partial.rs new file mode 100644 index 00000000..542de011 --- /dev/null +++ b/src/merkle/mmr/partial.rs @@ -0,0 +1,377 @@ +use crate::{ + hash::rpo::{Rpo256, RpoDigest}, + merkle::{ + mmr::{full::nodes_in_forest, leaf_to_corresponding_tree}, + InOrderIndex, MerklePath, MmrError, MmrPeaks, + }, + utils::collections::{BTreeMap, Vec}, +}; + +use super::MmrDelta; + +/// Partially materialized [Mmr], used to efficiently store and update the authentication paths for +/// a subset of the elements in a full [Mmr]. +/// +/// This structure store only the authentication path for a value, the value itself is stored +/// separately. +#[derive(Debug)] +pub struct PartialMmr { + /// The version of the [Mmr]. + /// + /// This value serves the following purposes: + /// + /// - The forest is a counter for the total number of elements in the [Mmr]. + /// - Since the [Mmr] is an append-only structure, every change to it causes a change to the + /// `forest`, so this value has a dual purpose as a version tag. + /// - The bits in the forest also corresponds to the count and size of every perfect binary + /// tree that composes the [Mmr] structure, which server to compute indexes and perform + /// validation. + pub(crate) forest: usize, + + /// The [Mmr] peaks. + /// + /// The peaks are used for two reasons: + /// + /// 1. It authenticates the addition of an element to the [PartialMmr], ensuring only valid + /// elements are tracked. + /// 2. During a [Mmr] update peaks can be merged by hashing the left and right hand sides. The + /// peaks are used as the left hand. + /// + /// 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. + pub(crate) peaks: Vec, + + /// Authentication nodes used to construct merkle paths for a subset of the [Mmr]'s leaves. + /// + /// This does not include the [Mmr]'s peaks nor the tracked nodes, only the elements required + /// to construct their authentication paths. This property is used to detect when elements can + /// be safely removed from, because they are no longer required to authenticate any element in + /// the [PartialMmr]. + /// + /// The elements in the [Mmr] are referenced using a in-order tree index. This indexing scheme + /// permits for easy computation of the relative nodes (left/right children, sibling, parent), + /// which is useful for traversal. The indexing is also stable, meaning that merges to the + /// trees in the [Mmr] can be represented without rewrites of the indexes. + pub(crate) nodes: BTreeMap, + + /// Flag indicating if the odd element should be tracked. + /// + /// This flag is necessary because the sibling of the odd doesn't exist yet, so it can not be + /// added into `nodes` to signal the value is being tracked. + pub(crate) track_latest: bool, +} + +impl PartialMmr { + // CONSTRUCTORS + // -------------------------------------------------------------------------------------------- + + /// Constructs a [PartialMmr] from the given [MmrPeaks]. + pub fn from_peaks(accumulator: MmrPeaks) -> Self { + let forest = accumulator.num_leaves(); + let peaks = accumulator.peaks().clone(); + let nodes = BTreeMap::new(); + let track_latest = false; + + Self { forest, peaks, nodes, track_latest } + } + + // ACCESSORS + // -------------------------------------------------------------------------------------------- + + // Gets the current `forest`. + // + // This value corresponds to the version of the [PartialMmr] and the number of leaves in it. + pub fn forest(&self) -> usize { + self.forest + } + + // Returns a reference to the current peaks in the [PartialMmr] + pub fn peaks(&self) -> &[RpoDigest] { + &self.peaks + } + + /// Return the authentication path of `leaf_pos`. + /// + /// This method returns an error if `leaf_pos` is too large, and `None` if the requested value + /// is not tracked. + /// + /// Note: `leaf_pos` corresponds to the position the [Mmr] and not on an individual tree. + pub fn get_merkle_path(&self, leaf_pos: usize) -> Result, MmrError> { + let tree_bit = leaf_to_corresponding_tree(leaf_pos, self.forest) + .ok_or(MmrError::InvalidPosition(leaf_pos))?; + let depth = tree_bit as usize; + + let mut nodes = Vec::with_capacity(depth); + let mut idx = InOrderIndex::from_leaf_pos(leaf_pos); + + while let Some(node) = self.nodes.get(&idx.sibling()) { + nodes.push(*node); + idx = idx.parent(); + } + + // If there are nodes then the path must be complete, otherwise it is a bug + debug_assert!(nodes.is_empty() || nodes.len() == depth); + + if nodes.len() != depth { + // The requested `leaf_pos` is not being tracked. + Ok(None) + } else { + Ok(Some(MerklePath::new(nodes))) + } + } + + // MODIFIERS + // -------------------------------------------------------------------------------------------- + + /// Add the authentication path represented by [MerklePath] if it is valid. + /// + /// The `index` refers to the global position of the leaf in the [Mmr], these are 0-indexed + /// values assigned in a strictly monotonic fashion as elements are inserted into the [Mmr], + /// this value corresponds to the values used in the [Mmr] structure. + /// + /// The `node` corresponds to the value at `index`, and `path` is the authentication path for + /// that element up to its corresponding Mmr peak. The `node` is only used to compute the root + /// from the authentication path to valid the data, only the authentication data is saved in + /// the structure. If the value is required it should be stored out-of-band. + pub fn add(&mut self, index: usize, node: RpoDigest, path: MerklePath) -> Result<(), MmrError> { + // Checks there is a tree with same depth as the authentication path, if not the path is + // invalid. + let tree = 1 << path.depth(); + if tree & self.forest == 0 { + return Err(MmrError::UnknownPeak); + }; + + if index + 1 == self.forest + && path.depth() == 0 + && self.peaks.last().map_or(false, |v| *v == node) + { + self.track_latest = true; + return Ok(()); + } + + // ignore the trees smaller than the target (these elements are position after the current + // target and don't affect the target index) + let target_forest = self.forest ^ (self.forest & (tree - 1)); + let peak_pos = (target_forest.count_ones() - 1) as usize; + + // translate from mmr index to merkle path + let path_idx = index - (target_forest ^ tree); + + // Compute the root of the authentication path, and check it matches the current version of + // the PartialMmr. + let computed = path.compute_root(path_idx as u64, node).map_err(MmrError::MerkleError)?; + if self.peaks[peak_pos] != computed { + return Err(MmrError::InvalidPeak); + } + + let mut idx = InOrderIndex::from_leaf_pos(index); + for node in path.nodes() { + self.nodes.insert(idx.sibling(), *node); + idx = idx.parent(); + } + + Ok(()) + } + + /// Remove a leaf of the [PartialMmr] and the unused nodes from the authentication path. + /// + /// Note: `leaf_pos` corresponds to the position the [Mmr] and not on an individual tree. + pub fn remove(&mut self, leaf_pos: usize) { + let mut idx = InOrderIndex::from_leaf_pos(leaf_pos); + + self.nodes.remove(&idx.sibling()); + + // `idx` represent the element that can be computed by the authentication path, because + // these elements can be computed they are not saved for the authentication of the current + // target. In other words, if the idx is present it was added for the authentication of + // another element, and no more elements should be removed otherwise it would remove that + // element's authentication data. + while !self.nodes.contains_key(&idx) { + idx = idx.parent(); + self.nodes.remove(&idx.sibling()); + } + } + + /// Applies updates to the [PartialMmr]. + pub fn update(&mut self, delta: MmrDelta) -> Result<(), MmrError> { + let forest = delta.forest; + let updates = &delta.data; + + if forest < self.forest { + return Err(MmrError::InvalidPeaks); + } + + if forest == self.forest { + if !updates.is_empty() { + return Err(MmrError::InvalidUpdate); + } + + return Ok(()); + } + + // find the tree merges + let changes = self.forest ^ forest; + let largest = 1 << changes.ilog2(); + let merges = self.forest & (largest - 1); + + debug_assert!( + !self.track_latest || (merges & 1) == 1, + "if there is an odd element, a merge is required" + ); + + // keeps track of how many data elements from the update have been consumed + let mut update_count = 0; + + if merges != 0 { + // starts at the smallest peak and follows the merged peaks + let mut peak_idx = forest_to_root_index(self.forest); + + // match order of the update data while applying it + self.peaks.reverse(); + + // set to true when the data is needed for authentication paths updates + let mut track = self.track_latest; + self.track_latest = false; + + let mut peak_count = 0; + let mut target = 1 << merges.trailing_zeros(); + let mut new = updates[0]; + update_count += 1; + + while target < largest { + // check if either the left or right subtrees have saved for authentication paths. + // If so, turn tracking on to update those paths. + if target != 1 && !track { + let left_child = peak_idx.left_child(); + let right_child = peak_idx.right_child(); + track = self.nodes.contains_key(&left_child) + | self.nodes.contains_key(&right_child); + } + + // update data only contains the nodes from the right subtrees, left nodes are + // either previously known peaks or computed values + let (left, right) = if target & merges != 0 { + let peak = self.peaks[peak_count]; + peak_count += 1; + (peak, new) + } else { + let update = updates[update_count]; + update_count += 1; + (new, update) + }; + + if track { + self.nodes.insert(peak_idx.sibling(), right); + } + + peak_idx = peak_idx.parent(); + new = Rpo256::merge(&[left, right]); + target <<= 1; + } + + debug_assert!(peak_count == (merges.count_ones() as usize)); + + // restore the peaks order + self.peaks.reverse(); + // remove the merged peaks + self.peaks.truncate(self.peaks.len() - peak_count); + // add the newly computed peak, the result of the merges + self.peaks.push(new); + } + + // The rest of the update data is composed of peaks. None of these elements can contain + // tracked elements because the peaks were unknown, and it is not possible to add elements + // for tacking without authenticating it to a peak. + self.peaks.extend_from_slice(&updates[update_count..]); + self.forest = forest; + + debug_assert!(self.peaks.len() == (self.forest.count_ones() as usize)); + + Ok(()) + } +} + +// CONVERSIONS +// ================================================================================================ + +impl From for PartialMmr { + fn from(peaks: MmrPeaks) -> Self { + Self::from_peaks(peaks) + } +} + +impl From for MmrPeaks { + fn from(partial_mmr: PartialMmr) -> Self { + // Safety: the [PartialMmr] maintains the constraints the number of true bits in the forest + // matches the number of peaks, as required by the [MmrPeaks] + MmrPeaks::new(partial_mmr.forest, partial_mmr.peaks).unwrap() + } +} + +impl From<&MmrPeaks> for PartialMmr { + fn from(peaks: &MmrPeaks) -> Self { + Self::from_peaks(peaks.clone()) + } +} + +impl From<&PartialMmr> for MmrPeaks { + fn from(partial_mmr: &PartialMmr) -> Self { + // Safety: the [PartialMmr] maintains the constraints the number of true bits in the forest + // matches the number of peaks, as required by the [MmrPeaks] + MmrPeaks::new(partial_mmr.forest, partial_mmr.peaks.clone()).unwrap() + } +} + +// UTILS +// ================================================================================================ + +/// Given the description of a `forest`, returns the index of the root element of the smallest tree +/// in it. +pub fn forest_to_root_index(forest: usize) -> InOrderIndex { + // Count total size of all trees in the forest. + let nodes = nodes_in_forest(forest); + + // Add the count for the parent nodes that separate each tree. These are allocated but + // currently empty, and correspond to the nodes that will be used once the trees are merged. + let open_trees = (forest.count_ones() - 1) as usize; + + // Remove the count of the right subtree of the target tree, target tree root index comes + // before the subtree for the in-order tree walk. + let right_subtree_count = ((1u32 << forest.trailing_zeros()) - 1) as usize; + + let idx = nodes + open_trees - right_subtree_count; + + InOrderIndex::new(idx.try_into().unwrap()) +} + +#[cfg(test)] +mod test { + use super::forest_to_root_index; + use crate::merkle::InOrderIndex; + + #[test] + fn test_forest_to_root_index() { + fn idx(pos: usize) -> InOrderIndex { + InOrderIndex::new(pos.try_into().unwrap()) + } + + // When there is a single tree in the forest, the index is equivalent to the number of + // leaves in that tree, which is `2^n`. + assert_eq!(forest_to_root_index(0b0001), idx(1)); + assert_eq!(forest_to_root_index(0b0010), idx(2)); + assert_eq!(forest_to_root_index(0b0100), idx(4)); + assert_eq!(forest_to_root_index(0b1000), idx(8)); + + assert_eq!(forest_to_root_index(0b0011), idx(5)); + assert_eq!(forest_to_root_index(0b0101), idx(9)); + assert_eq!(forest_to_root_index(0b1001), idx(17)); + assert_eq!(forest_to_root_index(0b0111), idx(13)); + assert_eq!(forest_to_root_index(0b1011), idx(21)); + assert_eq!(forest_to_root_index(0b1111), idx(29)); + + assert_eq!(forest_to_root_index(0b0110), idx(10)); + assert_eq!(forest_to_root_index(0b1010), idx(18)); + assert_eq!(forest_to_root_index(0b1100), idx(20)); + assert_eq!(forest_to_root_index(0b1110), idx(26)); + } +} diff --git a/src/merkle/mmr/proof.rs b/src/merkle/mmr/proof.rs index d9b4bcfb..8674c2a8 100644 --- a/src/merkle/mmr/proof.rs +++ b/src/merkle/mmr/proof.rs @@ -1,6 +1,6 @@ /// The representation of a single Merkle path. use super::super::MerklePath; -use super::full::{high_bitmask, leaf_to_corresponding_tree}; +use super::{full::high_bitmask, leaf_to_corresponding_tree}; #[derive(Debug, Clone, PartialEq)] #[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))] diff --git a/src/merkle/mmr/tests.rs b/src/merkle/mmr/tests.rs index 3cecaa49..9afa9a11 100644 --- a/src/merkle/mmr/tests.rs +++ b/src/merkle/mmr/tests.rs @@ -1,12 +1,12 @@ use super::{ super::{InnerNodeInfo, Vec}, bit::TrueBitPositionIterator, - full::{high_bitmask, leaf_to_corresponding_tree, nodes_in_forest}, - Mmr, MmrPeaks, Rpo256, + full::{high_bitmask, nodes_in_forest}, + leaf_to_corresponding_tree, Mmr, MmrPeaks, PartialMmr, Rpo256, }; use crate::{ hash::rpo::RpoDigest, - merkle::{int_to_node, MerklePath}, + merkle::{int_to_node, InOrderIndex, MerklePath, MmrProof}, Felt, Word, }; @@ -118,14 +118,14 @@ fn test_mmr_simple() { let mut postorder = Vec::new(); postorder.push(LEAVES[0]); postorder.push(LEAVES[1]); - postorder.push(Rpo256::merge(&[LEAVES[0], LEAVES[1]])); + postorder.push(merge(LEAVES[0], LEAVES[1])); postorder.push(LEAVES[2]); postorder.push(LEAVES[3]); - postorder.push(Rpo256::merge(&[LEAVES[2], LEAVES[3]])); - postorder.push(Rpo256::merge(&[postorder[2], postorder[5]])); + postorder.push(merge(LEAVES[2], LEAVES[3])); + postorder.push(merge(postorder[2], postorder[5])); postorder.push(LEAVES[4]); postorder.push(LEAVES[5]); - postorder.push(Rpo256::merge(&[LEAVES[4], LEAVES[5]])); + postorder.push(merge(LEAVES[4], LEAVES[5])); postorder.push(LEAVES[6]); let mut mmr = Mmr::new(); @@ -138,8 +138,8 @@ fn test_mmr_simple() { 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]]); + assert_eq!(acc.num_leaves(), 1); + assert_eq!(acc.peaks(), &[postorder[0]]); mmr.add(LEAVES[1]); assert_eq!(mmr.forest(), 2); @@ -147,8 +147,8 @@ fn test_mmr_simple() { 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]]); + assert_eq!(acc.num_leaves(), 2); + assert_eq!(acc.peaks(), &[postorder[2]]); mmr.add(LEAVES[2]); assert_eq!(mmr.forest(), 3); @@ -156,8 +156,8 @@ fn test_mmr_simple() { 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]]); + assert_eq!(acc.num_leaves(), 3); + assert_eq!(acc.peaks(), &[postorder[2], postorder[3]]); mmr.add(LEAVES[3]); assert_eq!(mmr.forest(), 4); @@ -165,8 +165,8 @@ fn test_mmr_simple() { 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]]); + assert_eq!(acc.num_leaves(), 4); + assert_eq!(acc.peaks(), &[postorder[6]]); mmr.add(LEAVES[4]); assert_eq!(mmr.forest(), 5); @@ -174,8 +174,8 @@ fn test_mmr_simple() { 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]]); + assert_eq!(acc.num_leaves(), 5); + assert_eq!(acc.peaks(), &[postorder[6], postorder[7]]); mmr.add(LEAVES[5]); assert_eq!(mmr.forest(), 6); @@ -183,8 +183,8 @@ fn test_mmr_simple() { 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]]); + assert_eq!(acc.num_leaves(), 6); + assert_eq!(acc.peaks(), &[postorder[6], postorder[9]]); mmr.add(LEAVES[6]); assert_eq!(mmr.forest(), 7); @@ -192,15 +192,15 @@ fn test_mmr_simple() { 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]]); + 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 = Rpo256::merge(&[LEAVES[0], LEAVES[1]]); - let h23 = Rpo256::merge(&[LEAVES[2], LEAVES[3]]); + let h01 = merge(LEAVES[0], LEAVES[1]); + let h23 = merge(LEAVES[2], LEAVES[3]); // node at pos 7 is the root assert!(mmr.open(7).is_err(), "Element 7 is not in the tree, result should be None"); @@ -293,6 +293,138 @@ fn test_mmr_open() { ); } +/// Tests the openings of a simple Mmr with a single tree of depth 8. +#[test] +fn test_mmr_open_eight() { + let leaves = [ + 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), + int_to_node(7), + ]; + let depth1 = [ + merge(leaves[0], leaves[1]), + merge(leaves[2], leaves[3]), + merge(leaves[4], leaves[5]), + merge(leaves[6], leaves[7]), + ]; + #[rustfmt::skip] + let depth2 = [ + merge(depth1[0], depth1[1]), + merge(depth1[2], depth1[3]), + ]; + let root = merge(depth2[0], depth2[1]); + + let forest = leaves.len(); + let mmr: Mmr = leaves.into(); + + let position = 0; + let proof = mmr.open(position).unwrap(); + let merkle_path: MerklePath = [leaves[1], depth1[1], depth2[1]].as_ref().into(); + assert_eq!(proof, MmrProof { forest, position, merkle_path }); + assert_eq!(proof.merkle_path.compute_root(position as u64, leaves[position]).unwrap(), root); + + let position = 1; + let proof = mmr.open(position).unwrap(); + let merkle_path: MerklePath = [leaves[0], depth1[1], depth2[1]].as_ref().into(); + assert_eq!(proof, MmrProof { forest, position, merkle_path }); + assert_eq!(proof.merkle_path.compute_root(position as u64, leaves[position]).unwrap(), root); + + let position = 2; + let proof = mmr.open(position).unwrap(); + let merkle_path: MerklePath = [leaves[3], depth1[0], depth2[1]].as_ref().into(); + assert_eq!(proof, MmrProof { forest, position, merkle_path }); + assert_eq!(proof.merkle_path.compute_root(position as u64, leaves[position]).unwrap(), root); + + let position = 3; + let proof = mmr.open(position).unwrap(); + let merkle_path: MerklePath = [leaves[2], depth1[0], depth2[1]].as_ref().into(); + assert_eq!(proof, MmrProof { forest, position, merkle_path }); + assert_eq!(proof.merkle_path.compute_root(position as u64, leaves[position]).unwrap(), root); + + let position = 4; + let proof = mmr.open(position).unwrap(); + let merkle_path: MerklePath = [leaves[5], depth1[3], depth2[0]].as_ref().into(); + assert_eq!(proof, MmrProof { forest, position, merkle_path }); + assert_eq!(proof.merkle_path.compute_root(position as u64, leaves[position]).unwrap(), root); + + let position = 5; + let proof = mmr.open(position).unwrap(); + let merkle_path: MerklePath = [leaves[4], depth1[3], depth2[0]].as_ref().into(); + assert_eq!(proof, MmrProof { forest, position, merkle_path }); + assert_eq!(proof.merkle_path.compute_root(position as u64, leaves[position]).unwrap(), root); + + let position = 6; + let proof = mmr.open(position).unwrap(); + let merkle_path: MerklePath = [leaves[7], depth1[2], depth2[0]].as_ref().into(); + assert_eq!(proof, MmrProof { forest, position, merkle_path }); + assert_eq!(proof.merkle_path.compute_root(position as u64, leaves[position]).unwrap(), root); + + let position = 7; + let proof = mmr.open(position).unwrap(); + let merkle_path: MerklePath = [leaves[6], depth1[2], depth2[0]].as_ref().into(); + assert_eq!(proof, MmrProof { forest, position, merkle_path }); + assert_eq!(proof.merkle_path.compute_root(position as u64, leaves[position]).unwrap(), root); +} + +/// Tests the openings of Mmr with a 3 trees of depths 4, 2, and 1. +#[test] +fn test_mmr_open_seven() { + let h01 = merge(LEAVES[0], LEAVES[1]); + let h23 = merge(LEAVES[2], LEAVES[3]); + let h0123 = merge(h01, h23); + let h45 = merge(LEAVES[4], LEAVES[5]); + + let forest = LEAVES.len(); + let mmr: Mmr = LEAVES.into(); + + let position = 0; + let proof = mmr.open(position).unwrap(); + let merkle_path: MerklePath = [LEAVES[1], h23].as_ref().into(); + assert_eq!(proof, MmrProof { forest, position, merkle_path }); + assert_eq!(proof.merkle_path.compute_root(0, LEAVES[0]).unwrap(), h0123); + + let position = 1; + let proof = mmr.open(position).unwrap(); + let merkle_path: MerklePath = [LEAVES[0], h23].as_ref().into(); + assert_eq!(proof, MmrProof { forest, position, merkle_path }); + assert_eq!(proof.merkle_path.compute_root(1, LEAVES[1]).unwrap(), h0123); + + let position = 2; + let proof = mmr.open(position).unwrap(); + let merkle_path: MerklePath = [LEAVES[3], h01].as_ref().into(); + assert_eq!(proof, MmrProof { forest, position, merkle_path }); + assert_eq!(proof.merkle_path.compute_root(2, LEAVES[2]).unwrap(), h0123); + + let position = 3; + let proof = mmr.open(position).unwrap(); + let merkle_path: MerklePath = [LEAVES[2], h01].as_ref().into(); + assert_eq!(proof, MmrProof { forest, position, merkle_path }); + assert_eq!(proof.merkle_path.compute_root(3, LEAVES[3]).unwrap(), h0123); + + let position = 4; + let proof = mmr.open(position).unwrap(); + let merkle_path: MerklePath = [LEAVES[5]].as_ref().into(); + assert_eq!(proof, MmrProof { forest, position, merkle_path }); + assert_eq!(proof.merkle_path.compute_root(0, LEAVES[4]).unwrap(), h45); + + let position = 5; + let proof = mmr.open(position).unwrap(); + let merkle_path: MerklePath = [LEAVES[4]].as_ref().into(); + assert_eq!(proof, MmrProof { forest, position, merkle_path }); + assert_eq!(proof.merkle_path.compute_root(1, LEAVES[5]).unwrap(), h45); + + let position = 6; + let proof = mmr.open(position).unwrap(); + let merkle_path: MerklePath = [].as_ref().into(); + assert_eq!(proof, MmrProof { forest, position, merkle_path }); + assert_eq!(proof.merkle_path.compute_root(0, LEAVES[6]).unwrap(), LEAVES[6]); +} + #[test] fn test_mmr_get() { let mmr: Mmr = LEAVES.into(); @@ -314,12 +446,13 @@ fn test_mmr_invariants() { 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, + 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(), + accumulator.num_leaves().count_ones() as usize, + accumulator.peaks().len(), "bits on leaves must match the number of peaks" ); @@ -418,10 +551,9 @@ fn test_mmr_peaks_hash_less_than_16() { for i in 0..16 { peaks.push(int_to_node(i)); - let accumulator = MmrPeaks { - num_leaves: (1 << peaks.len()) - 1, - peaks: peaks.clone(), - }; + + let num_leaves = (1 << peaks.len()) - 1; + let accumulator = MmrPeaks::new(num_leaves, peaks.clone()).unwrap(); // minimum length is 16 let mut expected_peaks = peaks.clone(); @@ -437,10 +569,8 @@ fn test_mmr_peaks_hash_less_than_16() { fn test_mmr_peaks_hash_odd() { let peaks: Vec<_> = (0..=17).map(int_to_node).collect(); - let accumulator = MmrPeaks { - num_leaves: (1 << peaks.len()) - 1, - peaks: peaks.clone(), - }; + let num_leaves = (1 << peaks.len()) - 1; + let accumulator = MmrPeaks::new(num_leaves, peaks.clone()).unwrap(); // odd length bigger than 16 is padded to the next even number let mut expected_peaks = peaks; @@ -451,6 +581,125 @@ fn test_mmr_peaks_hash_odd() { ); } +#[test] +fn test_mmr_updates() { + let mmr: Mmr = LEAVES.into(); + let acc = mmr.accumulator(); + + // original_forest can't have more elements + assert!( + mmr.updates(LEAVES.len() + 1).is_err(), + "Can not provide updates for a newer Mmr" + ); + + // if the number of elements is the same there is no change + assert!( + mmr.updates(LEAVES.len()).unwrap().data.is_empty(), + "There are no updates for the same Mmr version" + ); + + // missing the last element added, which is itself a tree peak + assert_eq!(mmr.updates(6).unwrap().data, vec![acc.peaks()[2]], "one peak"); + + // missing the sibling to complete the tree of depth 2, and the last element + assert_eq!( + mmr.updates(5).unwrap().data, + vec![LEAVES[5], acc.peaks()[2]], + "one sibling, one peak" + ); + + // missing the whole last two trees, only send the peaks + assert_eq!(mmr.updates(4).unwrap().data, vec![acc.peaks()[1], acc.peaks()[2]], "two peaks"); + + // missing the sibling to complete the first tree, and the two last trees + assert_eq!( + mmr.updates(3).unwrap().data, + vec![LEAVES[3], acc.peaks()[1], acc.peaks()[2]], + "one sibling, two peaks" + ); + + // missing half of the first tree, only send the computed element (not the leaves), and the new + // peaks + assert_eq!( + mmr.updates(2).unwrap().data, + vec![mmr.nodes[5], acc.peaks()[1], acc.peaks()[2]], + "one sibling, two peaks" + ); + + assert_eq!( + mmr.updates(1).unwrap().data, + vec![LEAVES[1], mmr.nodes[5], acc.peaks()[1], acc.peaks()[2]], + "one sibling, two peaks" + ); + + assert_eq!(&mmr.updates(0).unwrap().data, acc.peaks(), "all peaks"); +} + +#[test] +fn test_partial_mmr_simple() { + let mmr: Mmr = LEAVES.into(); + let acc = mmr.accumulator(); + let mut partial: PartialMmr = acc.clone().into(); + + // check initial state of the partial mmr + assert_eq!(partial.peaks(), acc.peaks()); + assert_eq!(partial.forest(), acc.num_leaves()); + assert_eq!(partial.forest(), LEAVES.len()); + assert_eq!(partial.peaks().len(), 3); + assert_eq!(partial.nodes.len(), 0); + + // check state after adding tracking one element + let proof1 = mmr.open(0).unwrap(); + let el1 = mmr.get(proof1.position).unwrap(); + partial.add(proof1.position, el1, proof1.merkle_path.clone()).unwrap(); + + // check the number of nodes increased by the number of nodes in the proof + assert_eq!(partial.nodes.len(), proof1.merkle_path.len()); + // check the values match + let idx = InOrderIndex::from_leaf_pos(proof1.position); + assert_eq!(partial.nodes[&idx.sibling()], proof1.merkle_path[0]); + let idx = idx.parent(); + assert_eq!(partial.nodes[&idx.sibling()], proof1.merkle_path[1]); + + let proof2 = mmr.open(1).unwrap(); + let el2 = mmr.get(proof2.position).unwrap(); + partial.add(proof2.position, el2, proof2.merkle_path.clone()).unwrap(); + + // check the number of nodes increased by a single element (the one that is not shared) + assert_eq!(partial.nodes.len(), 3); + // check the values match + let idx = InOrderIndex::from_leaf_pos(proof2.position); + assert_eq!(partial.nodes[&idx.sibling()], proof2.merkle_path[0]); + let idx = idx.parent(); + assert_eq!(partial.nodes[&idx.sibling()], proof2.merkle_path[1]); +} + +#[test] +fn test_partial_mmr_update_single() { + let mut full = Mmr::new(); + let zero = int_to_node(0); + full.add(zero); + let mut partial: PartialMmr = full.accumulator().into(); + + let proof = full.open(0).unwrap(); + partial.add(proof.position, zero, proof.merkle_path.clone()).unwrap(); + + for i in 1..100 { + let node = int_to_node(i); + full.add(node); + let updates = full.updates(partial.forest()).unwrap(); + partial.update(updates).unwrap(); + + assert_eq!(partial.forest(), full.forest()); + assert_eq!(partial.peaks(), full.accumulator().peaks()); + + let proof1 = full.open(i as usize).unwrap(); + partial.add(proof1.position, node, proof1.merkle_path.clone()).unwrap(); + let proof2 = partial.get_merkle_path(proof1.position).unwrap().unwrap(); + assert_eq!(proof1.merkle_path, proof2); + } +} + mod property_tests { use super::leaf_to_corresponding_tree; use proptest::prelude::*; @@ -471,10 +720,10 @@ mod property_tests { 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; + let tree_bit = leaf_to_corresponding_tree(pos, leaves).expect("pos is smaller than leaves, there should always be a corresponding tree"); + let mask = 1usize << tree_bit; - assert!(tree < usize::BITS, "the result must be a bit in usize"); + assert!(tree_bit < usize::BITS, "the result must be a bit in usize"); assert!(mask & leaves != 0, "the result should be a tree in leaves"); } } @@ -486,3 +735,8 @@ mod property_tests { fn digests_to_elements(digests: &[RpoDigest]) -> Vec { digests.iter().flat_map(Word::from).collect() } + +// short hand for the rpo hash, used to make test code more concise and easy to read +fn merge(l: RpoDigest, r: RpoDigest) -> RpoDigest { + Rpo256::merge(&[l, r]) +} diff --git a/src/merkle/mod.rs b/src/merkle/mod.rs index c4e43cac..10012d2c 100644 --- a/src/merkle/mod.rs +++ b/src/merkle/mod.rs @@ -29,7 +29,7 @@ mod tiered_smt; pub use tiered_smt::{TieredSmt, TieredSmtProof, TieredSmtProofError}; mod mmr; -pub use mmr::{Mmr, MmrPeaks, MmrProof}; +pub use mmr::{InOrderIndex, Mmr, MmrError, MmrPeaks, MmrProof, PartialMmr}; mod store; pub use store::{DefaultMerkleStore, MerkleStore, RecordingMerkleStore, StoreNode}; diff --git a/src/merkle/path.rs b/src/merkle/path.rs index 4535ba94..7ec9df4d 100644 --- a/src/merkle/path.rs +++ b/src/merkle/path.rs @@ -28,6 +28,11 @@ impl MerklePath { self.nodes.len() as u8 } + /// Returns a reference to the [MerklePath]'s nodes. + pub fn nodes(&self) -> &Vec { + &self.nodes + } + /// Computes the merkle root for this opening. pub fn compute_root(&self, index: u64, node: RpoDigest) -> Result { let mut index = NodeIndex::new(self.depth(), index)?; @@ -69,6 +74,9 @@ impl MerklePath { } } +// CONVERSIONS +// ================================================================================================ + impl From for Vec { fn from(path: MerklePath) -> Self { path.nodes @@ -81,6 +89,12 @@ impl From> for MerklePath { } } +impl From<&[RpoDigest]> for MerklePath { + fn from(path: &[RpoDigest]) -> Self { + Self::new(path.to_vec()) + } +} + impl Deref for MerklePath { // we use `Vec` here instead of slice so we can call vector mutation methods directly from the // merkle path (example: `Vec::remove`).