fix(SortitionModule): total stake update

Author: unknownunknown1

contracts/src/arbitration/KlerosCore.sol

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

contracts/src/arbitration/SortitionModule.sol

@@ -311,19 +311,27 @@ contract SortitionModule is ISortitionModule, Initializable, UUPSProxiable {
         // Current phase is Staking: set stakes.
         if (stakeIncrease) {
             pnkDeposit = stakeChange;
-            totalStaked += stakeChange;
         } else {
             pnkWithdrawal = stakeChange;
             uint256 possibleWithdrawal = juror.stakedPnk > juror.lockedPnk ? juror.stakedPnk - juror.lockedPnk : 0;
             if (pnkWithdrawal > possibleWithdrawal) {
                 // Ensure locked tokens remain in the contract. They can only be released during Execution.
                 pnkWithdrawal = possibleWithdrawal;
             }
-            totalStaked -= pnkWithdrawal;
         }
         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,

contracts/src/arbitration/interfaces/ISortitionModule.sol

@@ -57,6 +57,11 @@ 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,6 +130,11 @@ 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_Staking.t.sol

@@ -141,6 +141,32 @@ contract KlerosCore_StakingTest is KlerosCore_TestBase {
         );
     }
 
+    function test_setStake_totalStaked() public {
+        // Increase
+        vm.prank(staker1);
+        core.setStake(GENERAL_COURT, 4000);
+        vm.prank(staker1);
+        core.setStake(GENERAL_COURT, 5001);
+        vm.prank(staker2);
+        core.setStake(GENERAL_COURT, 1000);
+        vm.prank(staker2);
+        core.setStake(GENERAL_COURT, 1500);
+
+        assertEq(sortitionModule.totalStaked(), 6501, "Wrong totalStaked");
+
+        // Decrease
+        vm.prank(staker1);
+        core.setStake(GENERAL_COURT, 3000);
+        vm.prank(staker1);
+        core.setStake(GENERAL_COURT, 2500);
+        vm.prank(staker2);
+        core.setStake(GENERAL_COURT, 1400);
+        vm.prank(staker2);
+        core.setStake(GENERAL_COURT, 1200);
+
+        assertEq(sortitionModule.totalStaked(), 3700, "Wrong totalStaked");
+    }
+
     function test_setStake_maxStakePathCheck() public {
         uint256[] memory supportedDK = new uint256[](1);
         supportedDK[0] = DISPUTE_KIT_CLASSIC;