From 43c2aab9131a3de5556298f0805c36004e63f4d7 Mon Sep 17 00:00:00 2001
From: Mark Spanbroek <mark@spanbroek.net>
Date: Tue, 31 Oct 2023 13:40:07 +0100
Subject: [PATCH] move unmarshal into its own module called 'io'

Includes separate tests for unmarshalling that
were previously tested as part of the merkleRoot
calculation.

Includes tests for unmarshalling little endian
and big endian byte arrays.
---
 poseidon2.nim                     | 20 +--------
 poseidon2/io.nim                  | 18 ++++++++
 tests/poseidon2/testIo.nim        | 36 ++++++++++++++++
 tests/poseidon2/testPoseidon2.nim | 39 +++++++++++++++++
 tests/test.nim                    | 72 ++-----------------------------
 5 files changed, 97 insertions(+), 88 deletions(-)
 create mode 100644 poseidon2/io.nim
 create mode 100644 tests/poseidon2/testIo.nim
 create mode 100644 tests/poseidon2/testPoseidon2.nim

diff --git a/poseidon2.nim b/poseidon2.nim
index f750564..34a84cd 100644
--- a/poseidon2.nim
+++ b/poseidon2.nim
@@ -1,13 +1,10 @@
-
-# import std/sugar
-
 import
   constantine/math/arithmetic,
-  constantine/math/io/io_bigints,
   constantine/math/config/curves
 
 import poseidon2/types
 import poseidon2/roundconst
+import poseidon2/io
 
 #-------------------------------------------------------------------------------
 
@@ -100,20 +97,5 @@ 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/io.nim b/poseidon2/io.nim
new file mode 100644
index 0000000..307cc36
--- /dev/null
+++ b/poseidon2/io.nim
@@ -0,0 +1,18 @@
+import ./types
+import constantine/math/arithmetic
+import constantine/math/io/io_bigints
+
+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
diff --git a/tests/poseidon2/testIo.nim b/tests/poseidon2/testIo.nim
new file mode 100644
index 0000000..bcde36e
--- /dev/null
+++ b/tests/poseidon2/testIo.nim
@@ -0,0 +1,36 @@
+import std/unittest
+import std/sequtils
+import constantine/math/io/io_bigints
+import constantine/math/arithmetic
+import poseidon2/types
+import poseidon2/io
+
+suite "unmarshalling":
+
+  const F123456 = F.fromBig(B.fromHex("0x123456", bigEndian))
+
+  test "converts big endian bytes into field elements":
+    let bytes = [0x12'u8, 0x34'u8, 0x56'u8]
+    var padded: array[31, byte]
+    padded[^3..^1] = bytes
+    check bool(F.unmarshal(bytes, bigEndian)[0] == F123456)
+    check bool(F.unmarshal(padded, bigEndian)[0] == F123456)
+
+  test "converts little endian bytes into field elements":
+    let bytes = [0x56'u8, 0x34'u8, 0x12'u8]
+    var padded: array[31, byte]
+    padded[0..<3] = bytes
+    check bool(F.unmarshal(bytes, littleEndian)[0] == F123456)
+    check bool(F.unmarshal(padded, littleEndian)[0] == F123456)
+
+  test "converts every 31 bytes into a field element":
+    let bytes = toSeq 1'u8..80'u8
+    let element1 = F.fromBig(B.unmarshal(toSeq  1'u8..31'u8, bigEndian))
+    let element2 = F.fromBig(B.unmarshal(toSeq 32'u8..62'u8, bigEndian))
+    let element3 = F.fromBig(B.unmarshal(toSeq 63'u8..80'u8, bigEndian))
+    let elements = F.unmarshal(bytes, bigEndian)
+    check elements.len == 3
+    check bool(elements[0] == element1)
+    check bool(elements[1] == element2)
+    check bool(elements[2] == element3)
+
diff --git a/tests/poseidon2/testPoseidon2.nim b/tests/poseidon2/testPoseidon2.nim
new file mode 100644
index 0000000..f12290b
--- /dev/null
+++ b/tests/poseidon2/testPoseidon2.nim
@@ -0,0 +1,39 @@
+import std/unittest
+import std/math
+import std/sequtils
+
+import constantine/math/arithmetic
+import constantine/math/io/io_fields
+import constantine/math/io/io_bigints
+import constantine/math/config/curves
+
+import poseidon2/types
+import poseidon2
+
+suite "poseidon2":
+
+  test "permutation in place":
+    var x: F = toF(0)
+    var y: F = toF(1)
+    var z: F = toF(2)
+
+    permInplace(x, y, z)
+
+    check toDecimal(x) == "21882471761025344482456282050943515707267606647948403374880378562101343146243"
+    check toDecimal(y) == "09030699330013392132529464674294378792132780497765201297316864012141442630280"
+    check toDecimal(z) == "09137931384593657624554037900714196568304064431583163402259937475584578975855"
+
+  test "merkle root of field elements":
+    let m = 17
+    let n = 2^m
+    var xs: seq[F]
+    for i in 1..n:
+      xs.add( toF(i) )
+
+    let root = merkleRoot(xs)
+    check root.toHex == "0x1eabbb64b76d5aecd393601c4a01878450e23f45fe8b2748bb63a615351b11d1"
+
+  test "merkle root of bytes":
+    let bytes = toSeq 1'u8..80'u8
+    let root = merkleRoot(bytes)
+    check root.toHex == "0x2d2d6f27d20a9e5597e20a888a4deac0e0600f92f0b9691ace0fdfab2fc1f500"
diff --git a/tests/test.nim b/tests/test.nim
index b018a8a..46c8f83 100644
--- a/tests/test.nim
+++ b/tests/test.nim
@@ -1,70 +1,4 @@
-import std/unittest
-import std/math
-import std/strutils
-import std/sequtils
+import ./poseidon2/testPoseidon2
+import ./poseidon2/testIo
 
-import constantine/math/arithmetic
-import constantine/math/io/io_fields
-import constantine/math/io/io_bigints
-import constantine/math/config/curves
-
-import poseidon2/types
-import poseidon2
-
-suite "poseidon2":
-
-  test "permutation in place":
-    var x: F = toF(0)
-    var y: F = toF(1)
-    var z: F = toF(2)
-
-    permInplace(x, y, z)
-
-    check toDecimal(x) == "21882471761025344482456282050943515707267606647948403374880378562101343146243"
-    check toDecimal(y) == "09030699330013392132529464674294378792132780497765201297316864012141442630280"
-    check toDecimal(z) == "09137931384593657624554037900714196568304064431583163402259937475584578975855"
-
-  test "merkle root of field elements":
-    let m = 17
-    let n = 2^m
-    var xs: seq[F]
-    for i in 1..n:
-      xs.add( toF(i) )
-
-    let root = merkleRoot(xs)
-    check root.toHex == "0x1eabbb64b76d5aecd393601c4a01878450e23f45fe8b2748bb63a615351b11d1"
-
-  test "merkle root of bytes":
-    let bytes = toSeq 1'u8..80'u8
-    let root = merkleRoot(bytes)
-    check root.toHex == "0x2d2d6f27d20a9e5597e20a888a4deac0e0600f92f0b9691ace0fdfab2fc1f500"
-
-  test "merkle root converts every 31 bytes to a field element":
-    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
+{.warning[UnusedImport]: off.}