chore(rln-v2): use lazyImt for gas estimates

.gas-snapshot

@@ -1,17 +1,18 @@
-RlnTest:test__Constants() (gas: 8619)
-RlnTest:test__InvalidRegistration__DuplicateCommitment(uint256) (runs: 1000, μ: 144113, ~: 144113)
-RlnTest:test__InvalidRegistration__FullSet() (gas: 1433224)
-RlnTest:test__InvalidRegistration__InsufficientDeposit(uint256) (runs: 1000, μ: 17440, ~: 17440)
-RlnTest:test__InvalidRegistration__InvalidIdCommitment(uint256) (runs: 1000, μ: 17053, ~: 17058)
-RlnTest:test__InvalidRegistration__InvalidUserMessageLimit() (gas: 17055)
-RlnTest:test__InvalidRegistration__MaxUserMessageLimit() (gas: 17203)
-RlnTest:test__InvalidSlash__InvalidProof() (gas: 1081919)
-RlnTest:test__InvalidSlash__MemberNotRegistered(uint256) (runs: 1000, μ: 32521, ~: 32521)
-RlnTest:test__InvalidSlash__NoStake(uint256,address) (runs: 1000, μ: 319630, ~: 319682)
-RlnTest:test__InvalidSlash__ToRlnAddress() (gas: 151034)
-RlnTest:test__InvalidSlash__ToZeroAddress() (gas: 150939)
-RlnTest:test__InvalidWithdraw__InsufficientContractBalance() (gas: 145224)
+RlnTest:test__Constants() (gas: 8642)
+RlnTest:test__InvalidRegistration__DuplicateCommitment(uint256) (runs: 1000, μ: 218931, ~: 218931)
+RlnTest:test__InvalidRegistration__FullSet() (gas: 2083860)
+RlnTest:test__InvalidRegistration__InsufficientDeposit(uint256) (runs: 1000, μ: 17497, ~: 17497)
+RlnTest:test__InvalidRegistration__InvalidIdCommitment(uint256) (runs: 1000, μ: 17093, ~: 17097)
+RlnTest:test__InvalidRegistration__InvalidUserMessageLimit() (gas: 17094)
+RlnTest:test__InvalidRegistration__MaxUserMessageLimit() (gas: 17248)
+RlnTest:test__InvalidSlash__InvalidProof() (gas: 1411622)
+RlnTest:test__InvalidSlash__MemberNotRegistered(uint256) (runs: 1000, μ: 30343, ~: 30343)
+RlnTest:test__InvalidSlash__NoStake(uint256,address) (runs: 1000, μ: 379350, ~: 379347)
+RlnTest:test__InvalidSlash__ToRlnAddress() (gas: 225728)
+RlnTest:test__InvalidSlash__ToZeroAddress() (gas: 225633)
+RlnTest:test__InvalidWithdraw__InsufficientContractBalance() (gas: 205000)
 RlnTest:test__InvalidWithdraw__InsufficientWithdrawalBalance() (gas: 10538)
-RlnTest:test__ValidRegistration(uint256) (runs: 1000, μ: 135760, ~: 135760)
-RlnTest:test__ValidSlash(uint256,address) (runs: 1000, μ: 203402, ~: 203412)
-RlnTest:test__ValidWithdraw(address) (runs: 1000, μ: 202120, ~: 202108)
+RlnTest:test__ValidRegistration(uint256) (runs: 1000, μ: 210499, ~: 210499)
+RlnTest:test__ValidSlash(uint256,address) (runs: 1000, μ: 263197, ~: 263194)
+RlnTest:test__ValidWithdraw(address) (runs: 1000, μ: 261897, ~: 261885)
+RlnTest:test__root__kats() (gas: 397525)

foundry.toml

@@ -10,7 +10,7 @@
   gas_reports = ["*"]
   libs = ["lib"]
   optimizer = true
-  optimizer_runs = 10_000
+  optimizer_runs = 20_000
   out = "out"
   script = "script"
   solc = "0.8.19"

package.json

@@ -29,7 +29,7 @@
     "prettier:write": "prettier --write **/*.{json,md,yml} --ignore-path=.prettierignore"
   },
   "dependencies": {
-    "@zk-kit/imt.sol": "2.0.0-beta",
+    "@zk-kit/imt.sol": "https://gitpkg.now.sh/privacy-scaling-explorations/zk-kit/packages/imt.sol?0699fd1e5ad3683ae0090e0626f75d7834145500",
     "poseidon-solidity": "^0.0.5"
   }
 }

pnpm-lock.yaml

@@ -6,8 +6,8 @@ settings:
 
 dependencies:
   '@zk-kit/imt.sol':
-    specifier: 2.0.0-beta
-    version: 2.0.0-beta
+    specifier: https://gitpkg.now.sh/privacy-scaling-explorations/zk-kit/packages/imt.sol?0699fd1e5ad3683ae0090e0626f75d7834145500
+    version: '@gitpkg.now.sh/privacy-scaling-explorations/zk-kit/packages/imt.sol?0699fd1e5ad3683ae0090e0626f75d7834145500'
   poseidon-solidity:
     specifier: ^0.0.5
     version: 0.0.5
@@ -49,12 +49,6 @@ packages:
       antlr4ts: 0.5.0-alpha.4
     dev: true
 
