From 7ce2bc49d074a11d3bcf781caceb2ac332e6cb3e Mon Sep 17 00:00:00 2001
From: Balazs Komuves <bkomuves@gmail.com>
Date: Tue, 15 Oct 2024 21:46:56 +0200
Subject: [PATCH] add some marshalling functionality (digests to/from bytes)

---
 goldilocks_hash/marshal.nim             | 116 +++++++++++++++++++
 goldilocks_hash/types.nim               |  31 +++++-
 tests/goldilocks_hash/field/marshal.nim | 141 ++++++++++++++++++++++++
 tests/test.nim                          |   1 +
 4 files changed, 285 insertions(+), 4 deletions(-)
 create mode 100644 goldilocks_hash/marshal.nim
 create mode 100644 tests/goldilocks_hash/field/marshal.nim

diff --git a/goldilocks_hash/marshal.nim b/goldilocks_hash/marshal.nim
new file mode 100644
index 0000000..8ab1788
--- /dev/null
+++ b/goldilocks_hash/marshal.nim
@@ -0,0 +1,116 @@
+
+import std/bitops
+import types
+
+#-------------------------------------------------------------------------------
+
+func uint64ToBytesLE*(what: uint64): array[8, byte] =
+  var   bytes : array[8,byte]
+  var   x     : uint64 = what
+  const mask  : uint64 = 0xff
+  for i in 0..<8:
+    bytes[i] = byte(bitand(x,mask))
+    x = x shr 8
+  return bytes
+
+proc uint64ToBytesIntoLE[n: static int](what: uint64, tgt: var array[n,byte], ofs: int) =
+  var   x    : uint64 = what
+  const mask : uint64 = 0xff
+  for i in 0..<8:
+    tgt[ofs + i] = byte(bitand(x,mask))
+    x = x shr 8
+
+#---------------------------------------
+
+# simply store a hash digest (4 Goldilocks field elements) as 32 bytes,
+# encoding them as 64 bit little-endian integers
+func digestToBytes*(digest: Digest): array[32, byte] =
+  let arr   : F4 = fromDigest(digest)
+  var bytes : array[32,byte]
+  uint64ToBytesIntoLE[32]( fromF(arr[0]) , bytes , 0  )
+  uint64ToBytesIntoLE[32]( fromF(arr[1]) , bytes , 8  )
+  uint64ToBytesIntoLE[32]( fromF(arr[2]) , bytes , 16 )
+  uint64ToBytesIntoLE[32]( fromF(arr[3]) , bytes , 24 )
+  return bytes
+
+#-------------------------------------------------------------------------------
+
+func bytesToUint64FromLE(bytes: openarray[byte], ofs: int): uint64 =
+  assert( ofs+7 < bytes.len )
+  var x: uint64 = 0
+  for i in 0..<8:
+    x = x shl 8
+    x = bitor( x , uint64(bytes[ofs+7-i]) )
+  return x
+
+func bytesToUint64LE*(bytes: array[8, byte]): uint64 = bytesToUint64FromLE(bytes, 0)
+
+#-------------------------------------------------------------------------------
+
+func decodeBytesToDigestFrom(bytes: openarray[byte], ofs: int): Digest =
+  const mask : uint64 = 0x3fffffffffffffff'u64
+  let p = bytesToUint64FromLE( bytes , ofs +  0 )
+  let q = bytesToUint64FromLE( bytes , ofs +  7 )
+  let r = bytesToUint64FromLE( bytes , ofs + 15 )
+  let s = bytesToUint64FromLE( bytes , ofs + 23 )
+
+  let a = bitand( p       , mask )
+  let b = bitor(  q shr 6 , bitand(r , 0x0f) shl 58 )
+  let c = bitor(  r shr 4 , bitand(s , 0x03) shl 60 )
+  let d =         s shr 2
+
+  return mkDigestU64(a,b,c,d)   
+
+# takes 31 bytes (not 32!) and creates a unique Digest out of them
+# this is used when hashing byte sequences (also to be a drop-in replacement for the BN254 field)
+# what we do is we divide into 4 pieces of size 62 bits, and interpret
+# them as little-endian integers, and further interpret those as field elements
+func decodeBytesToDigest*(bytes: array[31, byte]): Digest = decodeBytesToDigestFrom( bytes, 0 )
+
+#-------------------------------------------------------------------------------
+
+# pad to a multiple of 31 bytes (with the 10* padding strategy) and decode into a sequence of digests
+func padAndDecodeBytesToDigest31*(bytes: openarray[byte]): seq[Digest] =
+  let m  = bytes.len
+  let n1 = m div 31
+  let n  = n1 + 1
+  let r  = n*31 - m      # 1 <= r <= 31
+  let q  = 31 - r        # 0 <= q <= 30
+
+  var ds : seq[Digest] = newSeq[Digest](n)
+
+  for i in 0..<n1:
+    ds[i] = decodeBytesToDigestFrom(bytes, 31*i)
+
+  var last : array[31,byte]
+  for i in 0..<q: last[i] = bytes[31*n1 + i]
+  last[q] = 0x01
+  ds[n1] = decodeBytesToDigest(last)
+
+  return ds
+
+#---------------------------------------
+
+# pad to a multiple of 62 bytes (with the 10* padding strategy) and decode into a sequence of digests
+func padAndDecodeBytesToDigest62*(bytes: openarray[byte]): seq[Digest] =
+  let m  = bytes.len
+  let n1 = m div 62
+  let n  = n1 + 1
+  let r  = n*62 - m      # 1 <= r <= 62
+  let q  = 62 - r        # 0 <= q <= 61
+
+  var ds : seq[Digest] = newSeq[Digest](2*n)
+
+  for i in 0..<2*n1:
+    ds[i] = decodeBytesToDigestFrom(bytes, 31*i)
+
+  var last : array[62,byte]
+  for i in 0..<q: last[i] = bytes[62*n1 + i]
+  last[q] = 0x01
+
+  ds[2*n1  ] = decodeBytesToDigestFrom(last, 0 )
+  ds[2*n1+1] = decodeBytesToDigestFrom(last, 31)
+
+  return ds
+
+#-------------------------------------------------------------------------------
diff --git a/goldilocks_hash/types.nim b/goldilocks_hash/types.nim
index 8a5e045..dad014d 100644
--- a/goldilocks_hash/types.nim
+++ b/goldilocks_hash/types.nim
@@ -34,11 +34,16 @@ type State*  = distinct F12
 func fromDigest* (x : Digest): F4  = return F4(x)
 func fromState * (x : State):  F12 = return F12(x)
 
