refactor: extracted distinct rewards and penalties execution functions

jaybuidl · jaybuidl · commit b7418c00b0fc · 2023-05-25T00:45:31.000+01:00
diff --git a/contracts/src/arbitration/KlerosCore.sol b/contracts/src/arbitration/KlerosCore.sol
@@ -657,103 +657,144 @@ contract KlerosCore is IArbitrator {
         uint256 numberOfVotesInRound = round.drawnJurors.length;
         uint256 coherentCount = disputeKit.getCoherentCount(_disputeID, _round); // Total number of jurors that are eligible to a reward in this round.
 
-        address account; // Address of the juror.
-        uint256 degreeOfCoherence; // [0, 1] value that determines how coherent the juror was in this round, in basis points.
-
         if (coherentCount == 0) {
             // We loop over the votes once as there are no rewards because it is not a tie and no one in this round is coherent with the final outcome.
             if (end > numberOfVotesInRound) end = numberOfVotesInRound;
         } else {
             // We loop over the votes twice, first to collect penalties, and second to distribute them as rewards along with arbitration fees.
             if (end > numberOfVotesInRound * 2) end = numberOfVotesInRound * 2;
         }
-
         for (uint256 i = round.repartitions; i < end; i++) {
             if (i < numberOfVotesInRound) {
-                // Penalty.
-                degreeOfCoherence = disputeKit.getDegreeOfCoherence(_disputeID, _round, i);
+                penaltiesInRoundCache = _executePenalties(
+                    _disputeID,
+                    _round,
+                    coherentCount,
+                    numberOfVotesInRound,
+                    penaltiesInRoundCache,
+                    i
+                );
+            } else {
+                _executeRewards(_disputeID, _round, coherentCount, numberOfVotesInRound, penaltiesInRoundCache, i);
+            }
+        }
+        if (round.penalties != penaltiesInRoundCache) {
+            round.penalties = penaltiesInRoundCache;
+        }
+        round.repartitions = end;
+    }
 
-                // Make sure the degree doesn't exceed 1, though it should be ensured by the dispute kit.
-                if (degreeOfCoherence > ALPHA_DIVISOR) {
-                    degreeOfCoherence = ALPHA_DIVISOR;
-                }
+    function _executePenalties(
+        uint256 _disputeID,
+        uint256 _round,
+        uint256 _coherentCount,
+        uint256 _numberOfVotesInRound,
+        uint256 _penaltiesInRound,
+        uint256 _repartition
+    ) internal returns (uint256) {
+        Dispute storage dispute = disputes[_disputeID];
+        Round storage round = dispute.rounds[_round];
+        IDisputeKit disputeKit = disputeKitNodes[round.disputeKitID].disputeKit;
 
-                // Fully coherent jurors won't be penalized.
-                uint256 penalty = (round.tokensAtStakePerJuror * (ALPHA_DIVISOR - degreeOfCoherence)) / ALPHA_DIVISOR;
-                penaltiesInRoundCache += penalty;
+        // [0, 1] value that determines how coherent the juror was in this round, in basis points.
+        uint256 degreeOfCoherence = disputeKit.getDegreeOfCoherence(_disputeID, _round, _repartition);
+        if (degreeOfCoherence > ALPHA_DIVISOR) {
+            // Make sure the degree doesn't exceed 1, though it should be ensured by the dispute kit.
+            degreeOfCoherence = ALPHA_DIVISOR;
+        }
 
-                account = round.drawnJurors[i];
-                jurors[account].lockedTokens[dispute.courtID] -= penalty; // Release this part of locked tokens.
+        // Fully coherent jurors won't be penalized.
+        uint256 penalty = (round.tokensAtStakePerJuror * (ALPHA_DIVISOR - degreeOfCoherence)) / ALPHA_DIVISOR;
+        _penaltiesInRound += penalty;
+
+        // Unlock the tokens affected by the penalty
+        address account = round.drawnJurors[_repartition];
+        jurors[account].lockedTokens[dispute.courtID] -= penalty;
+
+        // Apply the penalty to the staked tokens
+        if (jurors[account].stakedTokens[dispute.courtID] >= courts[dispute.courtID].minStake + penalty) {
+            // The juror still has enough staked token after penalty for this court.
+            uint256 newStake = jurors[account].stakedTokens[dispute.courtID] - penalty;
+            _setStakeForAccount(account, dispute.courtID, newStake, penalty);
+        } else if (jurors[account].stakedTokens[dispute.courtID] != 0) {
+            // The juror does not have enough staked tokens after penalty for this court, unstake them.
+            _setStakeForAccount(account, dispute.courtID, 0, penalty);
+        }
+        emit TokenAndETHShift(account, _disputeID, _round, degreeOfCoherence, -int256(penalty), 0);
 
-                // Can only update the stake if it is able to cover the minStake and penalty, otherwise unstake from the court.
-                if (jurors[account].stakedTokens[dispute.courtID] >= courts[dispute.courtID].minStake + penalty) {
-                    uint256 newStake = jurors[account].stakedTokens[dispute.courtID] - penalty;
-                    _setStakeForAccount(account, dispute.courtID, newStake, penalty);
-                } else if (jurors[account].stakedTokens[dispute.courtID] != 0) {
-                    _setStakeForAccount(account, dispute.courtID, 0, penalty);
-                }
+        if (!disputeKit.isVoteActive(_disputeID, _round, _repartition)) {
+            // The juror is inactive, unstake them.
+            sortitionModule.setJurorInactive(account);
+        }
+        if (_repartition == _numberOfVotesInRound - 1 && _coherentCount == 0) {
+            // No one was coherent, send the rewards to the governor.
+            payable(governor).send(round.totalFeesForJurors);
+            _safeTransfer(governor, _penaltiesInRound);
+            emit LeftoverRewardSent(_disputeID, _round, _penaltiesInRound, round.totalFeesForJurors);
+        }
+        return _penaltiesInRound;
+    }
 
