From 41466908996fe449127ae8ede728624402a095a7 Mon Sep 17 00:00:00 2001
From: Jaremy Creechley <creechley@gmail.com>
Date: Wed, 24 Apr 2024 13:14:41 +0300
Subject: [PATCH] Sets up environment for running benchmarks

---
 benchmarks/.gitignore           |   2 +
 benchmarks/ark_prover_bench.nim |  39 +++++--
 benchmarks/config.nims          |  12 +-
 benchmarks/create_circuits.nim  | 190 ++++++++++++++++++++++++++++----
 4 files changed, 210 insertions(+), 33 deletions(-)
 create mode 100644 benchmarks/.gitignore

diff --git a/benchmarks/.gitignore b/benchmarks/.gitignore
new file mode 100644
index 00000000..6f697152
--- /dev/null
+++ b/benchmarks/.gitignore
@@ -0,0 +1,2 @@
+ceremony
+circuit_bench_*
diff --git a/benchmarks/ark_prover_bench.nim b/benchmarks/ark_prover_bench.nim
index ac69f84e..d19601fe 100644
--- a/benchmarks/ark_prover_bench.nim
+++ b/benchmarks/ark_prover_bench.nim
@@ -23,20 +23,41 @@ import pkg/constantine/math/io/io_fields
 
 import codex/slots/backends/helpers
 
-proc setup() =
+import create_circuits
+
+proc setup(
+  circuitDir: string, name: string,
+) =
   let
     inputData = readFile("tests/circuits/fixtures/input.json")
     inputJson: JsonNode = !JsonNode.parse(inputData)
-    proofInput: ProofInputs[Poseidon2Hash] =
-      Poseidon2Hash.jsonToProofInput(inputJson)
+    proofInput: ProofInputs[Poseidon2Hash] = Poseidon2Hash.jsonToProofInput(inputJson)
 
-  let
-    datasetProof = Poseidon2Proof.init(
-                    proofInput.slotIndex,
-                    proofInput.nSlotsPerDataSet,
-                    proofInput.slotProof[0..<4]).tryGet
+  let datasetProof = Poseidon2Proof.init(
+    proofInput.slotIndex, proofInput.nSlotsPerDataSet, proofInput.slotProof[0 ..< 4]
+  ).tryGet
 
   let ver = datasetProof.verify(proofInput.slotRoot, proofInput.datasetRoot).tryGet
   echo "ver: ", ver
 
-setup()
\ No newline at end of file
+when isMainModule:
+  echo "Running benchmark"
+  # setup()
+  checkEnv()
+
+  let args = CircArgs(
+    depth: 32, # maximum depth of the slot tree 
+    maxslots: 256, # maximum number of slots  
+    cellsize: 2048, # cell size in bytes 
+    blocksize: 65536, # block size in bytes 
+    nsamples: 5, # number of samples to prove
+    entropy: 1234567, # external randomness
+    seed: 12345, # seed for creating fake data
+    nslots: 11, # number of slots in the dataset
+    index: 3, # which slot we prove (0..NSLOTS-1)
+    ncells: 512, # number of cells in this slot
+  )
+
+  createCircuit(args)
+
+  ## TODO: copy over testcircomcompat proving
diff --git a/benchmarks/config.nims b/benchmarks/config.nims
index 040aadba..805335f4 100644
--- a/benchmarks/config.nims
+++ b/benchmarks/config.nims
@@ -1,7 +1,11 @@
---path:".."
---path:"../tests"
---threads:on
---tlsEmulation:off
+--path:
+  ".."
+--path:
+  "../tests"
+--threads:
+  on
+--tlsEmulation:
+  off
 
 # when not defined(chronicles_log_level):
 #   --define:"chronicles_log_level:NONE" # compile all log statements
diff --git a/benchmarks/create_circuits.nim b/benchmarks/create_circuits.nim
index bd1d8c09..00c29d61 100644
--- a/benchmarks/create_circuits.nim
+++ b/benchmarks/create_circuits.nim
@@ -1,18 +1,17 @@
+import std/hashes
+import std/json
 import std/strutils
+import std/strformat
 import std/os
+import std/osproc
 