-func toDigest* (x : F4 ): Digest = Digest(x)
-func toState*  (x : F12): State  = State(x)
+func toDigest*   (x : F4 ): Digest = Digest(x)
+func toState*    (x : F12): State  = State(x)
 
-func mkDigestU64* (a,b,c,d: uint64): Digest = toDigest( [toF(a),toF(b),toF(c),toF(d) ] )
-func mkDigest*    (a,b,c,d: F     ): Digest = toDigest( [a,b,c,d] )
+func mkDigest*   (a,b,c,d: F): Digest = toDigest( [a,b,c,d] )
+
+func mkDigestU64*(a,b,c,d: uint64): Digest = toDigest( [toF(a),toF(b),toF(c),toF(d) ] )
+func toDigestU64*(x : array[4,uint64]): Digest = mkDigestU64( x[0], x[1], x[2], x[3] )
+
+func uint64ToDigest*(x: uint64): Digest = mkDigestU64( x,0,0,0 )
+func intToDigest*   (x: uint64): Digest = uint64ToDigest( uint64(x) )
 
 const zeroDigest* : Digest = mkDigestU64(0,0,0,0)
 
@@ -50,9 +55,27 @@ proc `$`*(x: Digest): string = return $(fromDigest(x))
 
 #-------------------------------------------------------------------------------
 
