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