From bf99ed5b914f4098a88170814591ab777fde910e Mon Sep 17 00:00:00 2001
From: andri lim <jangko128@gmail.com>
Date: Tue, 21 Mar 2023 09:30:50 +0700
Subject: [PATCH] Nim: better handling of trusted setup (#231)

---
 bindings/nim/kzg.nim               |  3 +++
 bindings/nim/kzg_ex.nim            | 16 +++++++++++++---
 bindings/nim/tests/test_kzg_ex.nim |  7 ++++++-
 3 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/bindings/nim/kzg.nim b/bindings/nim/kzg.nim
index 9f782bc..856af27 100644
--- a/bindings/nim/kzg.nim
+++ b/bindings/nim/kzg.nim
@@ -208,6 +208,9 @@ proc verifyProofs*(ctx: KzgCtx,
 template loadTrustedSetupFile*(input: File | string): untyped =
   loadTrustedSetup(input)
 
+template freeTrustedSetup*(ctx: KzgCtx) =
+  free_trusted_setup(ctx.val)
+  
 template blobToKzgCommitment*(ctx: KzgCtx,
                    blob: KzgBlob): untyped =
   toCommitment(ctx, blob)
diff --git a/bindings/nim/kzg_ex.nim b/bindings/nim/kzg_ex.nim
index 05f4ecd..04c1ccc 100644
--- a/bindings/nim/kzg_ex.nim
+++ b/bindings/nim/kzg_ex.nim
@@ -22,12 +22,16 @@ else:
 # Private helpers
 ##############################################################
 
-var gCtx: KzgCtx
+var gCtx = KzgCtx(nil)
 
 const
-  GlobalCtxErr = "kzg global context not loaded"
+  TrustedSetupNotLoadedErr = "Trusted setup not loaded."
+  TrustedSetupAlreadyLoadedErr =
+    "Trusted setup is already loaded. Free it before loading a new one."
 
 template setupCtx(body: untyped): untyped =
+  if not gCtx.isNil:
+    return err(TrustedSetupAlreadyLoadedErr)
   let res = body
   if res.isErr:
     return err(res.error)
@@ -37,7 +41,7 @@ template setupCtx(body: untyped): untyped =
 template verifyCtx(body: untyped): untyped =
   {.gcsafe.}:
     if gCtx.isNil:
-      return err(GlobalCtxErr)
+      return err(TrustedSetupNotLoadedErr)
     body
 
 ##############################################################
@@ -65,6 +69,12 @@ proc loadTrustedSetupFromString*(_: type Kzg,
   setupCtx:
     kzg.loadTrustedSetupFromString(input)
 
+proc freeTrustedSetup*(_: type Kzg): Result[void, string] =
+  verifyCtx:
+    gCtx.freeTrustedSetup()
+    gCtx = nil
+    ok()
+
 proc toCommitment*(blob: KzgBlob):
                     Result[KzgCommitment, string] {.gcsafe.} =
   verifyCtx:
diff --git a/bindings/nim/tests/test_kzg_ex.nim b/bindings/nim/tests/test_kzg_ex.nim
index 2a39b45..83a56e0 100644
--- a/bindings/nim/tests/test_kzg_ex.nim
+++ b/bindings/nim/tests/test_kzg_ex.nim
@@ -75,7 +75,8 @@ suite "verify proof (extended version)":
   test "template aliases":
     # no need to check return value
     # only test if those templates can be compiled succesfully
-    discard Kzg.loadTrustedSetupFile(trustedSetupFile)
+    check Kzg.freeTrustedSetup().isOk
+    check Kzg.loadTrustedSetupFile(trustedSetupFile).isOk
     discard blobToKzgCommitment(blob)
     let kp = computeKzgProof(blob, inputPoint)
     discard computeBlobKzgProof(blob, commitment)
@@ -83,3 +84,7 @@ suite "verify proof (extended version)":
     discard verifyBlobKzgProof(blob, commitment, proof)
     let kb = createKateBlobs(1)
     discard verifyBlobKzgProofBatch(kb.blobs, kb.kates, [kp.get.proof])
+
+  test "load trusted setup more than once":
+    let res = Kzg.loadTrustedSetupFromString(trustedSetup)
+    check res.isErr