test: rln subtree root onchain strategy

.gas-snapshot

@@ -1,17 +1,19 @@
-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__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__Constants() (gas: 8706)
+RlnTest:test__InvalidRegistration__DuplicateCommitment(uint256) (runs: 1000, μ: 149716, ~: 149716)
+RlnTest:test__InvalidRegistration__InsufficientDeposit(uint256) (runs: 1000, μ: 17510, ~: 17510)
+RlnTest:test__InvalidRegistration__InvalidIdCommitment(uint256) (runs: 1000, μ: 17099, ~: 17100)
+RlnTest:test__InvalidRegistration__InvalidUserMessageLimit() (gas: 17075)
+RlnTest:test__InvalidRegistration__MaxUserMessageLimit() (gas: 17251)
+RlnTest:test__InvalidSlash__InvalidProof() (gas: 1445036)
+RlnTest:test__InvalidSlash__MemberNotRegistered(uint256) (runs: 1000, μ: 30329, ~: 30329)
+RlnTest:test__InvalidSlash__NoStake(uint256,address) (runs: 1000, μ: 324171, ~: 324178)
+RlnTest:test__InvalidSlash__ToRlnAddress() (gas: 154267)
+RlnTest:test__InvalidSlash__ToZeroAddress() (gas: 154172)
+RlnTest:test__InvalidWithdraw__InsufficientContractBalance() (gas: 149684)
+RlnTest:test__InvalidWithdraw__InsufficientWithdrawalBalance() (gas: 10493)
+RlnTest:test__ValidRegistration(uint256) (runs: 1000, μ: 141272, ~: 141272)
+RlnTest:test__ValidSlash(uint256,address) (runs: 1000, μ: 207834, ~: 207841)
+RlnTest:test__ValidWithdraw(address) (runs: 1000, μ: 206639, ~: 206629)
+RlnTest:test__full_root() (gas: 197774266)
+RlnTest:test__kats__root() (gas: 4263258)
+RlnTest:test__root() (gas: 217521968)

src/BinaryIMTMemory.sol

