use super::{MmrDelta, MmrProof, Rpo256, RpoDigest}; use crate::{ merkle::{ mmr::{leaf_to_corresponding_tree, nodes_in_forest}, InOrderIndex, InnerNodeInfo, MerklePath, MmrError, MmrPeaks, }, utils::{ collections::{BTreeMap, BTreeSet, Vec}, vec, }, }; // PARTIAL MERKLE MOUNTAIN RANGE // ================================================================================================ /// Partially materialized Merkle Mountain Range (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, Clone, PartialEq, Eq)] 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(peaks: MmrPeaks) -> Self { let forest = peaks.num_leaves(); let peaks = peaks.peaks().to_vec(); let nodes = BTreeMap::new(); let track_latest = false; Self { forest, peaks, nodes, track_latest } } // PUBLIC ACCESSORS // -------------------------------------------------------------------------------------------- /// Returns the current `forest` of this [PartialMmr]. /// /// This value corresponds to the version of the [PartialMmr] and the number of leaves in the /// underlying MMR. pub fn forest(&self) -> usize { self.forest } /// Returns the number of leaves in the underlying MMR for this [PartialMmr]. pub fn num_leaves(&self) -> usize { self.forest } /// Returns the peaks of the MMR for this [PartialMmr]. pub fn peaks(&self) -> MmrPeaks { // expect() is OK here because the constructor ensures that MMR peaks can be constructed // correctly MmrPeaks::new(self.forest, self.peaks.clone()).expect("invalid MMR peaks") } /// 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. If the /// requested value is not tracked returns `None`. /// /// 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, MmrError> { let tree_bit = leaf_to_corresponding_tree(pos, self.forest).ok_or(MmrError::InvalidPosition(pos))?; let depth = tree_bit as usize; let mut nodes = Vec::with_capacity(depth); let mut idx = InOrderIndex::from_leaf_pos(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 `pos` is not being tracked. Ok(None) } else { Ok(Some(MmrProof { forest: self.forest, position: pos, merkle_path: MerklePath::new(nodes), })) } } // ITERATORS // -------------------------------------------------------------------------------------------- /// Returns an iterator nodes of all authentication paths of this [PartialMmr]. pub fn nodes(&self) -> impl Iterator { self.nodes.iter() } /// Returns an iterator over inner nodes of this [PartialMmr] for the specified leaves. /// /// The order of iteration is not defined. If a leaf is not presented in this partial MMR it /// is silently ignored. pub fn inner_nodes<'a, I: Iterator + 'a>( &'a self, mut leaves: I, ) -> impl Iterator + '_ { let stack = if let Some((pos, leaf)) = leaves.next() { let idx = InOrderIndex::from_leaf_pos(*pos); vec![(idx, *leaf)] } else { Vec::new() }; InnerNodeIterator { nodes: &self.nodes, leaves, stack, seen_nodes: BTreeSet::new(), } } // STATE MUTATORS // -------------------------------------------------------------------------------------------- /// 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 in 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 this [PartialMmr] and returns a vector of new authentication nodes /// inserted into the partial MMR. pub fn apply(&mut self, delta: MmrDelta) -> Result, MmrError> { if delta.forest < self.forest { return Err(MmrError::InvalidPeaks); } let mut inserted_nodes = Vec::new(); if delta.forest == self.forest { if !delta.data.is_empty() { return Err(MmrError::InvalidUpdate); } return Ok(inserted_nodes); } // find the tree merges let changes = self.forest ^ delta.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" ); // count the number elements needed to produce largest from the current state let (merge_count, new_peaks) = if merges != 0 { let depth = largest.trailing_zeros(); let skipped = merges.trailing_zeros(); let computed = merges.count_ones() - 1; let merge_count = depth - skipped - computed; let new_peaks = delta.forest & (largest - 1); (merge_count, new_peaks) } else { (0, changes) }; // verify the delta size if (delta.data.len() as u32) != merge_count + new_peaks.count_ones() { return Err(MmrError::InvalidUpdate); } // 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 = delta.data[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 { track = self.is_tracked_node(&peak_idx); } // 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]; let sibling_idx = peak_idx.sibling(); // if the sibling peak is tracked, add this peaks to the set of // authentication nodes if self.is_tracked_node(&sibling_idx) { self.nodes.insert(peak_idx, new); inserted_nodes.push((peak_idx, new)); } peak_count += 1; (peak, new) } else { let update = delta.data[update_count]; update_count += 1; (new, update) }; if track { let sibling_idx = peak_idx.sibling(); if peak_idx.is_left_child() { self.nodes.insert(sibling_idx, right); inserted_nodes.push((sibling_idx, right)); } else { self.nodes.insert(sibling_idx, left); inserted_nodes.push((sibling_idx, left)); } } 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(&delta.data[update_count..]); self.forest = delta.forest; debug_assert!(self.peaks.len() == (self.forest.count_ones() as usize)); Ok(inserted_nodes) } // HELPER METHODS // -------------------------------------------------------------------------------------------- /// Returns true if this [PartialMmr] tracks authentication path for the node at the specified /// index. fn is_tracked_node(&self, node_index: &InOrderIndex) -> bool { if node_index.is_leaf() { self.nodes.contains_key(&node_index.sibling()) } else { let left_child = node_index.left_child(); let right_child = node_index.right_child(); self.nodes.contains_key(&left_child) | self.nodes.contains_key(&right_child) } } } // 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() } } // ITERATORS // ================================================================================================ /// An iterator over every inner node of the [PartialMmr]. pub struct InnerNodeIterator<'a, I: Iterator> { nodes: &'a BTreeMap, leaves: I, stack: Vec<(InOrderIndex, RpoDigest)>, seen_nodes: BTreeSet, } impl<'a, I: Iterator> Iterator for InnerNodeIterator<'a, I> { type Item = InnerNodeInfo; fn next(&mut self) -> Option { while let Some((idx, node)) = self.stack.pop() { let parent_idx = idx.parent(); let new_node = self.seen_nodes.insert(parent_idx); // if we haven't seen this node's parent before, and the node has a sibling, return // the inner node defined by the parent of this node, and move up the branch if new_node { if let Some(sibling) = self.nodes.get(&idx.sibling()) { let (left, right) = if parent_idx.left_child() == idx { (node, *sibling) } else { (*sibling, node) }; let parent = Rpo256::merge(&[left, right]); let inner_node = InnerNodeInfo { value: parent, left, right }; self.stack.push((parent_idx, parent)); return Some(inner_node); } } // the previous leaf has been processed, try to process the next leaf if let Some((pos, leaf)) = self.leaves.next() { let idx = InOrderIndex::from_leaf_pos(*pos); self.stack.push((idx, *leaf)); } } None } } // UTILS // ================================================================================================ /// Given the description of a `forest`, returns the index of the root element of the smallest tree /// in it. 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()) } // TESTS // ================================================================================================ #[cfg(test)] mod tests { use super::{forest_to_root_index, BTreeSet, InOrderIndex, PartialMmr, RpoDigest, Vec}; use crate::merkle::{int_to_node, MerkleStore, Mmr, NodeIndex}; const LEAVES: [RpoDigest; 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_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)); } #[test] fn test_partial_mmr_apply_delta() { // build an MMR with 10 nodes (2 peaks) and a partial MMR based on it let mut mmr = Mmr::default(); (0..10).for_each(|i| mmr.add(int_to_node(i))); let mut partial_mmr: PartialMmr = mmr.peaks(mmr.forest()).unwrap().into(); // add authentication path for position 1 and 8 { let node = mmr.get(1).unwrap(); let proof = mmr.open(1, mmr.forest()).unwrap(); partial_mmr.add(1, node, &proof.merkle_path).unwrap(); } { let node = mmr.get(8).unwrap(); let proof = mmr.open(8, mmr.forest()).unwrap(); partial_mmr.add(8, node, &proof.merkle_path).unwrap(); } // add 2 more nodes into the MMR and validate apply_delta() (10..12).for_each(|i| mmr.add(int_to_node(i))); validate_apply_delta(&mmr, &mut partial_mmr); // add 1 more node to the MMR, validate apply_delta() and start tracking the node mmr.add(int_to_node(12)); validate_apply_delta(&mmr, &mut partial_mmr); { let node = mmr.get(12).unwrap(); let proof = mmr.open(12, mmr.forest()).unwrap(); partial_mmr.add(12, node, &proof.merkle_path).unwrap(); assert!(partial_mmr.track_latest); } // by this point we are tracking authentication paths for positions: 1, 8, and 12 // add 3 more nodes to the MMR (collapses to 1 peak) and validate apply_delta() (13..16).for_each(|i| mmr.add(int_to_node(i))); validate_apply_delta(&mmr, &mut partial_mmr); } fn validate_apply_delta(mmr: &Mmr, partial: &mut PartialMmr) { let tracked_leaves = partial .nodes .iter() .filter_map(|(index, _)| if index.is_leaf() { Some(index.sibling()) } else { None }) .collect::>(); let nodes_before = partial.nodes.clone(); // compute and apply delta let delta = mmr.get_delta(partial.forest(), mmr.forest()).unwrap(); let nodes_delta = partial.apply(delta).unwrap(); // new peaks were computed correctly assert_eq!(mmr.peaks(mmr.forest()).unwrap(), partial.peaks()); let mut expected_nodes = nodes_before; for (key, value) in nodes_delta { // nodes should not be duplicated assert!(expected_nodes.insert(key, value).is_none()); } // new nodes should be a combination of original nodes and delta assert_eq!(expected_nodes, partial.nodes); // make sure tracked leaves open to the same proofs as in the underlying MMR for index in tracked_leaves { let index_value: u64 = index.into(); let pos = index_value / 2; let proof1 = partial.open(pos as usize).unwrap().unwrap(); let proof2 = mmr.open(pos as usize, mmr.forest()).unwrap(); assert_eq!(proof1, proof2); } } #[test] fn test_partial_mmr_inner_nodes_iterator() { // build the MMR let mmr: Mmr = LEAVES.into(); let first_peak = mmr.peaks(mmr.forest).unwrap().peaks()[0]; // -- test single tree ---------------------------- // get path and node for position 1 let node1 = mmr.get(1).unwrap(); let proof1 = mmr.open(1, mmr.forest()).unwrap(); // create partial MMR and add authentication path to node at position 1 let mut partial_mmr: PartialMmr = mmr.peaks(mmr.forest()).unwrap().into(); partial_mmr.add(1, node1, &proof1.merkle_path).unwrap(); // empty iterator should have no nodes assert_eq!(partial_mmr.inner_nodes([].iter()).next(), None); // build Merkle store from authentication paths in partial MMR let mut store: MerkleStore = MerkleStore::new(); store.extend(partial_mmr.inner_nodes([(1, node1)].iter())); let index1 = NodeIndex::new(2, 1).unwrap(); let path1 = store.get_path(first_peak, index1).unwrap().path; assert_eq!(path1, proof1.merkle_path); // -- test no duplicates -------------------------- // build the partial MMR let mut partial_mmr: PartialMmr = mmr.peaks(mmr.forest()).unwrap().into(); let node0 = mmr.get(0).unwrap(); let proof0 = mmr.open(0, mmr.forest()).unwrap(); let node2 = mmr.get(2).unwrap(); let proof2 = mmr.open(2, mmr.forest()).unwrap(); partial_mmr.add(0, node0, &proof0.merkle_path).unwrap(); partial_mmr.add(1, node1, &proof1.merkle_path).unwrap(); partial_mmr.add(2, node2, &proof2.merkle_path).unwrap(); // make sure there are no duplicates let leaves = [(0, node0), (1, node1), (2, node2)]; let mut nodes = BTreeSet::new(); for node in partial_mmr.inner_nodes(leaves.iter()) { assert!(nodes.insert(node.value)); } // and also that the store is still be built correctly store.extend(partial_mmr.inner_nodes(leaves.iter())); let index0 = NodeIndex::new(2, 0).unwrap(); let index1 = NodeIndex::new(2, 1).unwrap(); let index2 = NodeIndex::new(2, 2).unwrap(); let path0 = store.get_path(first_peak, index0).unwrap().path; let path1 = store.get_path(first_peak, index1).unwrap().path; let path2 = store.get_path(first_peak, index2).unwrap().path; assert_eq!(path0, proof0.merkle_path); assert_eq!(path1, proof1.merkle_path); assert_eq!(path2, proof2.merkle_path); // -- test multiple trees ------------------------- // build the partial MMR let mut partial_mmr: PartialMmr = mmr.peaks(mmr.forest()).unwrap().into(); let node5 = mmr.get(5).unwrap(); let proof5 = mmr.open(5, mmr.forest()).unwrap(); partial_mmr.add(1, node1, &proof1.merkle_path).unwrap(); partial_mmr.add(5, node5, &proof5.merkle_path).unwrap(); // build Merkle store from authentication paths in partial MMR let mut store: MerkleStore = MerkleStore::new(); store.extend(partial_mmr.inner_nodes([(1, node1), (5, node5)].iter())); let index1 = NodeIndex::new(2, 1).unwrap(); let index5 = NodeIndex::new(1, 1).unwrap(); let second_peak = mmr.peaks(mmr.forest).unwrap().peaks()[1]; let path1 = store.get_path(first_peak, index1).unwrap().path; let path5 = store.get_path(second_peak, index5).unwrap().path; assert_eq!(path1, proof1.merkle_path); assert_eq!(path5, proof5.merkle_path); } }