PNK penalties applied to the sortition tree stakes

contracts/CHANGELOG.md

@@ -13,6 +13,12 @@ The format is based on [Common Changelog](https://common-changelog.org/).
 - **Breaking:** Rename the interface from `RNG` to `IRNG` ([#2054](https://github.com/kleros/kleros-v2/issues/2054))
 - **Breaking:** Remove the `_block` parameter from `IRNG.requestRandomness()` and `IRNG.receiveRandomness()`, not needed for the primary VRF-based RNG ([#2054](https://github.com/kleros/kleros-v2/issues/2054))
 - **Breaking:** Rename `governor` to `owner` in order to comply with the lightweight ownership standard [ERC-5313](https://eipsinsight.com/ercs/erc-5313) ([#2112](https://github.com/kleros/kleros-v2/issues/2112))
+- **Breaking:** Apply the penalties to the stakes in the Sortition Tree ([#2107](https://github.com/kleros/kleros-v2/issues/2107))
+- **Breaking:** Make `SortitionModule.getJurorBalance().stakedInCourt` include the penalties ([#2107](https://github.com/kleros/kleros-v2/issues/2107))
+- **Breaking:** Add a new field `drawnJurorFromCourtIDs` to the `Round` struct in `KlerosCoreBase` and `KlerosCoreUniversity` ([#2107](https://github.com/kleros/kleros-v2/issues/2107))
+- Make `IDisputeKit.draw()` and `ISortitionModule.draw()` return the court ID from which the juror was drawn ([#2107](https://github.com/kleros/kleros-v2/issues/2107))
+- Rename `SortitionModule.setJurorInactive()` to `SortitionModule.forcedUnstakeAllCourts()` ([#2107](https://github.com/kleros/kleros-v2/issues/2107))
+- Allow stake changes to by-pass delayed stakes when initiated by the SortitionModule by setting the `_noDelay` parameter to `true` in `SortitionModule.validateStake()` ([#2107](https://github.com/kleros/kleros-v2/issues/2107))
 - Make the primary VRF-based RNG fall back to `BlockhashRNG` if the VRF request is not fulfilled within a timeout ([#2054](https://github.com/kleros/kleros-v2/issues/2054))
 - Authenticate the calls to the RNGs to prevent 3rd parties from depleting the Chainlink VRF subscription funds ([#2054](https://github.com/kleros/kleros-v2/issues/2054))
 - Use `block.timestamp` rather than `block.number` for `BlockhashRNG` for better reliability on Arbitrum as block production is sporadic depending on network conditions. ([#2054](https://github.com/kleros/kleros-v2/issues/2054))

contracts/src/arbitration/KlerosCoreBase.sol

@@ -60,6 +60,7 @@ abstract contract KlerosCoreBase is IArbitratorV2, Initializable, UUPSProxiable
         uint256 repartitions; // A counter of reward repartitions made in this round.
         uint256 pnkPenalties; // The amount of PNKs collected from penalties in this round.
         address[] drawnJurors; // Addresses of the jurors that were drawn in this round.
+        uint96[] drawnJurorFromCourtIDs; // The courtIDs where the juror was drawn from, possibly their stake in a subcourt.
         uint256 sumFeeRewardPaid; // Total sum of arbitration fees paid to coherent jurors as a reward in this round.
         uint256 sumPnkRewardPaid; // Total sum of PNK paid to coherent jurors as a reward in this round.
         IERC20 feeToken; // The token used for paying fees in this round.
@@ -463,16 +464,16 @@ abstract contract KlerosCoreBase is IArbitratorV2, Initializable, UUPSProxiable
     /// @param _newStake The new stake.
     /// Note that the existing delayed stake will be nullified as non-relevant.
     function setStake(uint96 _courtID, uint256 _newStake) external virtual whenNotPaused {
-        _setStake(msg.sender, _courtID, _newStake, OnError.Revert);
+        _setStake(msg.sender, _courtID, _newStake, false, OnError.Revert);
     }
 
-    /// @dev Sets the stake of a specified account in a court, typically to apply a delayed stake or unstake inactive jurors.
+    /// @dev Sets the stake of a specified account in a court without delaying stake changes, typically to apply a delayed stake or unstake inactive jurors.
     /// @param _account The account whose stake is being set.
     /// @param _courtID The ID of the court.
     /// @param _newStake The new stake.
     function setStakeBySortitionModule(address _account, uint96 _courtID, uint256 _newStake) external {
         if (msg.sender != address(sortitionModule)) revert SortitionModuleOnly();
-        _setStake(_account, _courtID, _newStake, OnError.Return);
+        _setStake(_account, _courtID, _newStake, true, OnError.Return);
     }
 
     /// @dev Transfers PNK to the juror by SortitionModule.
@@ -606,13 +607,14 @@ abstract contract KlerosCoreBase is IArbitratorV2, Initializable, UUPSProxiable
         uint256 startIndex = round.drawIterations; // for gas: less storage reads
         uint256 i;
         while (i < _iterations && round.drawnJurors.length < round.nbVotes) {
-            address drawnAddress = disputeKit.draw(_disputeID, startIndex + i++);
+            (address drawnAddress, uint96 fromSubcourtID) = disputeKit.draw(_disputeID, startIndex + i++);
             if (drawnAddress == address(0)) {
                 continue;
             }
             sortitionModule.lockStake(drawnAddress, round.pnkAtStakePerJuror);
             emit Draw(drawnAddress, _disputeID, currentRound, round.drawnJurors.length);
             round.drawnJurors.push(drawnAddress);
+            round.drawnJurorFromCourtIDs.push(fromSubcourtID != 0 ? fromSubcourtID : dispute.courtID);
             if (round.drawnJurors.length == round.nbVotes) {
                 sortitionModule.postDrawHook(_disputeID, currentRound);
             }
@@ -775,7 +777,12 @@ abstract contract KlerosCoreBase is IArbitratorV2, Initializable, UUPSProxiable
         sortitionModule.unlockStake(account, penalty);
 
         // Apply the penalty to the staked PNKs.
-        (uint256 pnkBalance, uint256 availablePenalty) = sortitionModule.penalizeStake(account, penalty);
+        uint96 penalizedInCourtID = round.drawnJurorFromCourtIDs[_params.repartition];
+        (uint256 pnkBalance, uint256 newCourtStake, uint256 availablePenalty) = sortitionModule.setStakePenalty(
+            account,
+            penalizedInCourtID,
+            penalty
+        );
         _params.pnkPenaltiesInRound += availablePenalty;
         emit TokenAndETHShift(
             account,
@@ -786,10 +793,15 @@ abstract contract KlerosCoreBase is IArbitratorV2, Initializable, UUPSProxiable
             0,
             round.feeToken
         );
-        // Unstake the juror from all courts if he was inactive or his balance can't cover penalties anymore.
+
         if (pnkBalance == 0 || !disputeKit.isVoteActive(_params.disputeID, _params.round, _params.repartition)) {
-            sortitionModule.setJurorInactive(account);
+            // The juror is inactive or their balance is can't cover penalties anymore, unstake them from all courts.
+            sortitionModule.forcedUnstakeAllCourts(account);
+        } else if (newCourtStake < courts[penalizedInCourtID].minStake) {
+            // The juror's balance fell below the court minStake, unstake them from the court.
+            sortitionModule.forcedUnstake(account, penalizedInCourtID);
         }
+
         if (_params.repartition == _params.numberOfVotesInRound - 1 && _params.coherentCount == 0) {
             // No one was coherent, send the rewards to the owner.
             _transferFeeToken(round.feeToken, payable(owner), round.totalFeesForJurors);
@@ -1126,9 +1138,16 @@ abstract contract KlerosCoreBase is IArbitratorV2, Initializable, UUPSProxiable
     /// @param _account The account to set the stake for.
     /// @param _courtID The ID of the court to set the stake for.
     /// @param _newStake The new stake.
+    /// @param _noDelay True if the stake change should not be delayed.
     /// @param _onError Whether to revert or return false on error.
     /// @return Whether the stake was successfully set or not.
-    function _setStake(address _account, uint96 _courtID, uint256 _newStake, OnError _onError) internal returns (bool) {
+    function _setStake(
+        address _account,
+        uint96 _courtID,
+        uint256 _newStake,
+        bool _noDelay,
+        OnError _onError
+    ) internal returns (bool) {
         if (_courtID == FORKING_COURT || _courtID >= courts.length) {
             _stakingFailed(_onError, StakingResult.CannotStakeInThisCourt); // Staking directly into the forking court is not allowed.
             return false;
@@ -1140,7 +1159,8 @@ abstract contract KlerosCoreBase is IArbitratorV2, Initializable, UUPSProxiable
         (uint256 pnkDeposit, uint256 pnkWithdrawal, StakingResult stakingResult) = sortitionModule.validateStake(
             _account,
             _courtID,
-            _newStake
+            _newStake,
+            _noDelay
         );
         if (stakingResult != StakingResult.Successful && stakingResult != StakingResult.Delayed) {
             _stakingFailed(_onError, stakingResult);

contracts/src/arbitration/KlerosCoreNeo.sol

@@ -108,7 +108,7 @@ contract KlerosCoreNeo is KlerosCoreBase {
     /// Note that the existing delayed stake will be nullified as non-relevant.
     function setStake(uint96 _courtID, uint256 _newStake) external override whenNotPaused {
         if (jurorNft.balanceOf(msg.sender) == 0) revert NotEligibleForStaking();
-        super._setStake(msg.sender, _courtID, _newStake, OnError.Revert);
+        super._setStake(msg.sender, _courtID, _newStake, false, OnError.Revert);
     }
 
     // ************************************* //