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