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//! 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}