From 6dbf451f8a94e07793c96949723033ffe9d21301 Mon Sep 17 00:00:00 2001
From: Alex E <alex@creatoros.co>
Date: Wed, 15 Jan 2025 09:09:39 +0500
Subject: [PATCH] [FEAT] better init process for multithreaded environment

---
 gnark/libraries/core_test.go           | 48 ++++++++++++++++++++++----
 gnark/libraries/prover/impl/library.go |  6 ++++
 2 files changed, 48 insertions(+), 6 deletions(-)

diff --git a/gnark/libraries/core_test.go b/gnark/libraries/core_test.go
index d75a43a..924db54 100644
--- a/gnark/libraries/core_test.go
+++ b/gnark/libraries/core_test.go
@@ -15,6 +15,7 @@ import (
 	"math"
 	"math/big"
 	"os"
+	"sync"
 	"testing"
 
 	"github.com/consensys/gnark-crypto/ecc/bn254/twistededwards"
@@ -45,12 +46,47 @@ func init() {
 
 func TestInit(t *testing.T) {
 	assert := test.NewAssert(t)
-	assert.True(prover.InitAlgorithm(prover.CHACHA20, chachaKey, chachaR1CS))
-	assert.True(prover.InitAlgorithm(prover.AES_128, aes128Key, aes128r1cs))
-	assert.True(prover.InitAlgorithm(prover.AES_256, aes256Key, aes256r1cs))
-	assert.True(prover.InitAlgorithm(prover.CHACHA20_OPRF, chachaOprfKey, chachaOprfr1cs))
-	assert.True(prover.InitAlgorithm(prover.AES_128_OPRF, aes128OprfKey, aes128Oprfr1cs))
-	assert.True(prover.InitAlgorithm(prover.AES_256_OPRF, aes256OprfKey, aes256Oprfr1cs))
+
+	wg1 := &sync.WaitGroup{}
+	wg1.Add(1)
+
+	wg2 := &sync.WaitGroup{}
+	wg2.Add(24)
+
+	f := func(algorithmID uint8, provingKey []byte, r1csData []byte) {
+		go func() {
+			wg1.Wait()
+			assert.True(prover.InitAlgorithm(algorithmID, provingKey, r1csData))
+			wg2.Done()
+		}()
+	}
+
+	f(prover.CHACHA20, chachaKey, chachaR1CS)
+	f(prover.AES_128, aes128Key, aes128r1cs)
+	f(prover.AES_256, aes256Key, aes256r1cs)
+	f(prover.CHACHA20_OPRF, chachaOprfKey, chachaOprfr1cs)
+	f(prover.AES_128_OPRF, aes128OprfKey, aes128Oprfr1cs)
+	f(prover.AES_256_OPRF, aes256OprfKey, aes256Oprfr1cs)
+	f(prover.CHACHA20, chachaKey, chachaR1CS)
+	f(prover.AES_128, aes128Key, aes128r1cs)
+	f(prover.AES_256, aes256Key, aes256r1cs)
+	f(prover.CHACHA20_OPRF, chachaOprfKey, chachaOprfr1cs)
+	f(prover.AES_128_OPRF, aes128OprfKey, aes128Oprfr1cs)
+	f(prover.AES_256_OPRF, aes256OprfKey, aes256Oprfr1cs)
+	f(prover.CHACHA20, chachaKey, chachaR1CS)
+	f(prover.AES_128, aes128Key, aes128r1cs)
+	f(prover.AES_256, aes256Key, aes256r1cs)
+	f(prover.CHACHA20_OPRF, chachaOprfKey, chachaOprfr1cs)
+	f(prover.AES_128_OPRF, aes128OprfKey, aes128Oprfr1cs)
+	f(prover.AES_256_OPRF, aes256OprfKey, aes256Oprfr1cs)
+	f(prover.CHACHA20, chachaKey, chachaR1CS)
+	f(prover.AES_128, aes128Key, aes128r1cs)
+	f(prover.AES_256, aes256Key, aes256r1cs)
+	f(prover.CHACHA20_OPRF, chachaOprfKey, chachaOprfr1cs)
+	f(prover.AES_128_OPRF, aes128OprfKey, aes128Oprfr1cs)
+	f(prover.AES_256_OPRF, aes256OprfKey, aes256Oprfr1cs)
+	wg1.Done()
+	wg2.Wait()
 }
 
 func TestPanic(t *testing.T) {
diff --git a/gnark/libraries/prover/impl/library.go b/gnark/libraries/prover/impl/library.go
index 6e502aa..35a79fb 100644
--- a/gnark/libraries/prover/impl/library.go
+++ b/gnark/libraries/prover/impl/library.go
@@ -7,6 +7,7 @@ import (
 	"encoding/hex"
 	"encoding/json"
 	"fmt"
+	"sync"
 
 	"github.com/consensys/gnark-crypto/ecc"
 	"github.com/consensys/gnark/backend/groth16"
@@ -76,6 +77,7 @@ type ProverParams struct {
 	KeyHash     string
 	CircuitHash string
 	initDone    bool
+	initLock    sync.Mutex
 }
 
 func init() {
@@ -92,7 +94,10 @@ func InitAlgorithm(algorithmID uint8, provingKey []byte, r1csData []byte) (res b
 	}()
 	if alg, ok := algorithmNames[algorithmID]; ok {
 		proverParams := provers[alg]
+		proverParams.initLock.Lock()
+		defer proverParams.initLock.Unlock()
 		if proverParams.initDone {
+			fmt.Println("Algorithm already initialized", alg)
 			return true
 		}
 
@@ -129,6 +134,7 @@ func InitAlgorithm(algorithmID uint8, provingKey []byte, r1csData []byte) (res b
 
 		proverParams.SetParams(r1cs, pkey)
 		proverParams.initDone = true
+		fmt.Println("Initialized", alg)
 		return true
 	}
 	return false