chore: add the initial contract after few optimizations (#4)

* chore: add deps

* fix: fmt in package.json

* chore: add the initial contract after few optimizations

* fix: lint

* chore: more optimizations

* chore: adorno

* fix: test getCommitments too

* chore: add kats

* fix: adorno

* fix: install deps before running build

* ci: add pnpm install before other jobs

* fix: remove magic number

* fix: remove unused errors,events and hardcode depth

.gas-snapshot

@@ -1 +1 @@
-FooTest:test_Example() (gas: 8662)
+WakuRlnV2Test:test__ValidRegistration() (gas: 108661)

.github/workflows/ci.yml

@@ -59,6 +59,20 @@ jobs:
       - name: "Install Foundry"
         uses: "foundry-rs/foundry-toolchain@v1"
 
+      - name: "Install Pnpm"
+        uses: "pnpm/action-setup@v2"
+        with:
+          version: "8"
+
+      - name: "Install Node.js"
+        uses: "actions/setup-node@v3"
+        with:
+          cache: "pnpm"
+          node-version: "lts/*"
+
+      - name: "Install the Node.js dependencies"
+        run: "pnpm install"
+
       - name: "Build the contracts and print their size"
         run: "forge build --sizes"
 
@@ -79,6 +93,20 @@ jobs:
       - name: "Install Foundry"
         uses: "foundry-rs/foundry-toolchain@v1"
 
+      - name: "Install Pnpm"
+        uses: "pnpm/action-setup@v2"
+        with:
+          version: "8"
+
+      - name: "Install Node.js"
+        uses: "actions/setup-node@v3"
+        with:
+          cache: "pnpm"
+          node-version: "lts/*"
+
+      - name: "Install the Node.js dependencies"
+        run: "pnpm install"
+
       - name: "Show the Foundry config"
         run: "forge config"
 
@@ -108,6 +136,20 @@ jobs:
       - name: "Install Foundry"
         uses: "foundry-rs/foundry-toolchain@v1"
 
+      - name: "Install Pnpm"
+        uses: "pnpm/action-setup@v2"
+        with:
+          version: "8"
+
+      - name: "Install Node.js"
+        uses: "actions/setup-node@v3"
+        with:
+          cache: "pnpm"
+          node-version: "lts/*"
+
+      - name: "Install the Node.js dependencies"
+        run: "pnpm install"
+
       - name: "Generate the coverage report using the unit and the integration tests"
         run: 'forge coverage --match-path "test/**/*.sol" --report lcov'
 

script/Deploy.s.sol

@@ -1,13 +1,13 @@
 // SPDX-License-Identifier: UNLICENSED
 pragma solidity >=0.8.19 <=0.9.0;
 
-import { Foo } from "../src/Foo.sol";
+import { WakuRlnV2 } from "../src/WakuRlnV2.sol";
 import { BaseScript } from "./Base.s.sol";
 import { DeploymentConfig } from "./DeploymentConfig.s.sol";
 
 contract Deploy is BaseScript {
-    function run() public returns (Foo foo, DeploymentConfig deploymentConfig) {
+    function run() public returns (WakuRlnV2 w, DeploymentConfig deploymentConfig) {
         deploymentConfig = new DeploymentConfig(broadcaster);
-        foo = new Foo();
+        w = new WakuRlnV2(20);
     }
 }

src/Foo.sol

@@ -1,8 +0,0 @@
-// SPDX-License-Identifier: UNLICENSED
-pragma solidity >=0.8.19;
-
-contract Foo {
-    function id(uint256 value) external pure returns (uint256) {
-        return value;
-    }
-}

src/WakuRlnV2.sol

@@ -0,0 +1,141 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity 0.8.19;
+
+import { LazyIMT, LazyIMTData } from "@zk-kit/imt.sol/LazyIMT.sol";
+import { PoseidonT3 } from "poseidon-solidity/PoseidonT3.sol";
+
+/// The tree is full
+error FullTree();
+
+/// Member is already registered
+error DuplicateIdCommitment();
+
+/// Invalid idCommitment
+error InvalidIdCommitment(uint256 idCommitment);
+
+/// Invalid userMessageLimit
+error InvalidUserMessageLimit(uint32 messageLimit);
+
+/// Invalid pagination query
+error InvalidPaginationQuery(uint256 startIndex, uint256 endIndex);
+
+contract WakuRlnV2 {
+    /// @notice The Field
+    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;
+
+    /// @notice The max message limit per epoch
+    uint32 public immutable MAX_MESSAGE_LIMIT;
+
+    /// @notice The depth of the merkle tree
+    uint8 public constant DEPTH = 20;
+
+    /// @notice The size of the merkle tree, i.e 2^depth
+    uint32 public immutable SET_SIZE;
+
+    /// @notice The index of the next member to be registered
+    uint32 public idCommitmentIndex = 0;
+
+    struct MembershipInfo {
+        /// @notice the user message limit of each member
+        uint32 userMessageLimit;
+        /// @notice the index of the member in the set
+        uint32 index;
+    }
+
+    /// @notice the member metadata
+    mapping(uint256 => MembershipInfo) public memberInfo;
+
+    /// @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, uint32 userMessageLimit, uint32 index);
+
+    modifier onlyValidIdCommitment(uint256 idCommitment) {
+        if (!isValidCommitment(idCommitment)) revert InvalidIdCommitment(idCommitment);
+        _;
+    }
+
+    modifier onlyValidUserMessageLimit(uint32 messageLimit) {
+        if (messageLimit > MAX_MESSAGE_LIMIT) revert InvalidUserMessageLimit(messageLimit);
+        if (messageLimit == 0) revert InvalidUserMessageLimit(messageLimit);
+        _;
+    }
+
+    constructor(uint32 maxMessageLimit) {
+        MAX_MESSAGE_LIMIT = maxMessageLimit;
+        SET_SIZE = uint32(1 << DEPTH);
+        deployedBlockNumber = uint32(block.number);
+        LazyIMT.init(imtData, DEPTH);
+    }
+
+    function memberExists(uint256 idCommitment) public view returns (bool) {
+        MembershipInfo memory member = memberInfo[idCommitment];
+        return member.userMessageLimit > 0 && member.index >= 0;
+    }
+
+    /// Allows a user to register as a member
+    /// @param idCommitment The idCommitment of the member
+    /// @param userMessageLimit The message limit of the member
+    function register(
+        uint256 idCommitment,
+        uint32 userMessageLimit
+    )
+        external
+        onlyValidIdCommitment(idCommitment)
+        onlyValidUserMessageLimit(userMessageLimit)
+    {
+        _register(idCommitment, userMessageLimit);
+    }
+
+    /// Registers a member
+    /// @param idCommitment The idCommitment of the member
+    /// @param userMessageLimit The message limit of the member
+    function _register(uint256 idCommitment, uint32 userMessageLimit) internal {
+        if (memberExists(idCommitment)) revert DuplicateIdCommitment();
+        if (idCommitmentIndex >= SET_SIZE) revert FullTree();
+
+        uint256 rateCommitment = PoseidonT3.hash([idCommitment, userMessageLimit]);
+        MembershipInfo memory member = MembershipInfo({ userMessageLimit: userMessageLimit, index: idCommitmentIndex });
+        LazyIMT.insert(imtData, rateCommitment);
+        memberInfo[idCommitment] = member;
+
+        emit MemberRegistered(idCommitment, userMessageLimit, idCommitmentIndex);
+        idCommitmentIndex += 1;
+    }
+
+    function isValidCommitment(uint256 idCommitment) public pure returns (bool) {
+        return idCommitment != 0 && idCommitment < Q;
+    }
+
+    function indexToCommitment(uint32 index) public view returns (uint256) {
+        return imtData.elements[LazyIMT.indexForElement(0, index)];
+    }
+
+    function getCommitments(uint32 startIndex, uint32 endIndex) public view returns (uint256[] memory) {
+        if (startIndex >= endIndex) revert InvalidPaginationQuery(startIndex, endIndex);
+        if (endIndex > idCommitmentIndex) revert InvalidPaginationQuery(startIndex, endIndex);
+
+        uint256[] memory commitments = new uint256[](endIndex - startIndex);
+        for (uint32 i = startIndex; i < endIndex; i++) {
+            commitments[i - startIndex] = indexToCommitment(i);
+        }
+        return commitments;
+    }
+
+    function root() external view returns (uint256) {
+        return LazyIMT.root(imtData, DEPTH);
+    }
+
+    function merkleProofElements(uint40 index) public view returns (uint256[] memory) {
+        return LazyIMT.merkleProofElements(imtData, index, DEPTH);
+    }
+}

test/Foo.t.sol

@@ -1,26 +0,0 @@
-// SPDX-License-Identifier: UNLICENSED
-pragma solidity >=0.8.19 <0.9.0;
-
-import { Test, console } from "forge-std/Test.sol";
-
-import { Deploy } from "../script/Deploy.s.sol";
-import { DeploymentConfig } from "../script/DeploymentConfig.s.sol";
-import { Foo } from "../src/Foo.sol";
-
-contract FooTest is Test {
-    Foo internal foo;
-    DeploymentConfig internal deploymentConfig;
-
-    address internal deployer;
-
-    function setUp() public virtual {
-        Deploy deployment = new Deploy();
-        (foo, deploymentConfig) = deployment.run();
-    }
-
-    function test_Example() external {
-        console.log("Hello World");
-        uint256 x = 42;
-        assertEq(foo.id(x), x, "value mismatch");
-    }
-}

test/WakuRlnV2.t.sol

@@ -0,0 +1,48 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity >=0.8.19 <0.9.0;
+
+import { Test, console } from "forge-std/Test.sol";
+
+import { Deploy } from "../script/Deploy.s.sol";
+import { DeploymentConfig } from "../script/DeploymentConfig.s.sol";
+import { WakuRlnV2 } from "../src/WakuRlnV2.sol";
+import { PoseidonT3 } from "poseidon-solidity/PoseidonT3.sol";
+import { LazyIMT } from "@zk-kit/imt.sol/LazyIMT.sol";
+
+contract WakuRlnV2Test is Test {
+    WakuRlnV2 internal w;
+    DeploymentConfig internal deploymentConfig;
+
+    address internal deployer;
+
+    function setUp() public virtual {
+        Deploy deployment = new Deploy();
+        (w, deploymentConfig) = deployment.run();
+    }
+
+    function test__ValidRegistration() external {
+        vm.pauseGasMetering();
+        uint256 idCommitment = 2;
+        uint32 userMessageLimit = 2;
+        vm.resumeGasMetering();
+        w.register(idCommitment, userMessageLimit);
+        vm.pauseGasMetering();
+        assertEq(w.idCommitmentIndex(), 1);
+        assertEq(w.memberExists(idCommitment), true);
+        (uint32 fetchedUserMessageLimit, uint32 index) = w.memberInfo(idCommitment);
+        assertEq(fetchedUserMessageLimit, userMessageLimit);
+        assertEq(index, 0);
+        // kats from zerokit
+        uint256 rateCommitment =
+            4_699_387_056_273_519_054_140_667_386_511_343_037_709_699_938_246_587_880_795_929_666_834_307_503_001;
+        assertEq(w.indexToCommitment(0), rateCommitment);
+        uint256[] memory commitments = w.getCommitments(0, 1);
+        assertEq(commitments.length, 1);
+        assertEq(commitments[index], rateCommitment);
+        assertEq(
+            w.root(),
+            13_801_897_483_540_040_307_162_267_952_866_411_686_127_372_014_953_358_983_481_592_640_000_001_877_295
+        );
+        vm.resumeGasMetering();
+    }
+}