feat: make registries for gkr thread safe (#920)

Consensys · Nov 16, 2023 · 088a222 · 088a222
1 parent 835adfe
commit 088a222
Show file tree

Hide file tree

Showing 11 changed files with 248 additions and 40 deletions.
diff --git a/constraint/bls12-377/gkr.go b/constraint/bls12-377/gkr.go
diff --git a/constraint/bls12-381/gkr.go b/constraint/bls12-381/gkr.go
diff --git a/constraint/bls24-315/gkr.go b/constraint/bls24-315/gkr.go
diff --git a/constraint/bls24-317/gkr.go b/constraint/bls24-317/gkr.go
diff --git a/constraint/bn254/gkr.go b/constraint/bn254/gkr.go
diff --git a/constraint/bw6-633/gkr.go b/constraint/bw6-633/gkr.go
diff --git a/constraint/bw6-761/gkr.go b/constraint/bw6-761/gkr.go
diff --git a/internal/generator/backend/template/representations/gkr.go.tmpl b/internal/generator/backend/template/representations/gkr.go.tmpl
@@ -10,6 +10,7 @@ import (
 	"github.com/consensys/gnark/std/utils/algo_utils"
 	"hash"
 	"math/big"
+	"sync"
 )
 
 type GkrSolvingData struct {
@@ -146,9 +147,12 @@ func GkrProveHint(hashName string, data *GkrSolvingData) hint.Hint {
 			return b[:]
 		})
 
-		hsh := HashBuilderRegistry[hashName]()
+		hsh, err := GetHashBuilder(hashName)
+		if err != nil {
+			return err
+		}
 
-		proof, err := gkr.Prove(data.circuit, data.assignments, fiatshamir.WithHash(hsh, insBytes...), gkr.WithPool(&data.memoryPool), gkr.WithWorkers(data.workers))
+		proof, err := gkr.Prove(data.circuit, data.assignments, fiatshamir.WithHash(hsh(), insBytes...), gkr.WithPool(&data.memoryPool), gkr.WithWorkers(data.workers))
 		if err != nil {
 			return err
 		}
@@ -175,4 +179,23 @@ func GkrProveHint(hashName string, data *GkrSolvingData) hint.Hint {
 }
 
 // TODO: Move to gnark-crypto
-var HashBuilderRegistry = make(map[string]func() hash.Hash)
+var (
+	hashBuilderRegistry = make(map[string]func() hash.Hash)
+	hasBuilderLock sync.RWMutex
+)
+
+func RegisterHashBuilder(name string, builder func() hash.Hash) {
+	hasBuilderLock.Lock()
+	defer hasBuilderLock.Unlock()
+	hashBuilderRegistry[name] = builder
+}
+
+func GetHashBuilder(name string) (func() hash.Hash, error) {
+	hasBuilderLock.RLock()
+	defer hasBuilderLock.RUnlock()
+	builder, ok := hashBuilderRegistry[name]
+	if !ok {
+		return nil, fmt.Errorf("hash function not found")
+	}
+	return builder, nil
+}