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