feat(KC): instant staking

Author: unknownunknown1

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);

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);
  }
  }
 

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);
 }