miden-crypto/src/merkle/mmr/partial.rs
src/merkle/mmr/partial.rsRaw
1use super::{MmrDelta, MmrProof, Rpo256, RpoDigest};
2use crate::{
3 merkle::{
4 mmr::{leaf_to_corresponding_tree, nodes_in_forest},
5 InOrderIndex, InnerNodeInfo, MerklePath, MmrError, MmrPeaks,
6 },
7 utils::{
8 collections::{BTreeMap, BTreeSet, Vec},
9 vec,
10 },
11};
12
13// PARTIAL MERKLE MOUNTAIN RANGE
14// ================================================================================================
15/// Partially materialized Merkle Mountain Range (MMR), used to efficiently store and update the
16/// authentication paths for a subset of the elements in a full MMR.
17///
18/// This structure store only the authentication path for a value, the value itself is stored
19/// separately.
20#[derive(Debug, Clone, PartialEq, Eq)]
21pub struct PartialMmr {
22 /// The version of the MMR.
23 ///
24 /// This value serves the following purposes:
25 ///
26 /// - The forest is a counter for the total number of elements in the MMR.
27 /// - Since the MMR is an append-only structure, every change to it causes a change to the
28 /// `forest`, so this value has a dual purpose as a version tag.
29 /// - The bits in the forest also corresponds to the count and size of every perfect binary
30 /// tree that composes the MMR structure, which server to compute indexes and perform
31 /// validation.
32 pub(crate) forest: usize,
33
34 /// The MMR peaks.
35 ///
36 /// The peaks are used for two reasons:
37 ///
38 /// 1. It authenticates the addition of an element to the [PartialMmr], ensuring only valid
39 /// elements are tracked.
40 /// 2. During a MMR update peaks can be merged by hashing the left and right hand sides. The
41 /// peaks are used as the left hand.
42 ///
43 /// All the peaks of every tree in the MMR forest. The peaks are always ordered by number of
44 /// leaves, starting from the peak with most children, to the one with least.
45 pub(crate) peaks: Vec<RpoDigest>,
46
47 /// Authentication nodes used to construct merkle paths for a subset of the MMR's leaves.
48 ///
49 /// This does not include the MMR's peaks nor the tracked nodes, only the elements required
50 /// to construct their authentication paths. This property is used to detect when elements can
51 /// be safely removed from, because they are no longer required to authenticate any element in
52 /// the [PartialMmr].
53 ///
54 /// The elements in the MMR are referenced using a in-order tree index. This indexing scheme
55 /// permits for easy computation of the relative nodes (left/right children, sibling, parent),
56 /// which is useful for traversal. The indexing is also stable, meaning that merges to the
57 /// trees in the MMR can be represented without rewrites of the indexes.
58 pub(crate) nodes: BTreeMap<InOrderIndex, RpoDigest>,
59
60 /// Flag indicating if the odd element should be tracked.
61 ///
62 /// This flag is necessary because the sibling of the odd doesn't exist yet, so it can not be
63 /// added into `nodes` to signal the value is being tracked.
64 pub(crate) track_latest: bool,
65}
66
67impl PartialMmr {
68 // CONSTRUCTORS
69 // --------------------------------------------------------------------------------------------
70
71 /// Constructs a [PartialMmr] from the given [MmrPeaks].
72 pub fn from_peaks(peaks: MmrPeaks) -> Self {
73 let forest = peaks.num_leaves();
74 let peaks = peaks.peaks().to_vec();
75 let nodes = BTreeMap::new();
76 let track_latest = false;
77
78 Self { forest, peaks, nodes, track_latest }
79 }
80
81 // PUBLIC ACCESSORS
82 // --------------------------------------------------------------------------------------------
83
84 // Gets the current `forest`.
85 //
86 // This value corresponds to the version of the [PartialMmr] and the number of leaves in it.
87 pub fn forest(&self) -> usize {
88 self.forest
89 }
90
91 // Returns a reference to the current peaks in the [PartialMmr].
92 pub fn peaks(&self) -> MmrPeaks {
93 // expect() is OK here because the constructor ensures that MMR peaks can be constructed
94 // correctly
95 MmrPeaks::new(self.forest, self.peaks.clone()).expect("invalid MMR peaks")
96 }
97
98 /// Given a leaf position, returns the Merkle path to its corresponding peak. If the position
99 /// is greater-or-equal than the tree size an error is returned. If the requested value is not
100 /// tracked returns `None`.
101 ///
102 /// Note: The leaf position is the 0-indexed number corresponding to the order the leaves were
103 /// added, this corresponds to the MMR size _prior_ to adding the element. So the 1st element
104 /// has position 0, the second position 1, and so on.
105 pub fn open(&self, pos: usize) -> Result<Option<MmrProof>, MmrError> {
106 let tree_bit =
107 leaf_to_corresponding_tree(pos, self.forest).ok_or(MmrError::InvalidPosition(pos))?;
108 let depth = tree_bit as usize;
109
110 let mut nodes = Vec::with_capacity(depth);
111 let mut idx = InOrderIndex::from_leaf_pos(pos);
112
113 while let Some(node) = self.nodes.get(&idx.sibling()) {
114 nodes.push(*node);
115 idx = idx.parent();
116 }
117
118 // If there are nodes then the path must be complete, otherwise it is a bug
119 debug_assert!(nodes.is_empty() || nodes.len() == depth);
120
121 if nodes.len() != depth {
122 // The requested `pos` is not being tracked.
123 Ok(None)
124 } else {
125 Ok(Some(MmrProof {
126 forest: self.forest,
127 position: pos,
128 merkle_path: MerklePath::new(nodes),
129 }))
130 }
131 }
132
133 // ITERATORS
134 // --------------------------------------------------------------------------------------------
135
136 /// Returns an iterator over inner nodes of this [PartialMmr] for the specified leaves.
137 ///
138 /// The order of iteration is not defined. If a leaf is not presented in this partial MMR it
139 /// is silently ignored.
140 pub fn inner_nodes<'a, I: Iterator<Item = &'a (usize, RpoDigest)> + 'a>(
141 &'a self,
142 mut leaves: I,
143 ) -> impl Iterator<Item = InnerNodeInfo> + '_ {
144 let stack = if let Some((pos, leaf)) = leaves.next() {
145 let idx = InOrderIndex::from_leaf_pos(*pos);
146 vec![(idx, *leaf)]
147 } else {
148 Vec::new()
149 };
150
151 InnerNodeIterator {
152 nodes: &self.nodes,
153 leaves,
154 stack,
155 seen_nodes: BTreeSet::new(),
156 }
157 }
158
159 // STATE MUTATORS
160 // --------------------------------------------------------------------------------------------
161
162 /// Add the authentication path represented by [MerklePath] if it is valid.
163 ///
164 /// The `index` refers to the global position of the leaf in the MMR, these are 0-indexed
165 /// values assigned in a strictly monotonic fashion as elements are inserted into the MMR,
166 /// this value corresponds to the values used in the MMR structure.
167 ///
168 /// The `node` corresponds to the value at `index`, and `path` is the authentication path for
169 /// that element up to its corresponding Mmr peak. The `node` is only used to compute the root
170 /// from the authentication path to valid the data, only the authentication data is saved in
171 /// the structure. If the value is required it should be stored out-of-band.
172 pub fn add(
173 &mut self,
174 index: usize,
175 node: RpoDigest,
176 path: &MerklePath,
177 ) -> Result<(), MmrError> {
178 // Checks there is a tree with same depth as the authentication path, if not the path is
179 // invalid.
180 let tree = 1 << path.depth();
181 if tree & self.forest == 0 {
182 return Err(MmrError::UnknownPeak);
183 };
184
185 if index + 1 == self.forest
186 && path.depth() == 0
187 && self.peaks.last().map_or(false, |v| *v == node)
188 {
189 self.track_latest = true;
190 return Ok(());
191 }
192
193 // ignore the trees smaller than the target (these elements are position after the current
194 // target and don't affect the target index)
195 let target_forest = self.forest ^ (self.forest & (tree - 1));
196 let peak_pos = (target_forest.count_ones() - 1) as usize;
197
198 // translate from mmr index to merkle path
199 let path_idx = index - (target_forest ^ tree);
200
201 // Compute the root of the authentication path, and check it matches the current version of
202 // the PartialMmr.
203 let computed = path.compute_root(path_idx as u64, node).map_err(MmrError::MerkleError)?;
204 if self.peaks[peak_pos] != computed {
205 return Err(MmrError::InvalidPeak);
206 }
207
208 let mut idx = InOrderIndex::from_leaf_pos(index);
209 for node in path.nodes() {
210 self.nodes.insert(idx.sibling(), *node);
211 idx = idx.parent();
212 }
213
214 Ok(())
215 }
216
217 /// Remove a leaf of the [PartialMmr] and the unused nodes from the authentication path.
218 ///
219 /// Note: `leaf_pos` corresponds to the position in the MMR and not on an individual tree.
220 pub fn remove(&mut self, leaf_pos: usize) {
221 let mut idx = InOrderIndex::from_leaf_pos(leaf_pos);
222
223 self.nodes.remove(&idx.sibling());
224
225 // `idx` represent the element that can be computed by the authentication path, because
226 // these elements can be computed they are not saved for the authentication of the current
227 // target. In other words, if the idx is present it was added for the authentication of
228 // another element, and no more elements should be removed otherwise it would remove that
229 // element's authentication data.
230 while !self.nodes.contains_key(&idx) {
231 idx = idx.parent();
232 self.nodes.remove(&idx.sibling());
233 }
234 }
235
236 /// Applies updates to the [PartialMmr].
237 pub fn apply(&mut self, delta: MmrDelta) -> Result<(), MmrError> {
238 if delta.forest < self.forest {
239 return Err(MmrError::InvalidPeaks);
240 }
241
242 if delta.forest == self.forest {
243 if !delta.data.is_empty() {
244 return Err(MmrError::InvalidUpdate);
245 }
246
247 return Ok(());
248 }
249
250 // find the tree merges
251 let changes = self.forest ^ delta.forest;
252 let largest = 1 << changes.ilog2();
253 let merges = self.forest & (largest - 1);
254
255 debug_assert!(
256 !self.track_latest || (merges & 1) == 1,
257 "if there is an odd element, a merge is required"
258 );
259
260 // count the number elements needed to produce largest from the current state
261 let (merge_count, new_peaks) = if merges != 0 {
262 let depth = largest.trailing_zeros();
263 let skipped = merges.trailing_zeros();
264 let computed = merges.count_ones() - 1;
265 let merge_count = depth - skipped - computed;
266
267 let new_peaks = delta.forest & (largest - 1);
268
269 (merge_count, new_peaks)
270 } else {
271 (0, changes)
272 };
273
274 // verify the delta size
275 if (delta.data.len() as u32) != merge_count + new_peaks.count_ones() {
276 return Err(MmrError::InvalidUpdate);
277 }
278
279 // keeps track of how many data elements from the update have been consumed
280 let mut update_count = 0;
281
282 if merges != 0 {
283 // starts at the smallest peak and follows the merged peaks
284 let mut peak_idx = forest_to_root_index(self.forest);
285
286 // match order of the update data while applying it
287 self.peaks.reverse();
288
289 // set to true when the data is needed for authentication paths updates
290 let mut track = self.track_latest;
291 self.track_latest = false;
292
293 let mut peak_count = 0;
294 let mut target = 1 << merges.trailing_zeros();
295 let mut new = delta.data[0];
296 update_count += 1;
297
298 while target < largest {
299 // check if either the left or right subtrees have saved for authentication paths.
300 // If so, turn tracking on to update those paths.
301 if target != 1 && !track {
302 let left_child = peak_idx.left_child();
303 let right_child = peak_idx.right_child();
304 track = self.nodes.contains_key(&left_child)
305 | self.nodes.contains_key(&right_child);
306 }
307
308 // update data only contains the nodes from the right subtrees, left nodes are
309 // either previously known peaks or computed values
310 let (left, right) = if target & merges != 0 {
311 let peak = self.peaks[peak_count];
312 peak_count += 1;
313 (peak, new)
314 } else {
315 let update = delta.data[update_count];
316 update_count += 1;
317 (new, update)
318 };
319
320 if track {
321 self.nodes.insert(peak_idx.sibling(), right);
322 }
323
324 peak_idx = peak_idx.parent();
325 new = Rpo256::merge(&[left, right]);
326 target <<= 1;
327 }
328
329 debug_assert!(peak_count == (merges.count_ones() as usize));
330
331 // restore the peaks order
332 self.peaks.reverse();
333 // remove the merged peaks
334 self.peaks.truncate(self.peaks.len() - peak_count);
335 // add the newly computed peak, the result of the merges
336 self.peaks.push(new);
337 }
338
339 // The rest of the update data is composed of peaks. None of these elements can contain
340 // tracked elements because the peaks were unknown, and it is not possible to add elements
341 // for tacking without authenticating it to a peak.
342 self.peaks.extend_from_slice(&delta.data[update_count..]);
343 self.forest = delta.forest;
344
345 debug_assert!(self.peaks.len() == (self.forest.count_ones() as usize));
346
347 Ok(())
348 }
349}
350
351// CONVERSIONS
352// ================================================================================================
353
354impl From<MmrPeaks> for PartialMmr {
355 fn from(peaks: MmrPeaks) -> Self {
356 Self::from_peaks(peaks)
357 }
358}
359
360impl From<PartialMmr> for MmrPeaks {
361 fn from(partial_mmr: PartialMmr) -> Self {
362 // Safety: the [PartialMmr] maintains the constraints the number of true bits in the forest
363 // matches the number of peaks, as required by the [MmrPeaks]
364 MmrPeaks::new(partial_mmr.forest, partial_mmr.peaks).unwrap()
365 }
366}
367
368impl From<&MmrPeaks> for PartialMmr {
369 fn from(peaks: &MmrPeaks) -> Self {
370 Self::from_peaks(peaks.clone())
371 }
372}
373
374impl From<&PartialMmr> for MmrPeaks {
375 fn from(partial_mmr: &PartialMmr) -> Self {
376 // Safety: the [PartialMmr] maintains the constraints the number of true bits in the forest
377 // matches the number of peaks, as required by the [MmrPeaks]
378 MmrPeaks::new(partial_mmr.forest, partial_mmr.peaks.clone()).unwrap()
379 }
380}
381
382// ITERATORS
383// ================================================================================================
384
385/// An iterator over every inner node of the [PartialMmr].
386pub struct InnerNodeIterator<'a, I: Iterator<Item = &'a (usize, RpoDigest)>> {
387 nodes: &'a BTreeMap<InOrderIndex, RpoDigest>,
388 leaves: I,
389 stack: Vec<(InOrderIndex, RpoDigest)>,
390 seen_nodes: BTreeSet<InOrderIndex>,
391}
392
393impl<'a, I: Iterator<Item = &'a (usize, RpoDigest)>> Iterator for InnerNodeIterator<'a, I> {
394 type Item = InnerNodeInfo;
395
396 fn next(&mut self) -> Option<Self::Item> {
397 while let Some((idx, node)) = self.stack.pop() {
398 let parent_idx = idx.parent();
399 let new_node = self.seen_nodes.insert(parent_idx);
400
401 // if we haven't seen this node's parent before, and the node has a sibling, return
402 // the inner node defined by the parent of this node, and move up the branch
403 if new_node {
404 if let Some(sibling) = self.nodes.get(&idx.sibling()) {
405 let (left, right) = if parent_idx.left_child() == idx {
406 (node, *sibling)
407 } else {
408 (*sibling, node)
409 };
410 let parent = Rpo256::merge(&[left, right]);
411 let inner_node = InnerNodeInfo { value: parent, left, right };
412
413 self.stack.push((parent_idx, parent));
414 return Some(inner_node);
415 }
416 }
417
418 // the previous leaf has been processed, try to process the next leaf
419 if let Some((pos, leaf)) = self.leaves.next() {
420 let idx = InOrderIndex::from_leaf_pos(*pos);
421 self.stack.push((idx, *leaf));
422 }
423 }
424
425 None
426 }
427}
428
429// UTILS
430// ================================================================================================
431
432/// Given the description of a `forest`, returns the index of the root element of the smallest tree
433/// in it.
434pub fn forest_to_root_index(forest: usize) -> InOrderIndex {
435 // Count total size of all trees in the forest.
436 let nodes = nodes_in_forest(forest);
437
438 // Add the count for the parent nodes that separate each tree. These are allocated but
439 // currently empty, and correspond to the nodes that will be used once the trees are merged.
440 let open_trees = (forest.count_ones() - 1) as usize;
441
442 // Remove the count of the right subtree of the target tree, target tree root index comes
443 // before the subtree for the in-order tree walk.
444 let right_subtree_count = ((1u32 << forest.trailing_zeros()) - 1) as usize;
445
446 let idx = nodes + open_trees - right_subtree_count;
447
448 InOrderIndex::new(idx.try_into().unwrap())
449}
450
451// TESTS
452// ================================================================================================
453
454#[cfg(test)]
455mod test {
456 use super::{forest_to_root_index, BTreeSet, InOrderIndex, PartialMmr, RpoDigest};
457 use crate::merkle::{int_to_node, MerkleStore, Mmr, NodeIndex};
458
459 const LEAVES: [RpoDigest; 7] = [
460 int_to_node(0),
461 int_to_node(1),
462 int_to_node(2),
463 int_to_node(3),
464 int_to_node(4),
465 int_to_node(5),
466 int_to_node(6),
467 ];
468
469 #[test]
470 fn test_forest_to_root_index() {
471 fn idx(pos: usize) -> InOrderIndex {
472 InOrderIndex::new(pos.try_into().unwrap())
473 }
474
475 // When there is a single tree in the forest, the index is equivalent to the number of
476 // leaves in that tree, which is `2^n`.
477 assert_eq!(forest_to_root_index(0b0001), idx(1));
478 assert_eq!(forest_to_root_index(0b0010), idx(2));
479 assert_eq!(forest_to_root_index(0b0100), idx(4));
480 assert_eq!(forest_to_root_index(0b1000), idx(8));
481
482 assert_eq!(forest_to_root_index(0b0011), idx(5));
483 assert_eq!(forest_to_root_index(0b0101), idx(9));
484 assert_eq!(forest_to_root_index(0b1001), idx(17));
485 assert_eq!(forest_to_root_index(0b0111), idx(13));
486 assert_eq!(forest_to_root_index(0b1011), idx(21));
487 assert_eq!(forest_to_root_index(0b1111), idx(29));
488
489 assert_eq!(forest_to_root_index(0b0110), idx(10));
490 assert_eq!(forest_to_root_index(0b1010), idx(18));
491 assert_eq!(forest_to_root_index(0b1100), idx(20));
492 assert_eq!(forest_to_root_index(0b1110), idx(26));
493 }
494
495 #[test]
496 fn test_partial_mmr_inner_nodes_iterator() {
497 // build the MMR
498 let mmr: Mmr = LEAVES.into();
499 let first_peak = mmr.peaks(mmr.forest).unwrap().peaks()[0];
500
501 // -- test single tree ----------------------------
502
503 // get path and node for position 1
504 let node1 = mmr.get(1).unwrap();
505 let proof1 = mmr.open(1, mmr.forest()).unwrap();
506
507 // create partial MMR and add authentication path to node at position 1
508 let mut partial_mmr: PartialMmr = mmr.peaks(mmr.forest()).unwrap().into();
509 partial_mmr.add(1, node1, &proof1.merkle_path).unwrap();
510
511 // empty iterator should have no nodes
512 assert_eq!(partial_mmr.inner_nodes([].iter()).next(), None);
513
514 // build Merkle store from authentication paths in partial MMR
515 let mut store: MerkleStore = MerkleStore::new();
516 store.extend(partial_mmr.inner_nodes([(1, node1)].iter()));
517
518 let index1 = NodeIndex::new(2, 1).unwrap();
519 let path1 = store.get_path(first_peak, index1).unwrap().path;
520
521 assert_eq!(path1, proof1.merkle_path);
522
523 // -- test no duplicates --------------------------
524
525 // build the partial MMR
526 let mut partial_mmr: PartialMmr = mmr.peaks(mmr.forest()).unwrap().into();
527
528 let node0 = mmr.get(0).unwrap();
529 let proof0 = mmr.open(0, mmr.forest()).unwrap();
530
531 let node2 = mmr.get(2).unwrap();
532 let proof2 = mmr.open(2, mmr.forest()).unwrap();
533
534 partial_mmr.add(0, node0, &proof0.merkle_path).unwrap();
535 partial_mmr.add(1, node1, &proof1.merkle_path).unwrap();
536 partial_mmr.add(2, node2, &proof2.merkle_path).unwrap();
537
538 // make sure there are no duplicates
539 let leaves = [(0, node0), (1, node1), (2, node2)];
540 let mut nodes = BTreeSet::new();
541 for node in partial_mmr.inner_nodes(leaves.iter()) {
542 assert!(nodes.insert(node.value));
543 }
544
545 // and also that the store is still be built correctly
546 store.extend(partial_mmr.inner_nodes(leaves.iter()));
547
548 let index0 = NodeIndex::new(2, 0).unwrap();
549 let index1 = NodeIndex::new(2, 1).unwrap();
550 let index2 = NodeIndex::new(2, 2).unwrap();
551
552 let path0 = store.get_path(first_peak, index0).unwrap().path;
553 let path1 = store.get_path(first_peak, index1).unwrap().path;
554 let path2 = store.get_path(first_peak, index2).unwrap().path;
555
556 assert_eq!(path0, proof0.merkle_path);
557 assert_eq!(path1, proof1.merkle_path);
558 assert_eq!(path2, proof2.merkle_path);
559
560 // -- test multiple trees -------------------------
561
562 // build the partial MMR
563 let mut partial_mmr: PartialMmr = mmr.peaks(mmr.forest()).unwrap().into();
564
565 let node5 = mmr.get(5).unwrap();
566 let proof5 = mmr.open(5, mmr.forest()).unwrap();
567
568 partial_mmr.add(1, node1, &proof1.merkle_path).unwrap();
569 partial_mmr.add(5, node5, &proof5.merkle_path).unwrap();
570
571 // build Merkle store from authentication paths in partial MMR
572 let mut store: MerkleStore = MerkleStore::new();
573 store.extend(partial_mmr.inner_nodes([(1, node1), (5, node5)].iter()));
574
575 let index1 = NodeIndex::new(2, 1).unwrap();
576 let index5 = NodeIndex::new(1, 1).unwrap();
577
578 let second_peak = mmr.peaks(mmr.forest).unwrap().peaks()[1];
579
580 let path1 = store.get_path(first_peak, index1).unwrap().path;
581 let path5 = store.get_path(second_peak, index5).unwrap().path;
582
583 assert_eq!(path1, proof1.merkle_path);
584 assert_eq!(path5, proof5.merkle_path);
585 }
586}