-                // Unstake the juror if he lost due to inactivity.
-                if (!disputeKit.isVoteActive(_disputeID, _round, i)) {
-                    sortitionModule.setJurorInactive(account);
-                }
-                emit TokenAndETHShift(account, _disputeID, _round, degreeOfCoherence, -int256(penalty), 0);
-
-                if (i == numberOfVotesInRound - 1) {
-                    if (coherentCount == 0) {
-                        // No one was coherent. Send the rewards to governor.
-                        payable(governor).send(round.totalFeesForJurors);
-                        _safeTransfer(governor, penaltiesInRoundCache);
-                        emit LeftoverRewardSent(_disputeID, _round, penaltiesInRoundCache, round.totalFeesForJurors);
-                    }
-                }
-            } else {
-                // Reward.
-                degreeOfCoherence = disputeKit.getDegreeOfCoherence(_disputeID, _round, i % numberOfVotesInRound);
+    function _executeRewards(
+        uint256 _disputeID,
+        uint256 _round,
+        uint256 _coherentCount,
+        uint256 _numberOfVotesInRound,
+        uint256 _penaltiesInRound,
+        uint256 _repartition
+    ) internal {
+        Dispute storage dispute = disputes[_disputeID];
+        Round storage round = dispute.rounds[_round];
+        IDisputeKit disputeKit = disputeKitNodes[round.disputeKitID].disputeKit;
 
-                // Make sure the degree doesn't exceed 1, though it should be ensured by the dispute kit.
-                if (degreeOfCoherence > ALPHA_DIVISOR) {
-                    degreeOfCoherence = ALPHA_DIVISOR;
-                }
+        // [0, 1] value that determines how coherent the juror was in this round, in basis points.
+        uint256 degreeOfCoherence = disputeKit.getDegreeOfCoherence(
+            _disputeID,
+            _round,
+            _repartition % _numberOfVotesInRound
+        );
 
-                account = round.drawnJurors[i % numberOfVotesInRound];
+        // Make sure the degree doesn't exceed 1, though it should be ensured by the dispute kit.
+        if (degreeOfCoherence > ALPHA_DIVISOR) {
+            degreeOfCoherence = ALPHA_DIVISOR;
+        }
 
-                // Release the rest of the tokens of the juror for this round.
-                jurors[account].lockedTokens[dispute.courtID] -=
-                    (round.tokensAtStakePerJuror * degreeOfCoherence) /
-                    ALPHA_DIVISOR;
+        address account = round.drawnJurors[_repartition % _numberOfVotesInRound];
 
-                // Give back the locked tokens in case the juror fully unstaked earlier.
-                if (jurors[account].stakedTokens[dispute.courtID] == 0) {
-                    uint256 tokenLocked = (round.tokensAtStakePerJuror * degreeOfCoherence) / ALPHA_DIVISOR;
-                    _safeTransfer(account, tokenLocked);
-                }
+        // Release the rest of the tokens of the juror for this round.
+        jurors[account].lockedTokens[dispute.courtID] -=
+            (round.tokensAtStakePerJuror * degreeOfCoherence) /
+            ALPHA_DIVISOR;
 
-                uint256 tokenReward = ((penaltiesInRoundCache / coherentCount) * degreeOfCoherence) / ALPHA_DIVISOR;
-                round.sumTokenRewardPaid += tokenReward;
-                uint256 ethReward = ((round.totalFeesForJurors / coherentCount) * degreeOfCoherence) / ALPHA_DIVISOR;
-                round.sumRewardPaid += ethReward;
-                _safeTransfer(account, tokenReward);
-                payable(account).send(ethReward);
-                emit TokenAndETHShift(
-                    account,
-                    _disputeID,
-                    _round,
-                    degreeOfCoherence,
-                    int256(tokenReward),
-                    int256(ethReward)
-                );
+        // Give back the locked tokens in case the juror fully unstaked earlier.
+        if (jurors[account].stakedTokens[dispute.courtID] == 0) {
+            uint256 tokenLocked = (round.tokensAtStakePerJuror * degreeOfCoherence) / ALPHA_DIVISOR;
+            _safeTransfer(account, tokenLocked);
+        }
 
-                if (i == numberOfVotesInRound * 2 - 1) {
-                    _executeResidualRewards(_disputeID, _round, penaltiesInRoundCache); // distinct function to avoid a "stack too deep error"
+        // Transfer the rewards
+        uint256 tokenReward = ((_penaltiesInRound / _coherentCount) * degreeOfCoherence) / ALPHA_DIVISOR;
+        round.sumTokenRewardPaid += tokenReward;
+        uint256 ethReward = ((round.totalFeesForJurors / _coherentCount) * degreeOfCoherence) / ALPHA_DIVISOR;
+        round.sumRewardPaid += ethReward;
+        _safeTransfer(account, tokenReward);
+        payable(account).send(ethReward);
+        emit TokenAndETHShift(account, _disputeID, _round, degreeOfCoherence, int256(tokenReward), int256(ethReward));
+
+        // Transfer any residual rewards to the governor. It may happen due to partial coherence of the jurors.
+        if (_repartition == _numberOfVotesInRound * 2 - 1) {
+            uint256 leftoverReward = round.totalFeesForJurors - round.sumRewardPaid;
+            uint256 leftoverTokenReward = _penaltiesInRound - round.sumTokenRewardPaid;
+            if (leftoverReward != 0 || leftoverTokenReward != 0) {
+                if (leftoverReward != 0) {
+                    payable(governor).send(leftoverReward);
+                }
+                if (leftoverTokenReward != 0) {
+                    _safeTransfer(governor, leftoverTokenReward);
                 }
+                emit LeftoverRewardSent(_disputeID, _round, leftoverTokenReward, leftoverReward);
             }
         }
-
-        if (round.penalties != penaltiesInRoundCache) {
-            round.penalties = penaltiesInRoundCache;
-        }
-        round.repartitions = end;
     }
 
     /// @dev Executes a specified dispute's ruling. UNTRUSTED.
@@ -949,28 +990,6 @@ contract KlerosCore is IArbitrator {
         emit DisputeKitEnabled(_courtID, _disputeKitID, _enable);
     }
 
-    /// @dev Sends the residual rewards to the governor. It may happen due to partial coherence of the jurors.
-    /// @param _disputeID The ID of the dispute.
-    /// @param _round The round to execute.
-    /// @param _penaltiesInRound The total penalties in the round.
-    function _executeResidualRewards(uint256 _disputeID, uint256 _round, uint256 _penaltiesInRound) internal {
-        Dispute storage dispute = disputes[_disputeID];
-        Round storage round = dispute.rounds[_round];
-        uint256 leftoverReward;
-        uint256 leftoverTokenReward;
-        if (round.totalFeesForJurors > round.sumRewardPaid) {
-            leftoverReward = round.totalFeesForJurors - round.sumRewardPaid;
-            payable(governor).send(leftoverReward);
-        }
-        if (_penaltiesInRound > round.sumTokenRewardPaid) {
-            leftoverTokenReward = _penaltiesInRound - round.sumTokenRewardPaid;
-            _safeTransfer(governor, leftoverTokenReward);
-        }
-        if (leftoverReward != 0 || leftoverTokenReward != 0) {
-            emit LeftoverRewardSent(_disputeID, _round, leftoverTokenReward, leftoverReward);
-        }
-    }
-
     /// @dev Sets the specified juror's stake in a court.
     /// `O(n + p * log_k(j))` where
     /// `n` is the number of courts the juror has staked in,
