From 2100d6c861f3cc6cebb737d222427ffa96483a01 Mon Sep 17 00:00:00 2001
From: frisitano <francesco.risitano95@gmail.com>
Date: Tue, 11 Apr 2023 16:27:25 +0100
Subject: [PATCH] refactor(mmr): expose method to join mmr peaks in a vector
 and pad

---
 src/merkle/mmr/accumulator.rs | 58 +++++++++++++++++++++++------------
 src/merkle/mmr/mod.rs         |  2 +-
 2 files changed, 40 insertions(+), 20 deletions(-)

diff --git a/src/merkle/mmr/accumulator.rs b/src/merkle/mmr/accumulator.rs
index 610999a..ce0ee49 100644
--- a/src/merkle/mmr/accumulator.rs
+++ b/src/merkle/mmr/accumulator.rs
@@ -1,8 +1,4 @@
-use super::{
-    super::Vec,
-    super::{WORD_SIZE, ZERO},
-    MmrProof, Rpo256, Word,
-};
+use super::{super::Vec, super::ZERO, Felt, MmrProof, Rpo256, Word};
 
 #[derive(Debug, Clone, PartialEq)]
 pub struct MmrPeaks {
@@ -35,25 +31,49 @@ pub struct MmrPeaks {
 impl MmrPeaks {
     /// 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.
+    /// The procedure will:
+    /// - Flatten and pad the peaks to a vector of Felts.
+    /// - Hash the vector of Felts.
     pub fn hash_peaks(&self) -> Word {
-        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()
+        Rpo256::hash_elements(&self.flatten_and_pad_peaks()).into()
     }
 
     pub fn verify(&self, value: Word, opening: MmrProof) -> bool {
         let root = &self.peaks[opening.peak_index()];
         opening.merkle_path.verify(opening.relative_pos() as u64, value, root)
     }
+
+    /// Flattens and pads the peaks to make hashing inside of the Miden VM easier.
+    ///
+    /// The procedure will:
+    /// - Flatten the vector of Words into a vector of Felts.
+    /// - 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 flatten_and_pad_peaks(&self) -> Vec<Felt> {
+        let num_peaks = self.peaks.len();
+
+        // To achieve the padding rules above we calculate the length of the final vector.
+        // This is calculated as the number of field elements. Each peak is 4 field elements.
+        // The length is calculated as follows:
+        // - If there are less than 16 peaks, the data is padded to 16 peaks and as such requires
+        //   64 field elements.
+        // - If there are more than 16 peaks and the number of peaks is odd, the data is padded to
+        //   an even number of peaks and as such requires `(num_peaks + 1) * 4` field elements.
+        // - If there are more than 16 peaks and the number of peaks is even, the data is not padded
+        //   and as such requires `num_peaks * 4` field elements.
+        let len = if num_peaks < 16 {
+            64
+        } else if num_peaks % 2 == 1 {
+            (num_peaks + 1) * 4
+        } else {
+            num_peaks * 4
+        };
+
+        let mut elements = Vec::with_capacity(len);
+        elements.extend_from_slice(&self.peaks.as_slice().concat());
+        elements.resize(len, ZERO);
+        elements
+    }
 }
diff --git a/src/merkle/mmr/mod.rs b/src/merkle/mmr/mod.rs
index d8903ca..118bb12 100644
--- a/src/merkle/mmr/mod.rs
+++ b/src/merkle/mmr/mod.rs
@@ -6,7 +6,7 @@ mod proof;
 #[cfg(test)]
 mod tests;
 
-use super::{Rpo256, Word};
+use super::{Felt, Rpo256, Word};
 
 // REEXPORTS
 // ================================================================================================