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