feat(KC): instant staking

unknownunknown1 · unknownunknown1 · commit 17796c12c975 · 2023-07-05T18:02:09.000+10:00
diff --git a/contracts/src/arbitration/KlerosCore.sol b/contracts/src/arbitration/KlerosCore.sol
@@ -127,6 +127,7 @@ contract KlerosCore is IArbitratorV2 {
 
     event StakeSet(address indexed _address, uint256 _courtID, uint256 _amount);
     event StakeDelayed(address indexed _address, uint256 _courtID, uint256 _amount, uint256 _penalty);
+    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);
@@ -482,12 +483,28 @@ contract KlerosCore is IArbitratorV2 {
     /// @param _courtID The ID of the court.
     /// @param _stake The new stake.
     function setStake(uint96 _courtID, uint256 _stake) external {
-        if (!_setStakeForAccount(msg.sender, _courtID, _stake, 0)) revert StakingFailed();
+        if (!_setStakeForAccount(msg.sender, _courtID, _stake, 0, false)) revert StakingFailed();
     }
 
-    function setStakeBySortitionModule(address _account, uint96 _courtID, uint256 _stake, uint256 _penalty) external {
+    function withdrawDelayedStake(uint256 _delayedStakeIndex) external {
+        // All necessary checks are done in sortition module.
+        (uint256 amountToWithdraw, uint96 courtID) = sortitionModule.removeDelayedStake(_delayedStakeIndex, msg.sender);
+        if (jurors[msg.sender].stakedPnk[courtID] <= amountToWithdraw) {
+            amountToWithdraw = jurors[msg.sender].stakedPnk[courtID];
+        }
+        require(pinakion.safeTransfer(msg.sender, amountToWithdraw));
+        jurors[msg.sender].stakedPnk[courtID] -= amountToWithdraw;
+    }
+
+    function setStakeBySortitionModule(
+        address _account,
+        uint96 _courtID,
+        uint256 _stake,
+        uint256 _penalty,
+        bool _alreadyTransferred
+    ) external {
         if (msg.sender != address(sortitionModule)) revert WrongCaller();
-        _setStakeForAccount(_account, _courtID, _stake, _penalty);
+        _setStakeForAccount(_account, _courtID, _stake, _penalty, _alreadyTransferred);
     }
 
     /// @inheritdoc IArbitratorV2
@@ -768,10 +785,11 @@ contract KlerosCore is IArbitratorV2 {
         if (jurors[account].stakedPnk[dispute.courtID] >= courts[dispute.courtID].minStake + penalty) {
             // The juror still has enough staked PNKs after penalty for this court.
             uint256 newStake = jurors[account].stakedPnk[dispute.courtID] - penalty;
-            _setStakeForAccount(account, dispute.courtID, newStake, penalty);
+            // `alreadyTransferred` flag can be true only after manual stake increase, which can't happen during penalty.
+            _setStakeForAccount(account, dispute.courtID, newStake, penalty, false);
         } else if (jurors[account].stakedPnk[dispute.courtID] != 0) {
             // The juror does not have enough staked PNKs after penalty for this court, unstake them.
-            _setStakeForAccount(account, dispute.courtID, 0, penalty);
+            _setStakeForAccount(account, dispute.courtID, 0, penalty, false);
         }
         emit TokenAndETHShift(
             account,
@@ -1110,12 +1128,14 @@ contract KlerosCore is IArbitratorV2 {
     /// @param _courtID The ID of the court.
     /// @param _stake The new stake.
     /// @param _penalty Penalized amount won't be transferred back to juror when the stake is lowered.
+    /// @param _alreadyTransferred True if the tokens were already transferred. Only relevant for delayed stake execution.
     /// @return succeeded True if the call succeeded, false otherwise.
     function _setStakeForAccount(
         address _account,
         uint96 _courtID,
         uint256 _stake,
-        uint256 _penalty
+        uint256 _penalty,
+        bool _alreadyTransferred
     ) internal returns (bool succeeded) {
         if (_courtID == FORKING_COURT || _courtID > courts.length) return false;
 
@@ -1135,47 +1155,56 @@ contract KlerosCore is IArbitratorV2 {
             return true;
         }
 
-        uint256 transferredAmount;
-        if (_stake >= currentStake) {
-            transferredAmount = _stake - currentStake;
-            if (transferredAmount > 0) {
-                if (pinakion.safeTransferFrom(_account, address(this), transferredAmount)) {
-                    if (currentStake == 0) {
-                        juror.courtIDs.push(_courtID);
-                    }
-                } else {
-                    return false;
-                }
-            }
-        } else {
-            if (_stake == 0) {
-                // Keep locked PNKs in the contract and release them after dispute is executed.
-                transferredAmount = currentStake - juror.lockedPnk[_courtID] - _penalty;
+        // Don't transfer the tokens and only update the drawing chance if the transfer was already done.
+        if (!_alreadyTransferred) {
+            uint256 transferredAmount;
+            if (_stake >= currentStake) {
+                transferredAmount = _stake - currentStake;
                 if (transferredAmount > 0) {
-                    if (pinakion.safeTransfer(_account, transferredAmount)) {
-                        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;
-                            }
+                    if (pinakion.safeTransferFrom(_account, address(this), transferredAmount)) {
+                        if (currentStake == 0) {
+                            juror.courtIDs.push(_courtID);
                         }
                     } else {
                         return false;
                     }
                 }
             } else {
-                transferredAmount = currentStake - _stake - _penalty;
-                if (transferredAmount > 0) {
-                    if (!pinakion.safeTransfer(_account, transferredAmount)) {
-                        return false;
+                if (_stake == 0) {
+                    // Keep locked PNKs in the contract and release them after dispute is executed.
+                    transferredAmount = currentStake - juror.lockedPnk[_courtID] - _penalty;
+                    if (transferredAmount > 0) {
+                        if (pinakion.safeTransfer(_account, transferredAmount)) {
+                            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;
+                                }
+                            }
+                        } else {
+                            return false;
+                        }
+                    }
+                } else {
+                    transferredAmount = currentStake - _stake - _penalty;
+                    if (transferredAmount > 0) {
+                        if (!pinakion.safeTransfer(_account, transferredAmount)) {
+                            return false;
+                        }
                     }
                 }
             }
+
+            // Update juror's records.
+            juror.stakedPnk[_courtID] = _stake;
         }
 
-        // Update juror's records.
-        juror.stakedPnk[_courtID] = _stake;
+        // 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, _stake);
         emit StakeSet(_account, _courtID, _stake);
diff --git a/contracts/src/arbitration/SortitionModule.sol b/contracts/src/arbitration/SortitionModule.sol
@@ -37,6 +37,7 @@ contract SortitionModule is ISortitionModule {
         uint96 courtID; // The ID of the court.
         uint256 stake; // The new stake.
         uint256 penalty; // Penalty value, in case the stake was set during execution.
+        bool alreadyTransferred; // True if tokens were already transferred before delayed stake's execution.
     }
 
     // ************************************* //
@@ -185,17 +186,38 @@ contract SortitionModule is ISortitionModule {
 
         for (uint256 i = delayedStakeReadIndex; i < newDelayedStakeReadIndex; i++) {
             DelayedStake storage delayedStake = delayedStakes[i];
-            core.setStakeBySortitionModule(
-                delayedStake.account,
-                delayedStake.courtID,
-                delayedStake.stake,
-                delayedStake.penalty
-            );
-            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.penalty,
+                    delayedStake.alreadyTransferred
+                );
+                delete delayedStakes[i];
+            }
         }
         delayedStakeReadIndex = newDelayedStakeReadIndex;
     }
 
+    /// @dev Remove the delayed stake after its partial execution in order to return the tokens.
+    /// @param _index Index of the delayed stake to remove.
+    /// @param _sender Address that attempted removal.
+    /// @return stake Stake amount that was discarded.
+    /// @return courtID ID of the court related to delayed stake.
+    function removeDelayedStake(
+        uint256 _index,
+        address _sender
+    ) external override onlyByCore returns (uint256 stake, uint96 courtID) {
+        DelayedStake storage delayedStake = delayedStakes[_index];
+        require(delayedStake.account == _sender, "Can only remove your own stake");
+        require(delayedStake.alreadyTransferred, "No tokens to return");
+        stake = delayedStake.stake;
+        courtID = delayedStake.courtID;
+        delete delayedStakes[_index];
+    }
+
     function preStakeHook(
         address _account,
         uint96 _courtID,
@@ -208,13 +230,18 @@ contract SortitionModule is ISortitionModule {
             return preStakeHookResult.failed;
         } else {
             if (phase != Phase.staking) {
-                delayedStakes[++delayedStakeWriteIndex] = DelayedStake({
-                    account: _account,
-                    courtID: _courtID,
-                    stake: _stake,
-                    penalty: _penalty
-                });
-                return preStakeHookResult.delayed;
+                DelayedStake storage delayedStake = delayedStakes[++delayedStakeWriteIndex];
+                delayedStake.account = _account;
+                delayedStake.courtID = _courtID;
+                delayedStake.stake = _stake;
+                delayedStake.penalty = _penalty;
+                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;
@@ -264,7 +291,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, 0);
+            core.setStakeBySortitionModule(_account, courtIDs[j - 1], 0, 0, false);
         }
     }
 
diff --git a/contracts/src/arbitration/interfaces/ISortitionModule.sol b/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);
@@ -36,4 +37,6 @@ interface ISortitionModule {
     function createDisputeHook(uint256 _disputeID, uint256 _roundID) external;
 
     function postDrawHook(uint256 _disputeID, uint256 _roundID) external;
+
+    function removeDelayedStake(uint256 _index, address _sender) external returns (uint256 amount, uint96 courtID);
 }

Original file line number	Diff line number	Diff line change
`@@ -9,9 +9,10 @@ interface ISortitionModule {`
`9`	`9`	`}`
`10`	`10`
`11`	`11`	`enum preStakeHookResult {`
`12`		`- ok,`
`13`		`- delayed,`
`14`		`- failed`
	`12`	`+ ok, // Correct phase. All checks are passed.`
	`13`	`+ partiallyDelayed, // Wrong phase but stake is increased, so transfer the tokens without updating the drawing chance.`
	`14`	`+ delayed, // Wrong phase and stake is decreased. Delay the token transfer and drawing chance update.`
	`15`	`+ failed // Checks didn't pass. Do no changes.`
`15`	`16`	`}`
`16`	`17`
`17`	`18`	`event NewPhase(Phase _phase);`
`@@ -36,4 +37,6 @@ interface ISortitionModule {`
`36`	`37`	`function createDisputeHook(uint256 _disputeID, uint256 _roundID) external;`
`37`	`38`
`38`	`39`	`function postDrawHook(uint256 _disputeID, uint256 _roundID) external;`
	`40`	`+`
	`41`	`+ function removeDelayedStake(uint256 _index, address _sender) external returns (uint256 amount, uint96 courtID);`
`39`	`42`	`}`