refactor: submit onchain improvements

privacy-scaling-explorations · Nov 8, 2024 · 257b4fa · 257b4fa
1 parent 25d4d8b
commit 257b4fa
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 47 deletions.
diff --git a/packages/contracts/tasks/runner/initPoll.ts b/packages/contracts/tasks/runner/initPoll.ts
@@ -17,7 +17,7 @@ interface IInitPollParams {
 }
 
 /**
- * Command to merge signup and message queues of a MACI contract
+ * Command to initialize poll
  */
 task("initPoll", "Initialize poll")
   .addParam("poll", "The poll id", undefined, types.string)

diff --git a/packages/contracts/tasks/runner/submitOnChain.ts b/packages/contracts/tasks/runner/submitOnChain.ts
@@ -1,6 +1,5 @@
 /* eslint-disable no-await-in-loop */
 /* eslint-disable no-console */
-import { BaseContract } from "ethers";
 import { task, types } from "hardhat/config";
 import { ContractStorage, Deployment, Proof, Prover, type TallyData } from "maci-contracts";
 
@@ -14,10 +13,35 @@ import {
   type AccQueue,
   type MessageProcessor,
   type Tally,
-  type SimpleRegistry,
 } from "../../typechain-types";
 import { EContracts, ISubmitOnChainParams } from "../helpers/constants";
 
+/**
+ * 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
  */
@@ -57,13 +81,15 @@ task("submitOnChain", "Command to prove the result of a poll on-chain")
     } = await import("../../typechain-types");
 
     const maciContractAddress = storage.mustGetAddress(EContracts.MACI, network.name);
-    const maciContract = await deployment.getContract<MACI>({
-      name: EContracts.MACI,
-      address: maciContractAddress,
-      abi: MACIFactory.abi,
-    });
-    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,
+        abi: MACIFactory.abi,
+      }),
+      deployment.getContract<VkRegistry>({ name: EContracts.VkRegistry }),
+      deployment.getContract<Verifier>({ name: EContracts.Verifier }),
+    ]);
 
     const pollContracts = await maciContract.polls(poll);
     const pollContract = await deployment.getContract<Poll>({
@@ -72,39 +98,38 @@ task("submitOnChain", "Command to prove the result of a poll on-chain")
       abi: PollFactory.abi,
     });
 
-    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,
+          abi: TallyFactory.abi,
+        }),
+      ]);
     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,
-      abi: TallyFactory.abi,
-    });
-
     const data = {
       processProofs: [] as Proof[],
       tallyProofs: [] as Proof[],
@@ -114,22 +139,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,
@@ -144,18 +156,16 @@ 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);
 
     // submit the results with number participants to be taken from the registry
     const registryContractAddress = await pollContract.getRegistry();
 
-    const registryContract = new BaseContract(
-      registryContractAddress,
-      SimpleRegistryFactory.abi,
-      signer,
-    ) as SimpleRegistry;
+    const registryContract = SimpleRegistryFactory.connect(registryContractAddress, signer);
     const recipientCount = await registryContract.recipientCount();
 
     await prover.submitResults(tallyData, Number.parseInt(recipientCount.toString(), 10));