From 61a0764a614db9ef5017797220b3384d317e3f70 Mon Sep 17 00:00:00 2001
From: Bobbin Threadbare <bobbinth@protonmail.com>
Date: Mon, 5 Feb 2024 15:41:33 -0800
Subject: [PATCH] fix: peak index calculation in MmrProof

---
 src/merkle/mmr/proof.rs | 74 ++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 73 insertions(+), 1 deletion(-)

diff --git a/src/merkle/mmr/proof.rs b/src/merkle/mmr/proof.rs
index 8674c2a..1539550 100644
--- a/src/merkle/mmr/proof.rs
+++ b/src/merkle/mmr/proof.rs
@@ -2,6 +2,9 @@
 use super::super::MerklePath;
 use super::{full::high_bitmask, leaf_to_corresponding_tree};
 
+// MMR PROOF
+// ================================================================================================
+
 #[derive(Debug, Clone, PartialEq)]
 #[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
 pub struct MmrProof {
@@ -26,9 +29,78 @@ impl MmrProof {
         self.position - forest_before
     }
 
+    /// Returns index of the MMR peak against which the Merkle path in this proof can be verified.
     pub fn peak_index(&self) -> usize {
         let root = leaf_to_corresponding_tree(self.position, self.forest)
             .expect("position must be part of the forest");
-        (self.forest.count_ones() - root - 1) as usize
+        let smaller_peak_mask = 2_usize.pow(root) as usize - 1;
+        let num_smaller_peaks = (self.forest & smaller_peak_mask).count_ones();
+        (self.forest.count_ones() - num_smaller_peaks - 1) as usize
+    }
+}
+
+// TESTS
+// ================================================================================================
+
+#[cfg(test)]
+mod tests {
+    use super::{MerklePath, MmrProof};
+
+    #[test]
+    fn test_peak_index() {
+        // --- single peak forest ---------------------------------------------
+        let forest = 11;
+
+        // the first 4 leaves belong to peak 0
+        for position in 0..8 {
+            let proof = make_dummy_proof(forest, position);
+            assert_eq!(proof.peak_index(), 0);
+        }
+
+        // --- forest with non-consecutive peaks ------------------------------
+        let forest = 11;
+
+        // the first 8 leaves belong to peak 0
+        for position in 0..8 {
+            let proof = make_dummy_proof(forest, position);
+            assert_eq!(proof.peak_index(), 0);
+        }
+
+        // the next 2 leaves belong to peak 1
+        for position in 8..10 {
+            let proof = make_dummy_proof(forest, position);
+            assert_eq!(proof.peak_index(), 1);
+        }
+
+        // the last leaf is the peak 2
+        let proof = make_dummy_proof(forest, 10);
+        assert_eq!(proof.peak_index(), 2);
+
+        // --- forest with consecutive peaks ----------------------------------
+        let forest = 7;
+
+        // the first 4 leaves belong to peak 0
+        for position in 0..4 {
+            let proof = make_dummy_proof(forest, position);
+            assert_eq!(proof.peak_index(), 0);
+        }
+
+        // the next 2 leaves belong to peak 1
+        for position in 4..6 {
+            let proof = make_dummy_proof(forest, position);
+            assert_eq!(proof.peak_index(), 1);
+        }
+
+        // the last leaf is the peak 2
+        let proof = make_dummy_proof(forest, 6);
+        assert_eq!(proof.peak_index(), 2);
+    }
+
+    fn make_dummy_proof(forest: usize, position: usize) -> MmrProof {
+        MmrProof {
+            forest,
+            position,
+            merkle_path: MerklePath::default(),
+        }
     }
 }