Alternative Merkle tree proof fix for trees with duplicate leaf values

shared/src/main/scala/scorex/crypto/authds/merkle/MerkleTree.scala

@@ -7,16 +7,39 @@ import scorex.crypto.hash.{Digest, _}
 import scala.annotation.tailrec
 import scala.collection.mutable
 
-case class MerkleTree[D <: Digest](topNode: Node[D],
-                                   elementsHashIndex: Map[mutable.WrappedArray.ofByte, Int]) {
+/**
+  * Merkle tree constructed from leaves
+ */
+class MerkleTree[D <: Digest] private (leaves: Vector[Leaf[D]])(implicit val hf: CryptographicHash[D]) {
+
+  lazy val topNode: Node[D] = MerkleTree.calcTopNode(leaves)
 
   lazy val rootHash: D = topNode.hash
-  lazy val length: Int = elementsHashIndex.size
 
-  def proofByElement(element: Leaf[D]): Option[MerkleProof[D]] = proofByElementHash(element.hash)
+  lazy val length: Int = leaves.length
+
+  private lazy val leafHashes: Vector[D] = leaves.map(_.hash)
+
+  private lazy val hashToIndex: Map[mutable.WrappedArray.ofByte, Int] =
+    leaves.indices.foldLeft(Map.empty[mutable.WrappedArray.ofByte, Int]) { (acc, i) =>
+      val key = new mutable.WrappedArray.ofByte(leafHashes(i))
+      if (acc.contains(key)) acc else acc.updated(key, i)
+    }
+
+  def proofByElement(element: Leaf[D]): Option[MerkleProof[D]] = {
+    proofByElementHash(element.hash)
+  }
 
   def proofByElementHash(hash: D): Option[MerkleProof[D]] = {
-    elementsHashIndex.get(new mutable.WrappedArray.ofByte(hash)).flatMap(i => proofByIndex(i))
+    indexByElementHash(hash).flatMap(proofByIndex)
+  }
+
+  /**
+    * Position of the first leaf with the given hash, if any.
+    * Equal hashes are stored at several positions; only the first one is returned.
+    */
+  def indexByElementHash(hash: D): Option[Int] = {
+    hashToIndex.get(new mutable.WrappedArray.ofByte(hash))
   }
 
   def proofByIndex(index: Int): Option[MerkleProof[D]] = if (index >= 0 && index < length) {
@@ -45,7 +68,7 @@ case class MerkleTree[D <: Digest](topNode: Node[D],
     * @param indices - leaf indices
     * @return Optional BatchMerkleProof
     */
-  def proofByIndices(indices: Seq[Int])(implicit hf: CryptographicHash[D]): Option[BatchMerkleProof[D]] = {
+  def proofByIndices(indices: Seq[Int]): Option[BatchMerkleProof[D]] = {
 
     /**
       * Recursive function to build the multiproof
@@ -96,7 +119,7 @@ case class MerkleTree[D <: Digest](topNode: Node[D],
     }
 
     if (indices.nonEmpty && indices.forall(index => index >= 0 && index < length)) {
-      val hashes: Seq[Digest] = elementsHashIndex.toSeq.sortBy(_._2).map(_._1.toArray.asInstanceOf[Digest])
+      val hashes: Seq[Digest] = leafHashes
       val normalized_indices = indices.distinct.sorted
       val multiproof = loop(normalized_indices, hashes)
       Some(BatchMerkleProof(normalized_indices zip (normalized_indices map hashes.apply), multiproof))
@@ -144,13 +167,7 @@ object MerkleTree {
     */
   def apply[D <: Digest](payload: Seq[LeafData])
                         (implicit hf: CryptographicHash[D]): MerkleTree[D] = {
-    val leafs = payload.map(d => Leaf(d))
-    val elementsIndex: Map[mutable.WrappedArray.ofByte, Int] = leafs.indices.map { i =>
-      (new mutable.WrappedArray.ofByte(leafs(i).hash), i)
-    }.toMap
-    val topNode = calcTopNode[D](leafs)
-
-    MerkleTree(topNode, elementsIndex)
+    new MerkleTree(payload.map(d => Leaf(d)).toVector)(hf)
   }
 
   @tailrec

shared/src/test/scala/scorex/crypto/authds/merkle/MerkleTreeSpecification.scala

@@ -14,6 +14,14 @@ class MerkleTreeSpecification extends AnyPropSpec with ScalaCheckDrivenPropertyC
 
   private val LeafSize = 32
 
+  private def payloadsWithDuplicateLeaves: Seq[Seq[LeafData]] =
+    Seq(
+      Seq(1, 1),
+      Seq(1, 2, 1),
+      Seq(1, 2, 1, 3),
+      Seq(1, 2, 3, 4, 1)
+    ).map(_.map(value => LeafData @@ Array.fill[Byte](LeafSize)(value.toByte)))
+
   property("Proof generation by element") {
     forAll(smallInt) { (N: Int) =>
       whenever(N > 0) {
@@ -48,6 +56,63 @@ class MerkleTreeSpecification extends AnyPropSpec with ScalaCheckDrivenPropertyC
     }
   }
 
+  property("Proof by index remains bound to the requested position with duplicate leaves") {
+    val first = LeafData @@ Array.fill[Byte](LeafSize)(1: Byte)
+    val second = LeafData @@ Array.fill[Byte](LeafSize)(2: Byte)
+    val tree = MerkleTree(Seq(first, second, first))
+    val proof = tree.proofByIndex(0).get
+
+    proof.valid(tree.rootHash) shouldBe true
+    proof.leafData.sameElements(first) shouldBe true
+  }
+
+  property("Proof generation with duplicate leaves") {
+    payloadsWithDuplicateLeaves.foreach { data =>
+      val tree = MerkleTree(data)
+
+      withClue(s"leaf count ${data.length}: ") {
+        tree.length shouldBe data.length
+        data.indices.foreach { index =>
+          val proof = tree.proofByIndex(index).get
+          proof.leafData.sameElements(data(index)) shouldBe true
+          proof.valid(tree.rootHash) shouldBe true
+        }
+
+        Seq(data.head, data(1)).foreach { leafData =>
+          val leaf = Leaf(leafData)
+          Seq(tree.proofByElement(leaf), tree.proofByElementHash(leaf.hash)).foreach { proofOpt =>
+            val proof = proofOpt.get
+            proof.leafData.sameElements(leafData) shouldBe true
+            proof.valid(tree.rootHash) shouldBe true
+          }
+        }
+      }
+    }
+  }
+
+  property("Index by element hash returns the first occurrence with duplicate leaves") {
+    payloadsWithDuplicateLeaves.foreach { data =>
+      val tree = MerkleTree(data)
+      val hash = Leaf(data.head).hash
+
+      withClue(s"leaf count ${data.length}: ") {
+        tree.indexByElementHash(hash) shouldBe Some(0)
+        data.indices.foreach { index =>
+          tree.indexByElementHash(Leaf(data(index)).hash).isDefined shouldBe true
+        }
+      }
+    }
+  }
+
+  property("Index by element hash returns None for a missing hash") {
+    val d = (0 until 5).map(_ => LeafData @@ scorex.utils.Random.randomBytes(LeafSize))
+    val tree = MerkleTree(d)
+    val missing = Leaf(LeafData @@ scorex.utils.Random.randomBytes(LeafSize))
+
+    tree.indexByElementHash(missing.hash) shouldBe None
+    tree.proofByElementHash(missing.hash) shouldBe None
+  }
+
   property("Batch proof generation by indices") {
     val r = new Random()
     forAll(smallInt) { (N: Int) =>
@@ -69,6 +134,24 @@ class MerkleTreeSpecification extends AnyPropSpec with ScalaCheckDrivenPropertyC
     tree.proofByIndices(Seq(2,2,2,3,6,6,8,9,9)).get.valid(tree.rootHash) shouldBe true
   }
 
+  property("Batch proof generation with duplicate leaves") {
+    payloadsWithDuplicateLeaves.foreach { data =>
+      val tree = MerkleTree(data)
+      val duplicateIndex = data.lastIndexWhere(_.sameElements(data.head))
+      val requestedIndices = Seq(duplicateIndex, 0, duplicateIndex)
+      val expectedIndices = requestedIndices.distinct.sorted
+      val proof = tree.proofByIndices(requestedIndices).get
+
+      withClue(s"leaf count ${data.length}: ") {
+        proof.indices.map(_._1) shouldEqual expectedIndices
+        proof.indices.zip(expectedIndices).foreach { case ((_, hash), index) =>
+          hash.sameElements(Leaf(data(index)).hash) shouldBe true
+        }
+        proof.valid(tree.rootHash) shouldBe true
+      }
+    }
+  }
+
   property("Batch proof generation by negative indices") {
     val d = (0 until 5).map(_ => LeafData @@ scorex.utils.Random.randomBytes(LeafSize))
     val tree = MerkleTree(d)