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