From 2b95044921cd0bed72694d069e2f9ed8fa44e72e Mon Sep 17 00:00:00 2001
From: ctrlc03 <93448202+ctrlc03@users.noreply.github.com>
Date: Fri, 15 Dec 2023 17:31:17 +0000
Subject: [PATCH] feat(proofs): make preferred witness type in circuit's
 genProof function

Remove isArm check in circuits and instead accept a parameter to decide which witness type to use.
Made appropriate changes in tests and cli
---
 circuits/ts/index.ts          |  2 +-
 circuits/ts/proofs.ts         | 21 ++++++++++++++++++---
 circuits/ts/types.ts          |  1 +
 cli/tests/e2e.subsidy.test.ts |  3 +--
 cli/tests/e2e.test.ts         |  3 +--
 cli/tests/keyChange.test.ts   |  3 +--
 cli/tests/utils.ts            | 14 ++++++++++++--
 cli/ts/commands/genProofs.ts  |  3 +++
 8 files changed, 38 insertions(+), 12 deletions(-)

diff --git a/circuits/ts/index.ts b/circuits/ts/index.ts
index add0ed5816..90ba1ac8db 100644
--- a/circuits/ts/index.ts
+++ b/circuits/ts/index.ts
@@ -1,2 +1,2 @@
 export { genProof, verifyProof, extractVk } from "./proofs";
-export { isArm, cleanThreads } from "./utils";
+export { cleanThreads } from "./utils";
diff --git a/circuits/ts/proofs.ts b/circuits/ts/proofs.ts
index 55acb8a929..39f4b306c1 100644
--- a/circuits/ts/proofs.ts
+++ b/circuits/ts/proofs.ts
@@ -16,6 +16,7 @@ import { cleanThreads, isArm } from "./utils";
  * snark and a WASM witness
  * @param inputs - the inputs to the circuit
  * @param zkeyPath - the path to the zkey
+ * @param useWasm - whether we want to use the wasm witness or not
  * @param rapidsnarkExePath - the path to the rapidnsark binary
  * @param witnessExePath - the path to the compiled witness binary
  * @param wasmPath - the path to the wasm witness
