feat(sdk): allow to reuse inclusion proof for joinPoll

Author: ctrlc03

packages/sdk/ts/browser/index.ts

@@ -6,7 +6,14 @@ export * from "../trees";
 export * from "../vote";
 export * from "../maciKeys";
 export { getSignedupUserData, signup, hasUserSignedUp } from "../user/signup";
-export { getJoinedUserData, hasUserJoinedPoll } from "../user/utils";
+export {
+  getJoinedUserData,
+  hasUserJoinedPoll,
+  genMaciStateTree,
+  getPollJoiningCircuitEvents,
+  joiningCircuitInputs,
+  genMaciStateTreeWithEndKey,
+} from "../user/utils";
 
 export * from "./joinPoll";
 

packages/sdk/ts/browser/joinPoll.ts

@@ -3,27 +3,24 @@ import { MACI__factory as MACIFactory, Poll__factory as PollFactory } from "@mac
 import { poseidon } from "@maci-protocol/crypto";
 import { Keypair, PrivateKey } from "@maci-protocol/domainobjs";
 
-import type { IJoinPollArgs, IJoinPollData } from "../user/types";
+import type { IJoinPollBrowserArgs, IJoinPollData } from "../user/types";
+import type { TCircuitInputs } from "../utils/types";
 
-import {
-  getPollJoiningCircuitEvents,
-  getPollJoiningCircuitInputsFromStateFile,
-  hasUserJoinedPoll,
-} from "../user/utils";
+import { getPollJoiningCircuitEvents, hasUserJoinedPoll, joiningCircuitInputs } from "../user/utils";
 import { contractExists } from "../utils/contracts";
