Circuit compile and proving key gen enhancements

Signed-off-by: Jim Zhang <[email protected]>
hyperledger-labs · Aug 13, 2024 · f2c17c2 · f2c17c2
1 parent a7b9e39
commit f2c17c2
Show file tree

Hide file tree

Showing 5 changed files with 46 additions and 54 deletions.
diff --git a/.github/workflows/e2e.yaml b/.github/workflows/e2e.yaml
@@ -52,6 +52,7 @@ jobs:
 
       - name: Build circuits
         env:
+          CIRCUITS_ROOT: ${{ runner.temp }}/zeto-artifacts
           PROVING_KEYS_ROOT: ${{ runner.temp }}/zeto-artifacts
           PTAU_DOWNLOAD_PATH: ${{ runner.temp }}/zeto-artifacts
         working-directory: zkp/circuits
@@ -62,15 +63,15 @@ jobs:
       - name: Run golang e2e tests
         env:
           PROVING_KEYS_ROOT: ${{ runner.temp }}/zeto-artifacts
-          CIRCUITS_ROOT: ${{ github.workspace }}/zkp/js/lib
+          CIRCUITS_ROOT: ${{ runner.temp }}/zeto-artifacts
         working-directory: go-sdk
         run: |
           make e2e
 
       - name: Run js e2e tests
         env:
           PROVING_KEYS_ROOT: ${{ runner.temp }}/zeto-artifacts
-          CIRCUITS_ROOT: ${{ github.workspace }}/zkp/js/lib
+          CIRCUITS_ROOT: ${{ runner.temp }}/zeto-artifacts
         working-directory: zkp/js
         run: |
           npm install
@@ -79,7 +80,7 @@ jobs:
       - name: Run Zeto Tokens hardhat tests
         env:
           PROVING_KEYS_ROOT: ${{ runner.temp }}/zeto-artifacts
-          CIRCUITS_ROOT: ${{ github.workspace }}/zkp/js/lib
+          CIRCUITS_ROOT: ${{ runner.temp }}/zeto-artifacts
         working-directory: solidity
         run: |
           npm install

diff --git a/zkp/circuits/gen-config.json b/zkp/circuits/gen-config.json
@@ -39,10 +39,6 @@
     "ptau": "powersOfTau28_hez_final_16",
     "skipSolidityGenaration": false
   },
