From bc42846dc4cb185efc04aef173b3db10e8c54661 Mon Sep 17 00:00:00 2001
From: dan <themighty1@users.noreply.github.com>
Date: Thu, 11 Jan 2024 12:41:03 +0000
Subject: [PATCH] fix: sort indices correctly (#30)

Co-authored-by: themighty1 <you@example.com>
---
 src/merkle_proof.rs        |  6 ++++--
 tests/merkle_proof_test.rs | 29 +++++++++++++++++++++++++++++
 2 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/src/merkle_proof.rs b/src/merkle_proof.rs
index f8a0e49..285fff0 100644
--- a/src/merkle_proof.rs
+++ b/src/merkle_proof.rs
@@ -1,7 +1,7 @@
-use crate::prelude::*;
 use crate::{
     error::Error,
     partial_tree::PartialTree,
+    prelude::*,
     proof_serializers::{DirectHashesOrder, MerkleProofSerializer},
     utils, Hasher,
 };
@@ -222,8 +222,10 @@ impl<T: Hasher> MerkleProof<T> {
         // Sorting leaves by indexes in case they weren't sorted already
         leaf_tuples.sort_by(|(a, _), (b, _)| a.cmp(b));
         // Getting back _sorted_ indices
+        let (sorted_indices, _): (Vec<_>, Vec<_>) = leaf_tuples.iter().cloned().unzip();
+
         let proof_indices_by_layers =
-            utils::indices::proof_indices_by_layers(leaf_indices, total_leaves_count);
+            utils::indices::proof_indices_by_layers(&sorted_indices, total_leaves_count);
 
         // The next lines copy hashes from proof hashes and group them by layer index
         let mut proof_layers: Vec<Vec<(usize, T::Hash)>> = Vec::with_capacity(tree_depth + 1);
diff --git a/tests/merkle_proof_test.rs b/tests/merkle_proof_test.rs
index aa32042..a7ebf04 100644
--- a/tests/merkle_proof_test.rs
+++ b/tests/merkle_proof_test.rs
@@ -107,6 +107,35 @@ pub mod root {
             "Should return error not_enough_hashes_to_calculate_root"
         );
     }
+
+    // Expect to calculate the correct root even though the indices are unsorted
+    #[test]
+    fn should_sort_indices() {
+        let test_data = common::setup();
+        let leaf_hashes = &test_data.leaf_hashes;
+        let expected_root = test_data.expected_root_hex.clone();
+        let indices_to_prove = vec![0, 3];
+
+        let merkle_tree = MerkleTree::<Sha256>::from_leaves(&test_data.leaf_hashes);
+        let proof = merkle_tree.proof(&indices_to_prove);
+
+        // make indices unsorted
+        let indices_to_compute_root = vec![3, 0];
+        let leaves_to_compute_root: Vec<[u8; 32]> = indices_to_compute_root
+            .iter()
+            .map(|i| *leaf_hashes.get(*i).unwrap())
+            .collect();
+
+        let extracted_root = proof
+            .root_hex(
+                &indices_to_compute_root,
+                &leaves_to_compute_root,
+                test_data.leaf_values.len(),
+            )
+            .unwrap();
+
+        assert_eq!(extracted_root, expected_root);
+    }
 }
 
 pub mod to_bytes {