refactor: moved the juror stakes accounting from KlerosCore to SortitionModule

Author: jaybuidl

contracts/src/arbitration/KlerosCore.sol

@@ -70,13 +70,6 @@ contract KlerosCore is IArbitratorV2, UUPSProxiable, Initializable {
         uint256 drawIterations; // The number of iterations passed drawing the jurors for this round.
     }
 
-    struct Juror {
-        uint96[] courtIDs; // The IDs of courts where the juror's stake path ends. A stake path is a path from the general court to a court the juror directly staked in using `_setStake`.
-        uint256 stakedPnk; // The juror's total amount of tokens staked in subcourts. Reflects actual pnk balance.
-        uint256 lockedPnk; // The juror's total amount of tokens locked in disputes. Can reflect actual pnk balance when stakedPnk are fully withdrawn.
-        mapping(uint96 => uint256) stakedPnkByCourt; // The amount of PNKs the juror has staked in the court in the form `stakedPnkByCourt[courtID]`.
-    }
-
     // Workaround "stack too deep" errors
     struct ExecuteParams {
         uint256 disputeID; // The ID of the dispute to execute.
@@ -107,21 +100,12 @@ contract KlerosCore is IArbitratorV2, UUPSProxiable, Initializable {
     Court[] public courts; // The courts.
     IDisputeKit[] public disputeKits; // Array of dispute kits.
     Dispute[] public disputes; // The disputes.
-    mapping(address => Juror) internal jurors; // The jurors.
     mapping(IERC20 => CurrencyRate) public currencyRates; // The price of each token in ETH.
 
     // ************************************* //
     // *              Events               * //
     // ************************************* //
 
-    event StakeSet(address indexed _address, uint256 _courtID, uint256 _amount);
-    event StakeDelayedNotTransferred(address indexed _address, uint256 _courtID, uint256 _amount);
-    event StakeDelayedAlreadyTransferred(address indexed _address, uint256 _courtID, uint256 _amount);
-    event StakeDelayedAlreadyTransferredWithdrawn(
-        uint96 indexed _courtID,
-        address indexed _account,
-        uint256 _withdrawnAmount
-    );
     event NewPeriod(uint256 indexed _disputeID, Period _period);
     event AppealPossible(uint256 indexed _disputeID, IArbitrableV2 indexed _arbitrable);
     event AppealDecision(uint256 indexed _disputeID, IArbitrableV2 indexed _arbitrable);
@@ -464,7 +448,7 @@ 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 {
-        if (!_setStakeForAccount(msg.sender, _courtID, _newStake, false)) revert StakingFailed();
+        _setStake(msg.sender, _courtID, _newStake, false);
     }
 
     function setStakeBySortitionModule(
@@ -474,31 +458,7 @@ contract KlerosCore is IArbitratorV2, UUPSProxiable, Initializable {
         bool _alreadyTransferred
     ) external {
         if (msg.sender != address(sortitionModule)) revert SortitionModuleOnly();
-        _setStakeForAccount(_account, _courtID, _newStake, _alreadyTransferred);
-    }
-
-    function withdrawPartiallyDelayedStake(uint96 _courtID, address _juror, uint256 _amountToWithdraw) external {
-        if (msg.sender != address(sortitionModule)) revert SortitionModuleOnly();
-        uint256 actualAmount = _amountToWithdraw;
-        Juror storage juror = jurors[_juror];
-        if (juror.stakedPnk <= actualAmount) {
-            actualAmount = juror.stakedPnk;
-        }
-        require(pinakion.safeTransfer(_juror, actualAmount));
-        // StakePnk can become lower because of penalty, thus we adjust the amount for it. stakedPnkByCourt can't be penalized so subtract the default amount.
-        juror.stakedPnk -= actualAmount;
-        juror.stakedPnkByCourt[_courtID] -= _amountToWithdraw;
-        emit StakeDelayedAlreadyTransferredWithdrawn(_courtID, _juror, _amountToWithdraw);
-        // Note that if we don't delete court here it'll be duplicated after staking.
-        if (juror.stakedPnkByCourt[_courtID] == 0) {
-            for (uint256 i = juror.courtIDs.length; i > 0; i--) {
-                if (juror.courtIDs[i - 1] == _courtID) {
-                    juror.courtIDs[i - 1] = juror.courtIDs[juror.courtIDs.length - 1];
-                    juror.courtIDs.pop();
-                    break;
-                }
-            }
-        }
+        _setStake(_account, _courtID, _newStake, _alreadyTransferred);
     }
 
     /// @inheritdoc IArbitratorV2
@@ -625,7 +585,7 @@ contract KlerosCore is IArbitratorV2, UUPSProxiable, Initializable {
             if (drawnAddress == address(0)) {
                 continue;
             }
-            jurors[drawnAddress].lockedPnk += round.pnkAtStakePerJuror;
+            sortitionModule.lockStake(drawnAddress, round.pnkAtStakePerJuror);
             emit Draw(drawnAddress, _disputeID, currentRound, round.drawnJurors.length);
             round.drawnJurors.push(drawnAddress);
             if (round.drawnJurors.length == round.nbVotes) {
@@ -764,15 +724,10 @@ contract KlerosCore is IArbitratorV2, UUPSProxiable, Initializable {
 
         // Unlock the PNKs affected by the penalty
         address account = round.drawnJurors[_params.repartition];
-        jurors[account].lockedPnk -= penalty;
+        sortitionModule.unlockStake(account, penalty);
 
         // Apply the penalty to the staked PNKs.
-        // Note that lockedPnk will always cover penalty while stakedPnk can become lower after manual unstaking.
-        if (jurors[account].stakedPnk >= penalty) {
-            jurors[account].stakedPnk -= penalty;
-        } else {
-            jurors[account].stakedPnk = 0;
-        }
+        sortitionModule.penalizeStake(account, penalty);
         emit TokenAndETHShift(
             account,
             _params.disputeID,
@@ -831,10 +786,10 @@ contract KlerosCore is IArbitratorV2, UUPSProxiable, Initializable {
         uint256 pnkLocked = (round.pnkAtStakePerJuror * degreeOfCoherence) / ALPHA_DIVISOR;
 
         // Release the rest of the PNKs of the juror for this round.
-        jurors[account].lockedPnk -= pnkLocked;
+        sortitionModule.unlockStake(account, pnkLocked);
 
         // Give back the locked PNKs in case the juror fully unstaked earlier.
-        if (jurors[account].stakedPnk == 0) {
+        if (!sortitionModule.isJurorStaked(account)) {
             pinakion.safeTransfer(account, pnkLocked);
         }
 
@@ -980,17 +935,6 @@ contract KlerosCore is IArbitratorV2, UUPSProxiable, Initializable {
         return disputes[_disputeID].rounds.length;
     }
 
-    function getJurorBalance(
-        address _juror,
-        uint96 _courtID
-    ) external view returns (uint256 totalStaked, uint256 totalLocked, uint256 stakedInCourt, uint256 nbCourts) {
-        Juror storage juror = jurors[_juror];
-        totalStaked = juror.stakedPnk;
-        totalLocked = juror.lockedPnk;
-        stakedInCourt = juror.stakedPnkByCourt[_courtID];
-        nbCourts = juror.courtIDs.length;
-    }
-
     function isSupported(uint96 _courtID, uint256 _disputeKitID) external view returns (bool) {
         return courts[_courtID].supportedDisputeKits[_disputeKitID];
     }
@@ -1033,12 +977,6 @@ contract KlerosCore is IArbitratorV2, UUPSProxiable, Initializable {
         return disputeKits.length;
     }
 
-    /// @dev Gets the court identifiers where a specific `_juror` has staked.
-    /// @param _juror The address of the juror.
-    function getJurorCourtIDs(address _juror) public view returns (uint96[] memory) {
-        return jurors[_juror].courtIDs;
-    }
-
     function convertEthToTokenAmount(IERC20 _toToken, uint256 _amountInEth) public view returns (uint256) {
         return (_amountInEth * 10 ** currencyRates[_toToken].rateDecimals) / currencyRates[_toToken].rateInEth;
     }
@@ -1056,104 +994,34 @@ contract KlerosCore is IArbitratorV2, UUPSProxiable, Initializable {
         emit DisputeKitEnabled(_courtID, _disputeKitID, _enable);
     }
 
-    /// @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,
-    /// `p` is the depth of the court tree,
-    /// `k` is the minimum number of children per node of one of these courts' sortition sum tree,
-    /// and `j` is the maximum number of jurors that ever staked in one of these courts simultaneously.
-    /// @param _account The address of the juror.
-    /// @param _courtID The ID of the court.
-    /// @param _newStake The new stake.
-    /// @param _alreadyTransferred True if the tokens were already transferred from juror. Only relevant for delayed stakes.
-    /// @return succeeded True if the call succeeded, false otherwise.
-    function _setStakeForAccount(
+    function _setStake(
         address _account,
         uint96 _courtID,
         uint256 _newStake,
         bool _alreadyTransferred
-    ) internal returns (bool succeeded) {
-        if (_courtID == Constants.FORKING_COURT || _courtID > courts.length) return false;
-
-        Juror storage juror = jurors[_account];
-        uint256 currentStake = juror.stakedPnkByCourt[_courtID];
-
-        if (_newStake != 0) {
-            if (_newStake < courts[_courtID].minStake) return false;
-        } else if (currentStake == 0) {
-            return false;
+    ) internal returns (bool success) {
+        if (_courtID == Constants.FORKING_COURT || _courtID > courts.length) {
+            return false; // Staking directly into the forking court is not allowed.
         }
-
-        ISortitionModule.PreStakeHookResult result = sortitionModule.preStakeHook(_account, _courtID, _newStake);
-        if (result == ISortitionModule.PreStakeHookResult.failed) {
-            return false;
-        } else if (result == ISortitionModule.PreStakeHookResult.stakeDelayedNotTransferred) {
-            emit StakeDelayedNotTransferred(_account, _courtID, _newStake);
-            return true;
+        if (_newStake != 0 && _newStake < courts[_courtID].minStake) {
+            return false; // Staking less than the minimum stake is not allowed.
         }
-
-        uint256 transferredAmount;
-        if (_newStake >= currentStake) {
-            if (!_alreadyTransferred) {
-                // Stake increase
-                // When stakedPnk becomes lower than lockedPnk count the locked tokens in when transferring tokens from juror.
-                // (E.g. stakedPnk = 0, lockedPnk = 150) which can happen if the juror unstaked fully while having some tokens locked.
-                uint256 previouslyLocked = (juror.lockedPnk >= juror.stakedPnk) ? juror.lockedPnk - juror.stakedPnk : 0; // underflow guard
-                transferredAmount = (_newStake >= currentStake + previouslyLocked) // underflow guard
-                    ? _newStake - currentStake - previouslyLocked
-                    : 0;
-                if (transferredAmount > 0) {
-                    // Note we don't return false after incorrect transfer because when stake is increased the transfer is done immediately, thus it can't disrupt delayed stakes' queue.
-                    pinakion.safeTransferFrom(_account, address(this), transferredAmount);
-                }
-                if (currentStake == 0) {
-                    juror.courtIDs.push(_courtID);
-                }
-            }
-        } else {
-            // Note that stakes can be partially delayed only when stake is increased.
-            // Stake decrease: make sure locked tokens always stay in the contract. They can only be released during Execution.
-            if (juror.stakedPnk >= currentStake - _newStake + juror.lockedPnk) {
-                // We have enough pnk staked to afford withdrawal while keeping locked tokens.
-                transferredAmount = currentStake - _newStake;
-            } else if (juror.stakedPnk >= juror.lockedPnk) {
-                // Can't afford withdrawing the current stake fully. Take whatever is available while keeping locked tokens.
-                transferredAmount = juror.stakedPnk - juror.lockedPnk;
-            }
-            if (transferredAmount > 0) {
-                if (!pinakion.safeTransfer(_account, transferredAmount)) {
-                    return false;
-                }
-            }
-            if (_newStake == 0) {
-                for (uint256 i = juror.courtIDs.length; i > 0; i--) {
-                    if (juror.courtIDs[i - 1] == _courtID) {
-                        juror.courtIDs[i - 1] = juror.courtIDs[juror.courtIDs.length - 1];
-                        juror.courtIDs.pop();
-                        break;
-                    }
-                }
+        (uint256 pnkDeposit, uint256 pnkWithdrawal, bool sortitionSuccess) = sortitionModule.setStake(
+            _account,
+            _courtID,
+            _newStake,
+            _alreadyTransferred
+        );
+        if (pnkDeposit > 0 && pnkWithdrawal > 0) revert StakingFailed();
+        if (pnkDeposit > 0) {
+            // Note we don't return false after incorrect transfer because when stake is increased the transfer is done immediately, thus it can't disrupt delayed stakes' queue.
+            pinakion.safeTransferFrom(_account, address(this), pnkDeposit);
+        } else if (pnkWithdrawal > 0) {
+            if (!pinakion.safeTransfer(_account, pnkWithdrawal)) {
+                return false;
             }
         }
-
-        // Note that stakedPnk can become async with currentStake (e.g. after penalty).
-        // Also note that these values were already updated if the stake was only partially delayed.
-        if (!_alreadyTransferred) {
-            juror.stakedPnk = (juror.stakedPnk >= currentStake)
-                ? juror.stakedPnk - currentStake + _newStake
-                : _newStake;
-            juror.stakedPnkByCourt[_courtID] = _newStake;
-        }
-
-        // Transfer the tokens but don't update sortition module.
-        if (result == ISortitionModule.PreStakeHookResult.stakeDelayedAlreadyTransferred) {
-            emit StakeDelayedAlreadyTransferred(_account, _courtID, _newStake);
-            return true;
-        }
-
-        sortitionModule.setStake(_account, _courtID, _newStake);
-        emit StakeSet(_account, _courtID, _newStake);
-        return true;
+        return sortitionSuccess;
     }
 
     /// @dev Gets a court ID, the minimum number of jurors and an ID of a dispute kit from a specified extra data bytes array.