fix(SortitionModule): sync totalStaked with stakedPnk

Author: unknownunknown1

contracts/src/arbitration/KlerosCore.sol

@@ -1370,7 +1370,6 @@ contract KlerosCore is IArbitratorV2, Initializable, UUPSProxiable {
             }
         }
         sortitionModule.setStake(_account, _courtID, pnkDeposit, pnkWithdrawal, _newStake);
-        sortitionModule.updateTotalStake(pnkDeposit, pnkWithdrawal);
 
         return true;
     }

contracts/src/arbitration/SortitionModule.sol

@@ -322,16 +322,6 @@ contract SortitionModule is ISortitionModule, Initializable, UUPSProxiable {
         return (pnkDeposit, pnkWithdrawal, StakingResult.Successful);
     }
 
-    /// @inheritdoc ISortitionModule
-    function updateTotalStake(uint256 _pnkDeposit, uint256 _pnkWithdrawal) external override onlyByCore {
-        // Note that we don't update totalStake in setStake() function because it doesn't always change total (e.g. during rewards/penalties).
-        if (_pnkDeposit > 0) {
-            totalStaked += _pnkDeposit;
-        } else {
-            totalStaked -= _pnkWithdrawal;
-        }
-    }
-
     /// @inheritdoc ISortitionModule
     function setStake(
         address _account,
@@ -399,8 +389,10 @@ contract SortitionModule is ISortitionModule, Initializable, UUPSProxiable {
             }
             // Increase juror's balance by deposited amount.
             juror.stakedPnk += _pnkDeposit;
+            totalStaked += _pnkDeposit;
         } else {
             juror.stakedPnk -= _pnkWithdrawal;
+            totalStaked -= _pnkWithdrawal;
             if (_newStake == 0) {
                 // Cleanup
                 for (uint256 i = juror.courtIDs.length; i > 0; i--) {
@@ -468,6 +460,7 @@ contract SortitionModule is ISortitionModule, Initializable, UUPSProxiable {
         uint256 amount = getJurorLeftoverPNK(_account);
         if (amount == 0) revert NotEligibleForWithdrawal();
         jurors[_account].stakedPnk = 0;
+        totalStaked -= amount;
         core.transferBySortitionModule(_account, amount);
         emit LeftoverPNKWithdrawn(_account, amount);
     }

contracts/src/arbitration/interfaces/ISortitionModule.sol

@@ -57,11 +57,6 @@ interface ISortitionModule {
         bool _noDelay
     ) external returns (uint256 pnkDeposit, uint256 pnkWithdrawal, StakingResult stakingResult);
 
-    /// @notice Updates the total amount staked in all courts.
-    /// @param _pnkDeposit The amount of PNK that increases total stake.
-    /// @param _pnkWithdrawal The amount of PNK that decreases total stake.
-    function updateTotalStake(uint256 _pnkDeposit, uint256 _pnkWithdrawal) external;
-
     /// @notice Update the state of the stakes, called by KC at the end of setStake flow.
     ///
     /// @dev `O(n + p * log_k(j))` where

contracts/src/arbitration/university/SortitionModuleUniversity.sol

@@ -130,11 +130,6 @@ contract SortitionModuleUniversity is ISortitionModuleUniversity, UUPSProxiable,
         // NOP
     }
 
-    /// @inheritdoc ISortitionModule
-    function updateTotalStake(uint256 _pnkDeposit, uint256 _pnkWithdrawal) external override onlyByCore {
-        // NOP
-    }
-
     /// @inheritdoc ISortitionModule
     function createDisputeHook(uint256 /*_disputeID*/, uint256 /*_roundID*/) external override onlyByCore {
         disputesWithoutJurors++;

contracts/test/foundry/KlerosCore_Execution.t.sol

@@ -98,6 +98,9 @@ contract KlerosCore_ExecutionTest is KlerosCore_TestBase {
             "Wrong penalty coherence 2 vote ID"
         );
 
+        assertEq(pinakion.balanceOf(address(core)), 22000, "Wrong token balance of the core");
+        assertEq(sortitionModule.totalStaked(), 22000, "Total staked should be equal to the balance in this test");
+
         vm.expectEmit(true, true, true, true);
         emit SortitionModule.StakeLocked(staker1, 1000, true);
         vm.expectEmit(true, true, true, true);
@@ -143,9 +146,11 @@ contract KlerosCore_ExecutionTest is KlerosCore_TestBase {
         assertEq(staker1.balance, 0, "Wrong balance of the staker1");
         assertEq(staker2.balance, 0.09 ether, "Wrong balance of the staker2");
 
-        assertEq(pinakion.balanceOf(address(core)), 22000, "Wrong token balance of the core"); // Was 21500. 1000 was transferred to staker2
+        assertEq(pinakion.balanceOf(address(core)), 22000, "Token balance of the core shouldn't change after rewards");
+        assertEq(sortitionModule.totalStaked(), 22000, "Total staked shouldn't change after rewards");
+
         assertEq(pinakion.balanceOf(staker1), 999999999999998000, "Wrong token balance of staker1");
-        assertEq(pinakion.balanceOf(staker2), 999999999999980000, "Wrong token balance of staker2"); // 20k stake and 1k added as a reward, thus -19k from the default
+        assertEq(pinakion.balanceOf(staker2), 999999999999980000, "Wrong token balance of staker2");
     }
 
     function test_execute_NoCoherence() public {
@@ -479,13 +484,19 @@ contract KlerosCore_ExecutionTest is KlerosCore_TestBase {
         vm.prank(owner);
         core.transferBySortitionModule(staker1, 1000);
 
+        assertEq(pinakion.balanceOf(address(core)), 1000, "Wrong token balance of the core");
+        assertEq(sortitionModule.totalStaked(), 1000, "Wrong totalStaked before withdrawal");
+
         vm.expectEmit(true, true, true, true);
         emit SortitionModule.LeftoverPNKWithdrawn(staker1, 1000);
         sortitionModule.withdrawLeftoverPNK(staker1);
 
         (totalStaked, , , ) = sortitionModule.getJurorBalance(staker1, GENERAL_COURT);
         assertEq(totalStaked, 0, "Should be unstaked fully");
 
+        assertEq(pinakion.balanceOf(address(core)), 0, "Wrong token balance of the core");
+        assertEq(sortitionModule.totalStaked(), 0, "Wrong totalStaked after withdrawal");
+
         // Check that everything is withdrawn now
         assertEq(pinakion.balanceOf(address(core)), 0, "Core balance should be empty");
         assertEq(pinakion.balanceOf(staker1), 1 ether, "All PNK should be withdrawn");