From 7827e719c7475641cc334481d3395b770bb02ddc Mon Sep 17 00:00:00 2001
From: rymnc <43716372+rymnc@users.noreply.github.com>
Date: Thu, 3 Aug 2023 11:30:43 +0530
Subject: [PATCH] feat: waku rln registry contract to enable autosharding

---
 contracts/WakuRln.sol         |  6 ++-
 contracts/WakuRlnRegistry.sol | 57 +++++++++++++++++++++++
 docs/index.md                 | 88 ++++++++++++++++++++++++++++++++++-
 test/WakuRln.t.sol            |  2 +-
 4 files changed, 150 insertions(+), 3 deletions(-)
 create mode 100644 contracts/WakuRlnRegistry.sol

diff --git a/contracts/WakuRln.sol b/contracts/WakuRln.sol
index 8ce7891..765e408 100644
--- a/contracts/WakuRln.sol
+++ b/contracts/WakuRln.sol
@@ -8,7 +8,11 @@ import {Ownable} from "openzeppelin-contracts/contracts/access/Ownable.sol";
 error NotImplemented();
 
 contract WakuRln is Ownable, RlnBase {
-    constructor(address _poseidonHasher) Ownable() RlnBase(0, 20, _poseidonHasher, address(0)) {}
+    uint16 public immutable contractIndex;
+
+    constructor(address _poseidonHasher, uint16 _contractIndex) Ownable() RlnBase(0, 20, _poseidonHasher, address(0)) {
+        contractIndex = _contractIndex;
+    }
 
     /// Registers a member
     /// @param idCommitment The idCommitment of the member
diff --git a/contracts/WakuRlnRegistry.sol b/contracts/WakuRlnRegistry.sol
new file mode 100644
index 0000000..595a105
--- /dev/null
+++ b/contracts/WakuRlnRegistry.sol
@@ -0,0 +1,57 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.15;
+
+import {WakuRln} from "./WakuRln.sol";
+import {IPoseidonHasher} from "rln-contract/PoseidonHasher.sol";
+import {Ownable} from "openzeppelin-contracts/contracts/access/Ownable.sol";
+
+error StorageAlreadyExists(address storageAddress);
+error NoStorageContractAvailable();
+error FailedToRegister(string reason);
+
+contract WakuRlnRegistry is Ownable {
+    uint16 public currentStorageIndex;
+    mapping(uint16 => address) public storages;
+
+    uint16 public usingStorageIndex = 0;
+
+    IPoseidonHasher public immutable poseidonHasher;
+
+    event NewStorageContract(uint16 index, address storageAddress);
+
+    constructor(address _poseidonHasher) Ownable() {
+        poseidonHasher = IPoseidonHasher(_poseidonHasher);
+    }
+
+    function _insertIntoStorageMap(address storageAddress) internal {
+        storages[currentStorageIndex] = storageAddress;
+        emit NewStorageContract(currentStorageIndex, storageAddress);
+        currentStorageIndex += 1;
+    }
+
+    function registerStorage(address storageAddress) external onlyOwner {
+        if (storages[currentStorageIndex] != address(0)) revert StorageAlreadyExists(storageAddress);
+        _insertIntoStorageMap(storageAddress);
+    }
+
+    function newStorage() external onlyOwner {
+        WakuRln newStorageContract = new WakuRln(address(poseidonHasher), currentStorageIndex);
+        _insertIntoStorageMap(address(newStorageContract));
+    }
+
+    function register(uint256 commitment) external payable {
+        if (usingStorageIndex >= currentStorageIndex) revert NoStorageContractAvailable();
+
+        // iteratively check if the storage contract is full, and increment the usingStorageIndex if it is
+        while (true) {
+            try WakuRln(storages[usingStorageIndex]).register{value: msg.value}(commitment) {
+                break;
+            } catch Error(string memory reason) {
+                if (keccak256(abi.encodePacked(reason)) != keccak256(abi.encodePacked("FullTree()"))) {
+                    revert FailedToRegister(reason);
+                }
+                usingStorageIndex += 1;
+            }
+        }
+    }
+}
diff --git a/docs/index.md b/docs/index.md
index 042c5b2..4772849 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -8,10 +8,16 @@ error NotImplemented()
 
 ## WakuRln
 
+### contractIndex
+
+```solidity
+uint16 contractIndex
+```
+
 ### constructor
 
 ```solidity
-constructor(address _poseidonHasher) public
+constructor(address _poseidonHasher, uint16 _contractIndex) public
 ```
 
 ### \_register
@@ -85,3 +91,83 @@ function withdraw() external pure
 ```
 
 Allows a user to withdraw funds allocated to them upon slashing a member
+
+## StorageAlreadyExists
+
+```solidity
+error StorageAlreadyExists(address storageAddress)
+```
+
+## NoStorageContractAvailable
+
+```solidity
+error NoStorageContractAvailable()
+```
+
+## FailedToRegister
+
+```solidity
+error FailedToRegister(string reason)
+```
+
+## WakuRlnRegistry
+
+### currentStorageIndex
+
+```solidity
+uint16 currentStorageIndex
+```
+
+### storages
+
+```solidity
+mapping(uint16 => address) storages
+```
+
+### usingStorageIndex
+
+```solidity
+uint16 usingStorageIndex
+```
+
+### poseidonHasher
+
+```solidity
+contract IPoseidonHasher poseidonHasher
+```
+
+### NewStorageContract
+
+```solidity
+event NewStorageContract(uint16 index, address storageAddress)
+```
+
+### constructor
+
+```solidity
+constructor(address _poseidonHasher) public
+```
+
+### \_insertIntoStorageMap
+
+```solidity
+function _insertIntoStorageMap(address storageAddress) internal
+```
+
+### registerStorage
+
+```solidity
+function registerStorage(address storageAddress) external
+```
+
+### newStorage
+
+```solidity
+function newStorage() external
+```
+
+### register
+
+```solidity
+function register(uint256 commitment) external payable
+```
diff --git a/test/WakuRln.t.sol b/test/WakuRln.t.sol
index 27309af..4ac0ffb 100644
--- a/test/WakuRln.t.sol
+++ b/test/WakuRln.t.sol
@@ -33,7 +33,7 @@ contract WakuRlnTest is Test {
     /// @dev Setup the testing environment.
     function setUp() public {
         poseidon = new PoseidonHasher();
-        wakuRln = new WakuRln(address(poseidon));
+        wakuRln = new WakuRln(address(poseidon), 0);
     }
 
     /// @dev Ensure that you can hash a value.