From 4af85817bf7ca70df3c4d17252d6687f5af391a5 Mon Sep 17 00:00:00 2001
From: Mark Spanbroek <mark@spanbroek.net>
Date: Wed, 8 Nov 2023 13:08:33 +0100
Subject: [PATCH] Add streaming API for sponge

---
 poseidon2.nim                       | 66 ++--------------------------
 poseidon2/permutation.nim           | 18 ++++++++
 poseidon2/sponge.nim                | 67 +++++++++++++++++++++++++++++
 tests/poseidon2/testPermutation.nim | 22 ++++++++++
 tests/poseidon2/testPoseidon2.nim   | 53 -----------------------
 tests/poseidon2/testSponge.nim      | 50 +++++++++++++++++++++
 tests/test.nim                      |  2 +
 7 files changed, 162 insertions(+), 116 deletions(-)
 create mode 100644 poseidon2/permutation.nim
 create mode 100644 poseidon2/sponge.nim
 create mode 100644 tests/poseidon2/testPermutation.nim
 create mode 100644 tests/poseidon2/testSponge.nim

diff --git a/poseidon2.nim b/poseidon2.nim
index 1e3e568..d68ec1c 100644
--- a/poseidon2.nim
+++ b/poseidon2.nim
@@ -1,73 +1,13 @@
 import constantine/math/arithmetic
 
 import poseidon2/types
-import poseidon2/roundfun
+import poseidon2/permutation
 import poseidon2/io
+import poseidon2/sponge
 
+export sponge
 export toBytes
 
-#-------------------------------------------------------------------------------
-
-# the Poseidon2 permutation (mutable, in-place version)
-proc permInplace*(x, y, z : var F) =
-  linearLayer(x, y, z)
-  for j in 0..3:
-    externalRound(j, x, y, z)
-  for j in 0..55:
-    internalRound(j, x, y, z)
-  for j in 4..7:
-    externalRound(j, x, y, z)
-
-# the Poseidon2 permutation
-func perm*(xyz: S) : S =
-  var (x,y,z) = xyz
-  permInplace(x, y, z)
-  return (x,y,z)
-
-#-------------------------------------------------------------------------------
-
-# sponge with rate=1 (capacity=2)
-func spongeWithRate1*(xs: openArray[F]) : F =
-  var s0 : F = zero
-  var s1 : F = zero
-  var s2 : F = toF(0x0301) ; s2 += twoToThe64        # domain separation IV := (2^64 + 256*t + r)
-
-  for x in xs:
-    s0 += x
-    permInplace(s0,s1,s2)
-
-  # padding
-  s0 += one
-  permInplace(s0,s1,s2)
-  return s0
-
-# sponge with rate=2 (capacity=1)
-func spongeWithRate2*(xs: openArray[F]) : F =
-  let a = low(xs)
-  let b = high(xs)
-  let n = b-a+1
-  let halfn  : int  = n div 2
-
-  var s0 : F = zero
-  var s1 : F = zero
-  var s2 : F = toF(0x0302) ; s2 += twoToThe64        # domain separation IV := (2^64 + 256*t + r)
-
-  for i in 0..<halfn:
-    s0 += xs[a+2*i  ]
-    s1 += xs[a+2*i+1]
-    permInplace(s0,s1,s2)
-
-  if (2*halfn == n):
-    # padding even input
-    s0 += one
-    s1 += zero
-  else:
-    # padding odd input
-    s0 += xs[b]
-    s1 += one
-
-  permInplace(s0,s1,s2)
-  return s0
 
 #-------------------------------------------------------------------------------
 
