feat: allow to fetch state index from contracts (#2092)

packages/contracts/contracts/MACI.sol

@@ -63,36 +63,6 @@ contract MACI is IMACI, DomainObjs, Params, Hasher {
   /// For the N'th sign up, the state tree root will be stored at the index N
   uint256[] public stateRootsOnSignUp;
 
-  /// @notice A struct holding the addresses of poll, mp and tally
-  struct PollContracts {
-    address poll;
-    address messageProcessor;
-    address tally;
-  }
-  /// @notice A struct holding the params for poll deployment
-  struct DeployPollArgs {
-    /// @param duration How long should the Poll last for
-    uint256 duration;
-    /// @param treeDepths The depth of the Merkle trees
-    TreeDepths treeDepths;
-    /// @param messageBatchSize The message batch size
-    uint8 messageBatchSize;
-    /// @param coordinatorPubKey The coordinator's public key
-    PubKey coordinatorPubKey;
-    /// @param verifier The Verifier Contract
-    address verifier;
-    /// @param vkRegistry The VkRegistry Contract
-    address vkRegistry;
-    /// @param mode Voting mode
-    Mode mode;
-    /// @param gatekeeper The gatekeeper contract
-    address gatekeeper;
-    /// @param initialVoiceCreditProxy The initial voice credit proxy contract
-    address initialVoiceCreditProxy;
-    /// @param relayer The message relayer (optional)
-    address[] relayers;
-  }
-
   // Events
   event SignUp(uint256 _stateIndex, uint256 _timestamp, uint256 indexed _userPubKeyX, uint256 indexed _userPubKeyY);
   event DeployPoll(
@@ -139,14 +109,7 @@ contract MACI is IMACI, DomainObjs, Params, Hasher {
     if (hash2([uint256(1), uint256(1)]) == 0) revert PoseidonHashLibrariesNotLinked();
   }
 
-  /// @notice Allows any eligible user sign up. The sign-up gatekeeper should prevent
-  /// double sign-ups or ineligible users from doing so.  This function will
-  /// only succeed if the sign-up deadline has not passed.
-  /// @param _pubKey The user's desired public key.
-  /// @param _signUpGatekeeperData Data to pass to the sign-up gatekeeper's
-  ///     register() function. For instance, the POAPGatekeeper or
-  ///     SignUpTokenGatekeeper requires this value to be the ABI-encoded
-  ///     token ID.
+  /// @inheritdoc IMACI
   function signUp(PubKey memory _pubKey, bytes memory _signUpGatekeeperData) public virtual {
     // ensure we do not have more signups than what the circuits support
     if (leanIMTData.size >= maxSignups) revert TooManySignups();
@@ -170,9 +133,8 @@ contract MACI is IMACI, DomainObjs, Params, Hasher {
     emit SignUp(leanIMTData.size - 1, block.timestamp, _pubKey.x, _pubKey.y);
   }
 
-  /// @notice Deploy a new Poll contract.
-  /// @param args The deploy poll args
-  function deployPoll(DeployPollArgs calldata args) public virtual {
+  /// @inheritdoc IMACI
+  function deployPoll(DeployPollArgs calldata args) public virtual returns (PollContracts memory) {
     // cache the poll to a local variable so we can increment it
     uint256 pollId = nextPollId;
 
@@ -217,6 +179,8 @@ contract MACI is IMACI, DomainObjs, Params, Hasher {
     polls[pollId] = pollAddr;
 
     emit DeployPoll(pollId, args.coordinatorPubKey.x, args.coordinatorPubKey.y, args.mode);
+
+    return pollAddr;
   }
 
   /// @inheritdoc IMACI
@@ -241,4 +205,10 @@ contract MACI is IMACI, DomainObjs, Params, Hasher {
   function getStateRootOnIndexedSignUp(uint256 _index) external view returns (uint256 stateRoot) {
     stateRoot = stateRootsOnSignUp[_index];
   }
+
+  /// @inheritdoc IMACI
+  function getStateIndex(uint256 _pubKeyHash) external view returns (uint256 index) {
+    // need to subtract 1 because the index is 1 indexed due to 0 index reserved for deleted leaves
+    index = leanIMTData.leaves[_pubKeyHash] - 1;
+  }
 }

packages/contracts/contracts/Poll.sol

@@ -114,6 +114,7 @@ contract Poll is Params, Utilities, SnarkCommon, IPoll {
   error UserAlreadyJoined();
   error InvalidPollProof();
   error NotRelayer();
+  error StateLeafNotFound();
 
   event PublishMessage(Message _message, PubKey _encPubKey);
   event MergeState(uint256 indexed _stateRoot, uint256 indexed _numSignups);
@@ -319,13 +320,7 @@ contract Poll is Params, Utilities, SnarkCommon, IPoll {
     }
   }
 
-  /// @notice Join the poll for voting
-  /// @param _nullifier Hashed user's private key to check whether user has already voted
-  /// @param _pubKey Poll user's public key
-  /// @param _stateRootIndex Index of the MACI's stateRootOnSignUp for which the inclusion proof is generated
-  /// @param _proof The zk-SNARK proof
-  /// @param _signUpGatekeeperData Data to pass to the SignUpGatekeeper
-  /// @param _initialVoiceCreditProxyData Data to pass to the InitialVoiceCreditProxy
+  /// @inheritdoc IPoll
   function joinPoll(
     uint256 _nullifier,
     PubKey calldata _pubKey,
@@ -358,6 +353,7 @@ contract Poll is Params, Utilities, SnarkCommon, IPoll {
 
     // Store user in the pollStateTree
     uint256 stateLeaf = hashStateLeaf(StateLeaf(_pubKey, voiceCreditBalance, block.timestamp));
+
     uint256 stateRoot = InternalLazyIMT._insert(pollStateTree, stateLeaf);
 
     // Store the current state tree root in the array
@@ -487,4 +483,15 @@ contract Poll is Params, Utilities, SnarkCommon, IPoll {
   function getMaciContract() public view returns (IMACI maci) {
     return extContracts.maci;
   }
+
+  /// @inheritdoc IPoll
+  function getStateIndex(uint256 element) public view returns (uint40) {
+    for (uint40 i = 0; i <= pollStateTree.maxIndex; i++) {
+      if (pollStateTree.elements[i] == element) {
+        return i;
+      }
+    }
+
+    revert StateLeafNotFound();
+  }
 }

packages/contracts/contracts/interfaces/IMACI.sol

@@ -1,9 +1,42 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.28;
 
+import { Params } from "../utilities/Params.sol";
+import { DomainObjs } from "../utilities/DomainObjs.sol";
+
 /// @title IMACI
 /// @notice MACI interface
 interface IMACI {
+  /// @notice A struct holding the addresses of poll, mp and tally
+  struct PollContracts {
+    address poll;
+    address messageProcessor;
+    address tally;
+  }
+  /// @notice A struct holding the params for poll deployment
+  struct DeployPollArgs {
+    /// @param duration How long should the Poll last for
+    uint256 duration;
+    /// @param treeDepths The depth of the Merkle trees
+    Params.TreeDepths treeDepths;
+    /// @param messageBatchSize The message batch size
+    uint8 messageBatchSize;
+    /// @param coordinatorPubKey The coordinator's public key
+    DomainObjs.PubKey coordinatorPubKey;
+    /// @param verifier The Verifier Contract
+    address verifier;
+    /// @param vkRegistry The VkRegistry Contract
+    address vkRegistry;
+    /// @param mode Voting mode
+    DomainObjs.Mode mode;
+    /// @param gatekeeper The gatekeeper contract
+    address gatekeeper;
+    /// @param initialVoiceCreditProxy The initial voice credit proxy contract
+    address initialVoiceCreditProxy;
+    /// @param relayer The message relayer (optional)
+    address[] relayers;
+  }
+
   /// @notice Get the depth of the state tree
   /// @return The depth of the state tree
   function stateTreeDepth() external view returns (uint8);
@@ -12,6 +45,25 @@ interface IMACI {
   /// @return The Merkle root
   function getStateTreeRoot() external view returns (uint256);
 
+  /// @notice Get the index of a public key in the state tree
+  /// @param _pubKeyHash The hash of the public key
+  /// @return index The index of the public key in the state tree
+  function getStateIndex(uint256 _pubKeyHash) external view returns (uint256);
+
+  /// @notice Deploy a new Poll contract.
+  /// @param _pollArgs The deploy poll args
+  function deployPoll(DeployPollArgs memory _pollArgs) external returns (PollContracts memory);
+
+  /// @notice Allows any eligible user sign up. The sign-up gatekeeper should prevent
+  /// double sign-ups or ineligible users from doing so.  This function will
+  /// only succeed if the sign-up deadline has not passed.
+  /// @param _pubKey The user's desired public key.
+  /// @param _signUpGatekeeperData Data to pass to the sign-up gatekeeper's
+  ///     register() function. For instance, the POAPGatekeeper or
+  ///     SignUpTokenGatekeeper requires this value to be the ABI-encoded
+  ///     token ID.
+  function signUp(DomainObjs.PubKey memory _pubKey, bytes memory _signUpGatekeeperData) external;
+
   /// @notice Return the state root when the '_index' user signed up
   /// @param _index The serial number when the user signed up
   /// @return The Merkle root

packages/contracts/contracts/interfaces/IPoll.sol

@@ -7,7 +7,13 @@ import { IMACI } from "./IMACI.sol";
 /// @title IPoll
 /// @notice Poll interface
 interface IPoll {
-  /// @notice Join the poll
+  /// @notice Join the poll for voting
+  /// @param _nullifier Hashed user's private key to check whether user has already voted
+  /// @param _pubKey Poll user's public key
+  /// @param _stateRootIndex Index of the MACI's stateRootOnSignUp for which the inclusion proof is generated
+  /// @param _proof The zk-SNARK proof
+  /// @param _signUpGatekeeperData Data to pass to the SignUpGatekeeper
+  /// @param _initialVoiceCreditProxyData Data to pass to the InitialVoiceCreditProxy
   function joinPoll(
     uint256 _nullifier,
     DomainObjs.PubKey calldata _pubKey,
@@ -99,4 +105,9 @@ interface IPoll {
   /// @notice Get the external contracts
   /// @return maci The IMACI contract
   function getMaciContract() external view returns (IMACI maci);
+
+  /// @notice Get the index of a state leaf in the state tree
+  /// @param element The hash of thestate leaf
+  /// @return index The index of the state leaf in the state tree
+  function getStateIndex(uint256 element) external view returns (uint40);
 }

packages/contracts/tests/MACI.test.ts

@@ -187,6 +187,13 @@ describe("MACI", function test() {
     });
   });
 
+  describe("getStateIndex", () => {
+    it("should return the index of a state leaf", async () => {
+      const index = await maciContract.getStateIndex(users[0].pubKey.hash());
+      expect(index.toString()).to.eq("1");
+    });
+  });
+
   describe("Deploy a Poll", () => {
     let deployTime: number | undefined;
 

packages/contracts/tests/Poll.test.ts

@@ -5,7 +5,7 @@ import { AbiCoder, decodeBase58, encodeBase58, getBytes, hexlify, Signer, ZeroAd
 import { EthereumProvider } from "hardhat/types";
 import { MaciState, VOTE_OPTION_TREE_ARITY } from "maci-core";
 import { NOTHING_UP_MY_SLEEVE } from "maci-crypto";
-import { Keypair, Message, PCommand, PubKey } from "maci-domainobjs";
+import { Keypair, Message, PCommand, PubKey, StateLeaf } from "maci-domainobjs";
 
 import { EMode } from "../ts/constants";
 import { IVerifyingKeyStruct } from "../ts/types";
@@ -232,9 +232,13 @@ describe("Poll", () => {
 
         const expectedIndex = maciState.polls
           .get(pollId)
-          ?.joinPoll(BigInt(mockNullifier), keypair.pubKey, BigInt(voiceCredits), BigInt(timestamp));
+          ?.joinPoll(BigInt(mockNullifier), keypair.pubKey, voiceCredits, BigInt(timestamp));
 
         expect(index).to.eq(expectedIndex);
+
+        // get the index with getStateIndex
+        const stateLeaf = new StateLeaf(keypair.pubKey, voiceCredits, BigInt(timestamp));
+        expect(await pollContract.getStateIndex(stateLeaf.hash())).to.eq(index);
       }
     });