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