@@ -0,0 +1,120 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.19;
+
+import { PoseidonT3 } from "poseidon-solidity/PoseidonT3.sol";
+import "forge-std/console2.sol";
+
+// stripped down version of
+// solhint-disable-next-line max-line-length
+// https://github.com/privacy-scaling-explorations/zk-kit/blob/718a5c2fa0f6cd577cee3fd08373609ac985d3bb/packages/imt.sol/contracts/internal/InternalBinaryIMT.sol
+// that allows getting the root of the rln tree without expensive storage reads/writes
+struct BinaryIMTMemoryData {
+    uint256 root; // Root hash of the tree.
+    uint256 numberOfLeaves; // Number of leaves of the tree.
+    uint256 depth; // Depth of the tree.
+}
+
+/// @title In memory Incremental binary Merkle tree Root calculator
+/// @dev This helper library allows to calculate the root hash of the tree without using storage
+library BinaryIMTMemory {
+    uint8 public constant MAX_DEPTH = 20;
+    uint256 public constant SNARK_SCALAR_FIELD =
+        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;
+
+    uint256 public constant Z_0 = 0;
+    uint256 public constant Z_1 =
+        14_744_269_619_966_411_208_579_211_824_598_458_697_587_494_354_926_760_081_771_325_075_741_142_829_156;
+    uint256 public constant Z_2 =
+        7_423_237_065_226_347_324_353_380_772_367_382_631_490_014_989_348_495_481_811_164_164_159_255_474_657;
+    uint256 public constant Z_3 =
+        11_286_972_368_698_509_976_183_087_595_462_810_875_513_684_078_608_517_520_839_298_933_882_497_716_792;
+    uint256 public constant Z_4 =
+        3_607_627_140_608_796_879_659_380_071_776_844_901_612_302_623_152_076_817_094_415_224_584_923_813_162;
+    uint256 public constant Z_5 =
+        19_712_377_064_642_672_829_441_595_136_074_946_683_621_277_828_620_209_496_774_504_837_737_984_048_981;
+    uint256 public constant Z_6 =
+        20_775_607_673_010_627_194_014_556_968_476_266_066_927_294_572_720_319_469_184_847_051_418_138_353_016;
+    uint256 public constant Z_7 =
+        3_396_914_609_616_007_258_851_405_644_437_304_192_397_291_162_432_396_347_162_513_310_381_425_243_293;
+    uint256 public constant Z_8 =
+        21_551_820_661_461_729_022_865_262_380_882_070_649_935_529_853_313_286_572_328_683_688_269_863_701_601;
+    uint256 public constant Z_9 =
+        6_573_136_701_248_752_079_028_194_407_151_022_595_060_682_063_033_565_181_951_145_966_236_778_420_039;
+    uint256 public constant Z_10 =
+        12_413_880_268_183_407_374_852_357_075_976_609_371_175_688_755_676_981_206_018_884_971_008_854_919_922;
+    uint256 public constant Z_11 =
+        14_271_763_308_400_718_165_336_499_097_156_975_241_954_733_520_325_982_997_864_342_600_795_471_836_726;
+    uint256 public constant Z_12 =
+        20_066_985_985_293_572_387_227_381_049_700_832_219_069_292_839_614_107_140_851_619_262_827_735_677_018;
+    uint256 public constant Z_13 =
+        9_394_776_414_966_240_069_580_838_672_673_694_685_292_165_040_808_226_440_647_796_406_499_139_370_960;
+    uint256 public constant Z_14 =
+        11_331_146_992_410_411_304_059_858_900_317_123_658_895_005_918_277_453_009_197_229_807_340_014_528_524;
+    uint256 public constant Z_15 =
+        15_819_538_789_928_229_930_262_697_811_477_882_737_253_464_456_578_333_862_691_129_291_651_619_515_538;
+    uint256 public constant Z_16 =
+        19_217_088_683_336_594_659_449_020_493_828_377_907_203_207_941_212_636_669_271_704_950_158_751_593_251;
+    uint256 public constant Z_17 =
+        21_035_245_323_335_827_719_745_544_373_081_896_983_162_834_604_456_827_698_288_649_288_827_293_579_666;
+    uint256 public constant Z_18 =
+        6_939_770_416_153_240_137_322_503_476_966_641_397_417_391_950_902_474_480_970_945_462_551_409_848_591;
+    uint256 public constant Z_19 =
+        10_941_962_436_777_715_901_943_463_195_175_331_263_348_098_796_018_438_960_955_633_645_115_732_864_202;
+    uint256 public constant Z_20 =
+        15_019_797_232_609_675_441_998_260_052_101_280_400_536_945_603_062_888_308_240_081_994_073_687_793_470;
+
+    // solhint-disable-next-line code-complexity
+    function defaultZero(uint256 index) public pure returns (uint256) {
+        if (index == 0) return Z_0;
+        if (index == 1) return Z_1;
+        if (index == 2) return Z_2;
+        if (index == 3) return Z_3;
+        if (index == 4) return Z_4;
+        if (index == 5) return Z_5;
+        if (index == 6) return Z_6;
+        if (index == 7) return Z_7;
+        if (index == 8) return Z_8;
+        if (index == 9) return Z_9;
+        if (index == 10) return Z_10;
+        if (index == 11) return Z_11;
+        if (index == 12) return Z_12;
+        if (index == 13) return Z_13;
+        if (index == 14) return Z_14;
+        if (index == 15) return Z_15;
+        if (index == 16) return Z_16;
+        if (index == 17) return Z_17;
+        if (index == 18) return Z_18;
+        if (index == 19) return Z_19;
+        if (index == 20) return Z_20;
+        revert("IncrementalBinaryTree: defaultZero bad index");
+    }
+
+    function calcSubtreeRoot(BinaryIMTMemoryData memory self, uint256[] memory leaves) public pure returns (uint256) {
+        // instead of the depth being 0..19, we go from 10..19.
+        // this way, we preserve the subtree root
+        uint8 depth = 10;
+        uint256[2][] memory lastSubtrees = new uint256[2][](depth);
+        for (uint256 j = 0; j < leaves.length;) {
+            uint256 index = self.numberOfLeaves;
+            uint256 hash = leaves[j];
+            for (uint8 i = 0; i < depth;) {
+                if (index & 1 == 0) {
+                    lastSubtrees[i] = [hash, defaultZero(i + depth)];
+                } else {
+                    lastSubtrees[i][1] = hash;
+                }
+                hash = PoseidonT3.hash(lastSubtrees[i]);
+                index >>= 1;
+                unchecked {
+                    ++i;
+                }
+            }
+            self.root = hash;
+            self.numberOfLeaves += 1;
+            unchecked {
+                ++j;
+            }
+        }
+        return self.root;
+    }
+}