diff --git a/poseidon2/permutation.nim b/poseidon2/permutation.nim
new file mode 100644
index 0000000..6911f00
--- /dev/null
+++ b/poseidon2/permutation.nim
@@ -0,0 +1,18 @@
+import ./types
+import ./roundfun
+
+# the Poseidon2 permutation (mutable, in-place version)
+proc permInplace*(x, y, z : var F) =
+  linearLayer(x, y, z)
+  for j in 0..3:
+    externalRound(j, x, y, z)
+  for j in 0..55:
+    internalRound(j, x, y, z)
+  for j in 4..7:
+    externalRound(j, x, y, z)
+
+# the Poseidon2 permutation
+func perm*(xyz: S) : S =
+  var (x,y,z) = xyz
+  permInplace(x, y, z)
+  return (x,y,z)
diff --git a/poseidon2/sponge.nim b/poseidon2/sponge.nim
new file mode 100644
index 0000000..274f4b2
--- /dev/null
+++ b/poseidon2/sponge.nim
@@ -0,0 +1,67 @@
+import ./types
+import ./permutation
+import constantine/math/io/io_fields
+import constantine/math/arithmetic
+
+type
+  Sponge*[rate: static int] = object
+    s0: F
+    s1: F
+    s2: F
+    when rate == 2:
+      even: bool
+
+func init(sponge: var Sponge[1]) =
+  # domain separation IV := (2^64 + 256*t + r)
+  const IV = F.fromHex("0x10000000000000301")
+  sponge.s0 = zero
+  sponge.s1 = zero
+  sponge.s2 = IV
+
+func update*(sponge: var Sponge[1], element: F) =
+  sponge.s0 += element
+  permInPlace(sponge.s0, sponge.s1, sponge.s2)
+
+func finish*(sponge: var Sponge[1]): F =
+  # padding
+  sponge.s0 += one
+  permInPlace(sponge.s0, sponge.s1, sponge.s2)
+  return sponge.s0
+
+func init(sponge: var Sponge[2]) =
+  # domain separation IV := (2^64 + 256*t + r)
+  const IV = F.fromHex("0x10000000000000302")
+  sponge.s0 = zero
+  sponge.s1 = zero
+  sponge.s2 = IV
+  sponge.even = true
+
+func update*(sponge: var Sponge[2], element: F) =
+  if sponge.even:
+    sponge.s0 += element
+  else:
+    sponge.s1 += element
+    permInPlace(sponge.s0, sponge.s1, sponge.s2)
+  sponge.even = not sponge.even
+
+func finish*(sponge: var Sponge[2]): F =
+  if sponge.even:
+    # padding even input
+    sponge.s0 += one
+    sponge.s1 += zero
+  else:
+    # padding odd input
+    sponge.s1 += one
+  permInPlace(sponge.s0, sponge.s1, sponge.s2)
+  return sponge.s0
+
+func init*(_: type Sponge, rate: static int = 2): Sponge[rate] =
+  when rate notin {1, 2}:
+    {.error: "only rate 1 and 2 are supported".}
+  result.init
+
+func digest*(_: type Sponge, elements: openArray[F], rate: static int): F =
+  var sponge = Sponge.init(rate)
+  for element in elements:
+    sponge.update(element)
+  return sponge.finish()
diff --git a/tests/poseidon2/testPermutation.nim b/tests/poseidon2/testPermutation.nim
new file mode 100644
index 0000000..7d34100
--- /dev/null
+++ b/tests/poseidon2/testPermutation.nim
@@ -0,0 +1,22 @@
+import std/unittest
+
+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/permutation
+
+suite "permutation":
+
+  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"
diff --git a/tests/poseidon2/testPoseidon2.nim b/tests/poseidon2/testPoseidon2.nim
index d237506..9164653 100644
--- a/tests/poseidon2/testPoseidon2.nim
+++ b/tests/poseidon2/testPoseidon2.nim
@@ -10,61 +10,8 @@ import constantine/serialization/codecs
 import poseidon2/types
 import poseidon2
 
-#-------------------------------------------------------------------------------
-
-const expectedSpongeResultsRate1 : array[8, string] = 
-  [ "11474111961551684932675539562074905375756669035986300321099733737886849683321"
-  , "12075737409606154890751050839468327529267137715708285489737384891841319770833"
-  , "01607478768131843313297310704782442615640380643931196052095347138434114571392"
-  , "17583439011341576528906247721476731129932611848439423516301689821385840105693"
-  , "12983779044863516108508991186638610589212096523915590215701244866830295506005"
-  , "16646216251577650555646508049064625507758601195307236539843683725095763921505"
-  , "11914716034377431890952169039751213443286692885071871704776127977841051829452"
-  , "20798492850731331785912281726856492405884190236464781409482377236764537088662"
-  ]
-
-const expectedSpongeResultsRate2 : array[8, string] = 
-  [ "15335097698975718583905618186682475632756177170667436996250626760551196078076"
-  , "05101758095924000127790537496504070769319625501671400349336709520206095219618"
-  , "07306734450287348725566606192910189982345130476287345231433021147457815478255"
-  , "18511919414269811073023003336929505285555117419480831606637506641708579940507"
-  , "17917165106036607360653786499368288558581739128065811663709392730081030901634"
-  , "04630821736691665506072583795473163860465039714428126246168623896083265248907"
-  , "02020506076765964149531002674962673761843846094901604358961533722934321735239"
-  , "11732533243633999579592740965735640217427639382365959787508754341969556105663"
-  ]
-
-#-------------------------------------------------------------------------------
-
 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 "sponge with rate=1":
