miden-crypto/src/merkle/mmr/full.rs
src/merkle/mmr/full.rsRaw
1//! A fully materialized Merkle mountain range (MMR).
2//!
3//! A MMR is a forest structure, i.e. it is an ordered set of disjoint rooted trees. The trees are
4//! ordered by size, from the most to least number of leaves. Every tree is a perfect binary tree,
5//! meaning a tree has all its leaves at the same depth, and every inner node has a branch-factor
6//! of 2 with both children set.
7//!
8//! Additionally the structure only supports adding leaves to the right-most tree, the one with the
9//! least number of leaves. The structure preserves the invariant that each tree has different
10//! depths, i.e. as part of adding adding a new element to the forest the trees with same depth are
11//! merged, creating a new tree with depth d+1, this process is continued until the property is
12//! reestablished.
13use super::{
14 super::{InnerNodeInfo, MerklePath},
15 bit::TrueBitPositionIterator,
16 leaf_to_corresponding_tree, nodes_in_forest, MmrDelta, MmrError, MmrPeaks, MmrProof, Rpo256,
17 RpoDigest,
18};
19use crate::utils::collections::*;
20
21// MMR
22// ===============================================================================================
23
24/// A fully materialized Merkle Mountain Range, with every tree in the forest and all their
25/// elements.
26///
27/// Since this is a full representation of the MMR, elements are never removed and the MMR will
28/// grow roughly `O(2n)` in number of leaf elements.
29#[derive(Debug, Clone)]
30#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
31pub struct Mmr {
32 /// Refer to the `forest` method documentation for details of the semantics of this value.
33 pub(super) forest: usize,
34
35 /// Contains every element of the forest.
36 ///
37 /// The trees are in postorder sequential representation. This representation allows for all
38 /// the elements of every tree in the forest to be stored in the same sequential buffer. It
39 /// also means new elements can be added to the forest, and merging of trees is very cheap with
40 /// no need to copy elements.
41 pub(super) nodes: Vec<RpoDigest>,
42}
43
44impl Default for Mmr {
45 fn default() -> Self {
46 Self::new()
47 }
48}
49
50impl Mmr {
51 // CONSTRUCTORS
52 // ============================================================================================
53
54 /// Constructor for an empty `Mmr`.
55 pub fn new() -> Mmr {
56 Mmr { forest: 0, nodes: Vec::new() }
57 }
58
59 // ACCESSORS
60 // ============================================================================================
61
62 /// Returns the MMR forest representation.
63 ///
64 /// The forest value has the following interpretations:
65 /// - its value is the number of elements in the forest
66 /// - bit count corresponds to the number of trees in the forest
67 /// - each true bit position determines the depth of a tree in the forest
68 pub const fn forest(&self) -> usize {
69 self.forest
70 }
71
72 // FUNCTIONALITY
73 // ============================================================================================
74
75 /// Given a leaf position, returns the Merkle path to its corresponding peak. If the position
76 /// is greater-or-equal than the tree size an error is returned.
77 ///
78 /// Note: The leaf position is the 0-indexed number corresponding to the order the leaves were
79 /// added, this corresponds to the MMR size _prior_ to adding the element. So the 1st element
80 /// has position 0, the second position 1, and so on.
81 pub fn open(&self, pos: usize, target_forest: usize) -> Result<MmrProof, MmrError> {
82 // find the target tree responsible for the MMR position
83 let tree_bit =
84 leaf_to_corresponding_tree(pos, target_forest).ok_or(MmrError::InvalidPosition(pos))?;
85
86 // isolate the trees before the target
87 let forest_before = target_forest & high_bitmask(tree_bit + 1);
88 let index_offset = nodes_in_forest(forest_before);
89
90 // update the value position from global to the target tree
91 let relative_pos = pos - forest_before;
92
93 // collect the path and the final index of the target value
94 let (_, path) = self.collect_merkle_path_and_value(tree_bit, relative_pos, index_offset);
95
96 Ok(MmrProof {
97 forest: target_forest,
98 position: pos,
99 merkle_path: MerklePath::new(path),
100 })
101 }
102
103 /// Returns the leaf value at position `pos`.
104 ///
105 /// Note: The leaf position is the 0-indexed number corresponding to the order the leaves were
106 /// added, this corresponds to the MMR size _prior_ to adding the element. So the 1st element
107 /// has position 0, the second position 1, and so on.
108 pub fn get(&self, pos: usize) -> Result<RpoDigest, MmrError> {
109 // find the target tree responsible for the MMR position
110 let tree_bit =
111 leaf_to_corresponding_tree(pos, self.forest).ok_or(MmrError::InvalidPosition(pos))?;
112
113 // isolate the trees before the target
114 let forest_before = self.forest & high_bitmask(tree_bit + 1);
115 let index_offset = nodes_in_forest(forest_before);
116
117 // update the value position from global to the target tree
118 let relative_pos = pos - forest_before;
119
120 // collect the path and the final index of the target value
121 let (value, _) = self.collect_merkle_path_and_value(tree_bit, relative_pos, index_offset);
122
123 Ok(value)
124 }
125
126 /// Adds a new element to the MMR.
127 pub fn add(&mut self, el: RpoDigest) {
128 // Note: every node is also a tree of size 1, adding an element to the forest creates a new
129 // rooted-tree of size 1. This may temporarily break the invariant that every tree in the
130 // forest has different sizes, the loop below will eagerly merge trees of same size and
131 // restore the invariant.
132 self.nodes.push(el);
133
134 let mut left_offset = self.nodes.len().saturating_sub(2);
135 let mut right = el;
136 let mut left_tree = 1;
137 while self.forest & left_tree != 0 {
138 right = Rpo256::merge(&[self.nodes[left_offset], right]);
139 self.nodes.push(right);
140
141 left_offset = left_offset.saturating_sub(nodes_in_forest(left_tree));
142 left_tree <<= 1;
143 }
144
145 self.forest += 1;
146 }
147
148 /// Returns an peaks of the MMR for the version specified by `forest`.
149 pub fn peaks(&self, forest: usize) -> Result<MmrPeaks, MmrError> {
150 if forest > self.forest {
151 return Err(MmrError::InvalidPeaks);
152 }
153
154 let peaks: Vec<RpoDigest> = TrueBitPositionIterator::new(forest)
155 .rev()
156 .map(|bit| nodes_in_forest(1 << bit))
157 .scan(0, |offset, el| {
158 *offset += el;
159 Some(*offset)
160 })
161 .map(|offset| self.nodes[offset - 1])
162 .collect();
163
164 // Safety: the invariant is maintained by the [Mmr]
165 let peaks = MmrPeaks::new(forest, peaks).unwrap();
166
167 Ok(peaks)
168 }
169
170 /// Compute the required update to `original_forest`.
171 ///
172 /// The result is a packed sequence of the authentication elements required to update the trees
173 /// that have been merged together, followed by the new peaks of the [Mmr].
174 pub fn get_delta(&self, from_forest: usize, to_forest: usize) -> Result<MmrDelta, MmrError> {
175 if to_forest > self.forest || from_forest > to_forest {
176 return Err(MmrError::InvalidPeaks);
177 }
178
179 if from_forest == to_forest {
180 return Ok(MmrDelta { forest: to_forest, data: Vec::new() });
181 }
182
183 let mut result = Vec::new();
184
185 // Find the largest tree in this [Mmr] which is new to `from_forest`.
186 let candidate_trees = to_forest ^ from_forest;
187 let mut new_high = 1 << candidate_trees.ilog2();
188
189 // Collect authentication nodes used for tree merges
190 // ----------------------------------------------------------------------------------------
191
192 // Find the trees from `from_forest` that have been merged into `new_high`.
193 let mut merges = from_forest & (new_high - 1);
194
195 // Find the peaks that are common to `from_forest` and this [Mmr]
196 let common_trees = from_forest ^ merges;
197
198 if merges != 0 {
199 // Skip the smallest trees unknown to `from_forest`.
200 let mut target = 1 << merges.trailing_zeros();
201
202 // Collect siblings required to computed the merged tree's peak
203 while target < new_high {
204 // Computes the offset to the smallest know peak
205 // - common_trees: peaks unchanged in the current update, target comes after these.
206 // - merges: peaks that have not been merged so far, target comes after these.
207 // - target: tree from which to load the sibling. On the first iteration this is a
208 // value known by the partial mmr, on subsequent iterations this value is to be
209 // computed from the known peaks and provided authentication nodes.
210 let known = nodes_in_forest(common_trees | merges | target);
211 let sibling = nodes_in_forest(target);
212 result.push(self.nodes[known + sibling - 1]);
213
214 // Update the target and account for tree merges
215 target <<= 1;
216 while merges & target != 0 {
217 target <<= 1;
218 }
219 // Remove the merges done so far
220 merges ^= merges & (target - 1);
221 }
222 } else {
223 // The new high tree may not be the result of any merges, if it is smaller than all the
224 // trees of `from_forest`.
225 new_high = 0;
226 }
227
228 // Collect the new [Mmr] peaks
229 // ----------------------------------------------------------------------------------------
230
231 let mut new_peaks = to_forest ^ common_trees ^ new_high;
232 let old_peaks = to_forest ^ new_peaks;
233 let mut offset = nodes_in_forest(old_peaks);
234 while new_peaks != 0 {
235 let target = 1 << new_peaks.ilog2();
236 offset += nodes_in_forest(target);
237 result.push(self.nodes[offset - 1]);
238 new_peaks ^= target;
239 }
240
241 Ok(MmrDelta { forest: to_forest, data: result })
242 }
243
244 /// An iterator over inner nodes in the MMR. The order of iteration is unspecified.
245 pub fn inner_nodes(&self) -> MmrNodes {
246 MmrNodes {
247 mmr: self,
248 forest: 0,
249 last_right: 0,
250 index: 0,
251 }
252 }
253
254 // UTILITIES
255 // ============================================================================================
256
257 /// Internal function used to collect the Merkle path of a value.
258 ///
259 /// The arguments are relative to the target tree. To compute the opening of the second leaf
260 /// for a tree with depth 2 in the forest `0b110`:
261 ///
262 /// - `tree_bit`: Depth of the target tree, e.g. 2 for the smallest tree.
263 /// - `relative_pos`: 0-indexed leaf position in the target tree, e.g. 1 for the second leaf.
264 /// - `index_offset`: Node count prior to the target tree, e.g. 7 for the tree of depth 3.
265 fn collect_merkle_path_and_value(
266 &self,
267 tree_bit: u32,
268 relative_pos: usize,
269 index_offset: usize,
270 ) -> (RpoDigest, Vec<RpoDigest>) {
271 // see documentation of `leaf_to_corresponding_tree` for details
272 let tree_depth = (tree_bit + 1) as usize;
273 let mut path = Vec::with_capacity(tree_depth);
274
275 // The tree walk below goes from the root to the leaf, compute the root index to start
276 let mut forest_target = 1usize << tree_bit;
277 let mut index = nodes_in_forest(forest_target) - 1;
278
279 // Loop until the leaf is reached
280 while forest_target > 1 {
281 // Update the depth of the tree to correspond to a subtree
282 forest_target >>= 1;
283
284 // compute the indices of the right and left subtrees based on the post-order
285 let right_offset = index - 1;
286 let left_offset = right_offset - nodes_in_forest(forest_target);
287
288 let left_or_right = relative_pos & forest_target;
289 let sibling = if left_or_right != 0 {
290 // going down the right subtree, the right child becomes the new root
291 index = right_offset;
292 // and the left child is the authentication
293 self.nodes[index_offset + left_offset]
294 } else {
295 index = left_offset;
296 self.nodes[index_offset + right_offset]
297 };
298
299 path.push(sibling);
300 }
301
302 debug_assert!(path.len() == tree_depth - 1);
303
304 // the rest of the codebase has the elements going from leaf to root, adjust it here for
305 // easy of use/consistency sake
306 path.reverse();
307
308 let value = self.nodes[index_offset + index];
309 (value, path)
310 }
311}
312
313// CONVERSIONS
314// ================================================================================================
315
316impl<T> From<T> for Mmr
317where
318 T: IntoIterator<Item = RpoDigest>,
319{
320 fn from(values: T) -> Self {
321 let mut mmr = Mmr::new();
322 for v in values {
323 mmr.add(v)
324 }
325 mmr
326 }
327}
328
329// ITERATOR
330// ===============================================================================================
331
332/// Yields inner nodes of the [Mmr].
333pub struct MmrNodes<'a> {
334 /// [Mmr] being yielded, when its `forest` value is matched, the iterations is finished.
335 mmr: &'a Mmr,
336 /// Keeps track of the left nodes yielded so far waiting for a right pair, this matches the
337 /// semantics of the [Mmr]'s forest attribute, since that too works as a buffer of left nodes
338 /// waiting for a pair to be hashed together.
339 forest: usize,
340 /// Keeps track of the last right node yielded, after this value is set, the next iteration
341 /// will be its parent with its corresponding left node that has been yield already.
342 last_right: usize,
343 /// The current index in the `nodes` vector.
344 index: usize,
345}
346
347impl<'a> Iterator for MmrNodes<'a> {
348 type Item = InnerNodeInfo;
349
350 fn next(&mut self) -> Option<Self::Item> {
351 debug_assert!(self.last_right.count_ones() <= 1, "last_right tracks zero or one element");
352
353 // only parent nodes are emitted, remove the single node tree from the forest
354 let target = self.mmr.forest & (usize::MAX << 1);
355
356 if self.forest < target {
357 if self.last_right == 0 {
358 // yield the left leaf
359 debug_assert!(self.last_right == 0, "left must be before right");
360 self.forest |= 1;
361 self.index += 1;
362
363 // yield the right leaf
364 debug_assert!((self.forest & 1) == 1, "right must be after left");
365 self.last_right |= 1;
366 self.index += 1;
367 };
368
369 debug_assert!(
370 self.forest & self.last_right != 0,
371 "parent requires both a left and right",
372 );
373
374 // compute the number of nodes in the right tree, this is the offset to the
375 // previous left parent
376 let right_nodes = nodes_in_forest(self.last_right);
377 // the next parent position is one above the position of the pair
378 let parent = self.last_right << 1;
379
380 // the left node has been paired and the current parent yielded, removed it from the forest
381 self.forest ^= self.last_right;
382 if self.forest & parent == 0 {
383 // this iteration yielded the left parent node
384 debug_assert!(self.forest & 1 == 0, "next iteration yields a left leaf");
385 self.last_right = 0;
386 self.forest ^= parent;
387 } else {
388 // the left node of the parent level has been yielded already, this iteration
389 // was the right parent. Next iteration yields their parent.
390 self.last_right = parent;
391 }
392
393 // yields a parent
394 let value = self.mmr.nodes[self.index];
395 let right = self.mmr.nodes[self.index - 1];
396 let left = self.mmr.nodes[self.index - 1 - right_nodes];
397 self.index += 1;
398 let node = InnerNodeInfo { value, left, right };
399
400 Some(node)
401 } else {
402 None
403 }
404 }
405}
406
407// UTILITIES
408// ===============================================================================================
409
410/// Return a bitmask for the bits including and above the given position.
411pub(crate) const fn high_bitmask(bit: u32) -> usize {
412 if bit > usize::BITS - 1 {
413 0
414 } else {
415 usize::MAX << bit
416 }
417}