-import { TCircuitInputs } from "../utils/types";
 
 import { generateProofSnarkjs, formatProofForVerifierContract } from "./utils";
 
 /**
  * Join Poll user to the Poll contract
+ * @dev This version is optimised to work on browsers
+ * @dev It uses WASM + accepts already created inclusion proofs
  * @param {IJoinPollArgs} args - The arguments for the join poll command
  * @returns {IJoinPollData} The poll state index of the joined user and transaction hash
  */
 export const joinPoll = async ({
   maciAddress,
   privateKey,
-  stateFile,
   pollId,
   signer,
   startBlock,
@@ -33,7 +30,9 @@ export const joinPoll = async ({
   pollWasm,
   sgDataArg,
   ivcpDataArg,
-}: IJoinPollArgs): Promise<IJoinPollData> => {
+  inclusionProof,
+  useLatestStateIndex,
+}: IJoinPollBrowserArgs): Promise<IJoinPollData> => {
   const validContract = await contractExists(signer.provider!, maciAddress);
 
   if (!validContract) {
@@ -71,15 +70,12 @@ export const joinPoll = async ({
   // get the state index from the MACI contract
   const stateIndex = await maciContract.getStateIndex(userMaciPublicKey.hash()).catch(() => -1n);
 
+  const stateTreeDepth = await maciContract.stateTreeDepth();
+
   let circuitInputs: TCircuitInputs;
 
-  if (stateFile) {
-    circuitInputs = await getPollJoiningCircuitInputsFromStateFile({
-      stateFile,
-      pollId,
-      stateIndex,
-      userMaciPrivateKey,
-    });
+  if (inclusionProof) {
+    circuitInputs = joiningCircuitInputs(inclusionProof, stateTreeDepth, userMaciPrivateKey, userMaciPublicKey, pollId);
   } else {
     circuitInputs = await getPollJoiningCircuitEvents({
       maciContract,
@@ -93,7 +89,9 @@ export const joinPoll = async ({
     });
   }
 
-  const currentStateRootIndex = Number.parseInt((await maciContract.totalSignups()).toString(), 10) - 1;
+  const stateRootIndex = useLatestStateIndex
+    ? Number.parseInt((await maciContract.totalSignups()).toString(), 10) - 1
+    : stateIndex;
 
   // generate the proof for this batch
   const { proof } = await generateProofSnarkjs({
@@ -106,7 +104,7 @@ export const joinPoll = async ({
   const tx = await pollContract.joinPoll(
     nullifier,
     userMaciPublicKey.asContractParam(),
-    currentStateRootIndex,
+    stateRootIndex,
     formatProofForVerifierContract(proof),
     sgDataArg,
     ivcpDataArg,

packages/sdk/ts/trees/index.ts

@@ -1,2 +1,2 @@
-export { generateSignUpTree } from "./stateTree";
-export type { IGenerateSignUpTreeArgs, IGenerateSignUpTree } from "./types";
+export { generateSignUpTree, generateSignUpTreeWithEndKey } from "./stateTree";
+export type { IGenerateSignUpTreeArgs, IGenerateSignUpTree, IGenerateSignUpTreeWithEndKeyArgs } from "./types";

packages/sdk/ts/trees/stateTree.ts

@@ -4,9 +4,7 @@ import { hashLeanIMT, hashLeftRight, PAD_KEY_HASH } from "@maci-protocol/crypto"
 import { PublicKey } from "@maci-protocol/domainobjs";
 import { LeanIMT, LeanIMTHashFunction } from "@zk-kit/lean-imt";
 
-import assert from "assert";
-
-import type { IGenerateSignUpTreeArgs, IGenerateSignUpTree } from "./types";
+import type { IGenerateSignUpTreeArgs, IGenerateSignUpTree, IGenerateSignUpTreeWithEndKeyArgs } from "./types";
 
 import { sleep } from "../utils/utils";
 
@@ -44,7 +42,6 @@ export const generateSignUpTree = async ({
     // eslint-disable-next-line no-await-in-loop
     const signUpLogs = await maciContract.queryFilter(maciContract.filters.SignUp(), i, toBlock);
     signUpLogs.forEach((event) => {
-      assert(!!event);
       const publicKeyX = event.args._userPublicKeyX;
       const publicKeyY = event.args._userPublicKeyY;
 
@@ -65,3 +62,69 @@ export const generateSignUpTree = async ({
     publicKeys,
   };
 };
+
+/**
+ * Generate a State tree object from the events of a MACI smart contracts
+ * @param provider - the ethereum provider
+ * @param address - the address of the MACI contract
+ * @param fromBlock - the block number from which to start fetching events
+ * @param blocksPerRequest - the number of blocks to fetch in each request
+ * @param endBlock - the block number at which to stop fetching events
+ * @param sleepAmount - the amount of time to sleep between each request
+ * @returns State tree
+ */
+export const generateSignUpTreeWithEndKey = async ({
+  provider,
+  address,
+  fromBlock = 0,
+  blocksPerRequest = 50,
+  endBlock,
+  sleepAmount,
+  userPublicKey,
+}: IGenerateSignUpTreeWithEndKeyArgs): Promise<IGenerateSignUpTree> => {
+  const lastBlock = endBlock || (await provider.getBlockNumber());
+
+  const maciContract = MACIFactory.connect(address, provider);
+  const signUpTree = new LeanIMT(hashLeanIMT as LeanIMTHashFunction);
+  signUpTree.insert(PAD_KEY_HASH);
+  const publicKeys: PublicKey[] = [];
+
+  // Fetch event logs in batches (lastBlock inclusive)
+  for (let i = fromBlock; i <= lastBlock; i += blocksPerRequest + 1) {
+    // the last block batch will be either current iteration block + blockPerRequest
+    // or the end block if it is set
+    const toBlock = i + blocksPerRequest >= lastBlock ? lastBlock : i + blocksPerRequest;
+
+    // eslint-disable-next-line no-await-in-loop
+    const signUpLogs = await maciContract.queryFilter(maciContract.filters.SignUp(), i, toBlock);
+    // eslint-disable-next-line @typescript-eslint/prefer-for-of
+    for (let j = 0; j < signUpLogs.length; j += 1) {
+      const event = signUpLogs[j];
+      const publicKeyX = event.args._userPublicKeyX;
+      const publicKeyY = event.args._userPublicKeyY;
+
+      const publicKey = new PublicKey([publicKeyX, publicKeyY]);
+
+      publicKeys.push(publicKey);
+      signUpTree.insert(hashLeftRight(publicKeyX, publicKeyY));
+
+      // early return cause we found the user
+      if (publicKey.equals(userPublicKey)) {
+        return {
+          signUpTree,
+          publicKeys,
+        };
+      }
+    }
+
+    if (sleepAmount) {
+      // eslint-disable-next-line no-await-in-loop
+      await sleep(sleepAmount);
+    }
+  }
+
+  return {
+    signUpTree,
+    publicKeys,
+  };
+};

packages/sdk/ts/trees/types.ts

@@ -37,6 +37,16 @@ export interface IGenerateSignUpTreeArgs {
   sleepAmount?: number;
 }
 
+/**
+ * An interface that represents arguments of generation sign up tree which stops fetching at a given public key
+ */
+export interface IGenerateSignUpTreeWithEndKeyArgs extends IGenerateSignUpTreeArgs {
+  /**
+   * The public key of the user
+   */
+  userPublicKey: PublicKey;
+}
+
 /**
  * An interface that represents sign up tree and state leaves
  */

packages/sdk/ts/user/index.ts

@@ -1,6 +1,13 @@
 export { joinPoll } from "./joinPoll";
 export { getSignedupUserData, signup, hasUserSignedUp } from "./signup";
-export { getJoinedUserData, hasUserJoinedPoll } from "./utils";
+export {
+  getJoinedUserData,
+  hasUserJoinedPoll,
+  genMaciStateTree,
+  genMaciStateTreeWithEndKey,
+  getPollJoiningCircuitEvents,
+  joiningCircuitInputs,
+} from "./utils";
 export type {
   IJoinedUserArgs,
   IIsRegisteredUser,
@@ -18,4 +25,6 @@ export type {
   IParseSignupEventsArgs,
   ISignupData,
   IHasUserSignedUpArgs,
+  IGenMaciStateTreeArgs,
+  IGenMaciStateTreeWithEndKeyArgs,
 } from "./types";

packages/sdk/ts/user/types.ts

@@ -1,3 +1,5 @@
+import { LeanIMTMerkleProof } from "@zk-kit/lean-imt";
+
 import type { MACI, Poll } from "@maci-protocol/contracts/typechain-types";
 import type { PrivateKey, PublicKey } from "@maci-protocol/domainobjs";
 import type { Signer } from "ethers";
@@ -354,6 +356,21 @@ export interface IJoinPollArgs {
   ivcpDataArg: string;
 }
 
+/**
+ * Interface for the arguments to the joinPoll command for the browser
+ */
+export interface IJoinPollBrowserArgs extends IJoinPollArgs {
+  /**
+   * Whether to use of not the latest state index
+   */
+  useLatestStateIndex?: boolean;
+
+  /**
+   * The inclusion proof
+   */
+  inclusionProof?: LeanIMTMerkleProof;
+}
+
 /**
  * Interface for the return data to the joinPoll command
  */
@@ -401,6 +418,46 @@ export interface IIsNullifierOnChainArgs {
   signer: Signer;
 }
 
+/**
+ * Arguments for IGenMaciStateTreeArgs
+ */
+export interface IGenMaciStateTreeArgs {
+  /**
+   * The MACI contract
+   */
+  maciContract: MACI;
+
+  /**
+   * The signer
+   */
+  signer: Signer;
+
+  /**
+   * The start block
+   */
+  startBlock?: number;
+
+  /**
+   * The end block
+   */
+  endBlock?: number;
+
+  /**
+   * The blocks per batch
+   */
+  blocksPerBatch?: number;
+}
+
+/**
+ * Arguments for IGenMaciStateTreeWithEndKeyArgs
+ */
+export interface IGenMaciStateTreeWithEndKeyArgs extends IGenMaciStateTreeArgs {
+  /**
+   * The public key of the user
+   */
+  userPublicKey: PublicKey;
+}
+
 /**
  * Arguments for getPollJoiningCircuitEvents
  */