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