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::bit::TrueBitPositionIterator;
14use super::{super::Vec, MmrPeaks, MmrProof, Rpo256, Word};
15use crate::merkle::MerklePath;
16use core::fmt::{Display, Formatter};
17
18#[cfg(feature = "std")]
19use std::error::Error;
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.
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<Word>,
40}
41
42#[derive(Debug, PartialEq, Eq, Copy, Clone)]
43pub enum MmrError {
44 InvalidPosition(usize),
45}
46
47impl Display for MmrError {
48 fn fmt(&self, fmt: &mut Formatter<'_>) -> Result<(), core::fmt::Error> {
49 match self {
50 MmrError::InvalidPosition(pos) => write!(fmt, "Mmr does not contain position {pos}"),
51 }
52 }
53}
54
55#[cfg(feature = "std")]
56impl Error for MmrError {}
57
58impl Default for Mmr {
59 fn default() -> Self {
60 Self::new()
61 }
62}
63
64impl Mmr {
65 // CONSTRUCTORS
66 // ============================================================================================
67
68 /// Constructor for an empty `Mmr`.
69 pub fn new() -> Mmr {
70 Mmr {
71 forest: 0,
72 nodes: Vec::new(),
73 }
74 }
75
76 // ACCESSORS
77 // ============================================================================================
78
79 /// Returns the MMR forest representation.
80 ///
81 /// The forest value has the following interpretations:
82 /// - its value is the number of elements in the forest
83 /// - bit count corresponds to the number of trees in the forest
84 /// - each true bit position determines the depth of a tree in the forest
85 pub const fn forest(&self) -> usize {
86 self.forest
87 }
88
89 // FUNCTIONALITY
90 // ============================================================================================
91
92 /// Given a leaf position, returns the Merkle path to its corresponding peak. If the position
93 /// is greater-or-equal than the tree size an error is returned.
94 ///
95 /// Note: The leaf position is the 0-indexed number corresponding to the order the leaves were
96 /// added, this corresponds to the MMR size _prior_ to adding the element. So the 1st element
97 /// has position 0, the second position 1, and so on.
98 pub fn open(&self, pos: usize) -> Result<MmrProof, MmrError> {
99 // find the target tree responsible for the MMR position
100 let tree_bit =
101 leaf_to_corresponding_tree(pos, self.forest).ok_or(MmrError::InvalidPosition(pos))?;
102 let forest_target = 1usize << tree_bit;
103
104 // isolate the trees before the target
105 let forest_before = self.forest & high_bitmask(tree_bit + 1);
106 let index_offset = nodes_in_forest(forest_before);
107
108 // find the root
109 let index = nodes_in_forest(forest_target) - 1;
110
111 // update the value position from global to the target tree
112 let relative_pos = pos - forest_before;
113
114 // collect the path and the final index of the target value
115 let (_, path) =
116 self.collect_merkle_path_and_value(tree_bit, relative_pos, index_offset, index);
117
118 Ok(MmrProof {
119 forest: self.forest,
120 position: pos,
121 merkle_path: MerklePath::new(path),
122 })
123 }
124
125 /// Returns the leaf value at position `pos`.
126 ///
127 /// Note: The leaf position is the 0-indexed number corresponding to the order the leaves were
128 /// added, this corresponds to the MMR size _prior_ to adding the element. So the 1st element
129 /// has position 0, the second position 1, and so on.
130 pub fn get(&self, pos: usize) -> Result<Word, MmrError> {
131 // find the target tree responsible for the MMR position
132 let tree_bit =
133 leaf_to_corresponding_tree(pos, self.forest).ok_or(MmrError::InvalidPosition(pos))?;
134 let forest_target = 1usize << tree_bit;
135
136 // isolate the trees before the target
137 let forest_before = self.forest & high_bitmask(tree_bit + 1);
138 let index_offset = nodes_in_forest(forest_before);
139
140 // find the root
141 let index = nodes_in_forest(forest_target) - 1;
142
143 // update the value position from global to the target tree
144 let relative_pos = pos - forest_before;
145
146 // collect the path and the final index of the target value
147 let (value, _) =
148 self.collect_merkle_path_and_value(tree_bit, relative_pos, index_offset, index);
149
150 Ok(value)
151 }
152
153 /// Adds a new element to the MMR.
154 pub fn add(&mut self, el: Word) {
155 // Note: every node is also a tree of size 1, adding an element to the forest creates a new
156 // rooted-tree of size 1. This may temporarily break the invariant that every tree in the
157 // forest has different sizes, the loop below will eagerly merge trees of same size and
158 // restore the invariant.
159 self.nodes.push(el);
160
161 let mut left_offset = self.nodes.len().saturating_sub(2);
162 let mut right = el;
163 let mut left_tree = 1;
164 while self.forest & left_tree != 0 {
165 right = *Rpo256::merge(&[self.nodes[left_offset].into(), right.into()]);
166 self.nodes.push(right);
167
168 left_offset = left_offset.saturating_sub(nodes_in_forest(left_tree));
169 left_tree <<= 1;
170 }
171
172 self.forest += 1;
173 }
174
175 /// Returns an accumulator representing the current state of the MMMR.
176 pub fn accumulator(&self) -> MmrPeaks {
177 let peaks: Vec<Word> = TrueBitPositionIterator::new(self.forest)
178 .rev()
179 .map(|bit| nodes_in_forest(1 << bit))
180 .scan(0, |offset, el| {
181 *offset += el;
182 Some(*offset)
183 })
184 .map(|offset| self.nodes[offset - 1])
185 .collect();
186
187 MmrPeaks {
188 num_leaves: self.forest,
189 peaks,
190 }
191 }
192
193 // UTILITIES
194 // ============================================================================================
195
196 /// Internal function used to collect the Merkle path of a value.
197 fn collect_merkle_path_and_value(
198 &self,
199 tree_bit: u32,
200 relative_pos: usize,
201 index_offset: usize,
202 mut index: usize,
203 ) -> (Word, Vec<Word>) {
204 // collect the Merkle path
205 let mut tree_depth = tree_bit as usize;
206 let mut path = Vec::with_capacity(tree_depth + 1);
207 while tree_depth > 0 {
208 let bit = relative_pos & tree_depth;
209 let right_offset = index - 1;
210 let left_offset = right_offset - nodes_in_forest(tree_depth);
211
212 // Elements to the right have a higher position because they were
213 // added later. Therefore when the bit is true the node's path is
214 // to the right, and its sibling to the left.
215 let sibling = if bit != 0 {
216 index = right_offset;
217 self.nodes[index_offset + left_offset]
218 } else {
219 index = left_offset;
220 self.nodes[index_offset + right_offset]
221 };
222
223 tree_depth >>= 1;
224 path.push(sibling);
225 }
226
227 // the rest of the codebase has the elements going from leaf to root, adjust it here for
228 // easy of use/consistency sake
229 path.reverse();
230
231 let value = self.nodes[index_offset + index];
232 (value, path)
233 }
234}
235
236impl<T> From<T> for Mmr
237where
238 T: IntoIterator<Item = Word>,
239{
240 fn from(values: T) -> Self {
241 let mut mmr = Mmr::new();
242 for v in values {
243 mmr.add(v)
244 }
245 mmr
246 }
247}
248
249// UTILITIES
250// ===============================================================================================
251
252/// Given a 0-indexed leaf position and the current forest, return the tree number responsible for
253/// the position.
254///
255/// Note:
256/// The result is a tree position `p`, it has the following interpretations. $p+1$ is the depth of
257/// the tree, which corresponds to the size of a Merkle proof for that tree. $2^p$ is equal to the
258/// number of leaves in this particular tree. and $2^(p+1)-1$ corresponds to size of the tree.
259pub(crate) const fn leaf_to_corresponding_tree(pos: usize, forest: usize) -> Option<u32> {
260 if pos >= forest {
261 None
262 } else {
263 // - each bit in the forest is a unique tree and the bit position its power-of-two size
264 // - each tree owns a consecutive range of positions equal to its size from left-to-right
265 // - this means the first tree owns from `0` up to the `2^k_0` first positions, where `k_0`
266 // is the highest true bit position, the second tree from `2^k_0 + 1` up to `2^k_1` where
267 // `k_1` is the second higest bit, so on.
268 // - this means the highest bits work as a category marker, and the position is owned by
269 // the first tree which doesn't share a high bit with the position
270 let before = forest & pos;
271 let after = forest ^ before;
272 let tree = after.ilog2();
273
274 Some(tree)
275 }
276}
277
278/// Return a bitmask for the bits including and above the given position.
279pub(crate) const fn high_bitmask(bit: u32) -> usize {
280 if bit > usize::BITS - 1 {
281 0
282 } else {
283 usize::MAX << bit
284 }
285}
286
287/// Return the total number of nodes of a given forest
288///
289/// Panics:
290///
291/// This will panic if the forest has size greater than `usize::MAX / 2`
292pub(crate) const fn nodes_in_forest(forest: usize) -> usize {
293 // - the size of a perfect binary tree is $2^{k+1}-1$ or $2*2^k-1$
294 // - the forest represents the sum of $2^k$ so a single multiplication is necessary
295 // - the number of `-1` is the same as the number of trees, which is the same as the number
296 // bits set
297 let tree_count = forest.count_ones() as usize;
298 forest * 2 - tree_count
299}