Pad big endian bytes with 0's to the right

poseidon2.nim

@@ -98,4 +98,4 @@ func merkleRoot*(xs: openArray[F]) : F =
     return merkleRoot(ys)
 
 func merkleRoot*(bytes: openArray[byte]): F =
-  merkleRoot(F.unmarshal(bytes, littleEndian))
+  merkleRoot(seq[F].unmarshal(bytes, littleEndian))

poseidon2/io.nim

@@ -2,17 +2,32 @@ import ./types
 import constantine/math/arithmetic
 import constantine/math/io/io_bigints
 
+func padRight(source: openArray[byte], endian: static Endianness): array[32, byte] =
+  assert source.len <= 31
+  when endian == littleEndian:
+    copyMem(addr result[0], unsafeAddr source[0], source.len)
+  when endian == bigEndian:
+    copyMem(addr result[1], unsafeAddr source[0], source.len)
+
 func unmarshal*(
         _: type F,
         bytes: openArray[byte],
+        endian: static Endianness): F =
+  assert bytes.len <= 31
+  let padded = bytes.padRight(endian)
+  let bigint = B.unmarshal(padded, endian)
+  return F.fromBig(bigint)
+
+func unmarshal*(
+        _: type seq[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)
+  var chunkStart = 0
+  while chunkStart < bytes.len:
+    let chunkEnd = min(chunkStart + 31, bytes.len)
+    let element = F.unmarshal(bytes.toOpenArray(chunkStart, chunkEnd - 1), endian)
     elements.add(element)
-    i += chunkLen
+    chunkStart += chunkLen
   return elements

tests/poseidon2/testIo.nim

@@ -7,30 +7,49 @@ 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)
+    let bytes = toSeq 1'u8..31'u8
+    let paddedTo32 = @[0x00'u8] & bytes # most significant byte is not used
+    let expected = F.fromBig(B.unmarshal(paddedTo32, bigEndian))
+    let unmarshalled = F.unmarshal(bytes, bigEndian)
+    check bool(unmarshalled == expected)
 
   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)
+    let bytes = toSeq 1'u8..31'u8
+    let paddedTo32 = bytes & @[0x00'u8] # most significant byte is not used
+    let expected = F.fromBig(B.unmarshal(paddedTo32, littleEndian))
+    let unmarshalled = F.unmarshal(bytes, littleEndian)
+    check bool(unmarshalled == expected)
+
+  test "pads big endian bytes to the right with 0's":
+    let bytes = @[0x12'u8, 0x34, 0x56]
+    let paddedTo31 = bytes & 0x00'u8.repeat(31 - bytes.len)
+    let paddedTo32 = @[0x00'u8] & paddedTo31 # most significant byte is not used
+    let expected = F.fromBig(B.unmarshal(paddedTo32, bigEndian))
+    let unmarshalled = F.unmarshal(bytes, bigEndian)
+    check bool(unmarshalled == expected)
+
+  test "pads little endian bytes to the right with 0's":
+    let bytes = @[0x56'u8, 0x34, 0x12]
+    let paddedTo31 = bytes & 0x00'u8.repeat(31 - bytes.len)
+    let paddedTo32 = paddedTo31 & @[0x00'u8] # most significant byte is not used
+    let expected = F.fromBig(B.unmarshal(paddedTo32, littleEndian))
+    let unmarshalled = F.unmarshal(bytes, littleEndian)
+    check bool(unmarshalled == expected)
 
   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)
+    template checkConversion(endian) =
+      let bytes = toSeq 1'u8..80'u8
+      let padded = bytes & 0'u8.repeat(93 - bytes.len)
+      let expected1 = F.fromBig(B.unmarshal(padded[ 0..<31], endian))
+      let expected2 = F.fromBig(B.unmarshal(padded[31..<62], endian))
+      let expected3 = F.fromBig(B.unmarshal(padded[62..<93], endian))
+      let elements = seq[F].unmarshal(bytes, endian)
+      check elements.len == 3
+      check bool(elements[0] == expected1)
+      check bool(elements[1] == expected2)
+      check bool(elements[2] == expected3)
+
+    checkConversion(littleEndian)
+    checkConversion(bigEndian)