feat: automatically stake PNK rewards instead of transferring them

Author: jaybuidl

contracts/src/arbitration/KlerosCoreBase.sol

@@ -846,13 +846,20 @@ abstract contract KlerosCoreBase is IArbitratorV2, Initializable, UUPSProxiable
         // Release the rest of the PNKs of the juror for this round.
         sortitionModule.unlockStake(account, pnkLocked);
 
-        // Transfer the rewards
+        // Compute the rewards
         uint256 pnkReward = _applyCoherence(_params.pnkPenaltiesInRound / _params.coherentCount, pnkCoherence);
         round.sumPnkRewardPaid += pnkReward;
         uint256 feeReward = _applyCoherence(round.totalFeesForJurors / _params.coherentCount, feeCoherence);
         round.sumFeeRewardPaid += feeReward;
-        pinakion.safeTransfer(account, pnkReward);
+
+        // Transfer the fee reward
         _transferFeeToken(round.feeToken, payable(account), feeReward);
+
+        // Stake the PNK reward if possible, by-passes delayed stakes and other checks usually done by validateStake()
+        if (!sortitionModule.setStakeReward(account, dispute.courtID, pnkReward)) {
+            pinakion.safeTransfer(account, pnkReward);
+        }
+
         emit TokenAndETHShift(
             account,
             _params.disputeID,

contracts/src/arbitration/SortitionModuleBase.sol

@@ -306,6 +306,30 @@ abstract contract SortitionModuleBase is ISortitionModule, Initializable, UUPSPr
         _setStake(_account, _courtID, _pnkDeposit, _pnkWithdrawal, _newStake);
     }
 
+    /// @dev Update the state of the stakes with a PNK reward deposit, called by KC during rewards execution.
+    /// `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 _reward The amount of PNK to be deposited as a reward.
+    function setStakeReward(
+        address _account,
+        uint96 _courtID,
+        uint256 _reward
+    ) external override onlyByCore returns (bool success) {
+        if (_reward == 0) return true; // No reward to add.
+
+        uint256 currentStake = stakeOf(_account, _courtID);
+        if (currentStake == 0) return false; // Juror has been unstaked, don't increase their stake.
+
+        uint256 newStake = currentStake + _reward;
+        _setStake(_account, _courtID, _reward, 0, newStake);
+        return true;
+    }
+
     function _setStake(
         address _account,
         uint96 _courtID,

contracts/src/arbitration/interfaces/ISortitionModule.sol

@@ -29,6 +29,8 @@ interface ISortitionModule {
         uint256 _newStake
     ) external;
 
+    function setStakeReward(address _account, uint96 _courtID, uint256 _reward) external returns (bool success);
+
     function setJurorInactive(address _account) external;
 
     function lockStake(address _account, uint256 _relativeAmount) external;

contracts/src/arbitration/university/KlerosCoreUniversity.sol

@@ -840,19 +840,26 @@ contract KlerosCoreUniversity is IArbitratorV2, UUPSProxiable, Initializable {
         // Release the rest of the PNKs of the juror for this round.
         sortitionModule.unlockStake(account, pnkLocked);
 
-        // Transfer the rewards
+        // Compute the rewards
         uint256 pnkReward = ((_params.pnkPenaltiesInRound / _params.coherentCount) * pnkCoherence) / ONE_BASIS_POINT;
         round.sumPnkRewardPaid += pnkReward;
         uint256 feeReward = ((round.totalFeesForJurors / _params.coherentCount) * feeCoherence) / ONE_BASIS_POINT;
         round.sumFeeRewardPaid += feeReward;
-        pinakion.safeTransfer(account, pnkReward);
+
+        // Transfer the fee reward
         if (round.feeToken == NATIVE_CURRENCY) {
             // The dispute fees were paid in ETH
             payable(account).send(feeReward);
         } else {
             // The dispute fees were paid in ERC20
             round.feeToken.safeTransfer(account, feeReward);
         }
+
+        // Stake the PNK reward if possible, by-passes delayed stakes and other checks usually done by validateStake()
+        if (!sortitionModule.setStakeReward(account, dispute.courtID, pnkReward)) {
+            pinakion.safeTransfer(account, pnkReward);
+        }
+
         emit TokenAndETHShift(
             account,
             _params.disputeID,

contracts/src/arbitration/university/SortitionModuleUniversity.sol

@@ -190,6 +190,40 @@ contract SortitionModuleUniversity is ISortitionModuleUniversity, UUPSProxiable,
         uint256 _pnkWithdrawal,
         uint256 _newStake
     ) external override onlyByCore {
+        _setStake(_account, _courtID, _pnkDeposit, _pnkWithdrawal, _newStake);
+    }
+
+    /// @dev Update the state of the stakes with a PNK reward deposit, called by KC during rewards execution.
+    /// `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 _reward The amount of PNK to be deposited as a reward.
+    function setStakeReward(
+        address _account,
+        uint96 _courtID,
+        uint256 _reward
+    ) external override onlyByCore returns (bool success) {
+        if (_reward == 0) return true; // No reward to add.
+
+        uint256 currentStake = _stakeOf(_account, _courtID);
+        if (currentStake == 0) return false; // Juror has been unstaked, don't increase their stake.
+
+        uint256 newStake = currentStake + _reward;
+        _setStake(_account, _courtID, _reward, 0, newStake);
+        return true;
+    }
+
+    function _setStake(
+        address _account,
+        uint96 _courtID,
+        uint256 _pnkDeposit,
+        uint256 _pnkWithdrawal,
+        uint256 _newStake
+    ) internal {
         Juror storage juror = jurors[_account];
         uint256 currentStake = _stakeOf(_account, _courtID);
         if (_pnkDeposit > 0) {

contracts/test/foundry/KlerosCore.t.sol

@@ -2397,9 +2397,9 @@ contract KlerosCoreTest is Test {
         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)), 20500, "Wrong token balance of the core"); // Was 21500. 1000 was transferred to staker2
+        assertEq(pinakion.balanceOf(address(core)), 21500, "Wrong token balance of the core"); // Was 21500. 1000 was transferred to staker2
         assertEq(pinakion.balanceOf(staker1), 999999999999998500, "Wrong token balance of staker1");
-        assertEq(pinakion.balanceOf(staker2), 999999999999981000, "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"); // 20k stake and 1k added as a reward, thus -19k from the default
     }
 
     function test_execute_NoCoherence() public {