@@ -25,16 +26,30 @@ import { cleanThreads, isArm } from "./utils";
 export const genProof = async ({
   inputs,
   zkeyPath,
+  useWasm,
   rapidsnarkExePath,
   witnessExePath,
   wasmPath,
   silent = false,
 }: IGenProofOptions): Promise<FullProveResult> => {
-  // if we are running on an arm chip we can use snarkjs directly
-  if (isArm()) {
-    const { proof, publicSignals } = await groth16.fullProve(inputs, wasmPath!, zkeyPath);
+  // if we want to use a wasm witness we use snarkjs
+  if (useWasm) {
+    if (!wasmPath) {
+      throw new Error("wasmPath must be specified");
+    }
+
+    if (!fs.existsSync(wasmPath)) {
+      throw new Error(`wasmPath ${wasmPath} does not exist`);
+    }
+
+    const { proof, publicSignals } = await groth16.fullProve(inputs, wasmPath, zkeyPath);
     return { proof, publicSignals };
   }
+
+  if (isArm()) {
+    throw new Error("To use rapidnsnark you currently need to be running on an intel chip");
+  }
+
   // intel chip flow (use rapidnsark)
   // Create tmp directory
   const tmpPath = path.resolve(tmpdir(), `tmp-${Date.now()}`);
diff --git a/circuits/ts/types.ts b/circuits/ts/types.ts
index ae1c2f863e..620e58be77 100644
--- a/circuits/ts/types.ts
+++ b/circuits/ts/types.ts
@@ -6,6 +6,7 @@ import type { CircuitInputs } from "maci-core";
 export interface IGenProofOptions {
   inputs: CircuitInputs;
   zkeyPath: string;
+  useWasm?: boolean;
   rapidsnarkExePath?: string;
   witnessExePath?: string;
   wasmPath?: string;
diff --git a/cli/tests/e2e.subsidy.test.ts b/cli/tests/e2e.subsidy.test.ts
index 1041e0359e..6681af87fa 100644
--- a/cli/tests/e2e.subsidy.test.ts
+++ b/cli/tests/e2e.subsidy.test.ts
@@ -39,8 +39,7 @@ import {
   testTallyVotesWasmPath,
   testTallyVotesWitnessPath,
 } from "./constants";
-import { cleanSubsidy } from "./utils";
-import { isArm } from "maci-circuits";
+import { cleanSubsidy, isArm } from "./utils";
 import { genRandomSalt } from "maci-crypto";
 import { DeployedContracts, PollContracts } from "../ts/utils";
 
diff --git a/cli/tests/e2e.test.ts b/cli/tests/e2e.test.ts
index 46bb441a7b..e0357e6ab6 100644
--- a/cli/tests/e2e.test.ts
+++ b/cli/tests/e2e.test.ts
@@ -36,8 +36,7 @@ import {
   testTallyVotesWasmPath,
   testTallyVotesWitnessPath,
 } from "./constants";
-import { cleanVanilla } from "./utils";
-import { isArm } from "maci-circuits";
+import { cleanVanilla, isArm } from "./utils";
 import { DeployedContracts, PollContracts } from "../ts/utils";
 import { Keypair } from "maci-domainobjs";
 import { genRandomSalt } from "maci-crypto";
diff --git a/cli/tests/keyChange.test.ts b/cli/tests/keyChange.test.ts
index 5af2e45ab8..9f7cc2578d 100644
--- a/cli/tests/keyChange.test.ts
+++ b/cli/tests/keyChange.test.ts
@@ -1,4 +1,3 @@
-import { isArm } from "maci-circuits";
 import {
   deploy,
   deployPoll,
@@ -32,7 +31,7 @@ import {
   testTallyVotesWitnessPath,
 } from "./constants";
 import { Keypair } from "maci-domainobjs";
-import { cleanVanilla } from "./utils";
+import { cleanVanilla, isArm } from "./utils";
 import { readFileSync } from "fs";
 import { expect } from "chai";
 import { genRandomSalt } from "maci-crypto";
diff --git a/cli/tests/utils.ts b/cli/tests/utils.ts
index bcb0a40e3b..61e7e5d60f 100644
--- a/cli/tests/utils.ts
+++ b/cli/tests/utils.ts
@@ -1,5 +1,7 @@
 import { existsSync, readdirSync, rmSync } from "fs";
-import { join } from "path";
+import { arch } from "os";
+
+import path from "path";
 
 /**
  * Test utility to clean up the proofs directory
@@ -8,7 +10,7 @@ import { join } from "path";
 export const cleanVanilla = () => {
   const files = readdirSync("./proofs");
   for (const file of files) {
-    rmSync(join("./proofs", file));
+    rmSync(path.join("./proofs", file));
   }
   if (existsSync("./tally.json")) rmSync("./tally.json");
 };
@@ -21,3 +23,11 @@ export const cleanSubsidy = () => {
   cleanVanilla();
   if (existsSync("./subsidy.json")) rmSync("./subsidy.json");
 };
+
+/**
+ * Check if we are running on an arm chip
+ * @returns whether we are running on an arm chip
+ */
+export const isArm = (): boolean => {
+  return arch().includes("arm");
+};
diff --git a/cli/ts/commands/genProofs.ts b/cli/ts/commands/genProofs.ts
index ae30092135..138de5856b 100644
--- a/cli/ts/commands/genProofs.ts
+++ b/cli/ts/commands/genProofs.ts
@@ -228,6 +228,7 @@ export const genProofs = async (
       const r = await genProof({
         inputs: circuitInputs,
         zkeyPath: processZkey,
+        useWasm,
         rapidsnarkExePath: rapidsnark,
         witnessExePath: processWitgen,
         wasmPath: processWasm,
@@ -287,6 +288,7 @@ export const genProofs = async (
         const r = await genProof({
           inputs: subsidyCircuitInputs,
           zkeyPath: subsidyZkey,
+          useWasm,
           rapidsnarkExePath: rapidsnark,
           witnessExePath: subsidyWitgen,
           wasmPath: subsidyWasm,
@@ -350,6 +352,7 @@ export const genProofs = async (
       const r = await genProof({
         inputs: tallyCircuitInputs,
         zkeyPath: tallyZkey,
+        useWasm,
         rapidsnarkExePath: rapidsnark,
         witnessExePath: tallyWitgen,
         wasmPath: tallyWasm,