+func digestSeqToFeltSeq*( ds: seq[Digest] ): seq[F] = 
+  let n = ds.len
+  var output: seq[F] = newSeq[F]( 4*n )
+  for k in 0..<n:
+    let f4 = fromDigest( ds[k] )
+    let j = 4*k
+    output[j+0] = f4[0]
+    output[j+1] = f4[1]
+    output[j+2] = f4[2]
+    output[j+3] = f4[3]
+  return output
+
+#-------------------------------------------------------------------------------
+
 func numberOfBits*(T: type): int {.compileTime.} =
   when T is F:
     63
+#  elif T is uint32:
+#    32
+#  elif T is uint16:
+#    16
   elif T is byte:
     8
   elif T is bool:
diff --git a/tests/goldilocks_hash/field/marshal.nim b/tests/goldilocks_hash/field/marshal.nim
new file mode 100644
index 0000000..8411f63
--- /dev/null
+++ b/tests/goldilocks_hash/field/marshal.nim
@@ -0,0 +1,141 @@
+
+import std/sequtils
+import std/unittest
+
+import goldilocks_hash/types
+import goldilocks_hash/marshal
+
+#-------------------------------------------------------------------------------
+
+type U4 = array[4,uint64]
+
+func F4SeqToDigestSeq( inp : seq[F4] ): seq[Digest] = inp.map(toDigest   )
+func U4SeqToDigestSeq( inp : seq[U4] ): seq[Digest] = inp.map(toDigestU64)
+
+func intToByte(x: int): byte = byte(x)
+
+#-------------------------------------------------------------------------------
+
+suite "marshalling to/from Digests":
+
+  test "uint64ToBytesLE":
+    let word  : uint64        = 0x123456789abcdef3'u64
+    let bytes : array[8,byte] = [ 0xf3, 0xde, 0xbc, 0x9a, 0x78, 0x56, 0x34, 0x12 ]
+    check ( bytes == uint64ToBytesLE(word) );
+
+  test "bytesToUint64LE":
+    let word  : uint64        = 0x123456789abcdef3'u64
+    let bytes : array[8,byte] = [ 0xf3, 0xde, 0xbc, 0x9a, 0x78, 0x56, 0x34, 0x12 ]
+    check ( bytesToUint64LE(bytes) == word );
+
+  #-----------------
+  
+  test "digestToBytes":
+    let dig = mkDigestU64( 0x123456789abcdef3'u64
+                         , 0x43f5a2d9c2871a4b'u64
+                         , 0xb2d79201f9771e0e'u64
+                         , 0x2074c7946509c3a5'u64
+                         )
+    let bytes : array[32, uint8] = 
+          [ 0xf3, 0xde, 0xbc, 0x9a, 0x78, 0x56, 0x34, 0x12
+          , 0x4b, 0x1a, 0x87, 0xc2, 0xd9, 0xa2, 0xf5, 0x43
+          , 0x0e, 0x1e, 0x77, 0xf9, 0x01, 0x92, 0xd7, 0xb2
+          , 0xa5, 0xc3, 0x09, 0x65, 0x94, 0xc7, 0x74, 0x20
+          ]
+    check ( bytes == digestToBytes(dig) )
+
+  #-----------------
+
+  test "decodeBytesToDigest [1..31]":
+    let bytes : array[31,byte]  = [1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31]
+    let words : array[4,uint64] = [ 0x0807060504030201'u64 , 0x003c3834302c2824'u64 , 0x0171615141312111'u64 , 0x07c7874706c68646'u64 ]
+    check ( decodeBytesToDigest(bytes) == toDigestU64(words) )
+
+  test "decodeBytesToDigest [51..81]":
+    let bytes : array[31,byte]  = [51,52,53,54,55,56,57,58,59,60,61,62,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81]
+    let words : array[4,uint64] = [ 0x3a39383736353433'u64 , 0x090500fcf8f4f0ec'u64 , 0x2494847464544434'u64 , 0x145413d3935312d2'u64 ]
+    check ( decodeBytesToDigest(bytes) == toDigestU64(words) )
+
+  test "decodeBytesToDigest [171..201]":
+    let bytes : array[31,byte]  = [171,172,173,174,175,176,177,178,179,180,181,182,183,184,185,186,187,188,189,190,191,192,193,194,195,196,197,198,199,200,201]
+    let words : array[4,uint64] = [ 0x32b1b0afaeadacab'u64 , 0x2ae6e2dedad6d2ce'u64 , 0x2c1c0bfbebdbcbbb'u64 , 0x327231f1b17130f0'u64 ]
+    check ( decodeBytesToDigest(bytes) == toDigestU64(words) )
+
+  #-----------------
+
+  test "padAndDecodeBytes31 []":
+    check U4SeqToDigestSeq( @[ [1'u64, 0'u64, 0'u64, 0'u64] ]) == padAndDecodeBytesToDigest31( @[] )
+
+  test "padAndDecodeBytes31 [1..11]":
+    let f4s : seq[U4] = @[ [0x0807060504030201'u64,0x00000000042c2824'u64,0x0000000000000000'u64,0x0000000000000000'u64] ]
+    check U4SeqToDigestSeq(f4s) == padAndDecodeBytesToDigest31( (1..11).toSeq().map(intToByte) )
+
+  test "padAndDecodeBytes31 [1..30]":
+    let f4s : seq[U4] = @[ [0x0807060504030201'u64,0x003c3834302c2824'u64,0x0171615141312111'u64,0x0047874706c68646'u64] ]
+    check U4SeqToDigestSeq(f4s) == padAndDecodeBytesToDigest31( (1..30).toSeq().map(intToByte) )
+
+  test "padAndDecodeBytes31 [1..31]":
+    let f4s : seq[U4] = @[ [0x0807060504030201'u64,0x003c3834302c2824'u64,0x0171615141312111'u64,0x07c7874706c68646'u64]
+                         , [0x0000000000000001'u64,0x0000000000000000'u64,0x0000000000000000'u64,0x0000000000000000'u64]
+                         ]
+    check U4SeqToDigestSeq(f4s) == padAndDecodeBytesToDigest31( (1..31).toSeq().map(intToByte) )
+
+  test "padAndDecodeBytes31 [1..51]":
+    let f4s : seq[U4] = @[ [0x0807060504030201'u64,0x003c3834302c2824'u64,0x0171615141312111'u64,0x07c7874706c68646'u64]
+                         , [0x2726252423222120'u64,0x3cb8b4b0aca8a4a0'u64,0x0000001333231302'u64,0x0000000000000000'u64]
+                         ]
+    check U4SeqToDigestSeq(f4s) == padAndDecodeBytesToDigest31( (1..51).toSeq().map(intToByte) )
+
+  test "padAndDecodeBytes31 [1..71]":
+    let f4s : seq[U4] = @[ [0x0807060504030201'u64,0x003c3834302c2824'u64,0x0171615141312111'u64,0x07c7874706c68646'u64]
+                         , [0x2726252423222120'u64,0x3cb8b4b0aca8a4a0'u64,0x3363534333231302'u64,0x0f8f4f0ece8e4e0d'u64]
+                         , [0x064544434241403f'u64,0x000000000000051d'u64,0x0000000000000000'u64,0x0000000000000000'u64]
+                         ]
+    check U4SeqToDigestSeq(f4s) == padAndDecodeBytesToDigest31( (1..71).toSeq().map(intToByte) )
+
+  #-----------------
+
+  test "padAndDecodeBytes62 []":
+    check U4SeqToDigestSeq( @[ [1'u64, 0'u64, 0'u64, 0'u64] , [0'u64, 0'u64, 0'u64, 0'u64] ]) == padAndDecodeBytesToDigest62( @[] )
+
+  test "padAndDecodeBytes62 [1..11]":
+    let f4s : seq[U4] = @[ [0x0807060504030201'u64,0x00000000042c2824'u64,0x0000000000000000'u64,0x0000000000000000'u64]
+                         , [0x0000000000000000'u64,0x0000000000000000'u64,0x0000000000000000'u64,0x0000000000000000'u64] ]
+    check U4SeqToDigestSeq(f4s) == padAndDecodeBytesToDigest62( (1..11).toSeq().map(intToByte) )
+
+  test "padAndDecodeBytes62 [1..31]":
+    let f4s : seq[U4] = 
+               @[[0x0807060504030201'u64,0x003c3834302c2824'u64,0x0171615141312111'u64,0x07c7874706c68646'u64]
+                ,[0x0000000000000001'u64,0x0000000000000000'u64,0x0000000000000000'u64,0x0000000000000000'u64]]
+    check U4SeqToDigestSeq(f4s) == padAndDecodeBytesToDigest62( (1..31).toSeq().map(intToByte) )
+
+  test "padAndDecodeBytes62 [1..51]":
+    let f4s : seq[U4] = 
+               @[[0x0807060504030201'u64,0x003c3834302c2824'u64,0x0171615141312111'u64,0x07c7874706c68646'u64]
+                ,[0x2726252423222120'u64,0x3cb8b4b0aca8a4a0'u64,0x0000001333231302'u64,0x0000000000000000'u64]]
+    check U4SeqToDigestSeq(f4s) == padAndDecodeBytesToDigest62( (1..51).toSeq().map(intToByte) )
+
+  test "padAndDecodeBytes62 [1..61]":
+    let f4s : seq[U4] = 
+               @[[0x0807060504030201'u64,0x003c3834302c2824'u64,0x0171615141312111'u64,0x07c7874706c68646'u64]
+                ,[0x2726252423222120'u64,0x3cb8b4b0aca8a4a0'u64,0x3363534333231302'u64,0x004f4f0ece8e4e0d'u64]]
+    check U4SeqToDigestSeq(f4s) == padAndDecodeBytesToDigest62( (1..61).toSeq().map(intToByte) )
+
+  test "padAndDecodeBytes62 [1..62]":
+    let f4s : seq[U4] = 
+               @[[0x0807060504030201'u64,0x003c3834302c2824'u64,0x0171615141312111'u64,0x07c7874706c68646'u64]
+                ,[0x2726252423222120'u64,0x3cb8b4b0aca8a4a0'u64,0x3363534333231302'u64,0x0f8f4f0ece8e4e0d'u64]
+                ,[0x0000000000000001'u64,0x0000000000000000'u64,0x0000000000000000'u64,0x0000000000000000'u64]
+                ,[0x0000000000000000'u64,0x0000000000000000'u64,0x0000000000000000'u64,0x0000000000000000'u64]]
+    check U4SeqToDigestSeq(f4s) == padAndDecodeBytesToDigest62( (1..62).toSeq().map(intToByte) )
+
+  test "padAndDecodeBytes62 [1..71]":
+    let f4s : seq[U4] = 
+               @[[0x0807060504030201'u64,0x003c3834302c2824'u64,0x0171615141312111'u64,0x07c7874706c68646'u64]
+                ,[0x2726252423222120'u64,0x3cb8b4b0aca8a4a0'u64,0x3363534333231302'u64,0x0f8f4f0ece8e4e0d'u64]
+                ,[0x064544434241403f'u64,0x000000000000051d'u64,0x0000000000000000'u64,0x0000000000000000'u64]
+                ,[0x0000000000000000'u64,0x0000000000000000'u64,0x0000000000000000'u64,0x0000000000000000'u64]]
+    check U4SeqToDigestSeq(f4s) == padAndDecodeBytesToDigest62( (1..71).toSeq().map(intToByte) )
+
+##-------------------------------------------------------------------------------
+#
\ No newline at end of file
diff --git a/tests/test.nim b/tests/test.nim
index ebf1907..47d1ffc 100644
--- a/tests/test.nim
+++ b/tests/test.nim
@@ -1,5 +1,6 @@
 
 import ./goldilocks_hash/field/testField
+import ./goldilocks_hash/field/marshal
 
 import ./goldilocks_hash/poseidon2/testPermutation
 import ./goldilocks_hash/poseidon2/testCompress