-    for n in 0..7:
-      var xs: seq[F]
-      for i in 1..n:
-        xs.add( toF(i) )
-      let h = spongeWithRate1(xs)
-      check toDecimal(h) == expectedSpongeResultsRate1[n] 
-
-  test "sponge with rate=2":
-    for n in 0..7:
-      var xs: seq[F]
-      for i in 1..n:
-        xs.add( toF(i) )
-      let h = spongeWithRate2(xs)
-      check toDecimal(h) == expectedSpongeResultsRate2[n] 
-
   test "merkle root of field elements":
     let m = 17
     let n = 2^m
diff --git a/tests/poseidon2/testSponge.nim b/tests/poseidon2/testSponge.nim
new file mode 100644
index 0000000..266992a
--- /dev/null
+++ b/tests/poseidon2/testSponge.nim
@@ -0,0 +1,50 @@
+import std/unittest
+
+import constantine/math/io/io_fields
+
+import poseidon2/types
+import poseidon2
+
+#-------------------------------------------------------------------------------
+
+const expectedSpongeResultsRate1 : array[8, string] =
+  [ "11474111961551684932675539562074905375756669035986300321099733737886849683321"
+  , "12075737409606154890751050839468327529267137715708285489737384891841319770833"
+  , "01607478768131843313297310704782442615640380643931196052095347138434114571392"
+  , "17583439011341576528906247721476731129932611848439423516301689821385840105693"
+  , "12983779044863516108508991186638610589212096523915590215701244866830295506005"
+  , "16646216251577650555646508049064625507758601195307236539843683725095763921505"
+  , "11914716034377431890952169039751213443286692885071871704776127977841051829452"
+  , "20798492850731331785912281726856492405884190236464781409482377236764537088662"
+  ]
+
+const expectedSpongeResultsRate2 : array[8, string] =
+  [ "15335097698975718583905618186682475632756177170667436996250626760551196078076"
+  , "05101758095924000127790537496504070769319625501671400349336709520206095219618"
+  , "07306734450287348725566606192910189982345130476287345231433021147457815478255"
+  , "18511919414269811073023003336929505285555117419480831606637506641708579940507"
+  , "17917165106036607360653786499368288558581739128065811663709392730081030901634"
+  , "04630821736691665506072583795473163860465039714428126246168623896083265248907"
+  , "02020506076765964149531002674962673761843846094901604358961533722934321735239"
+  , "11732533243633999579592740965735640217427639382365959787508754341969556105663"
+  ]
+
+#-------------------------------------------------------------------------------
+
+suite "sponge":
+
+  test "sponge with rate=1":
+    for n in 0..7:
+      var xs: seq[F]
+      for i in 1..n:
+        xs.add( toF(i) )
+      let h = Sponge.digest(xs, rate = 1)
+      check toDecimal(h) == expectedSpongeResultsRate1[n]
+
+  test "sponge with rate=2":
+    for n in 0..7:
+      var xs: seq[F]
+      for i in 1..n:
+        xs.add( toF(i) )
+      let h = Sponge.digest(xs, rate = 2)
+      check toDecimal(h) == expectedSpongeResultsRate2[n]
diff --git a/tests/test.nim b/tests/test.nim
index 46c8f83..18adebc 100644
--- a/tests/test.nim
+++ b/tests/test.nim
@@ -1,3 +1,5 @@
+import ./poseidon2/testPermutation
+import ./poseidon2/testSponge
 import ./poseidon2/testPoseidon2
 import ./poseidon2/testIo