miden-crypto/src/merkle/mmr/partial.rs
src/merkle/mmr/partial.rsRaw
1use super::{MmrDelta, MmrProof, Rpo256, RpoDigest};
2use crate::{
3 merkle::{
4 mmr::{leaf_to_corresponding_tree, nodes_in_forest},
5 InOrderIndex, InnerNodeInfo, MerklePath, MmrError, MmrPeaks,
6 },
7 utils::{
8 collections::{BTreeMap, BTreeSet, Vec},
9 vec,
10 },
11};
12
13// PARTIAL MERKLE MOUNTAIN RANGE
14// ================================================================================================
15/// Partially materialized Merkle Mountain Range (MMR), used to efficiently store and update the
16/// authentication paths for a subset of the elements in a full MMR.
17///
18/// This structure store only the authentication path for a value, the value itself is stored
19/// separately.
20#[derive(Debug, Clone, PartialEq, Eq)]
21pub struct PartialMmr {
22 /// The version of the MMR.
23 ///
24 /// This value serves the following purposes:
25 ///
26 /// - The forest is a counter for the total number of elements in the MMR.
27 /// - Since the MMR is an append-only structure, every change to it causes a change to the
28 /// `forest`, so this value has a dual purpose as a version tag.
29 /// - The bits in the forest also corresponds to the count and size of every perfect binary
30 /// tree that composes the MMR structure, which server to compute indexes and perform
31 /// validation.
32 pub(crate) forest: usize,
33
34 /// The MMR peaks.
35 ///
36 /// The peaks are used for two reasons:
37 ///
38 /// 1. It authenticates the addition of an element to the [PartialMmr], ensuring only valid
39 /// elements are tracked.
40 /// 2. During a MMR update peaks can be merged by hashing the left and right hand sides. The
41 /// peaks are used as the left hand.
42 ///
43 /// All the peaks of every tree in the MMR forest. The peaks are always ordered by number of
44 /// leaves, starting from the peak with most children, to the one with least.
45 pub(crate) peaks: Vec<RpoDigest>,
46
47 /// Authentication nodes used to construct merkle paths for a subset of the MMR's leaves.
48 ///
49 /// This does not include the MMR's peaks nor the tracked nodes, only the elements required to
50 /// construct their authentication paths. This property is used to detect when elements can be
51 /// safely removed, because they are no longer required to authenticate any element in the
52 /// [PartialMmr].
53 ///
54 /// The elements in the MMR are referenced using a in-order tree index. This indexing scheme
55 /// permits for easy computation of the relative nodes (left/right children, sibling, parent),
56 /// which is useful for traversal. The indexing is also stable, meaning that merges to the
57 /// trees in the MMR can be represented without rewrites of the indexes.
58 pub(crate) nodes: BTreeMap<InOrderIndex, RpoDigest>,
59
60 /// Flag indicating if the odd element should be tracked.
61 ///
62 /// This flag is necessary because the sibling of the odd doesn't exist yet, so it can not be
63 /// added into `nodes` to signal the value is being tracked.
64 pub(crate) track_latest: bool,
65}
66
67impl PartialMmr {
68 // CONSTRUCTORS
69 // --------------------------------------------------------------------------------------------
70
71 /// Constructs a [PartialMmr] from the given [MmrPeaks].
72 pub fn from_peaks(peaks: MmrPeaks) -> Self {
73 let forest = peaks.num_leaves();
74 let peaks = peaks.peaks().to_vec();
75 let nodes = BTreeMap::new();
76 let track_latest = false;
77
78 Self { forest, peaks, nodes, track_latest }
79 }
80
81 // PUBLIC ACCESSORS
82 // --------------------------------------------------------------------------------------------
83
84 /// Returns the current `forest` of this [PartialMmr].
85 ///
86 /// This value corresponds to the version of the [PartialMmr] and the number of leaves in the
87 /// underlying MMR.
88 pub fn forest(&self) -> usize {
89 self.forest
90 }
91
92 /// Returns the number of leaves in the underlying MMR for this [PartialMmr].
93 pub fn num_leaves(&self) -> usize {
94 self.forest
95 }
96
97 /// Returns the peaks of the MMR for this [PartialMmr].
98 pub fn peaks(&self) -> MmrPeaks {
99 // expect() is OK here because the constructor ensures that MMR peaks can be constructed
100 // correctly
101 MmrPeaks::new(self.forest, self.peaks.clone()).expect("invalid MMR peaks")
102 }
103
104 /// Returns true if this partial MMR tracks an authentication path for the leaf at the
105 /// specified position.
106 pub fn is_tracked(&self, pos: usize) -> bool {
107 if pos >= self.forest {
108 return false;
109 } else if pos == self.forest - 1 && self.forest & 1 != 0 {
110 // if the number of leaves in the MMR is odd and the position is for the last leaf
111 // whether the leaf is tracked is defined by the `track_latest` flag
112 return self.track_latest;
113 }
114
115 let leaf_index = InOrderIndex::from_leaf_pos(pos);
116 self.is_tracked_node(&leaf_index)
117 }
118
119 /// Given a leaf position, returns the Merkle path to its corresponding peak, or None if this
120 /// partial MMR does not track an authentication paths for the specified leaf.
121 ///
122 /// Note: The leaf position is the 0-indexed number corresponding to the order the leaves were
123 /// added, this corresponds to the MMR size _prior_ to adding the element. So the 1st element
124 /// has position 0, the second position 1, and so on.
125 ///
126 /// # Errors
127 /// Returns an error if the specified position is greater-or-equal than the number of leaves
128 /// in the underlying MMR.
129 pub fn open(&self, pos: usize) -> Result<Option<MmrProof>, MmrError> {
130 let tree_bit =
131 leaf_to_corresponding_tree(pos, self.forest).ok_or(MmrError::InvalidPosition(pos))?;
132 let depth = tree_bit as usize;
133
134 let mut nodes = Vec::with_capacity(depth);
135 let mut idx = InOrderIndex::from_leaf_pos(pos);
136
137 while let Some(node) = self.nodes.get(&idx.sibling()) {
138 nodes.push(*node);
139 idx = idx.parent();
140 }
141
142 // If there are nodes then the path must be complete, otherwise it is a bug
143 debug_assert!(nodes.is_empty() || nodes.len() == depth);
144
145 if nodes.len() != depth {
146 // The requested `pos` is not being tracked.
147 Ok(None)
148 } else {
149 Ok(Some(MmrProof {
150 forest: self.forest,
151 position: pos,
152 merkle_path: MerklePath::new(nodes),
153 }))
154 }
155 }
156
157 // ITERATORS
158 // --------------------------------------------------------------------------------------------
159
160 /// Returns an iterator nodes of all authentication paths of this [PartialMmr].
161 pub fn nodes(&self) -> impl Iterator<Item = (&InOrderIndex, &RpoDigest)> {
162 self.nodes.iter()
163 }
164
165 /// Returns an iterator over inner nodes of this [PartialMmr] for the specified leaves.
166 ///
167 /// The order of iteration is not defined. If a leaf is not presented in this partial MMR it
168 /// is silently ignored.
169 pub fn inner_nodes<'a, I: Iterator<Item = (usize, RpoDigest)> + 'a>(
170 &'a self,
171 mut leaves: I,
172 ) -> impl Iterator<Item = InnerNodeInfo> + '_ {
173 let stack = if let Some((pos, leaf)) = leaves.next() {
174 let idx = InOrderIndex::from_leaf_pos(pos);
175 vec![(idx, leaf)]
176 } else {
177 Vec::new()
178 };
179
180 InnerNodeIterator {
181 nodes: &self.nodes,
182 leaves,
183 stack,
184 seen_nodes: BTreeSet::new(),
185 }
186 }
187
188 // STATE MUTATORS
189 // --------------------------------------------------------------------------------------------
190
191 /// Adds a new peak and optionally track it. Returns a vector of the authentication nodes
192 /// inserted into this [PartialMmr] as a result of this operation.
193 ///
194 /// When `track` is `true` the new leaf is tracked.
195 pub fn add(&mut self, leaf: RpoDigest, track: bool) -> Vec<(InOrderIndex, RpoDigest)> {
196 self.forest += 1;
197 let merges = self.forest.trailing_zeros() as usize;
198 let mut new_nodes = Vec::with_capacity(merges);
199
200 let peak = if merges == 0 {
201 self.track_latest = track;
202 leaf
203 } else {
204 let mut track_right = track;
205 let mut track_left = self.track_latest;
206
207 let mut right = leaf;
208 let mut right_idx = forest_to_rightmost_index(self.forest);
209
210 for _ in 0..merges {
211 let left = self.peaks.pop().expect("Missing peak");
212 let left_idx = right_idx.sibling();
213
214 if track_right {
215 let old = self.nodes.insert(left_idx, left);
216 new_nodes.push((left_idx, left));
217
218 debug_assert!(
219 old.is_none(),
220 "Idx {:?} already contained an element {:?}",
221 left_idx,
222 old
223 );
224 };
225 if track_left {
226 let old = self.nodes.insert(right_idx, right);
227 new_nodes.push((right_idx, right));
228
229 debug_assert!(
230 old.is_none(),
231 "Idx {:?} already contained an element {:?}",
232 right_idx,
233 old
234 );
235 };
236
237 // Update state for the next iteration.
238 // --------------------------------------------------------------------------------
239
240 // This layer is merged, go up one layer.
241 right_idx = right_idx.parent();
242
243 // Merge the current layer. The result is either the right element of the next
244 // merge, or a new peak.
245 right = Rpo256::merge(&[left, right]);
246
247 // This iteration merged the left and right nodes, the new value is always used as
248 // the next iteration's right node. Therefore the tracking flags of this iteration
249 // have to be merged into the right side only.
250 track_right = track_right || track_left;
251
252 // On the next iteration, a peak will be merged. If any of its children are tracked,
253 // then we have to track the left side
254 track_left = self.is_tracked_node(&right_idx.sibling());
255 }
256 right
257 };
258
259 self.peaks.push(peak);
260
261 new_nodes
262 }
263
264 /// Adds the authentication path represented by [MerklePath] if it is valid.
265 ///
266 /// The `leaf_pos` refers to the global position of the leaf in the MMR, these are 0-indexed
267 /// values assigned in a strictly monotonic fashion as elements are inserted into the MMR,
268 /// this value corresponds to the values used in the MMR structure.
269 ///
270 /// The `leaf` corresponds to the value at `leaf_pos`, and `path` is the authentication path for
271 /// that element up to its corresponding Mmr peak. The `leaf` is only used to compute the root
272 /// from the authentication path to valid the data, only the authentication data is saved in
273 /// the structure. If the value is required it should be stored out-of-band.
274 pub fn track(
275 &mut self,
276 leaf_pos: usize,
277 leaf: RpoDigest,
278 path: &MerklePath,
279 ) -> Result<(), MmrError> {
280 // Checks there is a tree with same depth as the authentication path, if not the path is
281 // invalid.
282 let tree = 1 << path.depth();
283 if tree & self.forest == 0 {
284 return Err(MmrError::UnknownPeak);
285 };
286
287 if leaf_pos + 1 == self.forest
288 && path.depth() == 0
289 && self.peaks.last().map_or(false, |v| *v == leaf)
290 {
291 self.track_latest = true;
292 return Ok(());
293 }
294
295 // ignore the trees smaller than the target (these elements are position after the current
296 // target and don't affect the target leaf_pos)
297 let target_forest = self.forest ^ (self.forest & (tree - 1));
298 let peak_pos = (target_forest.count_ones() - 1) as usize;
299
300 // translate from mmr leaf_pos to merkle path
301 let path_idx = leaf_pos - (target_forest ^ tree);
302
303 // Compute the root of the authentication path, and check it matches the current version of
304 // the PartialMmr.
305 let computed = path.compute_root(path_idx as u64, leaf).map_err(MmrError::MerkleError)?;
306 if self.peaks[peak_pos] != computed {
307 return Err(MmrError::InvalidPeak);
308 }
309
310 let mut idx = InOrderIndex::from_leaf_pos(leaf_pos);
311 for leaf in path.nodes() {
312 self.nodes.insert(idx.sibling(), *leaf);
313 idx = idx.parent();
314 }
315
316 Ok(())
317 }
318
319 /// Removes a leaf of the [PartialMmr] and the unused nodes from the authentication path.
320 ///
321 /// Note: `leaf_pos` corresponds to the position in the MMR and not on an individual tree.
322 pub fn untrack(&mut self, leaf_pos: usize) {
323 let mut idx = InOrderIndex::from_leaf_pos(leaf_pos);
324
325 self.nodes.remove(&idx.sibling());
326
327 // `idx` represent the element that can be computed by the authentication path, because
328 // these elements can be computed they are not saved for the authentication of the current
329 // target. In other words, if the idx is present it was added for the authentication of
330 // another element, and no more elements should be removed otherwise it would remove that
331 // element's authentication data.
332 while !self.nodes.contains_key(&idx) {
333 idx = idx.parent();
334 self.nodes.remove(&idx.sibling());
335 }
336 }
337
338 /// Applies updates to this [PartialMmr] and returns a vector of new authentication nodes
339 /// inserted into the partial MMR.
340 pub fn apply(&mut self, delta: MmrDelta) -> Result<Vec<(InOrderIndex, RpoDigest)>, MmrError> {
341 if delta.forest < self.forest {
342 return Err(MmrError::InvalidPeaks);
343 }
344
345 let mut inserted_nodes = Vec::new();
346
347 if delta.forest == self.forest {
348 if !delta.data.is_empty() {
349 return Err(MmrError::InvalidUpdate);
350 }
351
352 return Ok(inserted_nodes);
353 }
354
355 // find the tree merges
356 let changes = self.forest ^ delta.forest;
357 let largest = 1 << changes.ilog2();
358 let merges = self.forest & (largest - 1);
359
360 debug_assert!(
361 !self.track_latest || (merges & 1) == 1,
362 "if there is an odd element, a merge is required"
363 );
364
365 // count the number elements needed to produce largest from the current state
366 let (merge_count, new_peaks) = if merges != 0 {
367 let depth = largest.trailing_zeros();
368 let skipped = merges.trailing_zeros();
369 let computed = merges.count_ones() - 1;
370 let merge_count = depth - skipped - computed;
371
372 let new_peaks = delta.forest & (largest - 1);
373
374 (merge_count, new_peaks)
375 } else {
376 (0, changes)
377 };
378
379 // verify the delta size
380 if (delta.data.len() as u32) != merge_count + new_peaks.count_ones() {
381 return Err(MmrError::InvalidUpdate);
382 }
383
384 // keeps track of how many data elements from the update have been consumed
385 let mut update_count = 0;
386
387 if merges != 0 {
388 // starts at the smallest peak and follows the merged peaks
389 let mut peak_idx = forest_to_root_index(self.forest);
390
391 // match order of the update data while applying it
392 self.peaks.reverse();
393
394 // set to true when the data is needed for authentication paths updates
395 let mut track = self.track_latest;
396 self.track_latest = false;
397
398 let mut peak_count = 0;
399 let mut target = 1 << merges.trailing_zeros();
400 let mut new = delta.data[0];
401 update_count += 1;
402
403 while target < largest {
404 // check if either the left or right subtrees have saved for authentication paths.
405 // If so, turn tracking on to update those paths.
406 if target != 1 && !track {
407 track = self.is_tracked_node(&peak_idx);
408 }
409
410 // update data only contains the nodes from the right subtrees, left nodes are
411 // either previously known peaks or computed values
412 let (left, right) = if target & merges != 0 {
413 let peak = self.peaks[peak_count];
414 let sibling_idx = peak_idx.sibling();
415
416 // if the sibling peak is tracked, add this peaks to the set of
417 // authentication nodes
418 if self.is_tracked_node(&sibling_idx) {
419 self.nodes.insert(peak_idx, new);
420 inserted_nodes.push((peak_idx, new));
421 }
422 peak_count += 1;
423 (peak, new)
424 } else {
425 let update = delta.data[update_count];
426 update_count += 1;
427 (new, update)
428 };
429
430 if track {
431 let sibling_idx = peak_idx.sibling();
432 if peak_idx.is_left_child() {
433 self.nodes.insert(sibling_idx, right);
434 inserted_nodes.push((sibling_idx, right));
435 } else {
436 self.nodes.insert(sibling_idx, left);
437 inserted_nodes.push((sibling_idx, left));
438 }
439 }
440
441 peak_idx = peak_idx.parent();
442 new = Rpo256::merge(&[left, right]);
443 target <<= 1;
444 }
445
446 debug_assert!(peak_count == (merges.count_ones() as usize));
447
448 // restore the peaks order
449 self.peaks.reverse();
450 // remove the merged peaks
451 self.peaks.truncate(self.peaks.len() - peak_count);
452 // add the newly computed peak, the result of the merges
453 self.peaks.push(new);
454 }
455
456 // The rest of the update data is composed of peaks. None of these elements can contain
457 // tracked elements because the peaks were unknown, and it is not possible to add elements
458 // for tacking without authenticating it to a peak.
459 self.peaks.extend_from_slice(&delta.data[update_count..]);
460 self.forest = delta.forest;
461
462 debug_assert!(self.peaks.len() == (self.forest.count_ones() as usize));
463
464 Ok(inserted_nodes)
465 }
466
467 // HELPER METHODS
468 // --------------------------------------------------------------------------------------------
469
470 /// Returns true if this [PartialMmr] tracks authentication path for the node at the specified
471 /// index.
472 fn is_tracked_node(&self, node_index: &InOrderIndex) -> bool {
473 if node_index.is_leaf() {
474 self.nodes.contains_key(&node_index.sibling())
475 } else {
476 let left_child = node_index.left_child();
477 let right_child = node_index.right_child();
478 self.nodes.contains_key(&left_child) | self.nodes.contains_key(&right_child)
479 }
480 }
481}
482
483// CONVERSIONS
484// ================================================================================================
485
486impl From<MmrPeaks> for PartialMmr {
487 fn from(peaks: MmrPeaks) -> Self {
488 Self::from_peaks(peaks)
489 }
490}
491
492impl From<PartialMmr> for MmrPeaks {
493 fn from(partial_mmr: PartialMmr) -> Self {
494 // Safety: the [PartialMmr] maintains the constraints the number of true bits in the forest
495 // matches the number of peaks, as required by the [MmrPeaks]
496 MmrPeaks::new(partial_mmr.forest, partial_mmr.peaks).unwrap()
497 }
498}
499
500impl From<&MmrPeaks> for PartialMmr {
501 fn from(peaks: &MmrPeaks) -> Self {
502 Self::from_peaks(peaks.clone())
503 }
504}
505
506impl From<&PartialMmr> for MmrPeaks {
507 fn from(partial_mmr: &PartialMmr) -> Self {
508 // Safety: the [PartialMmr] maintains the constraints the number of true bits in the forest
509 // matches the number of peaks, as required by the [MmrPeaks]
510 MmrPeaks::new(partial_mmr.forest, partial_mmr.peaks.clone()).unwrap()
511 }
512}
513
514// ITERATORS
515// ================================================================================================
516
517/// An iterator over every inner node of the [PartialMmr].
518pub struct InnerNodeIterator<'a, I: Iterator<Item = (usize, RpoDigest)>> {
519 nodes: &'a BTreeMap<InOrderIndex, RpoDigest>,
520 leaves: I,
521 stack: Vec<(InOrderIndex, RpoDigest)>,
522 seen_nodes: BTreeSet<InOrderIndex>,
523}
524
525impl<'a, I: Iterator<Item = (usize, RpoDigest)>> Iterator for InnerNodeIterator<'a, I> {
526 type Item = InnerNodeInfo;
527
528 fn next(&mut self) -> Option<Self::Item> {
529 while let Some((idx, node)) = self.stack.pop() {
530 let parent_idx = idx.parent();
531 let new_node = self.seen_nodes.insert(parent_idx);
532
533 // if we haven't seen this node's parent before, and the node has a sibling, return
534 // the inner node defined by the parent of this node, and move up the branch
535 if new_node {
536 if let Some(sibling) = self.nodes.get(&idx.sibling()) {
537 let (left, right) = if parent_idx.left_child() == idx {
538 (node, *sibling)
539 } else {
540 (*sibling, node)
541 };
542 let parent = Rpo256::merge(&[left, right]);
543 let inner_node = InnerNodeInfo { value: parent, left, right };
544
545 self.stack.push((parent_idx, parent));
546 return Some(inner_node);
547 }
548 }
549
550 // the previous leaf has been processed, try to process the next leaf
551 if let Some((pos, leaf)) = self.leaves.next() {
552 let idx = InOrderIndex::from_leaf_pos(pos);
553 self.stack.push((idx, leaf));
554 }
555 }
556
557 None
558 }
559}
560
561// UTILS
562// ================================================================================================
563
564/// Given the description of a `forest`, returns the index of the root element of the smallest tree
565/// in it.
566fn forest_to_root_index(forest: usize) -> InOrderIndex {
567 // Count total size of all trees in the forest.
568 let nodes = nodes_in_forest(forest);
569
570 // Add the count for the parent nodes that separate each tree. These are allocated but
571 // currently empty, and correspond to the nodes that will be used once the trees are merged.
572 let open_trees = (forest.count_ones() - 1) as usize;
573
574 // Remove the count of the right subtree of the target tree, target tree root index comes
575 // before the subtree for the in-order tree walk.
576 let right_subtree_count = ((1u32 << forest.trailing_zeros()) - 1) as usize;
577
578 let idx = nodes + open_trees - right_subtree_count;
579
580 InOrderIndex::new(idx.try_into().unwrap())
581}
582
583/// Given the description of a `forest`, returns the index of the right most element.
584fn forest_to_rightmost_index(forest: usize) -> InOrderIndex {
585 // Count total size of all trees in the forest.
586 let nodes = nodes_in_forest(forest);
587
588 // Add the count for the parent nodes that separate each tree. These are allocated but
589 // currently empty, and correspond to the nodes that will be used once the trees are merged.
590 let open_trees = (forest.count_ones() - 1) as usize;
591
592 let idx = nodes + open_trees;
593
594 InOrderIndex::new(idx.try_into().unwrap())
595}
596
597// TESTS
598// ================================================================================================
599
600#[cfg(test)]
601mod tests {
602 use super::{
603 forest_to_rightmost_index, forest_to_root_index, BTreeSet, InOrderIndex, MmrPeaks,
604 PartialMmr, RpoDigest, Vec,
605 };
606 use crate::merkle::{int_to_node, MerkleStore, Mmr, NodeIndex};
607
608 const LEAVES: [RpoDigest; 7] = [
609 int_to_node(0),
610 int_to_node(1),
611 int_to_node(2),
612 int_to_node(3),
613 int_to_node(4),
614 int_to_node(5),
615 int_to_node(6),
616 ];
617
618 #[test]
619 fn test_forest_to_root_index() {
620 fn idx(pos: usize) -> InOrderIndex {
621 InOrderIndex::new(pos.try_into().unwrap())
622 }
623
624 // When there is a single tree in the forest, the index is equivalent to the number of
625 // leaves in that tree, which is `2^n`.
626 assert_eq!(forest_to_root_index(0b0001), idx(1));
627 assert_eq!(forest_to_root_index(0b0010), idx(2));
628 assert_eq!(forest_to_root_index(0b0100), idx(4));
629 assert_eq!(forest_to_root_index(0b1000), idx(8));
630
631 assert_eq!(forest_to_root_index(0b0011), idx(5));
632 assert_eq!(forest_to_root_index(0b0101), idx(9));
633 assert_eq!(forest_to_root_index(0b1001), idx(17));
634 assert_eq!(forest_to_root_index(0b0111), idx(13));
635 assert_eq!(forest_to_root_index(0b1011), idx(21));
636 assert_eq!(forest_to_root_index(0b1111), idx(29));
637
638 assert_eq!(forest_to_root_index(0b0110), idx(10));
639 assert_eq!(forest_to_root_index(0b1010), idx(18));
640 assert_eq!(forest_to_root_index(0b1100), idx(20));
641 assert_eq!(forest_to_root_index(0b1110), idx(26));
642 }
643
644 #[test]
645 fn test_forest_to_rightmost_index() {
646 fn idx(pos: usize) -> InOrderIndex {
647 InOrderIndex::new(pos.try_into().unwrap())
648 }
649
650 for forest in 1..256 {
651 assert!(forest_to_rightmost_index(forest).inner() % 2 == 1, "Leaves are always odd");
652 }
653
654 assert_eq!(forest_to_rightmost_index(0b0001), idx(1));
655 assert_eq!(forest_to_rightmost_index(0b0010), idx(3));
656 assert_eq!(forest_to_rightmost_index(0b0011), idx(5));
657 assert_eq!(forest_to_rightmost_index(0b0100), idx(7));
658 assert_eq!(forest_to_rightmost_index(0b0101), idx(9));
659 assert_eq!(forest_to_rightmost_index(0b0110), idx(11));
660 assert_eq!(forest_to_rightmost_index(0b0111), idx(13));
661 assert_eq!(forest_to_rightmost_index(0b1000), idx(15));
662 assert_eq!(forest_to_rightmost_index(0b1001), idx(17));
663 assert_eq!(forest_to_rightmost_index(0b1010), idx(19));
664 assert_eq!(forest_to_rightmost_index(0b1011), idx(21));
665 assert_eq!(forest_to_rightmost_index(0b1100), idx(23));
666 assert_eq!(forest_to_rightmost_index(0b1101), idx(25));
667 assert_eq!(forest_to_rightmost_index(0b1110), idx(27));
668 assert_eq!(forest_to_rightmost_index(0b1111), idx(29));
669 }
670
671 #[test]
672 fn test_partial_mmr_apply_delta() {
673 // build an MMR with 10 nodes (2 peaks) and a partial MMR based on it
674 let mut mmr = Mmr::default();
675 (0..10).for_each(|i| mmr.add(int_to_node(i)));
676 let mut partial_mmr: PartialMmr = mmr.peaks(mmr.forest()).unwrap().into();
677
678 // add authentication path for position 1 and 8
679 {
680 let node = mmr.get(1).unwrap();
681 let proof = mmr.open(1, mmr.forest()).unwrap();
682 partial_mmr.track(1, node, &proof.merkle_path).unwrap();
683 }
684
685 {
686 let node = mmr.get(8).unwrap();
687 let proof = mmr.open(8, mmr.forest()).unwrap();
688 partial_mmr.track(8, node, &proof.merkle_path).unwrap();
689 }
690
691 // add 2 more nodes into the MMR and validate apply_delta()
692 (10..12).for_each(|i| mmr.add(int_to_node(i)));
693 validate_apply_delta(&mmr, &mut partial_mmr);
694
695 // add 1 more node to the MMR, validate apply_delta() and start tracking the node
696 mmr.add(int_to_node(12));
697 validate_apply_delta(&mmr, &mut partial_mmr);
698 {
699 let node = mmr.get(12).unwrap();
700 let proof = mmr.open(12, mmr.forest()).unwrap();
701 partial_mmr.track(12, node, &proof.merkle_path).unwrap();
702 assert!(partial_mmr.track_latest);
703 }
704
705 // by this point we are tracking authentication paths for positions: 1, 8, and 12
706
707 // add 3 more nodes to the MMR (collapses to 1 peak) and validate apply_delta()
708 (13..16).for_each(|i| mmr.add(int_to_node(i)));
709 validate_apply_delta(&mmr, &mut partial_mmr);
710 }
711
712 fn validate_apply_delta(mmr: &Mmr, partial: &mut PartialMmr) {
713 let tracked_leaves = partial
714 .nodes
715 .iter()
716 .filter_map(|(index, _)| if index.is_leaf() { Some(index.sibling()) } else { None })
717 .collect::<Vec<_>>();
718 let nodes_before = partial.nodes.clone();
719
720 // compute and apply delta
721 let delta = mmr.get_delta(partial.forest(), mmr.forest()).unwrap();
722 let nodes_delta = partial.apply(delta).unwrap();
723
724 // new peaks were computed correctly
725 assert_eq!(mmr.peaks(mmr.forest()).unwrap(), partial.peaks());
726
727 let mut expected_nodes = nodes_before;
728 for (key, value) in nodes_delta {
729 // nodes should not be duplicated
730 assert!(expected_nodes.insert(key, value).is_none());
731 }
732
733 // new nodes should be a combination of original nodes and delta
734 assert_eq!(expected_nodes, partial.nodes);
735
736 // make sure tracked leaves open to the same proofs as in the underlying MMR
737 for index in tracked_leaves {
738 let index_value: u64 = index.into();
739 let pos = index_value / 2;
740 let proof1 = partial.open(pos as usize).unwrap().unwrap();
741 let proof2 = mmr.open(pos as usize, mmr.forest()).unwrap();
742 assert_eq!(proof1, proof2);
743 }
744 }
745
746 #[test]
747 fn test_partial_mmr_inner_nodes_iterator() {
748 // build the MMR
749 let mmr: Mmr = LEAVES.into();
750 let first_peak = mmr.peaks(mmr.forest).unwrap().peaks()[0];
751
752 // -- test single tree ----------------------------
753
754 // get path and node for position 1
755 let node1 = mmr.get(1).unwrap();
756 let proof1 = mmr.open(1, mmr.forest()).unwrap();
757
758 // create partial MMR and add authentication path to node at position 1
759 let mut partial_mmr: PartialMmr = mmr.peaks(mmr.forest()).unwrap().into();
760 partial_mmr.track(1, node1, &proof1.merkle_path).unwrap();
761
762 // empty iterator should have no nodes
763 assert_eq!(partial_mmr.inner_nodes([].iter().cloned()).next(), None);
764
765 // build Merkle store from authentication paths in partial MMR
766 let mut store: MerkleStore = MerkleStore::new();
767 store.extend(partial_mmr.inner_nodes([(1, node1)].iter().cloned()));
768
769 let index1 = NodeIndex::new(2, 1).unwrap();
770 let path1 = store.get_path(first_peak, index1).unwrap().path;
771
772 assert_eq!(path1, proof1.merkle_path);
773
774 // -- test no duplicates --------------------------
775
776 // build the partial MMR
777 let mut partial_mmr: PartialMmr = mmr.peaks(mmr.forest()).unwrap().into();
778
779 let node0 = mmr.get(0).unwrap();
780 let proof0 = mmr.open(0, mmr.forest()).unwrap();
781
782 let node2 = mmr.get(2).unwrap();
783 let proof2 = mmr.open(2, mmr.forest()).unwrap();
784
785 partial_mmr.track(0, node0, &proof0.merkle_path).unwrap();
786 partial_mmr.track(1, node1, &proof1.merkle_path).unwrap();
787 partial_mmr.track(2, node2, &proof2.merkle_path).unwrap();
788
789 // make sure there are no duplicates
790 let leaves = [(0, node0), (1, node1), (2, node2)];
791 let mut nodes = BTreeSet::new();
792 for node in partial_mmr.inner_nodes(leaves.iter().cloned()) {
793 assert!(nodes.insert(node.value));
794 }
795
796 // and also that the store is still be built correctly
797 store.extend(partial_mmr.inner_nodes(leaves.iter().cloned()));
798
799 let index0 = NodeIndex::new(2, 0).unwrap();
800 let index1 = NodeIndex::new(2, 1).unwrap();
801 let index2 = NodeIndex::new(2, 2).unwrap();
802
803 let path0 = store.get_path(first_peak, index0).unwrap().path;
804 let path1 = store.get_path(first_peak, index1).unwrap().path;
805 let path2 = store.get_path(first_peak, index2).unwrap().path;
806
807 assert_eq!(path0, proof0.merkle_path);
808 assert_eq!(path1, proof1.merkle_path);
809 assert_eq!(path2, proof2.merkle_path);
810
811 // -- test multiple trees -------------------------
812
813 // build the partial MMR
814 let mut partial_mmr: PartialMmr = mmr.peaks(mmr.forest()).unwrap().into();
815
816 let node5 = mmr.get(5).unwrap();
817 let proof5 = mmr.open(5, mmr.forest()).unwrap();
818
819 partial_mmr.track(1, node1, &proof1.merkle_path).unwrap();
820 partial_mmr.track(5, node5, &proof5.merkle_path).unwrap();
821
822 // build Merkle store from authentication paths in partial MMR
823 let mut store: MerkleStore = MerkleStore::new();
824 store.extend(partial_mmr.inner_nodes([(1, node1), (5, node5)].iter().cloned()));
825
826 let index1 = NodeIndex::new(2, 1).unwrap();
827 let index5 = NodeIndex::new(1, 1).unwrap();
828
829 let second_peak = mmr.peaks(mmr.forest).unwrap().peaks()[1];
830
831 let path1 = store.get_path(first_peak, index1).unwrap().path;
832 let path5 = store.get_path(second_peak, index5).unwrap().path;
833
834 assert_eq!(path1, proof1.merkle_path);
835 assert_eq!(path5, proof5.merkle_path);
836 }
837
838 #[test]
839 fn test_partial_mmr_add_without_track() {
840 let mut mmr = Mmr::default();
841 let empty_peaks = MmrPeaks::new(0, vec![]).unwrap();
842 let mut partial_mmr = PartialMmr::from_peaks(empty_peaks);
843
844 for el in (0..256).map(int_to_node) {
845 mmr.add(el);
846 partial_mmr.add(el, false);
847
848 let mmr_peaks = mmr.peaks(mmr.forest()).unwrap();
849 assert_eq!(mmr_peaks, partial_mmr.peaks());
850 assert_eq!(mmr.forest(), partial_mmr.forest());
851 }
852 }
853
854 #[test]
855 fn test_partial_mmr_add_with_track() {
856 let mut mmr = Mmr::default();
857 let empty_peaks = MmrPeaks::new(0, vec![]).unwrap();
858 let mut partial_mmr = PartialMmr::from_peaks(empty_peaks);
859
860 for i in 0..256 {
861 let el = int_to_node(i);
862 mmr.add(el);
863 partial_mmr.add(el, true);
864
865 let mmr_peaks = mmr.peaks(mmr.forest()).unwrap();
866 assert_eq!(mmr_peaks, partial_mmr.peaks());
867 assert_eq!(mmr.forest(), partial_mmr.forest());
868
869 for pos in 0..i {
870 let mmr_proof = mmr.open(pos as usize, mmr.forest()).unwrap();
871 let partialmmr_proof = partial_mmr.open(pos as usize).unwrap().unwrap();
872 assert_eq!(mmr_proof, partialmmr_proof);
873 }
874 }
875 }
876
877 #[test]
878 fn test_partial_mmr_add_existing_track() {
879 let mut mmr = Mmr::from((0..7).map(int_to_node));
880
881 // derive a partial Mmr from it which tracks authentication path to leaf 5
882 let mut partial_mmr = PartialMmr::from_peaks(mmr.peaks(mmr.forest()).unwrap());
883 let path_to_5 = mmr.open(5, mmr.forest()).unwrap().merkle_path;
884 let leaf_at_5 = mmr.get(5).unwrap();
885 partial_mmr.track(5, leaf_at_5, &path_to_5).unwrap();
886
887 // add a new leaf to both Mmr and partial Mmr
888 let leaf_at_7 = int_to_node(7);
889 mmr.add(leaf_at_7);
890 partial_mmr.add(leaf_at_7, false);
891
892 // the openings should be the same
893 assert_eq!(mmr.open(5, mmr.forest()).unwrap(), partial_mmr.open(5).unwrap().unwrap());
894 }
895}