From 7ba808630f2e1fa3231503b82429b6f88290b07b Mon Sep 17 00:00:00 2001
From: Balazs Komuves <bkomuves@gmail.com>
Date: Tue, 4 Mar 2025 15:05:35 +0100
Subject: [PATCH] add benchmarking iterated Poseidon2 permutation

---
 hash/cpu/bench/Poseidon2/Zikkurat/poseidon2.c | 44 ++++++++++++++++---
 1 file changed, 39 insertions(+), 5 deletions(-)

diff --git a/hash/cpu/bench/Poseidon2/Zikkurat/poseidon2.c b/hash/cpu/bench/Poseidon2/Zikkurat/poseidon2.c
index fbcfd5d..cd82478 100644
--- a/hash/cpu/bench/Poseidon2/Zikkurat/poseidon2.c
+++ b/hash/cpu/bench/Poseidon2/Zikkurat/poseidon2.c
@@ -295,6 +295,32 @@ void permutation_example() {
   print_std( "  w = ", tgt + 2*NLIMBS );
 }
 
+// -----------------------------------------------------------------------------
+// * single-threaded iterated permutation
+
+void permutation_iterated(int n) {
+  uint64_t state[3*NLIMBS];
+  uint64_t tgt  [3*NLIMBS];
+
+  bn128_r_mont_from_std( example_input_std            , state            );
+  bn128_r_mont_from_std( example_input_std +   NLIMBS , state +   NLIMBS );
+  bn128_r_mont_from_std( example_input_std + 2*NLIMBS , state + 2*NLIMBS );
+
+  for(int i=0; i<n; i++) {
+    poseidon_perm_mont( state, state );
+  }
+
+  bn128_r_mont_to_std( state            , tgt            );
+  bn128_r_mont_to_std( state +   NLIMBS , tgt +   NLIMBS );
+  bn128_r_mont_to_std( state + 2*NLIMBS , tgt + 2*NLIMBS );
+
+  printf("\npermutation of (0,1,2) iterated %d times in F^3:\n",n);
+  print_std( "  x = ", tgt            );
+  print_std( "  y = ", tgt +   NLIMBS );
+  print_std( "  z = ", tgt + 2*NLIMBS );
+
+}
+
 // -----------------------------------------------------------------------------
 // * single-threaded Merkle tree
 
@@ -394,7 +420,6 @@ int main( int argc, char *argv[] ) {
   
   int nThreads = 1;
   int mtDepth  = 16;
-
   switch(argc) {
 
     case 2:
@@ -406,20 +431,29 @@ int main( int argc, char *argv[] ) {
       nThreads = atoi(argv[2]);
       break;
 
+
     default:
       printf("usage:\n");
       printf("$ poseidon2 <merkle_depth>:\n");
       printf("$ poseidon2 <merkle_depth> <nthreads>:\n");
+      printf("$ poseidon2 <n_iters>:\n");
       exit(-1);
       break;
   }
 
-  printf("calculating Poseidon2 Merkle root with depth=%d on %d threads...\n",mtDepth,nThreads);
-  if (nThreads == 1) {
-    merkle_root_example(mtDepth);
+  if (mtDepth < 32) {
+    printf("calculating Poseidon2 Merkle root with depth=%d on %d threads...\n",mtDepth,nThreads);
+    if (nThreads == 1) {
+      merkle_root_example(mtDepth);
+    }
+    else {
+      merkle_root_multithread(mtDepth, nThreads);    
+    }
   }
   else {
-    merkle_root_multithread(mtDepth, nThreads);    
+    int nIters = mtDepth;
+    printf("calculating Poseidon2 permutation iterated %d times on 1 threads...\n",nIters);
+    permutation_iterated(nIters); 
   }
 
   printf("\n");