feat(joinpoll): use genSignUpTree instead of genMaciStateFromContract

cli/ts/commands/joinPoll.ts

@@ -1,17 +1,98 @@
 import { extractVk, genProof, verifyProof } from "maci-circuits";
-import { formatProofForVerifierContract, genMaciStateFromContract } from "maci-contracts";
+import { formatProofForVerifierContract, genSignUpTree, IGenSignUpTree } from "maci-contracts";
 import { MACI__factory as MACIFactory, Poll__factory as PollFactory } from "maci-contracts/typechain-types";
-import { CircuitInputs, IJsonMaciState, MaciState } from "maci-core";
-import { poseidon } from "maci-crypto";
+import { CircuitInputs, IJsonMaciState, MaciState, IPollJoiningCircuitInputs } from "maci-core";
+import { poseidon, sha256Hash, stringifyBigInts } from "maci-crypto";
 import { Keypair, PrivKey, PubKey } from "maci-domainobjs";
 
+import assert from "assert";
 import fs from "fs";
 
 import type { IJoinPollArgs, IJoinedUserArgs, IParsePollJoinEventsArgs } from "../utils/interfaces";
 
 import { contractExists, logError, logYellow, info, logGreen, success } from "../utils";
 import { banner } from "../utils/banner";
 
+/**
+ * Create circuit input for pollJoining
+ * @param signUpData Sign up tree and state leaves
+ * @param stateTreeDepth Maci state tree depth
+ * @param maciPrivKey User's private key for signing up
+ * @param stateLeafIndex Index where the user is stored in the state leaves
+ * @param credits Credits for voting
+ * @param pollPrivKey Poll's private key for the poll joining
+ * @param pollPubKey Poll's public key for the poll joining
+ * @returns stringified circuit inputs
+ */
+const joiningCircuitInputs = (
+  signUpData: IGenSignUpTree,
+  stateTreeDepth: bigint,
+  maciPrivKey: PrivKey,
+  stateLeafIndex: bigint,
+  credits: bigint,
+  pollPrivKey: PrivKey,
+  pollPubKey: PubKey,
+): IPollJoiningCircuitInputs => {
+  // Get the state leaf on the index position
+  const { signUpTree: stateTree, stateLeaves } = signUpData;
+  const stateLeaf = stateLeaves[Number(stateLeafIndex)];
+  const { pubKey, voiceCreditBalance, timestamp } = stateLeaf;
+  const pubKeyX = pubKey.asArray()[0];
+  const pubKeyY = pubKey.asArray()[1];
+  const stateLeafArray = [pubKeyX, pubKeyY, voiceCreditBalance, timestamp];
+  const pollPubKeyArray = pollPubKey.asArray();
+
+  assert(credits <= voiceCreditBalance, "Credits must be lower than signed up credits");
+
+  // calculate the path elements for the state tree given the original state tree
+  const { siblings, index } = stateTree.generateProof(Number(stateLeafIndex));
+  const depth = siblings.length;
+
+  // The index must be converted to a list of indices, 1 for each tree level.
+  // The circuit tree depth is this.stateTreeDepth, so the number of siblings must be this.stateTreeDepth,
+  // even if the tree depth is actually 3. The missing siblings can be set to 0, as they
+  // won't be used to calculate the root in the circuit.
+  const indices: bigint[] = [];
+
+  for (let i = 0; i < stateTreeDepth; i += 1) {
+    // eslint-disable-next-line no-bitwise
+    indices.push(BigInt((index >> i) & 1));
+
+    if (i >= depth) {
+      siblings[i] = BigInt(0);
+    }
+  }
+
+  // Create nullifier from private key
+  const inputNullifier = BigInt(maciPrivKey.asCircuitInputs());
+  const nullifier = poseidon([inputNullifier]);
+
+  // Get pll state tree's root
+  const stateRoot = stateTree.root;
+
+  // Convert actualStateTreeDepth to BigInt
+  const actualStateTreeDepth = BigInt(stateTree.depth);
+
+  // Calculate public input hash from nullifier, credits and current root
+  const inputHash = sha256Hash([nullifier, credits, stateRoot, pollPubKeyArray[0], pollPubKeyArray[1]]);
+
+  const circuitInputs = {
+    privKey: maciPrivKey.asCircuitInputs(),
+    pollPrivKey: pollPrivKey.asCircuitInputs(),
+    pollPubKey: pollPubKey.asCircuitInputs(),
+    stateLeaf: stateLeafArray,
+    siblings,
+    indices,
+    nullifier,
+    credits,
+    stateRoot,
+    actualStateTreeDepth,
+    inputHash,
+  };
+
+  return stringifyBigInts(circuitInputs) as unknown as IPollJoiningCircuitInputs;
+};
+
 export const joinPoll = async ({
   maciAddress,
   privateKey,
@@ -33,7 +114,6 @@ export const joinPoll = async ({
   quiet = true,
 }: IJoinPollArgs): Promise<number> => {
   banner(quiet);
-  const userSideOnly = true;
 
   if (!(await contractExists(signer.provider!, maciAddress))) {
     logError("MACI contract does not exist");
@@ -70,6 +150,9 @@ export const joinPoll = async ({
   const pollContract = PollFactory.connect(pollAddress, signer);
 
   let maciState: MaciState | undefined;
+  let signUpData: IGenSignUpTree | undefined;
+  let currentStateRootIndex: number;
+  let circuitInputs: CircuitInputs;
   if (stateFile) {
     try {
       const file = await fs.promises.readFile(stateFile);
@@ -78,62 +161,65 @@ export const joinPoll = async ({
     } catch (error) {
       logError((error as Error).message);
     }
+    const poll = maciState!.polls.get(pollId)!;
+
+    if (poll.hasJoined(nullifier)) {
+      throw new Error("User the given nullifier has already joined");
+    }
+
+    currentStateRootIndex = poll.maciStateRef.numSignUps - 1;
+
+    poll.updatePoll(BigInt(maciState!.stateLeaves.length));
+
+    circuitInputs = poll.joiningCircuitInputs({
+      maciPrivKey: userMaciPrivKey,
+      stateLeafIndex: stateIndex,
+      credits: newVoiceCreditBalance,
+      pollPrivKey: pollPrivKeyDeserialized,
+      pollPubKey,
+    }) as unknown as CircuitInputs;
   } else {
     // build an off-chain representation of the MACI contract using data in the contract storage
-    const [defaultStartBlockSignup, defaultStartBlockPoll, stateRoot, numSignups] = await Promise.all([
+    const [defaultStartBlockSignup, defaultStartBlockPoll, stateTreeDepth, numSignups] = await Promise.all([
       maciContract.queryFilter(maciContract.filters.SignUp(), startBlock).then((events) => events[0]?.blockNumber ?? 0),
       maciContract
         .queryFilter(maciContract.filters.DeployPoll(), startBlock)
         .then((events) => events[0]?.blockNumber ?? 0),
-      maciContract.getStateTreeRoot(),
+      maciContract.stateTreeDepth(),
       maciContract.numSignUps(),
     ]);
     const defaultStartBlock = Math.min(defaultStartBlockPoll, defaultStartBlockSignup);
     let fromBlock = startBlock ? Number(startBlock) : defaultStartBlock;
 
-    const defaultEndBlock = await Promise.all([
-      pollContract
-        .queryFilter(pollContract.filters.MergeMaciState(stateRoot, numSignups), fromBlock)
-        .then((events) => events[events.length - 1]?.blockNumber),
-    ]).then((blocks) => Math.max(...blocks));
-
     if (transactionHash) {
       const tx = await signer.provider!.getTransaction(transactionHash);
       fromBlock = tx?.blockNumber ?? defaultStartBlock;
     }
 
     logYellow(quiet, info(`starting to fetch logs from block ${fromBlock}`));
-    // TODO: create genPollStateTree ?
-    maciState = await genMaciStateFromContract(
-      signer.provider!,
-      await maciContract.getAddress(),
-      new Keypair(), // Not important in this context
-      pollId,
+
+    signUpData = await genSignUpTree({
+      provider: signer.provider!,
+      address: await maciContract.getAddress(),
+      blocksPerRequest: blocksPerBatch || 50,
       fromBlock,
-      blocksPerBatch,
-      endBlock || defaultEndBlock,
-      0,
-      userSideOnly,
-    );
-  }
+      endBlock,
+      sleepAmount: 0,
+    });
 
-  const poll = maciState!.polls.get(pollId)!;
+    currentStateRootIndex = Number(numSignups) - 1;
 
-  if (poll.hasJoined(nullifier)) {
-    throw new Error("User the given nullifier has already joined");
+    circuitInputs = joiningCircuitInputs(
+      signUpData,
+      stateTreeDepth,
+      userMaciPrivKey,
+      stateIndex,
+      newVoiceCreditBalance,
+      pollPrivKeyDeserialized,
+      pollPubKey,
+    ) as unknown as CircuitInputs;
   }
 
-  const currentStateRootIndex = poll.maciStateRef.numSignUps - 1;
-  poll.updatePoll(BigInt(maciState!.stateLeaves.length));
-
-  const circuitInputs = poll.joiningCircuitInputs({
-    maciPrivKey: userMaciPrivKey,
-    stateLeafIndex: stateIndex,
-    credits: newVoiceCreditBalance,
-    pollPrivKey: pollPrivKeyDeserialized,
-    pollPubKey,
-  }) as unknown as CircuitInputs;
-
   const pollVk = await extractVk(pollJoiningZkey);
 
   try {

contracts/package.json

@@ -100,6 +100,7 @@
     "@nomicfoundation/hardhat-toolbox": "^5.0.0",
     "@openzeppelin/contracts": "^5.0.2",
     "@zk-kit/imt.sol": "2.0.0-beta.12",
+    "@zk-kit/lean-imt": "^2.1.0",
     "circomlibjs": "^0.1.7",
     "ethers": "^6.13.1",
     "hardhat": "^2.22.4",

contracts/ts/genMaciState.ts

@@ -32,7 +32,6 @@ export const genMaciStateFromContract = async (
   blocksPerRequest = 50,
   endBlock: number | undefined = undefined,
   sleepAmount: number | undefined = undefined,
-  userSideOnly: boolean | undefined = undefined,
 ): Promise<MaciState> => {
   // ensure the pollId is valid
   assert(pollId >= 0);
@@ -124,10 +123,8 @@ export const genMaciStateFromContract = async (
       pollContract.messageBatchSize(),
     ]);
 
-  if (!userSideOnly) {
-    assert(coordinatorPubKeyOnChain[0].toString() === coordinatorKeypair.pubKey.rawPubKey[0].toString());
-    assert(coordinatorPubKeyOnChain[1].toString() === coordinatorKeypair.pubKey.rawPubKey[1].toString());
-  }
+  assert(coordinatorPubKeyOnChain[0].toString() === coordinatorKeypair.pubKey.rawPubKey[0].toString());
+  assert(coordinatorPubKeyOnChain[1].toString() === coordinatorKeypair.pubKey.rawPubKey[1].toString());
 
   const maxVoteOptions = Number(onChainMaxVoteOptions);
 

contracts/ts/genSignUpTree.ts

@@ -1,12 +1,13 @@
-import { STATE_TREE_ARITY, STATE_TREE_DEPTH } from "maci-core/build/ts/utils/constants";
-import { hash2, IncrementalQuinTree } from "maci-crypto";
-import { blankStateLeafHash } from "maci-domainobjs";
+/* eslint-disable no-underscore-dangle */
+import { LeanIMT, LeanIMTHashFunction } from "@zk-kit/lean-imt";
+import { hashLeanIMT } from "maci-crypto";
+import { PubKey, StateLeaf, blankStateLeaf, blankStateLeafHash } from "maci-domainobjs";
 
 import { assert } from "console";
 
 import { MACI__factory as MACIFactory } from "../typechain-types";
 
-import { IGenSignUpTreeArgs } from "./types";
+import { IGenSignUpTreeArgs, IGenSignUpTree } from "./types";
 import { sleep } from "./utils";
 
 /**
@@ -26,12 +27,13 @@ export const genSignUpTree = async ({
   blocksPerRequest = 50,
   endBlock,
   sleepAmount,
-}: IGenSignUpTreeArgs): Promise<IncrementalQuinTree> => {
+}: IGenSignUpTreeArgs): Promise<IGenSignUpTree> => {
   const lastBlock = endBlock || (await provider.getBlockNumber());
 
   const maciContract = MACIFactory.connect(address, provider);
-  const signUpTree = new IncrementalQuinTree(STATE_TREE_DEPTH, blankStateLeafHash, STATE_TREE_ARITY, hash2);
+  const signUpTree = new LeanIMT(hashLeanIMT as LeanIMTHashFunction);
   signUpTree.insert(blankStateLeafHash);
+  const stateLeaves: StateLeaf[] = [blankStateLeaf];
 
   // Fetch event logs in batches (lastBlock inclusive)
   for (let i = fromBlock; i <= lastBlock; i += blocksPerRequest + 1) {
@@ -45,7 +47,15 @@ export const genSignUpTree = async ({
     ] = await Promise.all([maciContract.queryFilter(maciContract.filters.SignUp(), i, toBlock)]);
     signUpLogs.forEach((event) => {
       assert(!!event);
-      // eslint-disable-next-line no-underscore-dangle
+      const pubKeyX = event.args._userPubKeyX;
+      const pubKeyY = event.args._userPubKeyY;
+      const voiceCreditBalance = event.args._voiceCreditBalance;
+      const timestamp = event.args._timestamp;
+
+      const pubKey = new PubKey([pubKeyX, pubKeyY]);
+      const stateLeaf = new StateLeaf(pubKey, voiceCreditBalance, timestamp);
+
+      stateLeaves.push(stateLeaf);
       signUpTree.insert(event.args._stateLeaf);
     });
 
@@ -54,5 +64,8 @@ export const genSignUpTree = async ({
       await sleep(sleepAmount);
     }
   }
-  return signUpTree;
+  return {
+    signUpTree,
+    stateLeaves,
+  };
 };

contracts/ts/index.ts

@@ -25,5 +25,5 @@ export { Prover } from "../tasks/helpers/Prover";
 export { EContracts } from "../tasks/helpers/types";
 export { linkPoseidonLibraries } from "../tasks/helpers/abi";
 
-export type { IVerifyingKeyStruct, SnarkProof, Groth16Proof, Proof } from "./types";
+export type { IVerifyingKeyStruct, SnarkProof, Groth16Proof, Proof, IGenSignUpTree } from "./types";
 export * from "../typechain-types";

contracts/ts/types.ts

@@ -1,3 +1,5 @@
+import { LeanIMT } from "@zk-kit/lean-imt";
+
 import type {
   ConstantInitialVoiceCreditProxy,
   FreeForAllGatekeeper,
@@ -12,7 +14,7 @@ import type {
 } from "../typechain-types";
 import type { BigNumberish, Provider, Signer } from "ethers";
 import type { CircuitInputs } from "maci-core";
-import type { Message, PubKey } from "maci-domainobjs";
+import type { Message, PubKey, StateLeaf } from "maci-domainobjs";
 import type { PublicSignals } from "snarkjs";
 
 /**
@@ -171,6 +173,9 @@ export interface IDeployedMaci {
   };
 }
 
+/**
+ * An interface that represents arguments of generation sign up tree and state leaves
+ */
 export interface IGenSignUpTreeArgs {
   /**
    * The etherum provider
@@ -202,3 +207,18 @@ export interface IGenSignUpTreeArgs {
    */
   sleepAmount?: number;
 }
+
+/**
+ * An interface that represents sign up tree and state leaves
+ */
+export interface IGenSignUpTree {
+  /**
+   * Sign up tree
+   */
+  signUpTree: LeanIMT;
+
+  /**
+   * State leaves
+   */
+  stateLeaves: StateLeaf[];
+}

core/ts/Poll.ts

@@ -429,14 +429,14 @@ export class Poll implements IPoll {
   };
 
   /**
+   * Create circuit input for pollJoining
    * @param maciPrivKey User's private key for signing up
    * @param stateLeafIndex Index where the user is stored in the state leaves
    * @param credits Credits for voting
    * @param pollPrivKey Poll's private key for the poll joining
    * @param pollPubKey Poll's public key for the poll joining
    * @returns stringified circuit inputs
    */
-
   joiningCircuitInputs = ({
     maciPrivKey,
     stateLeafIndex,