From 06bd13b2bd891afca3aaf7038c03d35affed8604 Mon Sep 17 00:00:00 2001
From: Balazs Komuves <bkomuves@gmail.com>
Date: Thu, 3 Oct 2024 00:32:49 +0200
Subject: [PATCH] a small optimization in the poseidon2 hash

---
 cbits/goldilocks.c             | 20 +++++++++++++++-----
 goldilocks_hash/goldilocks.nim |  1 -
 2 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/cbits/goldilocks.c b/cbits/goldilocks.c
index 822b2d9..49619df 100644
--- a/cbits/goldilocks.c
+++ b/cbits/goldilocks.c
@@ -17,6 +17,11 @@ uint64_t goldilocks_add(uint64_t x, uint64_t y) {
   return ( (z >= GOLDILOCKS_PRIME) || (z<x) ) ? (z - GOLDILOCKS_PRIME) : z;
 }
 
+uint64_t goldilocks_add_to_uint64(uint64_t x, uint64_t y) {
+  uint64_t z = x + y;
+  return (z<x) ? (z - GOLDILOCKS_PRIME) : z;
+}
+
 uint64_t goldilocks_sub(uint64_t x, uint64_t y) {
   uint64_t z = x - y;
   return (z > x) ? (z + GOLDILOCKS_PRIME) : z;
@@ -80,6 +85,11 @@ uint64_t goldilocks_mul(uint64_t x, uint64_t y) {
   return goldilocks_rdc(z); 
 }
 
+uint64_t goldilocks_mul_to_uint64(uint64_t x, uint64_t y) {
+  __uint128_t z = (__uint128_t)x * (__uint128_t)y;
+  return goldilocks_rdc_to_uint64(z); 
+}
+
 uint64_t goldilocks_mul_add128(uint64_t x, uint64_t y, __uint128_t z) {
   __uint128_t w = (__uint128_t)x * (__uint128_t)y + z;
   return goldilocks_rdc(w); 
@@ -147,11 +157,11 @@ from <https://github.com/HorizenLabs/poseidon2/blob/main/plain_implementations/s
 */
 
 uint64_t goldilocks_poseidon2_sbox(uint64_t x0, uint64_t rc) {
-  uint64_t x  = goldilocks_add( x0 , rc );
-  uint64_t x2 = goldilocks_sqr( x       );
-  uint64_t x4 = goldilocks_sqr( x2      );
-  uint64_t x6 = goldilocks_mul( x4 , x2 );
-  uint64_t x7 = goldilocks_mul( x6 , x  );
+  uint64_t x  = goldilocks_add_to_uint64( x0 , rc );
+  uint64_t x2 = goldilocks_mul_to_uint64( x  , x  );
+  uint64_t x4 = goldilocks_mul_to_uint64( x2 , x2 );
+  uint64_t x6 = goldilocks_mul_to_uint64( x4 , x2 );
+  uint64_t x7 = goldilocks_mul          ( x6 , x  );
   return x7;
 }
 
diff --git a/goldilocks_hash/goldilocks.nim b/goldilocks_hash/goldilocks.nim
index 77a50e3..37c0111 100644
--- a/goldilocks_hash/goldilocks.nim
+++ b/goldilocks_hash/goldilocks.nim
@@ -10,4 +10,3 @@ func `*`* (x, y: F): F {. header: "../cbits/goldilocks.h", importc: "goldilocks_
 proc `+=`* (x: var F, y: F) = x = x + y
 proc `-=`* (x: var F, y: F) = x = x - y
 proc `*=`* (x: var F, y: F) = x = x * y
-