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