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