From 25878c88cf7d1fd61f51aed761fc99e1d696a225 Mon Sep 17 00:00:00 2001
From: Mark Spanbroek <mark@spanbroek.net>
Date: Thu, 26 Oct 2023 14:39:27 +0200
Subject: [PATCH] Determine merkle root of byte arrays

---
 poseidon2.nim       | 21 +++++++++++++++++----
 poseidon2/types.nim | 16 ++++++++--------
 tests/test.nim      | 34 +++++++++++++++++++++++++++++++++-
 3 files changed, 58 insertions(+), 13 deletions(-)

diff --git a/poseidon2.nim b/poseidon2.nim
index 3bef275..f750564 100644
--- a/poseidon2.nim
+++ b/poseidon2.nim
@@ -3,7 +3,7 @@
 
 import
   constantine/math/arithmetic,
-  constantine/math/io/io_fields,
+  constantine/math/io/io_bigints,
   constantine/math/config/curves
 
 import poseidon2/types
@@ -100,7 +100,20 @@ proc merkleRoot*(xs: openArray[F]) : F =
 
     return merkleRoot(ys)
 
-#-------------------------------------------------------------------------------
-
-
+proc unmarshal(
+        _: type F,
+        bytes: openArray[byte],
+        endian: static Endianness): seq[F] =
+  const chunkLen = 31
+  var elements: seq[F]
+  var i = 0
+  while i < bytes.len:
+    let chunk = bytes[i..<min(i + chunkLen, bytes.len)]
+    let bigint = B.unmarshal(chunk, endian)
+    let element = F.fromBig(bigint)
+    elements.add(element)
+    i += chunkLen
+  return elements
 
+proc merkleRoot*(bytes: openArray[byte]): F =
+  merkleRoot(F.unmarshal(bytes, littleEndian))
diff --git a/poseidon2/types.nim b/poseidon2/types.nim
index 4e4837a..a14d541 100644
--- a/poseidon2/types.nim
+++ b/poseidon2/types.nim
@@ -1,6 +1,7 @@
 
 import
   constantine/math/io/io_fields,
+  constantine/math/io/io_bigints,
   constantine/math/arithmetic,
   constantine/math/config/curves
 
@@ -22,15 +23,14 @@ func toF*(a: int) : F =
   fromInt(y, a);
   return y
 
-func hexToF*(s : string) : F =
-  var y : F
-  fromHex(y, s)
-  return y
+func hexToF*(s : string, endian: static Endianness = bigEndian) : F =
+  let bigint = B.fromHex(s, endian)
+  return F.fromBig(bigint)
 
-func arrayFromHex*[N]( inp: array[N, string]) : array[N, F] =
+func arrayFromHex*[N](
+    inp: array[N, string],
+    endian: static Endianness = bigEndian) : array[N, F] =
   var tmp : array[N, F]
   for i in low(inp)..high(inp):
-    tmp[i] = hexToF( inp[i] )
+    tmp[i] = hexToF(inp[i], endian)
   return tmp
-
-#-------------------------------------------------------------------------------
diff --git a/tests/test.nim b/tests/test.nim
index 49d1691..19147a2 100644
--- a/tests/test.nim
+++ b/tests/test.nim
@@ -1,7 +1,10 @@
 import std/unittest
 import std/math
+import std/strutils
 
+import constantine/math/arithmetic
 import constantine/math/io/io_fields
+import constantine/math/io/io_bigints
 import constantine/math/config/curves
 
 import poseidon2/types
@@ -20,7 +23,7 @@ suite "poseidon2":
     check toDecimal(y) == "09030699330013392132529464674294378792132780497765201297316864012141442630280"
     check toDecimal(z) == "09137931384593657624554037900714196568304064431583163402259937475584578975855"
 
-  test "merkle root":
+  test "merkle root of field elements":
     let m = 17
     let n = 2^m
     var xs: seq[F]
@@ -30,3 +33,32 @@ suite "poseidon2":
     let root = merkleRoot(xs)
     check toHex(root) == "0x1eabbb64b76d5aecd393601c4a01878450e23f45fe8b2748bb63a615351b11d1"
 
+  test "merkle root of bytes":
+    let hex = [ # 31 bytes per field element
+      "000102030405060708090A0B0C0D0E0F101112131415161718191A1B1C1D1E",
+      "1F202122232425262728292A2B2C2D2E2F303132333435363738393A3B3C3D",
+      "3E3F404142434445464748494A4B4C4D4E4F505152535455565758595A5B5C"
+    ]
+    let padded = [ # padded to 32 bytes
+      hex[0] & "00",
+      hex[1] & "00",
+      hex[2] & "00"
+    ]
+    let bytes = cast[seq[byte]](parseHexStr(hex.join()))
+    let elements = arrayFromHex(padded, littleEndian)
+    check merkleRoot(bytes).toHex == merkleRoot(elements).toHex
+
+  test "merkle root of bytes whose length is not a not multiple of 31":
+    let hex = [ # 31 bytes per field element
+      "000102030405060708090A0B0C0D0E0F101112131415161718191A1B1C1D1E",
+      "1F202122232425262728292A2B2C2D2E2F303132333435363738393A3B3C3D",
+      "3E3F404142434445464748494A4B4C4D4E4F505152535455565758"
+    ]
+    let padded = [ # padded to 32 bytes
+      hex[0] & "00",
+      hex[1] & "00",
+      hex[2] & "0000000000"
+    ]
+    let bytes = cast[seq[byte]](parseHexStr(hex.join()))
+    let elements = arrayFromHex(padded, littleEndian)
+    check merkleRoot(bytes).toHex == merkleRoot(elements).toHex