Merge pull request #127 from 0xPolygonMiden/hacka-optimized-peak-hash

mmr: optimized peak hash for Miden VM

src/merkle/mmr/accumulator.rs

@@ -1,4 +1,8 @@
-use super::{super::Vec, MmrProof, Rpo256, Word};
+use super::{
+    super::Vec,
+    super::{WORD_SIZE, ZERO},
+    MmrProof, Rpo256, Word,
+};
 
 #[derive(Debug, Clone, PartialEq)]
 pub struct MmrPeaks {
@@ -8,18 +12,17 @@ pub struct MmrPeaks {
     /// the MMR has a power-of-two number of leaves there is a single peak.
     ///
     /// Every tree in the MMR forest has a distinct power-of-two size, this means only the right
-    /// most tree can have an odd number of elements (1). Additionally this means that the bits in
+    /// most tree can have an odd number of elements (e.g. `1`). Additionally this means that the bits in
     /// `num_leaves` conveniently encode the size of each individual tree.
     ///
     /// Examples:
     ///
-    ///    Example 1: With 5 leaves, the binary 0b101. The number of set bits is equal the number
+    ///    - With 5 leaves, the binary `0b101`. The number of set bits is equal the number
-    ///    of peaks, in this case there are 2 peaks. The 0-indexed least-significant position of
+    ///      of peaks, in this case there are 2 peaks. The 0-indexed least-significant position of
-    ///    the bit determines the number of elements of a tree, so the rightmost tree has 2**0
+    ///      the bit determines the number of elements of a tree, so the rightmost tree has `2**0`
-    ///    elements and the left most has 2**2.
+    ///      elements and the left most has `2**2`.
-    ///
+    ///    - With 12 leaves, the binary is `0b1100`, this case also has 2 peaks, the
-    ///    Example 2: With 12 leaves, the binary is 0b1100, this case also has 2 peaks, the
+    ///      leftmost tree has `2**3=8` elements, and the right most has `2**2=4` elements.
-    ///    leftmost tree has 2**3=8 elements, and the right most has 2**2=4 elements.
     pub num_leaves: usize,
 
     /// All the peaks of every tree in the MMR forest. The peaks are always ordered by number of
@@ -30,9 +33,23 @@ pub struct MmrPeaks {
 }
 
 impl MmrPeaks {
-    /// Hashes the peaks sequentially, compacting it to a single digest
+    /// Hashes the peaks.
+    ///
+    /// The hashing is optimized to work with the Miden VM, the procedure will:
+    ///
+    /// - Pad the peaks with ZERO to an even number of words, this removes the need to handle RPO padding.
+    /// - Pad the peaks to a minimum length of 16 words, which reduces the constant cost of
+    ///   hashing.
     pub fn hash_peaks(&self) -> Word {
-        Rpo256::hash_elements(&self.peaks.as_slice().concat()).into()
+        let mut copy = self.peaks.clone();
+
+        if copy.len() < 16 {
+            copy.resize(16, [ZERO; WORD_SIZE])
+        } else if copy.len() % 2 == 1 {
+            copy.push([ZERO; WORD_SIZE])
+        }
+
+        Rpo256::hash_elements(&copy.as_slice().concat()).into()
     }
 
     pub fn verify(&self, value: Word, opening: MmrProof) -> bool {

src/merkle/mmr/full.rs

@@ -174,7 +174,7 @@ impl Mmr {
         self.forest += 1;
     }
 
-    /// Returns an accumulator representing the current state of the MMMR.
+    /// Returns an accumulator representing the current state of the MMR.
     pub fn accumulator(&self) -> MmrPeaks {
         let peaks: Vec<Word> = TrueBitPositionIterator::new(self.forest)
             .rev()
@@ -192,7 +192,7 @@ impl Mmr {
         }
     }
 
-    /// An iterator over inner nodes in the [Mmm]. The order of iteration is unspecified.
+    /// An iterator over inner nodes in the MMR. The order of iteration is unspecified.
     pub fn inner_nodes(&self) -> MmrNodes {
         MmrNodes {
             mmr: self,

src/merkle/mmr/tests.rs

@@ -1,8 +1,8 @@
 use super::bit::TrueBitPositionIterator;
 use super::full::{high_bitmask, leaf_to_corresponding_tree, nodes_in_forest};
 use super::{
-    super::{InnerNodeInfo, Vec},
+    super::{InnerNodeInfo, Vec, WORD_SIZE, ZERO},
-    Mmr, Rpo256, Word,
+    Mmr, MmrPeaks, Rpo256, Word,
 };
 use crate::merkle::{int_to_node, MerklePath};
 
@@ -448,6 +448,66 @@ fn test_mmr_inner_nodes() {
     assert_eq!(postorder, nodes);
 }
 
+#[test]
+fn test_mmr_hash_peaks() {
+    let mmr: Mmr = LEAVES.into();
+    let peaks = mmr.accumulator();
+
+    let first_peak = *Rpo256::merge(&[
+        Rpo256::hash_elements(&[LEAVES[0], LEAVES[1]].concat()),
+        Rpo256::hash_elements(&[LEAVES[2], LEAVES[3]].concat()),
+    ]);
+    let second_peak = *Rpo256::hash_elements(&[LEAVES[4], LEAVES[5]].concat());
+    let third_peak = LEAVES[6];
+
+    // minimum length is 16
+    let mut expected_peaks = [first_peak, second_peak, third_peak].to_vec();
+    expected_peaks.resize(16, [ZERO; WORD_SIZE]);
+    assert_eq!(
+        peaks.hash_peaks(),
+        *Rpo256::hash_elements(&expected_peaks.as_slice().concat())
+    );
+}
+
+#[test]
+fn test_mmr_peaks_hash_less_than_16() {
+    let mut peaks = Vec::new();
+
+    for i in 0..16 {
+        peaks.push(int_to_node(i));
+        let accumulator = MmrPeaks {
+            num_leaves: (1 << peaks.len()) - 1,
+            peaks: peaks.clone(),
+        };
+
+        // minimum length is 16
+        let mut expected_peaks = peaks.clone();
+        expected_peaks.resize(16, [ZERO; WORD_SIZE]);
+        assert_eq!(
+            accumulator.hash_peaks(),
+            *Rpo256::hash_elements(&expected_peaks.as_slice().concat())
+        );
+    }
+}
+
+#[test]
+fn test_mmr_peaks_hash_odd() {
+    let peaks: Vec<_> = (0..=17).map(|i| int_to_node(i)).collect();
+
+    let accumulator = MmrPeaks {
+        num_leaves: (1 << peaks.len()) - 1,
+        peaks: peaks.clone(),
+    };
+
+    // odd length bigger than 16 is padded to the next even nubmer
+    let mut expected_peaks = peaks.clone();
+    expected_peaks.resize(18, [ZERO; WORD_SIZE]);
+    assert_eq!(
+        accumulator.hash_peaks(),
+        *Rpo256::hash_elements(&expected_peaks.as_slice().concat())
+    );
+}
+
 mod property_tests {
     use super::leaf_to_corresponding_tree;
     use proptest::prelude::*;