Merge pull request #1924 from privacy-scaling-explorations/refactor/s…

…ubmit-on-chain refactor(contracts): improve submit onchain task
privacy-scaling-explorations · Nov 20, 2024 · 20bff7a · 20bff7a
2 parents 6bdb342 + cf71ccc
commit 20bff7a
Show file tree

Hide file tree

Showing 3 changed files with 61 additions and 37 deletions.
diff --git a/packages/contracts/package.json b/packages/contracts/package.json
@@ -172,7 +172,7 @@
     "@openzeppelin/merkle-tree": "^1.0.7",
     "circomlibjs": "^0.1.7",
     "ethers": "^6.13.4",
-    "hardhat": "^2.22.8",
+    "hardhat": "^2.22.15",
     "lowdb": "^1.0.0",
     "maci-circuits": "^2.4.0",
     "maci-core": "^2.4.0",

diff --git a/packages/contracts/tasks/runner/submitOnChain.ts b/packages/contracts/tasks/runner/submitOnChain.ts
@@ -12,6 +12,32 @@ import { Deployment } from "../helpers/Deployment";
 import { Prover } from "../helpers/Prover";
 import { EContracts, TallyData, type ISubmitOnChainParams } from "../helpers/types";
 
+/**
+ * Interface that represents read proofs arguments
+ */
+interface IReadProofsArgs {
+  files: string[];
+  folder: string;
+  type: "tally" | "process";
+}
+
+/**
+ * Read and parse proofs
+ *
+ * @param args - read proofs arguments
+ * @returns proofs
+ */
+async function readProofs({ files, folder, type }: IReadProofsArgs): Promise<Proof[]> {
+  return Promise.all(
+    files
+      .filter((f) => f.startsWith(`${type}_`) && f.endsWith(".json"))
+      .sort()
+      .map(async (file) =>
+        fs.promises.readFile(`${folder}/${file}`, "utf8").then((result) => JSON.parse(result) as Proof),
+      ),
+  );
+}
+
 /**
  * Prove hardhat task for submitting proofs on-chain as well as uploading tally results
  */
@@ -22,6 +48,8 @@ task("submitOnChain", "Command to prove the result of a poll on-chain")
   .setAction(async ({ outputDir, poll, tallyFile }: ISubmitOnChainParams, hre) => {
     const deployment = Deployment.getInstance();
     deployment.setHre(hre);
+    deployment.setContractNames(EContracts);
+
     const storage = ContractStorage.getInstance();
     // if we do not have the output directory just create it
     const isOutputDirExists = fs.existsSync(outputDir);
@@ -41,45 +69,52 @@ task("submitOnChain", "Command to prove the result of a poll on-chain")
     console.log("Start balance: ", Number(startBalance / 10n ** 12n) / 1e6);
 
     const maciContractAddress = storage.mustGetAddress(EContracts.MACI, network.name);
-    const maciContract = await deployment.getContract<MACI>({ name: EContracts.MACI, address: maciContractAddress });
-    const vkRegistryContract = await deployment.getContract<VkRegistry>({ name: EContracts.VkRegistry });
-    const verifierContract = await deployment.getContract<Verifier>({ name: EContracts.Verifier });
+    const [maciContract, vkRegistryContract, verifierContract] = await Promise.all([
+      deployment.getContract<MACI>({
+        name: EContracts.MACI,
+        address: maciContractAddress,
+      }),
+      deployment.getContract<VkRegistry>({ name: EContracts.VkRegistry }),
+      deployment.getContract<Verifier>({ name: EContracts.Verifier }),
+    ]);
 
     const pollContracts = await maciContract.polls(poll);
-    const pollContract = await deployment.getContract<Poll>({ name: EContracts.Poll, address: pollContracts.poll });
+    const pollContract = await deployment.getContract<Poll>({
+      name: EContracts.Poll,
+      address: pollContracts.poll,
+    });
 
-    const [, messageAqContractAddress] = await pollContract.extContracts();
+    const [[, messageAqContractAddress], isStateAqMerged, messageTreeDepth, mpContract, tallyContract] =
+      await Promise.all([
+        pollContract.extContracts(),
+        pollContract.stateMerged(),
+        pollContract.treeDepths().then((depths) => Number(depths[2])),
+        deployment.getContract<MessageProcessor>({
+          name: EContracts.MessageProcessor,
+          address: pollContracts.messageProcessor,
+        }),
+        deployment.getContract<Tally>({
+          name: EContracts.Tally,
+          address: pollContracts.tally,
+        }),
+      ]);
     const messageAqContract = await deployment.getContract<AccQueue>({
       name: EContracts.AccQueue,
       address: messageAqContractAddress,
     });
-    const isStateAqMerged = await pollContract.stateMerged();
 
     // Check that the state and message trees have been merged for at least the first poll
     if (!isStateAqMerged && poll.toString() === "0") {
       throw new Error("The state tree has not been merged yet. Please use the mergeSignups subcommand to do so.");
     }
 
-    const messageTreeDepth = await pollContract.treeDepths().then((depths) => Number(depths[2]));
-
     // check that the main root is set
     const mainRoot = await messageAqContract.getMainRoot(messageTreeDepth.toString());
 
     if (mainRoot.toString() === "0") {
       throw new Error("The message tree has not been merged yet. Please use the mergeMessages subcommand to do so.");
     }
 
-    const mpContract = await deployment.getContract<MessageProcessor>({
-      name: EContracts.MessageProcessor,
-      address: pollContracts.messageProcessor,
-    });
-
-    // get the tally contract based on the useQuadraticVoting flag
-    const tallyContract = await deployment.getContract<Tally>({
-      name: EContracts.Tally,
-      address: pollContracts.tally,
-    });
-
     const data = {
       processProofs: [] as Proof[],
       tallyProofs: [] as Proof[],
@@ -89,22 +124,9 @@ task("submitOnChain", "Command to prove the result of a poll on-chain")
     const files = await fs.promises.readdir(outputDir);
 
     // Read process proofs
-    const processProofFiles = files.filter((f) => f.startsWith("process_") && f.endsWith(".json"));
-    await Promise.all(
-      processProofFiles.sort().map(async (file) => {
-        const proofData = JSON.parse(await fs.promises.readFile(`${outputDir}/${file}`, "utf8")) as Proof;
-        data.processProofs.push(proofData);
-      }),
-    );
-
+    data.processProofs = await readProofs({ files, folder: outputDir, type: "process" });
     // Read tally proofs
-    const tallyProofFiles = files.filter((f) => f.startsWith("tally_") && f.endsWith(".json"));
-    await Promise.all(
-      tallyProofFiles.sort().map(async (file) => {
-        const proofData = JSON.parse(await fs.promises.readFile(`${outputDir}/${file}`, "utf8")) as Proof;
-        data.tallyProofs.push(proofData);
-      }),
-    );
+    data.tallyProofs = await readProofs({ files, folder: outputDir, type: "tally" });
 
     const prover = new Prover({
       maciContract,
@@ -119,7 +141,9 @@ task("submitOnChain", "Command to prove the result of a poll on-chain")
     await prover.proveMessageProcessing(data.processProofs);
 
     // read tally data
-    const tallyData = JSON.parse(await fs.promises.readFile(tallyFile, "utf8")) as unknown as TallyData;
+    const tallyData = await fs.promises
+      .readFile(tallyFile, "utf8")
+      .then((result) => JSON.parse(result) as unknown as TallyData);
 
     await prover.proveTally(data.tallyProofs);
 

diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml