refactor: simplification by moving KC._deleteDelayedStake() inside SM.preStakeHook()

Author: jaybuidl

contracts/src/arbitration/KlerosCore.sol

@@ -464,7 +464,6 @@ contract KlerosCore is IArbitratorV2, UUPSProxiable, Initializable {
     /// @param _newStake The new stake.
     /// Note that the existing delayed stake will be nullified as non-relevant.
     function setStake(uint96 _courtID, uint256 _newStake) external {
-        _deleteDelayedStake(_courtID);
         if (!_setStakeForAccount(msg.sender, _courtID, _newStake, false)) revert StakingFailed();
     }
 
@@ -475,10 +474,6 @@ contract KlerosCore is IArbitratorV2, UUPSProxiable, Initializable {
         bool _alreadyTransferred
     ) external {
         if (msg.sender != address(sortitionModule)) revert SortitionModuleOnly();
-        // Always nullify the latest delayed stake before setting a new value.
-        // Note that we call _deleteDelayedStake() here too because the one in `setStake` can be bypassed
-        // if the stake was updated automatically during `execute` (e.g. when unstaking inactive juror).
-        _deleteDelayedStake(_courtID);
         _setStakeForAccount(_account, _courtID, _newStake, _alreadyTransferred);
     }
 
@@ -1061,12 +1056,6 @@ contract KlerosCore is IArbitratorV2, UUPSProxiable, Initializable {
         emit DisputeKitEnabled(_courtID, _disputeKitID, _enable);
     }
 
-    /// @dev Removes the latest delayed stake if there is any.
-    /// @param _courtID The ID of the court.
-    function _deleteDelayedStake(uint96 _courtID) private {
-        sortitionModule.deleteDelayedStake(_courtID, msg.sender);
-    }
-
     /// @dev Sets the specified juror's stake in a court.
     /// `O(n + p * log_k(j))` where
     /// `n` is the number of courts the juror has staked in,

contracts/src/arbitration/SortitionModule.sol

@@ -218,30 +218,13 @@ contract SortitionModule is ISortitionModule, UUPSProxiable, Initializable {
         delayedStakeReadIndex = newDelayedStakeReadIndex;
     }
 
-    /// @dev Checks if there is already a delayed stake. In this case consider it irrelevant and remove it.
-    /// @param _courtID ID of the court.
-    /// @param _juror Juror whose stake to check.
-    function deleteDelayedStake(uint96 _courtID, address _juror) external override onlyByCore {
-        uint256 latestIndex = latestDelayedStakeIndex[_juror][_courtID];
-        if (latestIndex != 0) {
-            DelayedStake storage delayedStake = delayedStakes[latestIndex];
-            if (delayedStake.alreadyTransferred) {
-                bytes32 stakePathID = _accountAndCourtIDToStakePathID(_juror, _courtID);
-                // Sortition stake represents the stake value that was last updated during Staking phase.
-                uint256 sortitionStake = stakeOf(bytes32(uint256(_courtID)), stakePathID);
-                // Withdraw the tokens that were added with the latest delayed stake.
-                core.withdrawPartiallyDelayedStake(_courtID, _juror, delayedStake.stake - sortitionStake);
-            }
-            delete delayedStakes[latestIndex];
-            delete latestDelayedStakeIndex[_juror][_courtID];
-        }
-    }
-
     function preStakeHook(
         address _account,
         uint96 _courtID,
         uint256 _stake
     ) external override onlyByCore returns (PreStakeHookResult) {
+        _deleteDelayedStake(_courtID, _account);
+
         (, , uint256 currentStake, uint256 nbCourts) = core.getJurorBalance(_account, _courtID);
         if (currentStake == 0 && nbCourts >= MAX_STAKE_PATHS) {
             // Prevent staking beyond MAX_STAKE_PATHS but unstaking is always allowed.
@@ -268,6 +251,25 @@ contract SortitionModule is ISortitionModule, UUPSProxiable, Initializable {
         return PreStakeHookResult.ok;
     }
 
+    /// @dev Checks if there is already a delayed stake. In this case consider it irrelevant and remove it.
+    /// @param _courtID ID of the court.
+    /// @param _juror Juror whose stake to check.
+    function _deleteDelayedStake(uint96 _courtID, address _juror) internal {
+        uint256 latestIndex = latestDelayedStakeIndex[_juror][_courtID];
+        if (latestIndex != 0) {
+            DelayedStake storage delayedStake = delayedStakes[latestIndex];
+            if (delayedStake.alreadyTransferred) {
+                bytes32 stakePathID = _accountAndCourtIDToStakePathID(_juror, _courtID);
+                // Sortition stake represents the stake value that was last updated during Staking phase.
+                uint256 sortitionStake = stakeOf(bytes32(uint256(_courtID)), stakePathID);
+                // Withdraw the tokens that were added with the latest delayed stake.
+                core.withdrawPartiallyDelayedStake(_courtID, _juror, delayedStake.stake - sortitionStake);
+            }
+            delete delayedStakes[latestIndex];
+            delete latestDelayedStakeIndex[_juror][_courtID];
+        }
+    }
+
     function createDisputeHook(uint256 /*_disputeID*/, uint256 /*_roundID*/) external override onlyByCore {
         disputesWithoutJurors++;
     }

contracts/src/arbitration/interfaces/ISortitionModule.sol

@@ -32,6 +32,4 @@ interface ISortitionModule {
     function createDisputeHook(uint256 _disputeID, uint256 _roundID) external;
 
     function postDrawHook(uint256 _disputeID, uint256 _roundID) external;
-
-    function deleteDelayedStake(uint96 _courtID, address _juror) external;
 }