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