-  /@zk-kit/imt.sol@2.0.0-beta:
-    resolution: {integrity: sha512-bH7RvI5WHAEswUwPspUY582O2+71xbYv5aL+DM4xkaA0GdMyMLUwf5c1yJ4wrt46hp07iXCXJsLXdtLNsTnvZw==}
-    dependencies:
-      poseidon-solidity: 0.0.5
-    dev: false
-
   /ajv@6.12.6:
     resolution: {integrity: sha512-j3fVLgvTo527anyYyJOGTYJbG+vnnQYvE0m5mmkc1TK+nxAppkCLMIL0aZ4dblVCNoGShhm+kzE4ZUykBoMg4g==}
     dependencies:
@@ -468,3 +462,11 @@ packages:
   /wrappy@1.0.2:
     resolution: {integrity: sha512-l4Sp/DRseor9wL6EvV2+TuQn63dMkPjZ/sp9XkghTEbV9KlPS1xUsZ3u7/IQO4wxtcFB4bgpQPRcR3QCvezPcQ==}
     dev: true
+
+  '@gitpkg.now.sh/privacy-scaling-explorations/zk-kit/packages/imt.sol?0699fd1e5ad3683ae0090e0626f75d7834145500':
+    resolution: {tarball: https://gitpkg.now.sh/privacy-scaling-explorations/zk-kit/packages/imt.sol?0699fd1e5ad3683ae0090e0626f75d7834145500}
+    name: imt.sol
+    version: 0.0.0
+    dependencies:
+      poseidon-solidity: 0.0.5
+    dev: false

remappings.txt

@@ -1,3 +1,3 @@
 forge-std/=lib/forge-std/src/
-@zk-kit/imt.sol/=node_modules/@zk-kit/imt.sol/
-poseidon-solidity/=node_modules/poseidon-solidity/
+@zk-kit/imt.sol/=node_modules/@zk-kit/imt.sol/contracts/
+poseidon-solidity/=node_modules/poseidon-solidity/

src/RlnBase.sol

@@ -3,6 +3,8 @@
 pragma solidity ^0.8.19;
 
 import { IVerifier } from "./IVerifier.sol";
+import { LazyIMT, LazyIMTData } from "@zk-kit/imt.sol/LazyIMT.sol";
+import { PoseidonT3 } from "poseidon-solidity/PoseidonT3.sol";
 
 /// The tree is full
 error FullTree();
@@ -92,10 +94,13 @@ abstract contract RlnBase {
     /// @notice the deployed block number
     uint32 public immutable deployedBlockNumber;
 
+    /// @notice the stored imt data
+    LazyIMTData public imtData;
     /// Emitted when a new member is added to the set
     /// @param idCommitment The idCommitment of the member
     /// @param userMessageLimit the user message limit of the member
     /// @param index The index of the member in the set
+
     event MemberRegistered(uint256 idCommitment, uint256 userMessageLimit, uint256 index);
 
     /// Emitted when a member is removed from the set
@@ -121,6 +126,7 @@ abstract contract RlnBase {
         SET_SIZE = 1 << depth;
         verifier = IVerifier(_verifier);
         deployedBlockNumber = uint32(block.number);
+        LazyIMT.init(imtData, 20);
     }
 
     /// Returns the deposit amount required to register as a member
@@ -161,6 +167,8 @@ abstract contract RlnBase {
 
         members[idCommitment] = idCommitmentIndex;
         indexToCommitment[idCommitmentIndex] = idCommitment;
+        uint256 rateCommitment = PoseidonT3.hash([idCommitment, userMessageLimit]);
+        LazyIMT.insert(imtData, rateCommitment);
         memberExists[idCommitment] = true;
         stakedAmounts[idCommitment] = stake;
         userMessageLimits[idCommitment] = userMessageLimit;
@@ -215,7 +223,8 @@ abstract contract RlnBase {
         memberExists[idCommitment] = false;
         stakedAmounts[idCommitment] = 0;
         userMessageLimits[idCommitment] = 0;
-
+        // cast to uint40 is required and perhaps we should use the same later
+        LazyIMT.update(imtData, 0, uint40(index));
         // refund deposit
         withdrawalBalance[receiver] += amountToTransfer;
 
@@ -278,4 +287,12 @@ abstract contract RlnBase {
         }
         return commitments;
     }
+
+    function root() external view returns (uint256) {
+        return LazyIMT.root(imtData, 20);
+    }
+
+    function merkleProofElements(uint40 index) public view returns (uint256[] memory) {
+        return LazyIMT.merkleProofElements(imtData, index, 20);
+    }
 }

test/Rln.t.sol

@@ -218,4 +218,15 @@ contract RlnTest is Test {
         rln.withdraw();
         assertEq(rln.withdrawalBalance(to), 0);
     }
+
+    function test__root__kats() public {
+        vm.pauseGasMetering();
+        rln.register{ value: MEMBERSHIP_DEPOSIT }(10, 1);
+        vm.resumeGasMetering();
+        uint256 root = rln.root();
+        assertEq(
+            root,
+            5_310_040_806_297_387_474_753_844_093_785_987_393_201_115_008_355_232_131_367_722_760_020_080_961_576
+        );
+    }
 }