-  "check_smt_proof": {
-    "ptau": "powersOfTau28_hez_final_16",
-    "skipSolidityGenaration": false
-  },
   "check_nullifiers": {
     "ptau": "powersOfTau28_hez_final_11",
     "skipSolidityGenaration": true

diff --git a/zkp/circuits/gen.js b/zkp/circuits/gen.js
@@ -3,20 +3,29 @@ const path = require('path');
 const { exec } = require('child_process');
 const { promisify } = require('util');
 const axios = require('axios');
+const yargs = require('yargs/yargs');
+const { hideBin } = require('yargs/helpers');
+const argv = yargs(hideBin(process.argv)).argv;
 
+const circuitsRoot = process.env.CIRCUITS_ROOT;
 const provingKeysRoot = process.env.PROVING_KEYS_ROOT;
 const ptauDownload = process.env.PTAU_DOWNLOAD_PATH;
-const specificCircuit = process.argv[2];
+const specificCircuits = argv.c;
+const compileOnly = argv.compileOnly;
 const parallelLimit = parseInt(process.env.GEN_CONCURRENCY, 10) || 10; // Default to compile 10 circuits in parallel
 
 // check env vars
+if (!circuitsRoot) {
+  console.error('Error: CIRCUITS_ROOT is not set.');
+  process.exit(1);
+}
 
-if (!provingKeysRoot) {
+if (!compileOnly && !provingKeysRoot) {
   console.error('Error: PROVING_KEYS_ROOT is not set.');
   process.exit(1);
 }
 
-if (!ptauDownload) {
+if (!compileOnly && !ptauDownload) {
   console.error('Error: PTAU_DOWNLOAD_PATH is not set.');
   process.exit(1);
 }
@@ -54,15 +63,12 @@ const processCircuit = async (circuit, ptau, skipSolidityGenaration) => {
     return;
   }
 
-  if (!fs.existsSync(ptauFile)) {
+  if (!compileOnly && !fs.existsSync(ptauFile)) {
     log(circuit, `PTAU file does not exist, downloading: ${ptauFile}`);
     try {
-      const response = await axios.get(
-        `https://storage.googleapis.com/zkevm/ptau/${ptau}.ptau`,
-        {
-          responseType: 'stream',
-        }
-      );
+      const response = await axios.get(`https://storage.googleapis.com/zkevm/ptau/${ptau}.ptau`, {
+        responseType: 'stream',
+      });
       response.data.pipe(fs.createWriteStream(ptauFile));
       await new Promise((resolve, reject) => {
         response.data.on('end', resolve);
@@ -75,70 +81,54 @@ const processCircuit = async (circuit, ptau, skipSolidityGenaration) => {
   }
 
   log(circuit, `Compiling circuit`);
-  await execAsync(`circom ${circomInput} --output ../js/lib --sym --wasm`);
+  await execAsync(`circom ${circomInput} --output ${circuitsRoot} --sym --wasm`);
+  if (compileOnly) {
+    return;
+  }
+
   await execAsync(`circom ${circomInput} --output ${provingKeysRoot} --r1cs`);
 
   log(circuit, `Generating test proving key with ${ptau}`);
-  await execAsync(
-    `snarkjs groth16 setup ${path.join(
-      provingKeysRoot,
-      `${circuit}.r1cs`
-    )} ${ptauFile} ${zkeyOutput}`
-  );
+  await execAsync(`snarkjs groth16 setup ${path.join(provingKeysRoot, `${circuit}.r1cs`)} ${ptauFile} ${zkeyOutput}`);
 
   log(circuit, `Generating verification key`);
-  await execAsync(
-    `snarkjs zkey export verificationkey ${zkeyOutput} ${path.join(
-      provingKeysRoot,
-      `${circuit}-vkey.json`
-    )}`
-  );
+  await execAsync(`snarkjs zkey export verificationkey ${zkeyOutput} ${path.join(provingKeysRoot, `${circuit}-vkey.json`)}`);
 
   if (skipSolidityGenaration) {
     log(circuit, `Skipping solidity verifier generation`);
     return;
   }
 
   log(circuit, `Generating solidity verifier`);
-  const solidityFile = path.join(
-    '..',
-    '..',
-    'solidity',
-    'contracts',
-    'lib',
-    `verifier_${circuit}.sol`
-  );
-  await execAsync(
-    `snarkjs zkey export solidityverifier ${zkeyOutput} ${solidityFile}`
-  );
+  const solidityFile = path.join('..', '..', 'solidity', 'contracts', 'lib', `verifier_${circuit}.sol`);
+  await execAsync(`snarkjs zkey export solidityverifier ${zkeyOutput} ${solidityFile}`);
 
   log(circuit, `Modifying the contract name in the Solidity file`);
   const camelCaseCircuitName = toCamelCase(circuit);
   const solidityFileTmp = `${solidityFile}.tmp`;
 
   const fileContent = fs.readFileSync(solidityFile, 'utf8');
-  const updatedContent = fileContent.replace(
-    ' Groth16Verifier ',
-    ` Groth16Verifier_${camelCaseCircuitName} `
-  );
+  const updatedContent = fileContent.replace(' Groth16Verifier ', ` Groth16Verifier_${camelCaseCircuitName} `);
   fs.writeFileSync(solidityFileTmp, updatedContent);
   fs.renameSync(solidityFileTmp, solidityFile);
 };
 
 const run = async () => {
-  if (specificCircuit) {
-    // if a specific circuit is provided, check it's in the map
-    if (!circuits[specificCircuit]) {
-      console.error(`Error: Unknown circuit: ${specificCircuit}`);
-      process.exit(1);
+  if (specificCircuits) {
+    // if specific circuits are provided, check it's in the map
+    for (const circuit of specificCircuits) {
+      if (!circuits[circuit]) {
+        console.error(`Error: Unknown circuit: ${circuit}`);
+        process.exit(1);
+      }
     }
   }
 
   const circuitsArray = Object.entries(circuits);
   const activePromises = new Set();
 
   for (const [circuit, { ptau, skipSolidityGenaration }] of circuitsArray) {
-    if (specificCircuit && circuit !== specificCircuit) {
+    if (specificCircuits && !specificCircuits.includes(circuit)) {
       continue;
     }
 

diff --git a/zkp/circuits/package.json b/zkp/circuits/package.json
@@ -4,7 +4,8 @@
   "description": "Circom library for ZK circuits for use with Zeto based privacy-preserving tokens",
   "license": "Apache-2.0",
   "dependencies": {
-    "circomlib": "^2.0.5"
+    "circomlib": "^2.0.5",
+    "yargs": "^17.7.2"
   },
   "scripts": {
     "gen": "node gen.js"

diff --git a/zkp/js/index.js b/zkp/js/index.js
@@ -23,8 +23,12 @@ function loadCircuit(type) {
   if (!type) {
     throw new Error('The circuit name must be provided');
   }
-  const WitnessCalculator = require(`./lib/${type}_js/witness_calculator.js`);
-  const buffer = readFileSync(path.join(__dirname, `./lib/${type}_js/${type}.wasm`));
+  const circuitsRoot = process.env.CIRCUITS_ROOT;
+  if (!circuitsRoot) {
+    throw new Error('CIRCUITS_ROOT is not set');
+  }
+  const WitnessCalculator = require(path.join(circuitsRoot, `${type}_js/witness_calculator.js`));
+  const buffer = readFileSync(path.join(circuitsRoot, `${type}_js/${type}.wasm`));
   return WitnessCalculator(buffer);
 }