feat(KC): instant staking

Author: unknownunknown1

contracts/src/arbitration/KlerosCore.sol

@@ -129,6 +129,7 @@ contract KlerosCore is IArbitratorV2 {
 
     event StakeSet(address indexed _address, uint256 _courtID, uint256 _amount);
     event StakeDelayed(address indexed _address, uint256 _courtID, uint256 _amount);
+    event StakePartiallyDelayed(address indexed _address, uint256 _courtID, uint256 _amount);
     event NewPeriod(uint256 indexed _disputeID, Period _period);
     event AppealPossible(uint256 indexed _disputeID, IArbitrableV2 indexed _arbitrable);
     event AppealDecision(uint256 indexed _disputeID, IArbitrableV2 indexed _arbitrable);
@@ -187,6 +188,7 @@ contract KlerosCore is IArbitratorV2 {
         uint256 _feeAmount,
         IERC20 _feeToken
     );
+    event PartiallyDelayedStakeWithdrawn(uint96 indexed _courtID, address indexed _account, uint256 _withdrawnAmount);
 
     // ************************************* //
     // *        Function Modifiers         * //
@@ -484,13 +486,54 @@ contract KlerosCore is IArbitratorV2 {
     /// @dev Sets the caller's stake in a court.
     /// @param _courtID The ID of the court.
     /// @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)) revert StakingFailed();
+        removeDelayedStake(_courtID);
+        if (!_setStakeForAccount(msg.sender, _courtID, _newStake, false)) revert StakingFailed();
     }
 
-    function setStakeBySortitionModule(address _account, uint96 _courtID, uint256 _newStake) external {
+    /// @dev Removes the latest delayed stake if there is any.
+    /// @param _courtID The ID of the court.
+    function removeDelayedStake(uint96 _courtID) public {
+        sortitionModule.checkExistingDelayedStake(_courtID, msg.sender);
+    }
+
+    function withdrawPartiallyDelayedStake(uint96 _courtID, address _juror, uint256 _amountToWithdraw) external {
+        if (msg.sender != address(sortitionModule)) revert WrongCaller();
+        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 PartiallyDelayedStakeWithdrawn(_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;
+                }
+            }
+        }
+    }
+
+    function setStakeBySortitionModule(
+        address _account,
+        uint96 _courtID,
+        uint256 _stake,
+        bool _alreadyTransferred
+    ) external {
         if (msg.sender != address(sortitionModule)) revert WrongCaller();
-        _setStakeForAccount(_account, _courtID, _newStake);
+        // Always nullify the latest delayed stake before setting a new value.
+        // Note that we check the delayed stake here too because the check in `setStake` can be bypassed
+        // if the stake was updated automatically during `execute` (e.g. when unstaking inactive juror).
+        removeDelayedStake(_courtID);
+        _setStakeForAccount(_account, _courtID, _newStake, _alreadyTransferred);
     }
 
     /// @inheritdoc IArbitratorV2
