1use super::{
 2    super::{RpoDigest, Vec, ZERO},
 3    Felt, MmrProof, Rpo256, Word,
 4};
 5
 6#[derive(Debug, Clone, PartialEq)]
 7pub struct MmrPeaks {
 8    /// The number of leaves is used to differentiate accumulators that have the same number of
 9    /// peaks. This happens because the number of peaks goes up-and-down as the structure is used
10    /// causing existing trees to be merged and new ones to be created. As an example, every time
11    /// the MMR has a power-of-two number of leaves there is a single peak.
12    ///
13    /// Every tree in the MMR forest has a distinct power-of-two size, this means only the right
14    /// most tree can have an odd number of elements (e.g. `1`). Additionally this means that the bits in
15    /// `num_leaves` conveniently encode the size of each individual tree.
16    ///
17    /// Examples:
18    ///
19    ///    - With 5 leaves, the binary `0b101`. The number of set bits is equal the number
20    ///      of peaks, in this case there are 2 peaks. The 0-indexed least-significant position of
21    ///      the bit determines the number of elements of a tree, so the rightmost tree has `2**0`
22    ///      elements and the left most has `2**2`.
23    ///    - With 12 leaves, the binary is `0b1100`, this case also has 2 peaks, the
24    ///      leftmost tree has `2**3=8` elements, and the right most has `2**2=4` elements.
25    pub num_leaves: usize,
26
27    /// All the peaks of every tree in the MMR forest. The peaks are always ordered by number of
28    /// leaves, starting from the peak with most children, to the one with least.
29    ///
30    /// Invariant: The length of `peaks` must be equal to the number of true bits in `num_leaves`.
31    pub peaks: Vec<RpoDigest>,
32}
33
34impl MmrPeaks {
35    /// Hashes the peaks.
36    ///
37    /// The procedure will:
38    /// - Flatten and pad the peaks to a vector of Felts.
39    /// - Hash the vector of Felts.
40    pub fn hash_peaks(&self) -> Word {
41        Rpo256::hash_elements(&self.flatten_and_pad_peaks()).into()
42    }
43
44    pub fn verify(&self, value: RpoDigest, opening: MmrProof) -> bool {
45        let root = &self.peaks[opening.peak_index()];
46        opening.merkle_path.verify(opening.relative_pos() as u64, value, root)
47    }
48
49    /// Flattens and pads the peaks to make hashing inside of the Miden VM easier.
50    ///
51    /// The procedure will:
52    /// - Flatten the vector of Words into a vector of Felts.
53    /// - Pad the peaks with ZERO to an even number of words, this removes the need to handle RPO
54    ///   padding.
55    /// - Pad the peaks to a minimum length of 16 words, which reduces the constant cost of
56    ///   hashing.
57    pub fn flatten_and_pad_peaks(&self) -> Vec<Felt> {
58        let num_peaks = self.peaks.len();
59
60        // To achieve the padding rules above we calculate the length of the final vector.
61        // This is calculated as the number of field elements. Each peak is 4 field elements.
62        // The length is calculated as follows:
63        // - If there are less than 16 peaks, the data is padded to 16 peaks and as such requires
64        //   64 field elements.
65        // - If there are more than 16 peaks and the number of peaks is odd, the data is padded to
66        //   an even number of peaks and as such requires `(num_peaks + 1) * 4` field elements.
67        // - If there are more than 16 peaks and the number of peaks is even, the data is not padded
68        //   and as such requires `num_peaks * 4` field elements.
69        let len = if num_peaks < 16 {
70            64
71        } else if num_peaks % 2 == 1 {
72            (num_peaks + 1) * 4
73        } else {
74            num_peaks * 4
75        };
76
77        let mut elements = Vec::with_capacity(len);
78        elements.extend_from_slice(
79            &self
80                .peaks
81                .as_slice()
82                .iter()
83                .map(|digest| digest.into())
84                .collect::<Vec<Word>>()
85                .concat(),
86        );
87        elements.resize(len, ZERO);
88        elements
89    }
90}