From f5835d5fdc65f02db0248b50989ab9c2efd16c87 Mon Sep 17 00:00:00 2001
From: Balazs Komuves <bkomuves@gmail.com>
Date: Tue, 7 Nov 2023 14:54:10 +0100
Subject: [PATCH] add domain separation for the sponge construction

---
 poseidon2.nim                     |  4 ++--
 poseidon2/types.nim               |  2 ++
 tests/poseidon2/testPoseidon2.nim | 36 ++++++++++++++-----------------
 3 files changed, 20 insertions(+), 22 deletions(-)

diff --git a/poseidon2.nim b/poseidon2.nim
index e8f7f77..a36ba24 100644
--- a/poseidon2.nim
+++ b/poseidon2.nim
@@ -34,7 +34,7 @@ func spongeWithRate1*(xs: openArray[F]) : F =
 
   var s0 : F = zero
   var s1 : F = zero
-  var s2 : F = zero
+  var s2 : F = toF(0x0301) ; s2 += twoToThe64        # domain separation IV := (2^64 + 256*t + r)
 
   for x in xs:
     s0 += x
@@ -54,7 +54,7 @@ func spongeWithRate2*(xs: openArray[F]) : F =
 
   var s0 : F = zero
   var s1 : F = zero
-  var s2 : 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  ]
diff --git a/poseidon2/types.nim b/poseidon2/types.nim
index 8ee1f7e..ba72715 100644
--- a/poseidon2/types.nim
+++ b/poseidon2/types.nim
@@ -29,6 +29,8 @@ func toF*(a: int) : F =
 const zero* : F = getZero()
 const one*  : F = fromHex(F,"0x01")     # note: `fromUint()` does not work at compile time
 
+const twoToThe64* : F = fromHex(F,"0x10000000000000000")
+
 #-------------------------------------------------------------------------------
 
 func hexToF*(s : string, endian: static Endianness = bigEndian) : F =
diff --git a/tests/poseidon2/testPoseidon2.nim b/tests/poseidon2/testPoseidon2.nim
index 26de78f..f093162 100644
--- a/tests/poseidon2/testPoseidon2.nim
+++ b/tests/poseidon2/testPoseidon2.nim
@@ -13,29 +13,25 @@ import poseidon2
 #-------------------------------------------------------------------------------
 
 const expectedSpongeResultsRate1 : array[8, string] = 
-  [ "12363515589665961836680709257448433057869762330741639517836048636244832188495"
-  , "10755250120808789043370150604836786069442045362641800439807384337872752972068"
-  , "04842014531366721455661330916203255410159059117951668762867230544004815370337"
-  , "13502515636936876459766686836354199651004594178376827739246669803080321705927"
-  , "19312121576697000598919845239663673946550934099828684806027699882665482322097"
-  , "21509595983900483103260021285060939918324350560398732346653142062765920502059"
-  , "11892726572958426459775026381831352388154613015696290329810000571844227402585"
-  , "10284126944232604349630438079200913190801781418325975675236599364113149409058"
+  [ "11474111961551684932675539562074905375756669035986300321099733737886849683321"
+  , "12075737409606154890751050839468327529267137715708285489737384891841319770833"
+  , "01607478768131843313297310704782442615640380643931196052095347138434114571392"
+  , "17583439011341576528906247721476731129932611848439423516301689821385840105693"
+  , "12983779044863516108508991186638610589212096523915590215701244866830295506005"
+  , "16646216251577650555646508049064625507758601195307236539843683725095763921505"
+  , "11914716034377431890952169039751213443286692885071871704776127977841051829452"
+  , "20798492850731331785912281726856492405884190236464781409482377236764537088662"
   ]
 
-# TODO: add domain separation between rate=1 and rate=2, so that the empty input
-# gives different results. But this has to be done in all the other Poseidon2 libraries
-# too (circom, Haskell, C...)
-
 const expectedSpongeResultsRate2 : array[8, string] = 
-  [ "12363515589665961836680709257448433057869762330741639517836048636244832188495"
-  , "00899009032366875286186953183805404053380636995610127460025486428583509745414"
-  , "16500906802543951227422597869354004883060519121579073949799015758201044544012"
-  , "05275430613748165078459451567241807462288293965310307668712900802458919462965"
-  , "13763559248248167400098483085605230840597893317332127197498651878933380690961"
-  , "14871143128308815290845020646262475973102494373985615216162863857354721038367"
-  , "02746725081632011689597680224823496636241961292066939394613880404914874634920"
-  , "02290144245981244996669076598332792758523446545263085369617640761875376727694"
+  [ "15335097698975718583905618186682475632756177170667436996250626760551196078076"
+  , "05101758095924000127790537496504070769319625501671400349336709520206095219618"
+  , "07306734450287348725566606192910189982345130476287345231433021147457815478255"
+  , "18511919414269811073023003336929505285555117419480831606637506641708579940507"
+  , "17917165106036607360653786499368288558581739128065811663709392730081030901634"
+  , "04630821736691665506072583795473163860465039714428126246168623896083265248907"
+  , "02020506076765964149531002674962673761843846094901604358961533722934321735239"
+  , "11732533243633999579592740965735640217427639382365959787508754341969556105663"
   ]
 
 #-------------------------------------------------------------------------------