Add remove delegations function to proxy batch on mainnet and mainnet…

… delegation on base (#383)

contracts/scripts/deployments/DeployProxyBatchDelegation.s.sol

@@ -57,6 +57,14 @@ contract DeployProxyBatchDelegation is Deployer {
       mainnetDelegation = _getMainnetDelegation();
     }
 
+    if (riverToken == address(0)) {
+      revert("DeployProxyBatchDelegation: River token not deployed");
+    }
+
+    if (claimers == address(0)) {
+      revert("DeployProxyBatchDelegation: Claimers not deployed");
+    }
+
     vm.broadcast(deployer);
     return
       address(

contracts/scripts/deployments/facets/DeployMainnetDelegation.s.sol

@@ -23,6 +23,7 @@ contract DeployMainnetDelegation is FacetHelper, Deployer {
     addSelector(MainnetDelegation.setBatchAuthorizedClaimers.selector);
     addSelector(MainnetDelegation.getProxyDelegation.selector);
     addSelector(MainnetDelegation.getMessenger.selector);
+    addSelector(MainnetDelegation.removeDelegations.selector);
   }
 
   function initializer() public pure override returns (bytes4) {

contracts/src/tokens/river/base/delegation/IMainnetDelegation.sol

@@ -81,6 +81,12 @@ interface IMainnetDelegation is IMainnetDelegationBase {
     uint256 quantity
   ) external;
 
+  /**
+   * @notice Remove delegation of a delegator
+   * @param delegators The delegator address
+   */
+  function removeDelegations(address[] memory delegators) external;
+
   /**
    * @notice Get delegation of a delegator
    * @param delegator The delegator address

contracts/src/tokens/river/base/delegation/MainnetDelegation.sol

@@ -57,6 +57,15 @@ contract MainnetDelegation is
     return address(_getProxyDelegation());
   }
 
+  // =============================================================
+  //                       Remove Delegators
+  // =============================================================
+  function removeDelegations(
+    address[] calldata delegators
+  ) external onlyCrossDomainMessenger {
+    _removeDelegations(delegators);
+  }
+
   // =============================================================
   //                  Batch Authorized Claimers
   // =============================================================
@@ -80,6 +89,7 @@ contract MainnetDelegation is
     uint256[] calldata quantities
   ) external onlyCrossDomainMessenger {
     uint256 delegatorsLen = delegators.length;
+
     for (uint256 i; i < delegatorsLen; i++) {
       _replaceDelegation(
         delegators[i],
@@ -98,6 +108,7 @@ contract MainnetDelegation is
   }
 
   /// @inheritdoc IMainnetDelegation
+  /// @notice deprecated
   function setDelegation(
     address delegator,
     address operator,

contracts/src/tokens/river/base/delegation/MainnetDelegationBase.sol

@@ -14,6 +14,25 @@ import {MainnetDelegationStorage} from "./MainnetDelegationStorage.sol";
 abstract contract MainnetDelegationBase is IMainnetDelegationBase {
   using EnumerableSet for EnumerableSet.AddressSet;
 
+  function _removeDelegations(address[] memory delegators) internal {
+    MainnetDelegationStorage.Layout storage ds = MainnetDelegationStorage
+      .layout();
+
+    for (uint256 i = 0; i < delegators.length; i++) {
+      address delegator = delegators[i];
+      address operator = ds.delegationByDelegator[delegator].operator;
+
+      ds.delegatorsByOperator[operator].remove(delegator);
+      delete ds.delegationByDelegator[delegator];
+      delete ds.delegationByDelegator[delegator].operator;
+      delete ds.delegationByDelegator[delegator].quantity;
+      delete ds.delegationByDelegator[delegator].delegationTime;
+      delete ds.delegationByDelegator[delegator].delegator;
+
+      ds.delegators.remove(delegator);
+    }
+  }
+
   function _replaceDelegation(
     address delegator,
     address claimer,
@@ -22,6 +41,10 @@ abstract contract MainnetDelegationBase is IMainnetDelegationBase {
   ) internal {
     MainnetDelegationStorage.Layout storage ds = MainnetDelegationStorage
       .layout();
+
+    // Remove the current delegator
+    ds.delegators.remove(delegator);
+
     Delegation storage delegation = ds.delegationByDelegator[delegator];
     address currentClaimer = ds.claimerByDelegator[delegator];
 
@@ -37,13 +60,16 @@ abstract contract MainnetDelegationBase is IMainnetDelegationBase {
     }
 
     // Set the new delegation
+    ds.delegators.add(delegator);
     ds.delegatorsByOperator[operator].add(delegator);
-    ds.delegationByDelegator[delegator] = Delegation(
-      operator,
-      quantity,
-      delegator,
-      block.timestamp
-    );
+
+    if (delegation.operator != operator || delegation.quantity != quantity) {
+      delegation.delegationTime = block.timestamp;
+    }
+
+    delegation.operator = operator;
+    delegation.quantity = quantity;
+    delegation.delegator = delegator;
 
     // Update the claimer if it has changed
     if (claimer != currentClaimer) {

contracts/src/tokens/river/base/delegation/MainnetDelegationStorage.sol

@@ -24,6 +24,7 @@ library MainnetDelegationStorage {
     ICrossDomainMessenger messenger;
     mapping(address claimer => EnumerableSet.AddressSet delegators) delegatorsByAuthorizedClaimer;
     address proxyDelegation;
+    EnumerableSet.AddressSet delegators;
   }
 
   function layout() internal pure returns (Layout storage s) {

contracts/src/tokens/river/mainnet/delegation/ProxyBatchDelegation.sol

@@ -32,6 +32,21 @@ contract ProxyBatchDelegation is IProxyBatchDelegation {
     TARGET = _target;
   }
 
+  function removeDelegators() external {
+    address[] memory delegators = new address[](2);
+    delegators[0] = 0x204f1aA5B666d0eAc07228D3065a461e92AC399c;
+    delegators[1] = 0x3541F646d321CACbc0fF4A7cCcB583E8B6E413da;
+
+    ICrossDomainMessenger(MESSENGER).sendMessage(
+      TARGET,
+      abi.encodeWithSelector(
+        IMainnetDelegation.removeDelegations.selector,
+        delegators
+      ),
+      200_000
+    );
+  }
+
   function sendAuthorizedClaimers() external {
     address[] memory delegators = rvr.getDelegators();
     address[] memory authorizedClaimers = new address[](delegators.length);

contracts/test/fork/ForkMainnetDelegation.t.sol

@@ -5,6 +5,7 @@ pragma solidity ^0.8.19;
 import {TestUtils} from "contracts/test/utils/TestUtils.sol";
 
 //interfaces
+import {IMainnetDelegationBase} from "contracts/src/tokens/river/base/delegation/IMainnetDelegation.sol";
 
 //libraries
 
@@ -13,7 +14,7 @@ import {MockMessenger} from "contracts/test/mocks/MockMessenger.sol";
 import {MainnetDelegation} from "contracts/src/tokens/river/base/delegation/MainnetDelegation.sol";
 
 // Base
-contract ForkMainnetDelegationTest is TestUtils {
+contract ForkMainnetDelegationTest is TestUtils, IMainnetDelegationBase {
   address baseRegistry = 0x7c0422b31401C936172C897802CF0373B35B7698;
 
   MainnetDelegation internal mainnetDelegation;
@@ -22,6 +23,38 @@ contract ForkMainnetDelegationTest is TestUtils {
     mainnetDelegation = MainnetDelegation(baseRegistry);
   }
 
+  function test_removeDelegators() external onlyForked {
+    address getMessenger = mainnetDelegation.getMessenger();
+    address getProxyDelegation = mainnetDelegation.getProxyDelegation();
+
+    MockMessenger mockMessenger = new MockMessenger();
+    vm.etch(getMessenger, address(mockMessenger).code);
+    MockMessenger(getMessenger).setXDomainMessageSender(getProxyDelegation);
+
+    address delegatorAddress = 0x204f1aA5B666d0eAc07228D3065a461e92AC399c;
+
+    Delegation memory delegator = mainnetDelegation.getDelegationByDelegator(
+      delegatorAddress
+    );
+
+    assertEq(delegator.delegator, delegatorAddress);
+
+    vm.prank(address(getMessenger));
+    (bool success, ) = baseRegistry.call{gas: 200_000}(
+      // solhint-disable-next-line max-line-length
+      hex"012ad9da00000000000000000000000000000000000000000000000000000000000000200000000000000000000000000000000000000000000000000000000000000002000000000000000000000000204f1aa5b666d0eac07228d3065a461e92ac399c0000000000000000000000003541f646d321cacbc0ff4a7cccb583e8b6e413da"
+    );
+
+    assertTrue(success);
+
+    delegator = mainnetDelegation.getDelegationByDelegator(delegatorAddress);
+
+    assertEq(delegator.operator, address(0));
+    assertEq(delegator.quantity, 0);
+    assertEq(delegator.delegationTime, 0);
+    assertEq(delegator.delegator, address(0));
+  }
+
   function test_setBatchAuthorizedClaimers() external onlyForked {
     address getMessenger = mainnetDelegation.getMessenger();
     address getProxyDelegation = mainnetDelegation.getProxyDelegation();

contracts/test/fork/ForkProxyBatchDelegation.t.sol

@@ -33,4 +33,9 @@ contract ForkProxyBatchDelegationTest is TestUtils {
     vm.prank(_randomAddress());
     proxyBatchDelegation.sendAuthorizedClaimers();
   }
+
+  function test_removeDelegators() external onlyForked {
+    vm.prank(_randomAddress());
+    proxyBatchDelegation.removeDelegators();
+  }
 }

contracts/test/tokens/delegation/ProxyBatchDelegation.t.sol

@@ -113,6 +113,50 @@ contract ProxyBatchDelegationTest is BaseSetup, IMainnetDelegationBase {
     );
   }
 
+  function test_removeDelegators() external {
+    address operator = 0x09285F179a9bA06CEBA12DeCd1755Ac6942A8cf4;
+
+    address[] memory delegators = new address[](2);
+    delegators[0] = 0x204f1aA5B666d0eAc07228D3065a461e92AC399c;
+    delegators[1] = 0x3541F646d321CACbc0fF4A7cCcB583E8B6E413da;
+
+    // given delegators have tokens
+    vm.startPrank(vault);
+    rvr.transfer(delegators[0], tokens);
+    rvr.transfer(delegators[1], tokens);
+    vm.stopPrank();
+
+    // given delegators have delegated tokens
+    vm.prank(delegators[0]);
+    rvr.delegate(operator);
+
+    vm.prank(delegators[1]);
+    rvr.delegate(operator);
+
+    vm.prank(_randomAddress());
+    proxyDelegation.sendDelegators();
+
+    Delegation memory delegator = delegation.getDelegationByDelegator(
+      delegators[0]
+    );
+
+    assertEq(rvr.delegates(delegators[0]), delegator.operator);
+
+    vm.prank(_randomAddress());
+    proxyDelegation.removeDelegators();
+
+    delegator = delegation.getDelegationByDelegator(delegators[0]);
+
+    assertEq(delegator.operator, address(0));
+    assertEq(delegator.quantity, 0);
+    assertEq(delegator.delegator, address(0));
+
+    Delegation[] memory mainneDelegators = delegation
+      .getMainnetDelegationsByOperator(operator);
+
+    assertEq(mainneDelegators.length, 0);
+  }
+
   function _getRandomValue(
     address[] memory addresses
   ) internal view returns (address) {

packages/generated/addresses/ethereum_mainnet/authorizedClaimers.json

@@ -0,0 +1 @@
+{"address": "0x0bEe55b52d01C4D5d4D0cfcE1d6e0baE6722db05"}