-proc createCircuits() =
-  let cmds = """
-    ${NIMCLI_DIR}/cli $CLI_ARGS -v --circom=${CIRCUIT_MAIN}.circom --output=input.json
-    circom --r1cs --wasm --O2 -l${CIRCUIT_DIR} ${CIRCUIT_MAIN}.circom
-    NODE_OPTIONS="--max-old-space-size=8192" snarkjs groth16 setup ${CIRCUIT_MAIN}.r1cs $PTAU_PATH ${CIRCUIT_MAIN}_0000.zkey
-    echo "some_entropy_75289v3b7rcawcsyiur" | NODE_OPTIONS="--max-old-space-size=8192" snarkjs zkey contribute ${CIRCUIT_MAIN}_0000.zkey ${CIRCUIT_MAIN}_0001.zkey --name="1st Contributor Name"
-    """.splitLines()
-
-    # rm ${CIRCUIT_MAIN}_0000.zkey
-    # mv ${CIRCUIT_MAIN}_0001.zkey ${CIRCUIT_MAIN}.zkey
-
-let nimCircuitRefFl = "vendor"/"codex-storage-proofs-circuits"/"reference"/"nim"/"proof_input"/"cli"
+template withDir(dir: string, blk: untyped) =
+  let prev = getCurrentDir()
+  try:
+    setCurrentDir(dir)
+    `blk`
+  finally:
+    setCurrentDir(prev)
 
 proc setProjDir(prev = getCurrentDir()): string =
   if not "codex.nimble".fileExists():
@@ -25,14 +24,165 @@ proc setProjDir(prev = getCurrentDir()): string =
   else:
     getCurrentDir()
 
-let codexProjDir = setProjDir()
-echo "\n\nFound project dir: ", codexProjDir
+var
+  nimCircuitCli =
+    "vendor" / "codex-storage-proofs-circuits" / "reference" / "nim" / "proof_input" /
+    "cli"
+  circuitDirIncludes = "vendor" / "codex-storage-proofs-circuits" / "circuit"
+  ptauDefPath = "benchmarks" / "ceremony" / "powersOfTau28_hez_final_21.ptau"
+  ptauDefUrl =
+    "https://storage.googleapis.com/zkevm/ptau/powersOfTau28_hez_final_21.ptau"
+  codexProjDir = ""
 
-if not nimCircuitRefFl.fileExists:
-  echo "Nim Circuit reference cli not found: ", nimCircuitRefFl
-  echo "Building Circuit reference cli..."
-  setCurrentDir(nimCircuitRefFl.parentDir)
-  let res = execShellCmd("nimble build -d:release --styleCheck:off cli")
+proc checkEnv*() =
+  codexProjDir = setProjDir()
+  echo "\n\nFound project dir: ", codexProjDir
 
