From 2274a03e339f213361453590b54917bbfd0a0c31 Mon Sep 17 00:00:00 2001
From: Justin Traglia <95511699+jtraglia@users.noreply.github.com>
Date: Mon, 24 Jul 2023 18:13:34 +0200
Subject: [PATCH] crypto/kzg4844: do lazy init in all ckzg funcs (#27679)

* crypto/kzg4844: remove unnecessary init call & fix typo

* Fix kzg4844 tests/benchmarks

* Make init lazy & revert changes to tests
---
 crypto/kzg4844/kzg4844.go          |  2 +-
 crypto/kzg4844/kzg4844_ckzg_cgo.go |  8 ++++++++
 crypto/kzg4844/kzg4844_test.go     | 12 ++++++++++--
 3 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/crypto/kzg4844/kzg4844.go b/crypto/kzg4844/kzg4844.go
index 3bd814159..5969d1c2c 100644
--- a/crypto/kzg4844/kzg4844.go
+++ b/crypto/kzg4844/kzg4844.go
@@ -54,7 +54,7 @@ func UseCKZG(use bool) error {
 	useCKZG.Store(use)
 
 	// Initializing the library can take 2-4 seconds - and can potentially crash
-	// on CKZG and non-ADX CPUs - so might as well so it now and don't wait until
+	// on CKZG and non-ADX CPUs - so might as well do it now and don't wait until
 	// a crypto operation is actually needed live.
 	if use {
 		ckzgIniter.Do(ckzgInit)
diff --git a/crypto/kzg4844/kzg4844_ckzg_cgo.go b/crypto/kzg4844/kzg4844_ckzg_cgo.go
index d62ca3fad..540028569 100644
--- a/crypto/kzg4844/kzg4844_ckzg_cgo.go
+++ b/crypto/kzg4844/kzg4844_ckzg_cgo.go
@@ -74,6 +74,8 @@ func ckzgBlobToCommitment(blob Blob) (Commitment, error) {
 // ckzgComputeProof computes the KZG proof at the given point for the polynomial
 // represented by the blob.
 func ckzgComputeProof(blob Blob, point Point) (Proof, Claim, error) {
+	ckzgIniter.Do(ckzgInit)
+
 	proof, claim, err := ckzg4844.ComputeKZGProof((ckzg4844.Blob)(blob), (ckzg4844.Bytes32)(point))
 	if err != nil {
 		return Proof{}, Claim{}, err
@@ -84,6 +86,8 @@ func ckzgComputeProof(blob Blob, point Point) (Proof, Claim, error) {
 // ckzgVerifyProof verifies the KZG proof that the polynomial represented by the blob
 // evaluated at the given point is the claimed value.
 func ckzgVerifyProof(commitment Commitment, point Point, claim Claim, proof Proof) error {
+	ckzgIniter.Do(ckzgInit)
+
 	valid, err := ckzg4844.VerifyKZGProof((ckzg4844.Bytes48)(commitment), (ckzg4844.Bytes32)(point), (ckzg4844.Bytes32)(claim), (ckzg4844.Bytes48)(proof))
 	if err != nil {
 		return err
@@ -99,6 +103,8 @@ func ckzgVerifyProof(commitment Commitment, point Point, claim Claim, proof Proo
 //
 // This method does not verify that the commitment is correct with respect to blob.
 func ckzgComputeBlobProof(blob Blob, commitment Commitment) (Proof, error) {
+	ckzgIniter.Do(ckzgInit)
+
 	proof, err := ckzg4844.ComputeBlobKZGProof((ckzg4844.Blob)(blob), (ckzg4844.Bytes48)(commitment))
 	if err != nil {
 		return Proof{}, err
@@ -108,6 +114,8 @@ func ckzgComputeBlobProof(blob Blob, commitment Commitment) (Proof, error) {
 
 // ckzgVerifyBlobProof verifies that the blob data corresponds to the provided commitment.
 func ckzgVerifyBlobProof(blob Blob, commitment Commitment, proof Proof) error {
+	ckzgIniter.Do(ckzgInit)
+
 	valid, err := ckzg4844.VerifyBlobKZGProof((ckzg4844.Blob)(blob), (ckzg4844.Bytes48)(commitment), (ckzg4844.Bytes48)(proof))
 	if err != nil {
 		return err
diff --git a/crypto/kzg4844/kzg4844_test.go b/crypto/kzg4844/kzg4844_test.go
index 0d35667a9..fae8a7a76 100644
--- a/crypto/kzg4844/kzg4844_test.go
+++ b/crypto/kzg4844/kzg4844_test.go
@@ -47,7 +47,6 @@ func randBlob() Blob {
 
 func TestCKZGWithPoint(t *testing.T)  { testKZGWithPoint(t, true) }
 func TestGoKZGWithPoint(t *testing.T) { testKZGWithPoint(t, false) }
-
 func testKZGWithPoint(t *testing.T, ckzg bool) {
 	if ckzg && !ckzgAvailable {
 		t.Skip("CKZG unavailable in this test build")
@@ -73,7 +72,6 @@ func testKZGWithPoint(t *testing.T, ckzg bool) {
 
 func TestCKZGWithBlob(t *testing.T)  { testKZGWithBlob(t, true) }
 func TestGoKZGWithBlob(t *testing.T) { testKZGWithBlob(t, false) }
-
 func testKZGWithBlob(t *testing.T, ckzg bool) {
 	if ckzg && !ckzgAvailable {
 		t.Skip("CKZG unavailable in this test build")
@@ -106,6 +104,8 @@ func benchmarkBlobToCommitment(b *testing.B, ckzg bool) {
 	useCKZG.Store(ckzg)
 
 	blob := randBlob()
+
+	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
 		BlobToCommitment(blob)
 	}
@@ -124,6 +124,8 @@ func benchmarkComputeProof(b *testing.B, ckzg bool) {
 		blob  = randBlob()
 		point = randFieldElement()
 	)
+
+	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
 		ComputeProof(blob, point)
 	}
@@ -144,6 +146,8 @@ func benchmarkVerifyProof(b *testing.B, ckzg bool) {
 		commitment, _   = BlobToCommitment(blob)
 		proof, claim, _ = ComputeProof(blob, point)
 	)
+
+	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
 		VerifyProof(commitment, point, claim, proof)
 	}
@@ -162,6 +166,8 @@ func benchmarkComputeBlobProof(b *testing.B, ckzg bool) {
 		blob          = randBlob()
 		commitment, _ = BlobToCommitment(blob)
 	)
+
+	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
 		ComputeBlobProof(blob, commitment)
 	}
@@ -181,6 +187,8 @@ func benchmarkVerifyBlobProof(b *testing.B, ckzg bool) {
 		commitment, _ = BlobToCommitment(blob)
 		proof, _      = ComputeBlobProof(blob, commitment)
 	)
+
+	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
 		VerifyBlobProof(blob, commitment, proof)
 	}