From c65b8e388294763790c9ebc01552edbab2d31594 Mon Sep 17 00:00:00 2001
From: Filip Krawczyk <kotmawszy@gmail.com>
Date: Fri, 22 Mar 2024 20:04:30 +0100
Subject: [PATCH] MMR append change loops to pop from span

---
 src/data_structures/mmr/mmr.cairo   | 32 +++++++++++++----------------
 src/data_structures/mmr/utils.cairo |  4 ++++
 2 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/src/data_structures/mmr/mmr.cairo b/src/data_structures/mmr/mmr.cairo
index 5a8240d..6da2f5b 100644
--- a/src/data_structures/mmr/mmr.cairo
+++ b/src/data_structures/mmr/mmr.cairo
@@ -49,32 +49,28 @@ impl MMRImpl of MMRTrait {
 
         self.last_pos += 1;
 
+        // TODO: why number of new peaks is equal to number of trailing ones
         let new_peaks_count = trailing_ones(leaf_count);
-        let mut new_peak = hash;
-        let mut i = 0;
+        let mut preserved_peaks = peaks.slice(0, peaks_count - new_peaks_count);
+        let mut merged_peaks = peaks.slice(peaks_count - new_peaks_count, new_peaks_count);
 
+        let mut last_peak = hash;
         loop {
-            if i == new_peaks_count {
-                break ();
-            }
-
-            new_peak = PoseidonHasher::hash_double(*peaks.at(peaks.len() - i - 1), new_peak);
-
-            i += 1;
+            match merged_peaks.pop_back() {
+                Option::Some(x) => { last_peak = PoseidonHasher::hash_double(*x, last_peak); },
+                Option::None => { break; }
+            };
             self.last_pos += 1;
         };
 
         let mut new_peaks = ArrayTrait::new();
-        let mut i = 0;
         loop {
-            if i == peaks_count - new_peaks_count {
-                break ();
-            }
-            new_peaks.append(*peaks.at(i));
-
-            i += 1;
+            match preserved_peaks.pop_front() {
+                Option::Some(x) => { new_peaks.append(*x); },
+                Option::None => { break; }
+            };
         };
-        new_peaks.append(new_peak);
+        new_peaks.append(last_peak);
 
         let new_root = compute_root(self.last_pos.into(), new_peaks.span());
         self.root = new_root;
@@ -98,7 +94,7 @@ impl MMRImpl of MMRTrait {
         if !peaks.valid(*self.last_pos, *self.root) {
             return Result::Err('Invalid peaks');
         }
-        
+
         let (peak_index, peak_height) = get_peak_info(*self.last_pos, index);
 
         if proof.len() != peak_height {
diff --git a/src/data_structures/mmr/utils.cairo b/src/data_structures/mmr/utils.cairo
index 966f620..aa66ac7 100644
--- a/src/data_structures/mmr/utils.cairo
+++ b/src/data_structures/mmr/utils.cairo
@@ -45,6 +45,8 @@ fn count_ones(n: u32) -> u32 {
 // @notice Convert a leaf index to an Merkle Mountain Range tree index
 // @param n The leaf index
 // @return The MMR index
+// Explanation of why this formula is correct
+// https://mmr.herodotus.dev/mmr-size-vs-leaf-count#leaf-count-to-mmr-size-algorithm
 fn leaf_index_to_mmr_index(n: u32) -> u32 {
     2 * n - 1 - count_ones(n - 1)
 }
@@ -52,6 +54,8 @@ fn leaf_index_to_mmr_index(n: u32) -> u32 {
 // @notice Convert a Merkle Mountain Range tree size to number of leaves
 // @param n MMR size
 // @result Number of leaves
+// Explanation of why this algorithm is correct
+// https://mmr.herodotus.dev/mmr-size-vs-leaf-count#mmr-size-to-leaf-count-algorithm
 fn mmr_size_to_leaf_count(n: u32) -> u32 {
     let mut mmr_size = n;
     let bits = bit_length(mmr_size + 1);