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