From 2040166c5856b3ea3b0690c2d0030fd7535a5f85 Mon Sep 17 00:00:00 2001
From: Mark Spanbroek <mark@spanbroek.net>
Date: Thu, 9 Nov 2023 10:20:10 +0100
Subject: [PATCH] Fix merkle root of odd number of elements

---
 poseidon2.nim                     | 17 ++++++++---------
 tests/poseidon2/testPoseidon2.nim | 14 ++++++++++++--
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/poseidon2.nim b/poseidon2.nim
index 261879b..50ac201 100644
--- a/poseidon2.nim
+++ b/poseidon2.nim
@@ -85,17 +85,16 @@ func merkleRoot*(xs: openArray[F]) : F =
     let n      : int  = 2*halfn
     let isOdd : bool = (n != m)
 
-    var ys : seq[F] = newSeq[F](halfn)
-
+    var ys : seq[F]
     if not isOdd:
-      for i in 0..<halfn:
-        ys[i] = compress( xs[a+2*i], xs[a+2*i+1] )
-
+      ys = newSeq[F](halfn)
     else:
-      for i in 0..<halfn-1:
-        ys[i] = compress( xs[a+2*i], xs[a+2*i+1] )
-      # and the last one:
-      ys[halfn-1] = compress( xs[a+n-2], zero )
+      ys = newSeq[F](halfn+1)
+
+    for i in 0..<halfn:
+      ys[i] = compress( xs[a+2*i], xs[a+2*i+1] )
+    if isOdd:
+      ys[halfn] = compress( xs[n], zero )
 
     return merkleRoot(ys)
 
diff --git a/tests/poseidon2/testPoseidon2.nim b/tests/poseidon2/testPoseidon2.nim
index 70e6ea4..729f956 100644
--- a/tests/poseidon2/testPoseidon2.nim
+++ b/tests/poseidon2/testPoseidon2.nim
@@ -33,12 +33,22 @@ suite "poseidon2":
     let root = merkleRoot(xs)
     check root.toHex(littleEndian) == "0xd1111b3515a663bb48278bfe453fe2508487014a1c6093d3ec5a6db764bbab1e"
 
+  test "merkle root of even elements":
+    let elements = toSeq(1..4).mapIt(toF(it))
+    let expected = compress(compress(1.toF, 2.toF), compress(3.toF, 4.toF))
+    check bool(merkleRoot(elements) == expected)
+
+  test "merkle root of odd elements":
+    let elements = toSeq(1..3).mapIt(toF(it))
+    let expected = compress(compress(1.toF, 2.toF), compress(3.toF, 0.toF))
+    check bool(merkleRoot(elements) == expected)
+
   test "merkle root of bytes":
     let bytes = toSeq 1'u8..80'u8
     let root = merkleRoot(bytes)
-    check root.toHex(littleEndian) == "0x00f5c12fabdf0fce1a69b9f0920f60e0c0ea4d8a880ae297559e0ad2276f2d2d"
+    check root.toHex(littleEndian) == "0xa1dffa3f60d166283d60396023d95a1d7996d119e5290fe31131e7c6a7a27817"
 
   test "merkle root of bytes converted to bytes":
     let bytes = toSeq 1'u8..80'u8
     let rootAsBytes = merkleRoot(bytes).toBytes()
-    check rootAsBytes.toHex == "0x00f5c12fabdf0fce1a69b9f0920f60e0c0ea4d8a880ae297559e0ad2276f2d2d"
+    check rootAsBytes.toHex == "0xa1dffa3f60d166283d60396023d95a1d7996d119e5290fe31131e7c6a7a27817"