@@ -1081,11 +1124,17 @@ contract KlerosCore is IArbitratorV2 {
     /// @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(
+        
         address _account,
+       
         uint96 _courtID,
+       
         uint256 _newStake
+    ,
+        bool _alreadyTransferred
     ) internal returns (bool succeeded) {
         if (_courtID == FORKING_COURT || _courtID > courts.length) return false;
 
@@ -1108,22 +1157,24 @@ contract KlerosCore is IArbitratorV2 {
 
         uint256 transferredAmount;
         if (_newStake >= currentStake) {
-            // Stake increase
+            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) {
-                if (!pinakion.safeTransferFrom(_account, address(this), transferredAmount)) {
-                    return false;
+                // (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);
                 }
-            }
-            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.
@@ -1149,8 +1200,17 @@ contract KlerosCore is IArbitratorV2 {
         }
 
         // Note that stakedPnk can become async with currentStake (e.g. after penalty).
-        juror.stakedPnk = (juror.stakedPnk >= currentStake) ? juror.stakedPnk - currentStake + _newStake : _newStake;
-        juror.stakedPnkByCourt[_courtID] = _newStake;
+        // 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.partiallyDelayed) {
+            emit StakePartiallyDelayed(_account, _courtID, _stake);
+            return true;
+        }
 
         sortitionModule.setStake(_account, _courtID, _newStake);
         emit StakeSet(_account, _courtID, _newStake);

contracts/src/arbitration/SortitionModule.sol

@@ -36,6 +36,7 @@ contract SortitionModule is ISortitionModule {
         address account; // The address of the juror.
         uint96 courtID; // The ID of the court.
         uint256 stake; // The new stake.
+        bool alreadyTransferred; // True if tokens were already transferred before delayed stake's execution.
     }
 
     // ************************************* //
@@ -60,6 +61,7 @@ contract SortitionModule is ISortitionModule {
     uint256 public delayedStakeReadIndex = 1; // The index of the next `delayedStake` item that should be processed. Starts at 1 because 0 index is skipped.
     mapping(bytes32 => SortitionSumTree) sortitionSumTrees; // The mapping trees by keys.
     mapping(uint256 => DelayedStake) public delayedStakes; // Stores the stakes that were changed during Drawing phase, to update them when the phase is switched to Staking.
+    mapping(address => mapping(uint96 => uint256)) public latestDelayedStakeIndex; // Maps the juror to its latest delayed stake. If there is already a delayed stake for this juror then it'll be replaced. latestDelayedStakeIndex[juror][courtID].
 
     // ************************************* //
     // *        Function Modifiers         * //
@@ -184,12 +186,40 @@ contract SortitionModule is ISortitionModule {
 
         for (uint256 i = delayedStakeReadIndex; i < newDelayedStakeReadIndex; i++) {
             DelayedStake storage delayedStake = delayedStakes[i];
-            core.setStakeBySortitionModule(delayedStake.account, delayedStake.courtID, delayedStake.stake);
-            delete delayedStakes[i];
+            // Delayed stake could've been manually removed already. In this case simply move on to the next item.
+            if (delayedStake.account != address(0)) {
+                core.setStakeBySortitionModule(
+                    delayedStake.account,
+                    delayedStake.courtID,
+                    delayedStake.stake,
+                    delayedStake.alreadyTransferred
+                );
+                delete latestDelayedStakeIndex[delayedStake.account][delayedStake.courtID];
+                delete delayedStakes[i];
+            }
         }
         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 checkExistingDelayedStake(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,
@@ -201,12 +231,18 @@ contract SortitionModule is ISortitionModule {
             return preStakeHookResult.failed;
         } else {
             if (phase != Phase.staking) {
-                delayedStakes[++delayedStakeWriteIndex] = DelayedStake({
-                    account: _account,
-                    courtID: _courtID,
-                    stake: _stake
-                });
-                return preStakeHookResult.delayed;
+                DelayedStake storage delayedStake = delayedStakes[++delayedStakeWriteIndex];
+                delayedStake.account = _account;
+                delayedStake.courtID = _courtID;
+                delayedStake.stake = _stake;
+                latestDelayedStakeIndex[_account][_courtID] = delayedStakeWriteIndex;
+                if (_stake > currentStake) {
+                    // Actual token transfer is done right after this hook.
+                    delayedStake.alreadyTransferred = true;
+                    return preStakeHookResult.partiallyDelayed;
+                } else {
+                    return preStakeHookResult.delayed;
+                }
             }
         }
         return preStakeHookResult.ok;
@@ -256,7 +292,7 @@ contract SortitionModule is ISortitionModule {
     function setJurorInactive(address _account) external override onlyByCore {
         uint96[] memory courtIDs = core.getJurorCourtIDs(_account);
         for (uint256 j = courtIDs.length; j > 0; j--) {
-            core.setStakeBySortitionModule(_account, courtIDs[j - 1], 0);
+            core.setStakeBySortitionModule(_account, courtIDs[j - 1], 0, false);
         }
     }
 
@@ -469,4 +505,12 @@ contract SortitionModule is ISortitionModule {
             stakePathID := mload(ptr)
         }
     }
+
+    function stakeOf(bytes32 _key, bytes32 _ID) public view returns (uint256 value) {
+        SortitionSumTree storage tree = sortitionSumTrees[_key];
+        uint treeIndex = tree.IDsToNodeIndexes[_ID];
+
+        if (treeIndex == 0) value = 0;
+        else value = tree.nodes[treeIndex];
+    }
 }

contracts/src/arbitration/interfaces/ISortitionModule.sol

@@ -9,9 +9,10 @@ interface ISortitionModule {
     }
 
     enum preStakeHookResult {
-        ok,
-        delayed,
-        failed
+        ok, // Correct phase. All checks are passed.
+        partiallyDelayed, // Wrong phase but stake is increased, so transfer the tokens without updating the drawing chance.
+        delayed, // Wrong phase and stake is decreased. Delay the token transfer and drawing chance update.
+        failed // Checks didn't pass. Do no changes.
     }
 
     event NewPhase(Phase _phase);
@@ -31,4 +32,6 @@ interface ISortitionModule {
     function createDisputeHook(uint256 _disputeID, uint256 _roundID) external;
 
     function postDrawHook(uint256 _disputeID, uint256 _roundID) external;
+
+    function checkExistingDelayedStake(uint96 _courtID, address _juror) external;
 }