+  let snarkjs = findExe("snarkjs")
+  if snarkjs == "":
+    echo dedent"""
+    ERROR: must install snarkjs first
 
-echo "huzzah"
+      npm install -g snarkjs@latest
+    """
+
+  let circom = findExe("circom")
+  if circom == "":
+    echo dedent"""
+    ERROR: must install circom first
+
+      git clone https://github.com/iden3/circom.git
+      cargo install --path circom
+    """
+
+  if snarkjs == "" or circom == "":
+    quit 2
+
+  echo "Found SnarkJS: ", snarkjs
+  echo "Found Circom: ", circom
+
+  if not nimCircuitCli.fileExists:
+    echo "Nim Circuit reference cli not found: ", nimCircuitCli
+    echo "Building Circuit reference cli...\n"
+    withDir(nimCircuitCli.parentDir):
+      discard execShellCmd("nimble build -d:release --styleCheck:off cli")
+    echo "CWD: ", getCurrentDir()
+    assert nimCircuitCli.fileExists()
+
+  nimCircuitCli = nimCircuitCli.absolutePath()
+  echo "Found NimCircuitCli: ", nimCircuitCli
+
+  circuitDirIncludes = circuitDirIncludes.absolutePath
+  echo "Found Circuit Path: ", circuitDirIncludes
+
+  ptauDefPath = ptauDefPath.absolutePath
+  echo "Found PTAU file: ", ptauDefPath
+
+type CircArgs* = object
+  depth*: int
+  maxslots*: int
+  cellsize*: int
+  blocksize*: int
+  nsamples*: int
+  entropy*: int
+  seed*: int
+  nslots*: int
+  ncells*: int
+  index*: int
+
+proc downloadPtau*(ptauPath, ptauUrl: string) =
+  if not ptauPath.fileExists:
+    echo "Ceremony file not found, downloading..."
+    echo "PTAU url: ", ptauUrl
+    createDir(ptauPath.parentDir)
+    withDir(ptauPath.parentDir):
+      discard execShellCmd(fmt"curl -LOC - {ptauDefUrl}")
+
+proc getCircuitBenchPath*(args: CircArgs): string =
+  absolutePath("benchmarks/circuit_bench_" & $cast[uint](hash(args)))
+
+proc createCircuit*(
+    args: CircArgs,
+    name = "proof_main",
+    circBenchDir = getCircuitBenchPath(args),
+    circuitDirIncludes = circuitDirIncludes,
+    ptauPath = ptauDefPath,
+    ptauUrl = ptauDefUrl,
+    someEntropy = "some_entropy_75289v3b7rcawcsyiur",
+) =
+  ## Generates all the files needed for to run a proof circuit. Downloads the PTAU file if needed.
+  ## 
+  let circdir = circBenchDir
+
+  downloadPtau(ptauPath, ptauUrl)
+
+  echo "Creating circuit dir: ", circdir
+  createDir(circdir)
+  withDir circdir:
+    writeFile("circuit_params.json", pretty(%*args))
+    let
+      zkey = circdir / fmt"{name}.zkey"
+      wasm = circdir / fmt"{name}.wasm"
+      r1cs = circdir / fmt"{name}.r1cs"
+
+    var cliCmd = nimCircuitCli
+    for f, v in fieldPairs(args):
+      cliCmd &= " --" & f & "=" & $v
+
+    if not "input.json".fileExists:
+      echo "Generating Circom Files..."
+      # ${NIMCLI_DIR}/cli $CLI_ARGS -v --circom=${CIRCUIT_MAIN}.circom --output=input.json
+      cliCmd &= fmt" -v --circom={name}.circom --output=input.json"
+      echo "CWD: ", getCurrentDir()
+      echo "CLI_CMD: ", cliCmd
+
+      let cliRes = execShellCmd(cliCmd)
+      echo "RES: ", cliRes
+      assert cliRes == 0
+
+    if not wasm.fileExists or not r1cs.fileExists:
+      let cmd = fmt"circom --r1cs --wasm --O2 -l{circuitDirIncludes} {name}.circom"
+      echo "CMD: ", cmd
+      let cmdRes = execShellCmd(cmd)
+      echo "RES: ", cmdRes
+      assert cmdRes == 0
+      moveFile(fmt"{name}_js" / fmt"{name}.wasm", fmt"{name}.wasm")
+    echo "Found wasm: ", wasm
+    echo "Found r1cs: ", r1cs
+
+    if not zkey.fileExists:
+      echo "Zkey not found, generating..."
+      putEnv("NODE_OPTIONS", "--max-old-space-size=8192")
+      discard execShellCmd("echo $NODE_OPTIONS")
+      if not fmt"{name}_0000.zkey".fileExists:
+        let cmd = fmt"snarkjs groth16 setup {r1cs} {ptauPath} {name}_0000.zkey"
+        echo "CMD: ", cmd
+        let cmdRes = execShellCmd(cmd)
+        assert cmdRes == 0
+        echo fmt"Generated {name}_0000.zkey"
+
+      let cmd =
+        fmt"snarkjs zkey contribute {name}_0000.zkey {name}_0001.zkey --name='1st Contributor Name'"
+      echo "CMD: ", cmd
+      let cmdRes = execCmdEx(cmd, options = {}, input = someEntropy & "\n")
+      assert cmdRes.exitCode == 0
+
+      moveFile(fmt"{name}_0001.zkey", fmt"{name}.zkey")
+      removeFile(fmt"{name}_0000.zkey")
+
+when isMainModule:
+  checkEnv()
+
+  let args = CircArgs(
+    depth: 32, # maximum depth of the slot tree 
+    maxslots: 256, # maximum number of slots  
+    cellsize: 2048, # cell size in bytes 
+    blocksize: 65536, # block size in bytes 
+    nsamples: 5, # number of samples to prove
+    entropy: 1234567, # external randomness
+    seed: 12345, # seed for creating fake data
+    nslots: 11, # number of slots in the dataset
+    index: 3, # which slot we prove (0..NSLOTS-1)
+    ncells: 512, # number of cells in this slot
+  )
+  createCircuit(args)