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"
   }
 }

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();
@@ -51,27 +53,27 @@
         21_888_242_871_839_275_222_246_405_745_257_275_088_548_364_400_416_034_343_698_204_186_575_808_495_617;
 
     /// @notice The max message limit per epoch
     uint256 public immutable MAX_MESSAGE_LIMIT;
 
     /// @notice The deposit amount required to register as a member
     uint256 public immutable MEMBERSHIP_DEPOSIT;
 
     /// @notice The depth of the merkle tree
     uint256 public immutable DEPTH;
 
     /// @notice The size of the merkle tree, i.e 2^depth
     uint256 public immutable SET_SIZE;
 
     /// @notice The index of the next member to be registered
     uint256 public idCommitmentIndex = 0;
 
     /// @notice The amount of eth staked by each member
     /// maps from idCommitment to the amount staked
     mapping(uint256 => uint256) public stakedAmounts;
 
     /// @notice The membership status of each member
     /// maps from idCommitment to their index in the set
     mapping(uint256 => uint256) public members;
 
     /// @notice the user message limit of each member
     /// maps from idCommitment to their user message limit
@@ -92,10 +94,13 @@
     /// @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 @@
         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 @@
 
         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 @@
         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 @@
         }
         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,14 @@ 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
+        );
+    }
 }