src/RlnBase.sol

@@ -3,6 +3,9 @@
 pragma solidity ^0.8.19;
 
 import { IVerifier } from "./IVerifier.sol";
+import { BinaryIMTMemory, BinaryIMTMemoryData } from "./BinaryIMTMemory.sol";
+import { PoseidonT3 } from "poseidon-solidity/PoseidonT3.sol";
+import "forge-std/console2.sol";
 
 /// The tree is full
 error FullTree();
@@ -50,6 +53,19 @@
     uint256 public constant Q =
         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;
 
+    /// To ensure that the roots are accesible at each level, we shard the tree into subtrees
+    /// The leaves of the subtrees are stored in the leaves array.
+
+    /// @notice The max number of leaves in each subtree
+    uint256 public immutable SUBTREE_SIZE = 1024;
+    /// @notice the number of subtrees
+    uint256 public immutable SUBTREE_0_LENGTH = 1024;
+
+    mapping(uint256 => mapping(uint256 => uint256)) public leaves;
+    mapping(uint256 => mapping(uint256 => uint256)) public leavesIndex;
+    mapping(uint256 => bool) public memberExists;
+    mapping(uint256 => uint256) public leavesSet;
+
     /// @notice The max message limit per epoch
     uint256 public immutable MAX_MESSAGE_LIMIT;
 
@@ -62,27 +78,16 @@
     /// @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;
+    uint256 public currentShardIndex = 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
     mapping(uint256 => uint256) public userMessageLimits;
 
-    /// @notice the index to commitment mapping
-    mapping(uint256 => uint256) public indexToCommitment;
-
-    /// @notice The membership status of each member
-    mapping(uint256 => bool) public memberExists;
-
     /// @notice The balance of each user that can be withdrawn
     mapping(address => uint256) public withdrawalBalance;
 
@@ -96,12 +101,12 @@
     /// @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);
+    event MemberRegistered(uint256 shardIndex, uint256 idCommitment, uint256 userMessageLimit, uint256 index);
 
     /// Emitted when a member is removed from the set
     /// @param idCommitment The idCommitment of the member
     /// @param index The index of the member in the set
