Finish epoch processing split

packages/protocol/contracts-0.8/common/EpochManager.sol

@@ -60,6 +60,8 @@ contract EpochManager is
   mapping(uint256 => Epoch) internal epochs;
   mapping(address => uint256) public validatorPendingPayments;
 
+  uint256 public toProcessGroups = 0;
+
   /**
    * @notice Event emited when epochProcessing has begun.
    * @param epochNumber The epoch number that is being processed.
@@ -201,6 +203,69 @@ contract EpochManager is
     emit EpochProcessingStarted(currentEpochNumber);
   }
 
+  function setToProcessGroups() external {
+    require(isOnEpochProcess(), "Epoch process is not started");
+
+    EpochProcessState storage _epochProcessing = epochProcessing;
+
+    IValidators validators = getValidators();
+    IElection election = getElection();
+    IScoreReader scoreReader = getScoreReader();
+    for (uint i = 0; i < elected.length; i++) {
+      address group = validators.getValidatorsGroup(elected[i]);
+      if (processedGroups[group] == 0) {
+        toProcessGroups++;
+        uint256 groupScore = scoreReader.getGroupScore(group);
+        // We need to precompute epoch rewards for each group since computation depends on total active votes for all groups.
+        uint256 epochRewards = election.getGroupEpochRewardsBasedOnScore(
+          group,
+          _epochProcessing.totalRewardsVoter,
+          groupScore
+        );
+        processedGroups[group] = epochRewards == 0 ? type(uint256).max : epochRewards;
+      }
+    }
+  }
+
+  function processGroup(address group, address lesser, address greater) external {
+    require(isOnEpochProcess(), "Epoch process is not started");
+    require(toProcessGroups > 0, "no more groups to process");
+
+    uint256 epochRewards = processedGroups[group];
+    // checks that group is actually from elected group
+    require(epochRewards > 0, "group not from current elected set");
+    IElection election = getElection();
+
+    if (epochRewards != type(uint256).max) {
+      election.distributeEpochRewards(group, epochRewards, lesser, greater);
+    }
+
+    delete processedGroups[group];
+    toProcessGroups--;
+
+    if (toProcessGroups == 0) {
+      getCeloUnreleasedTreasury().release(
+        registry.getAddressForOrDie(GOVERNANCE_REGISTRY_ID),
+        epochProcessing.totalRewardsCommunity
+      );
+      getCeloUnreleasedTreasury().release(
+        getEpochRewards().carbonOffsettingPartner(),
+        epochProcessing.totalRewardsCarbonFund
+      );
+      // run elections
+      elected = election.electValidatorAccounts();
+      epochProcessing.status = EpochProcessStatus.NotStarted;
+
+      // finalize epoch
+      // last block should be the block before and timestamp from previous block
+      epochs[currentEpochNumber].lastBlock = block.number - 1;
+      // start new epoch
+      currentEpochNumber++;
+      epochs[currentEpochNumber].firstBlock = block.number;
+      epochs[currentEpochNumber].startTimestamp = block.timestamp;
+    }
+  }
+
   /**
    * @notice Finishes processing an epoch and releasing funds to the beneficiaries.
    * @param groups List of validator groups to be processed.
@@ -213,6 +278,7 @@ contract EpochManager is
     address[] calldata greaters
   ) external virtual nonReentrant {
     require(isOnEpochProcess(), "Epoch process is not started");
+    require(toProcessGroups == 0, "not all groups processed");
     // finalize epoch
     // last block should be the block before and timestamp from previous block
     epochs[currentEpochNumber].lastBlock = block.number - 1;
@@ -223,14 +289,14 @@ contract EpochManager is
 
     EpochProcessState storage _epochProcessing = epochProcessing;
 
-    uint256 toProcessGroups = 0;
+    uint256 _toProcessGroups = 0;
     IValidators validators = getValidators();
     IElection election = getElection();
     IScoreReader scoreReader = getScoreReader();
     for (uint i = 0; i < elected.length; i++) {
       address group = validators.getValidatorsGroup(elected[i]);
       if (processedGroups[group] == 0) {
-        toProcessGroups++;
+        _toProcessGroups++;
         uint256 groupScore = scoreReader.getGroupScore(group);
         // We need to precompute epoch rewards for each group since computation depends on total active votes for all groups.
         uint256 epochRewards = election.getGroupEpochRewardsBasedOnScore(
@@ -243,7 +309,7 @@ contract EpochManager is
       delete elected[i];
     }
 
-    require(toProcessGroups == groups.length, "number of groups does not match");
+    require(_toProcessGroups == groups.length, "number of groups does not match");
 
     for (uint i = 0; i < groups.length; i++) {
       uint256 epochRewards = processedGroups[groups[i]];

packages/protocol/contracts-0.8/common/test/MockEpochManager.sol

@@ -60,6 +60,9 @@ contract MockEpochManager is IEpochManager {
     address[] calldata greaters
   ) external {}
 
+  function setToProcessGroups() external {}
+  function processGroup(address group, address lesser, address greater) external {}
+
   function setIsTimeForNextEpoch(bool _isTime) external {
     _isTimeForNextEpoch = _isTime;
   }

packages/protocol/contracts/common/interfaces/IEpochManager.sol

@@ -13,6 +13,8 @@ interface IEpochManager {
     address[] calldata lessers,
     address[] calldata greaters
   ) external;
+  function setToProcessGroups() external;
+  function processGroup(address group, address lesser, address greater) external;
   function sendValidatorPayment(address) external;
   function getCurrentEpoch() external view returns (uint256, uint256, uint256, uint256);
   function getCurrentEpochNumber() external view returns (uint256);

packages/protocol/test-sol/devchain/e2e/common/EpochManager.t.sol

@@ -745,3 +745,143 @@ contract E2E_GasTest2_FinishNextEpochProcess is E2E_GasTest_Setup {
     console.log("elected count2: ", epochManager.getElected().length);
   }
 }
+
+contract E2E_FinishNextEpochProcess_Split is E2E_GasTest_Setup {
+  using EnumerableSet for EnumerableSet.AddressSet;
+
+  function setUp() public override {
+    super.setUp();
+
+    activateValidators();
+    whenL2(vm);
+
+    vm.prank(epochManagerEnabler);
+    epochManager.initializeSystem(1, 1, firstElected);
+
+    validatorsArray = getValidators().getRegisteredValidators();
+    groups = getValidators().getRegisteredValidatorGroups();
+
+    vm.startPrank(scoreManager.owner());
+    scoreManager.setGroupScore(groups[0], groupScore[0]);
+    scoreManager.setGroupScore(groups[1], groupScore[1]);
+    scoreManager.setGroupScore(groups[2], groupScore[2]);
+
+    scoreManager.setValidatorScore(validatorsArray[0], validatorScore[0]);
+    scoreManager.setValidatorScore(validatorsArray[1], validatorScore[1]);
+    scoreManager.setValidatorScore(validatorsArray[2], validatorScore[2]);
+    scoreManager.setValidatorScore(validatorsArray[3], validatorScore[3]);
+    scoreManager.setValidatorScore(validatorsArray[4], validatorScore[4]);
+    scoreManager.setValidatorScore(validatorsArray[5], validatorScore[5]);
+
+    vm.stopPrank();
+
+    timeTravel(vm, epochDuration + 1);
+    epochManager.startNextEpochProcess();
+
+    address[] memory lessers;
+    address[] memory greaters;
+    address[] memory groupsEligible;
+    GroupWithVotes[] memory groupWithVotes;
+    uint256[] memory groupActiveBalances;
+    (lessers, greaters, groupWithVotes) = getLessersAndGreaters(groups);
+
+    uint256 currentEpoch = epochManager.getCurrentEpochNumber();
+    address[] memory currentlyElected = epochManager.getElected();
+    for (uint256 i = 0; i < currentlyElected.length; i++) {
+      originalyElected.add(currentlyElected[i]);
+    }
+
+    // wait some time before finishing
+    timeTravel(vm, epochDuration / 2);
+    blockTravel(vm, 100);
+
+    epochManager.setToProcessGroups();
+    for (uint256 i = 0; i < groups.length; i++) {
+      epochManager.processGroup(groups[i], lessers[i], greaters[i]);
+    }
+
+    assertEq(currentEpoch + 1, epochManager.getCurrentEpochNumber());
+
+    address[] memory newlyElected = epochManager.getElected();
+
+    for (uint256 i = 0; i < currentlyElected.length; i++) {
+      assertEq(originalyElected.contains(currentlyElected[i]), true);
+    }
+
+    timeTravel(vm, epochDuration + 1);
+    epochManager.startNextEpochProcess();
+
+    // wait some time before finishing
+    timeTravel(vm, epochDuration / 2);
+    blockTravel(vm, 100);
+
+    (lessers, greaters, groupWithVotes) = getLessersAndGreaters(groups);
+    epochManager.setToProcessGroups();
+    for (uint256 i = 0; i < groups.length; i++) {
+      epochManager.processGroup(groups[i], lessers[i], greaters[i]);
+    }
+    // epochManager.finishNextEpochProcess(groups, lessers, greaters);
+    assertGroupWithVotes(groupWithVotes);
+
+    assertEq(currentEpoch + 2, epochManager.getCurrentEpochNumber());
+
+    address[] memory newlyElected2 = epochManager.getElected();
+
+    for (uint256 i = 0; i < currentlyElected.length; i++) {
+      assertEq(originalyElected.contains(newlyElected2[i]), true);
+    }
+    uint256 validatorGroupCount = 60;
+    uint256 validatorPerGroupCount = 2;
+
+    for (uint256 i = 0; i < validatorGroupCount; i++) {
+      (address newValidatorGroup, address newValidator) = registerNewValidatorGroupWithValidator(
+        i,
+        validatorPerGroupCount
+      );
+    }
+
+    timeTravel(vm, epochDuration + 1);
+    epochManager.startNextEpochProcess();
+
+    timeTravel(vm, epochDuration / 2);
+    blockTravel(vm, 100);
+
+    (lessers, greaters, groupWithVotes) = getLessersAndGreaters(groups);
+
+    epochManager.setToProcessGroups();
+    for (uint256 i = 0; i < groups.length; i++) {
+      epochManager.processGroup(groups[i], lessers[i], greaters[i]);
+    }
+
+    activateValidators();
+
+    timeTravel(vm, epochDuration + 1);
+    epochManager.startNextEpochProcess();
+
+    groups = getCurrentlyElectedGroups();
+
+    timeTravel(vm, epochDuration / 2);
+    blockTravel(vm, 100);
+  }
+
+  /**
+    * @notice Test the gas used by finishNextEpochProcess
+    This test is trying to measure gas used by finishNextEpochProcess in a real life worst case. We have 126 validators and 123 groups.
+    There are two main loops in the function, one for calculating rewards and the other for updating the elected validators.
+    FinishNextEpochProcess is called twice, first time with going from 6 -> 110 validators which consumes approx. 6M gas and the second time with going from 110 -> 110 validators which consumes approx. 19M gas. 
+     */
+  function test_shouldFinishNextEpochProcessing_GasTest_Split() public {
+    address[] memory lessers;
+    address[] memory greaters;
+    GroupWithVotes[] memory groupWithVotes;
+    (lessers, greaters, groupWithVotes) = getLessersAndGreaters(groups);
+    epochManager.setToProcessGroups();
+
+    for (uint256 i = 0; i < groups.length; i++) {
+      uint256 gasLeftBefore1 = gasleft();
+      epochManager.processGroup(groups[i], lessers[i], greaters[i]);
+      uint256 gasLeftAfter1 = gasleft();
+      console.log("processGroup gas used: ", gasLeftBefore1 - gasLeftAfter1);
+    }
+  }
+}