-    event MemberWithdrawn(uint256 idCommitment, uint256 index);
+    event MemberWithdrawn(uint256 shardIndex, uint256 idCommitment, uint256 index);
 
     modifier onlyValidIdCommitment(uint256 idCommitment) {
         if (!isValidCommitment(idCommitment)) revert InvalidIdCommitment(idCommitment);
@@ -157,16 +162,23 @@
     /// @param stake The amount of eth staked by the member
     function _register(uint256 idCommitment, uint256 userMessageLimit, uint256 stake) internal virtual {
         if (memberExists[idCommitment]) revert DuplicateIdCommitment();
-        if (idCommitmentIndex >= SET_SIZE) revert FullTree();
+        if (currentShardIndex == SUBTREE_0_LENGTH - 1 && leavesSet[currentShardIndex] == SUBTREE_SIZE - 1) {
+            revert FullTree();
+        }
+        uint256 index = leavesSet[currentShardIndex];
+
+        if (index == SUBTREE_SIZE - 1) {
+            currentShardIndex += 1;
+        }
 
-        members[idCommitment] = idCommitmentIndex;
-        indexToCommitment[idCommitmentIndex] = idCommitment;
+        leaves[currentShardIndex][idCommitment] = index;
+        leavesIndex[currentShardIndex][index] = idCommitment;
         memberExists[idCommitment] = true;
         stakedAmounts[idCommitment] = stake;
         userMessageLimits[idCommitment] = userMessageLimit;
 
-        emit MemberRegistered(idCommitment, userMessageLimit, idCommitmentIndex);
-        idCommitmentIndex += 1;
+        emit MemberRegistered(currentShardIndex, idCommitment, userMessageLimit, index);
+        leavesSet[currentShardIndex] += 1;
     }
 
     /// @dev Inheriting contracts MUST override this function
@@ -175,6 +187,7 @@
     /// @dev Allows a user to slash a member
     /// @param idCommitment The idCommitment of the member
     function slash(
+        uint256 shardIndex,
         uint256 idCommitment,
         address payable receiver,
         uint256[8] calldata proof
@@ -184,14 +197,22 @@
         onlyValidIdCommitment(idCommitment)
     {
         _validateSlash(idCommitment, receiver, proof);
-        _slash(idCommitment, receiver, proof);
+        _slash(shardIndex, idCommitment, receiver, proof);
     }
 
     /// @dev Slashes a member by removing them from the set, and adding their
     /// stake to the receiver's available withdrawal balance
     /// @param idCommitment The idCommitment of the member
     /// @param receiver The address to receive the funds
-    function _slash(uint256 idCommitment, address payable receiver, uint256[8] calldata proof) internal virtual {
+    function _slash(
+        uint256 shardIndex,
+        uint256 idCommitment,
+        address payable receiver,
+        uint256[8] calldata proof
+    )
+        internal
+        virtual
+    {
         if (receiver == address(this) || receiver == address(0)) {
             revert InvalidReceiverAddress(receiver);
         }
@@ -209,17 +230,17 @@
         uint256 amountToTransfer = stakedAmounts[idCommitment];
 
         // delete member
-        uint256 index = members[idCommitment];
-        members[idCommitment] = 0;
-        indexToCommitment[index] = 0;
+        uint256 index = leaves[shardIndex][idCommitment];
+        leaves[shardIndex][idCommitment] = 0;
+        leavesIndex[shardIndex][index] = 0;
         memberExists[idCommitment] = false;
         stakedAmounts[idCommitment] = 0;
         userMessageLimits[idCommitment] = 0;
 
         // refund deposit
         withdrawalBalance[receiver] += amountToTransfer;
 
-        emit MemberWithdrawn(idCommitment, index);
+        emit MemberWithdrawn(shardIndex, idCommitment, index);
     }
 
     function _validateSlash(
@@ -268,14 +289,48 @@
         );
     }
 
-    function getCommitments(uint256 startIndex, uint256 endIndex) public view returns (uint256[] memory) {
-        if (startIndex >= endIndex) revert InvalidPaginationQuery(startIndex, endIndex);
-        if (endIndex > idCommitmentIndex) revert InvalidPaginationQuery(startIndex, endIndex);
+    function getCommitments(uint256 shardIndex) public view returns (uint256[] memory) {
+        if (shardIndex >= SUBTREE_0_LENGTH) {
+            revert InvalidPaginationQuery(shardIndex, SUBTREE_0_LENGTH);
+        }
 
-        uint256[] memory commitments = new uint256[](endIndex - startIndex);
-        for (uint256 i = startIndex; i < endIndex; i++) {
-            commitments[i - startIndex] = indexToCommitment[i];
+        uint256 endIndex = leavesSet[shardIndex];
+        if (endIndex == 0) {
+            return new uint256[](0);
+        }
+        uint256[] memory commitments = new uint256[](endIndex);
+        for (uint256 i = 0; i < endIndex; i++) {
+            commitments[i] = leavesIndex[shardIndex][i];
         }
         return commitments;
     }
+
+    function root(uint256 shardIndex) public view returns (uint256) {
+        BinaryIMTMemoryData memory imtData;
+        uint256 idCommitmentIndex = leavesSet[shardIndex];
+        uint256[] memory calcLeaves = new uint256[](1024);
+
+        if (idCommitmentIndex == 0) {
+            return BinaryIMTMemory.calcSubtreeRoot(imtData, calcLeaves);
+        }
+        for (uint256 i = 0; i < idCommitmentIndex; i++) {
+            uint256 idCommitment = leavesIndex[shardIndex][i];
+            uint256 userMessageLimit = userMessageLimits[idCommitment];
+            calcLeaves[i] = PoseidonT3.hash([idCommitment, userMessageLimit]);
+        }
+        return BinaryIMTMemory.calcSubtreeRoot(imtData, calcLeaves);
+    }
+
+    function fullRoot(uint256[] calldata roots) public pure returns (uint256) {
+        BinaryIMTMemoryData memory imtData;
+        return BinaryIMTMemory.calcSubtreeRoot(imtData, roots);
+    }
+
+    function getLeafAtShard(uint256 shardIndex, uint256 index) public view returns (uint256) {
+        return leavesIndex[shardIndex][index];
+    }
+
+    function getLeafIndex(uint256 shardIndex, uint256 idCommitment) public view returns (uint256) {
+        return leaves[shardIndex][idCommitment];
+    }
 }