PNK penalties applied to the sortition tree stakes

contracts/hardhat.config.ts

@@ -31,7 +31,7 @@ const config: HardhatUserConfig = {
           viaIR: process.env.VIA_IR !== "false", // Defaults to true
           optimizer: {
             enabled: true,
-            runs: 10000,
+            runs: 2000,
           },
           outputSelection: {
             "*": {

contracts/src/arbitration/KlerosCoreBase.sol

@@ -60,6 +60,7 @@ abstract contract KlerosCoreBase is IArbitratorV2, Initializable, UUPSProxiable
         uint256 repartitions; // A counter of reward repartitions made in this round.
         uint256 pnkPenalties; // The amount of PNKs collected from penalties in this round.
         address[] drawnJurors; // Addresses of the jurors that were drawn in this round.
+        uint96[] drawnJurorFromCourtIDs; // The courtIDs where the juror was drawn from, possibly their stake in a subcourt.
         uint256 sumFeeRewardPaid; // Total sum of arbitration fees paid to coherent jurors as a reward in this round.
         uint256 sumPnkRewardPaid; // Total sum of PNK paid to coherent jurors as a reward in this round.
         IERC20 feeToken; // The token used for paying fees in this round.
@@ -610,13 +611,14 @@ abstract contract KlerosCoreBase is IArbitratorV2, Initializable, UUPSProxiable
         uint256 startIndex = round.drawIterations; // for gas: less storage reads
         uint256 i;
         while (i < _iterations && round.drawnJurors.length < round.nbVotes) {
-            address drawnAddress = disputeKit.draw(_disputeID, startIndex + i++);
+            (address drawnAddress, uint96 fromSubcourtID) = disputeKit.draw(_disputeID, startIndex + i++);
             if (drawnAddress == address(0)) {
                 continue;
             }
             sortitionModule.lockStake(drawnAddress, round.pnkAtStakePerJuror);
             emit Draw(drawnAddress, _disputeID, currentRound, round.drawnJurors.length);
             round.drawnJurors.push(drawnAddress);
+            round.drawnJurorFromCourtIDs.push(fromSubcourtID != 0 ? fromSubcourtID : dispute.courtID);
             if (round.drawnJurors.length == round.nbVotes) {
                 sortitionModule.postDrawHook(_disputeID, currentRound);
             }
@@ -786,7 +788,12 @@ abstract contract KlerosCoreBase is IArbitratorV2, Initializable, UUPSProxiable
         sortitionModule.unlockStake(account, penalty);
 
         // Apply the penalty to the staked PNKs.
-        (uint256 pnkBalance, uint256 availablePenalty) = sortitionModule.penalizeStake(account, penalty);
+        uint96 penalizedInCourtID = round.drawnJurorFromCourtIDs[_params.repartition];
+        (uint256 pnkBalance, uint256 availablePenalty) = sortitionModule.setStakePenalty(
+            account,
+            penalizedInCourtID,
+            penalty
+        );
         _params.pnkPenaltiesInRound += availablePenalty;
         emit TokenAndETHShift(
             account,

contracts/src/arbitration/SortitionModuleBase.sol

@@ -17,14 +17,20 @@ abstract contract SortitionModuleBase is ISortitionModule, Initializable, UUPSPr
     // *         Enums / Structs           * //
     // ************************************* //
 
+    struct SubCourtStakes {
+        uint256 totalStakedInCourts;
+        uint96[MAX_STAKE_PATHS] courtIDs;
+        uint256[MAX_STAKE_PATHS] stakedInCourts;
+    }
+
     struct SortitionSumTree {
         uint256 K; // The maximum number of children per node.
-        // We use this to keep track of vacant positions in the tree after removing a leaf. This is for keeping the tree as balanced as possible without spending gas on moving nodes around.
-        uint256[] stack;
-        uint256[] nodes;
+        uint256[] stack; // We use this to keep track of vacant positions in the tree after removing a leaf. This is for keeping the tree as balanced as possible without spending gas on moving nodes around.
+        uint256[] nodes; // The tree nodes.
         // Two-way mapping of IDs to node indexes. Note that node index 0 is reserved for the root node, and means the ID does not have a node.
-        mapping(bytes32 => uint256) IDsToNodeIndexes;
-        mapping(uint256 => bytes32) nodeIndexesToIDs;
+        mapping(bytes32 stakePathID => uint256 nodeIndex) IDsToNodeIndexes;
+        mapping(uint256 nodeIndex => bytes32 stakePathID) nodeIndexesToIDs;
+        mapping(bytes32 stakePathID => SubCourtStakes subcourtStakes) IDsToSubCourtStakes;
     }
 
     struct DelayedStake {
@@ -36,7 +42,7 @@ abstract contract SortitionModuleBase is ISortitionModule, Initializable, UUPSPr
 
     struct Juror {
         uint96[] courtIDs; // The IDs of courts where the juror's stake path ends. A stake path is a path from the general court to a court the juror directly staked in using `_setStake`.
-        uint256 stakedPnk; // The juror's total amount of tokens staked in subcourts. Reflects actual pnk balance.
+        uint256 stakedPnk; // The juror's total amount of tokens staked in subcourts. PNK balance including locked PNK and penalty deductions.
         uint256 lockedPnk; // The juror's total amount of tokens locked in disputes.
     }
 
@@ -306,6 +312,28 @@ abstract contract SortitionModuleBase is ISortitionModule, Initializable, UUPSPr
         _setStake(_account, _courtID, _pnkDeposit, _pnkWithdrawal, _newStake);
     }
 
+    function setStakePenalty(
+        address _account,
+        uint96 _courtID,
+        uint256 _penalty
+    ) external override onlyByCore returns (uint256 pnkBalance, uint256 availablePenalty) {
+        Juror storage juror = jurors[_account];
+        availablePenalty = _penalty;
+        if (juror.stakedPnk < _penalty) {
+            availablePenalty = juror.stakedPnk;
+        }
+
+        if (availablePenalty == 0) return (juror.stakedPnk, 0); // No penalty to apply.
+
+        uint256 currentStake = stakeOf(_account, _courtID);
+        uint256 newStake = 0;
+        if (currentStake >= availablePenalty) {
+            newStake = currentStake - availablePenalty;
+        }
+        _setStake(_account, _courtID, 0, availablePenalty, newStake);
+        pnkBalance = juror.stakedPnk; // updated by _setStake()
+    }
+
     function _setStake(
         address _account,
         uint96 _courtID,
@@ -339,12 +367,14 @@ abstract contract SortitionModuleBase is ISortitionModule, Initializable, UUPSPr
         bytes32 stakePathID = _accountAndCourtIDToStakePathID(_account, _courtID);
         bool finished = false;
         uint96 currentCourtID = _courtID;
+        uint96 fromSubCourtID = 0; // 0 means it is not coming from a subcourt.
         while (!finished) {
             // Tokens are also implicitly staked in parent courts through sortition module to increase the chance of being drawn.
-            _set(bytes32(uint256(currentCourtID)), _newStake, stakePathID);
+            _set(bytes32(uint256(currentCourtID)), _newStake, stakePathID, fromSubCourtID);
             if (currentCourtID == GENERAL_COURT) {
                 finished = true;
             } else {
+                fromSubCourtID = currentCourtID;
                 (currentCourtID, , , , , , ) = core.courts(currentCourtID); // Get the parent court.
             }
         }
@@ -367,25 +397,6 @@ abstract contract SortitionModuleBase is ISortitionModule, Initializable, UUPSPr
         }
     }
 
-    function penalizeStake(
-        address _account,
-        uint256 _relativeAmount
-    ) external override onlyByCore returns (uint256 pnkBalance, uint256 availablePenalty) {
-        Juror storage juror = jurors[_account];
-        uint256 stakedPnk = juror.stakedPnk;
-
-        if (stakedPnk >= _relativeAmount) {
-            availablePenalty = _relativeAmount;
-            juror.stakedPnk -= _relativeAmount;
-        } else {
-            availablePenalty = stakedPnk;
-            juror.stakedPnk = 0;
-        }
-
-        pnkBalance = juror.stakedPnk;
-        return (pnkBalance, availablePenalty);
-    }
-
     /// @dev Unstakes the inactive juror from all courts.
     /// `O(n * (p * log_k(j)) )` where
     /// `n` is the number of courts the juror has staked in,
@@ -433,12 +444,12 @@ abstract contract SortitionModuleBase is ISortitionModule, Initializable, UUPSPr
         bytes32 _key,
         uint256 _coreDisputeID,
         uint256 _nonce
-    ) public view override returns (address drawnAddress) {
+    ) public view override returns (address drawnAddress, uint96 fromSubcourtID) {
         if (phase != Phase.drawing) revert NotDrawingPhase();
         SortitionSumTree storage tree = sortitionSumTrees[_key];
 
         if (tree.nodes[0] == 0) {
-            return address(0); // No jurors staked.
+            return (address(0), 0); // No jurors staked.
         }
 
         uint256 currentDrawnNumber = uint256(keccak256(abi.encodePacked(randomNumber, _coreDisputeID, _nonce))) %
@@ -462,7 +473,33 @@ abstract contract SortitionModuleBase is ISortitionModule, Initializable, UUPSPr
                 }
             }
         }
-        drawnAddress = _stakePathIDToAccount(tree.nodeIndexesToIDs[treeIndex]);
+
+        bytes32 stakePathID = tree.nodeIndexesToIDs[treeIndex];
+        drawnAddress = _stakePathIDToAccount(stakePathID);
+
+        // Identify which subcourt was selected based on currentDrawnNumber
+        SubCourtStakes storage subcourtStakes = tree.IDsToSubCourtStakes[stakePathID];
+
+        // The current court stake is the node value minus all subcourt stakes
+        uint256 currentCourtStake = 0;
+        if (tree.nodes[treeIndex] > subcourtStakes.totalStakedInCourts) {
+            currentCourtStake = tree.nodes[treeIndex] - subcourtStakes.totalStakedInCourts;
+        }
+
+        // Check if the drawn number falls within current court range
+        if (currentDrawnNumber >= currentCourtStake) {
+            // Find which subcourt range contains the drawn number
+            uint256 accumulatedStake = currentCourtStake;
+            for (uint256 i = 0; i < MAX_STAKE_PATHS; i++) {
+                if (subcourtStakes.stakedInCourts[i] > 0) {
+                    if (currentDrawnNumber < accumulatedStake + subcourtStakes.stakedInCourts[i]) {
+                        fromSubcourtID = subcourtStakes.courtIDs[i];
+                        break;
+                    }
+                    accumulatedStake += subcourtStakes.stakedInCourts[i];
+                }
+            }
+        }
     }
 
     /// @dev Get the stake of a juror in a court.
@@ -487,6 +524,13 @@ abstract contract SortitionModuleBase is ISortitionModule, Initializable, UUPSPr
         return tree.nodes[treeIndex];
     }
 
+    /// @dev Gets the balance of a juror in a court.
+    /// @param _juror The address of the juror.
+    /// @param _courtID The ID of the court.
+    /// @return totalStaked The total amount of tokens staked including locked tokens and penalty deductions. Equivalent to the effective stake in the General court.
+    /// @return totalLocked The total amount of tokens locked in disputes.
+    /// @return stakedInCourt The amount of tokens staked in the specified court including locked tokens and penalty deductions.
+    /// @return nbCourts The number of courts the juror has directly staked in.
     function getJurorBalance(
         address _juror,
         uint96 _courtID
@@ -546,6 +590,41 @@ abstract contract SortitionModuleBase is ISortitionModule, Initializable, UUPSPr
         }
     }
 
+    /// @dev Update the subcourt stakes.
+    /// @param _subcourtStakes The subcourt stakes.
+    /// @param _value The value to update.
+    /// @param _fromSubCourtID The ID of the subcourt from which the stake is coming from, or 0 if the stake does not come from a subcourt.
+    /// `O(1)` complexity with `MAX_STAKE_PATHS` a small constant.
+    function _updateSubCourtStakes(
+        SubCourtStakes storage _subcourtStakes,
+        uint256 _value,
+        uint96 _fromSubCourtID
+    ) internal {
+        // Update existing stake item if found
+        for (uint256 i = 0; i < MAX_STAKE_PATHS; i++) {
+            if (_subcourtStakes.courtIDs[i] == _fromSubCourtID) {
+                if (_value == 0) {
+                    delete _subcourtStakes.courtIDs[i];
+                    delete _subcourtStakes.stakedInCourts[i];
+                } else {
+                    _subcourtStakes.totalStakedInCourts += _value;
+                    _subcourtStakes.totalStakedInCourts -= _subcourtStakes.stakedInCourts[i];
+                    _subcourtStakes.stakedInCourts[i] = _value;
+                }
+                return;
+            }
+        }
+        // Not found so add a new stake item
+        for (uint256 i = 0; i < MAX_STAKE_PATHS; i++) {
+            if (_subcourtStakes.courtIDs[i] == 0) {
+                _subcourtStakes.courtIDs[i] = _fromSubCourtID;
+                _subcourtStakes.totalStakedInCourts += _value;
+                _subcourtStakes.stakedInCourts[i] = _value;
+                return;
+            }
+        }
+    }
+
     /// @dev Retrieves a juror's address from the stake path ID.
     /// @param _stakePathID The stake path ID to unpack.
     /// @return account The account.
@@ -579,10 +658,11 @@ abstract contract SortitionModuleBase is ISortitionModule, Initializable, UUPSPr
     /// @param _key The key of the tree.
     /// @param _value The new value.
     /// @param _ID The ID of the value.
+    /// @param _fromSubCourtID The ID of the subcourt from which the stake is coming from, or 0 if the stake does not come from a subcourt.
     /// `O(log_k(n))` where
     /// `k` is the maximum number of children per node in the tree,
     ///  and `n` is the maximum number of nodes ever appended.
-    function _set(bytes32 _key, uint256 _value, bytes32 _ID) internal {
+    function _set(bytes32 _key, uint256 _value, bytes32 _ID, uint96 _fromSubCourtID) internal {
         SortitionSumTree storage tree = sortitionSumTrees[_key];
         uint256 treeIndex = tree.IDsToNodeIndexes[_ID];
 
@@ -652,6 +732,8 @@ abstract contract SortitionModuleBase is ISortitionModule, Initializable, UUPSPr
                 _updateParents(_key, treeIndex, plusOrMinus, plusOrMinusValue);
             }
         }
+
+        _updateSubCourtStakes(tree.IDsToSubCourtStakes[_ID], _value, _fromSubCourtID);
     }
 
     /// @dev Packs an account and a court ID into a stake path ID.

contracts/src/arbitration/dispute-kits/DisputeKitClassicBase.sol

@@ -228,7 +228,7 @@ abstract contract DisputeKitClassicBase is IDisputeKit, Initializable, UUPSProxi
     function draw(
         uint256 _coreDisputeID,
         uint256 _nonce
-    ) external override onlyByCore notJumped(_coreDisputeID) returns (address drawnAddress) {
+    ) external override onlyByCore notJumped(_coreDisputeID) returns (address drawnAddress, uint96 fromSubcourtID) {
         uint256 localDisputeID = coreDisputeIDToLocal[_coreDisputeID];
         Dispute storage dispute = disputes[localDisputeID];
         uint256 localRoundID = dispute.rounds.length - 1;
@@ -238,10 +238,10 @@ abstract contract DisputeKitClassicBase is IDisputeKit, Initializable, UUPSProxi
         (uint96 courtID, , , , ) = core.disputes(_coreDisputeID);
         bytes32 key = bytes32(uint256(courtID)); // Get the ID of the tree.
 
-        drawnAddress = sortitionModule.draw(key, _coreDisputeID, _nonce);
+        (drawnAddress, fromSubcourtID) = sortitionModule.draw(key, _coreDisputeID, _nonce);
         if (drawnAddress == address(0)) {
             // Sortition can return 0 address if no one has staked yet.
-            return drawnAddress;
+            return (drawnAddress, fromSubcourtID);
         }
 
         if (_postDrawCheck(round, _coreDisputeID, drawnAddress)) {

contracts/src/arbitration/interfaces/IDisputeKit.sol

@@ -48,7 +48,10 @@ interface IDisputeKit {
     /// @param _coreDisputeID The ID of the dispute in Kleros Core, not in the Dispute Kit.
     /// @param _nonce Nonce.
     /// @return drawnAddress The drawn address.
-    function draw(uint256 _coreDisputeID, uint256 _nonce) external returns (address drawnAddress);
+    function draw(
+        uint256 _coreDisputeID,
+        uint256 _nonce
+    ) external returns (address drawnAddress, uint96 fromSubcourtID);
 
     // ************************************* //
     // *           Public Views            * //

contracts/src/arbitration/interfaces/ISortitionModule.sol

@@ -29,20 +29,25 @@ interface ISortitionModule {
         uint256 _newStake
     ) external;
 
+    function setStakePenalty(
+        address _account,
+        uint96 _courtID,
+        uint256 _penalty
+    ) external returns (uint256 pnkBalance, uint256 availablePenalty);
+
     function setJurorInactive(address _account) external;
 
     function lockStake(address _account, uint256 _relativeAmount) external;
 
     function unlockStake(address _account, uint256 _relativeAmount) external;
 
-    function penalizeStake(
-        address _account,
-        uint256 _relativeAmount
-    ) external returns (uint256 pnkBalance, uint256 availablePenalty);
-
     function notifyRandomNumber(uint256 _drawnNumber) external;
 
-    function draw(bytes32 _court, uint256 _coreDisputeID, uint256 _nonce) external view returns (address);
+    function draw(
+        bytes32 _court,
+        uint256 _coreDisputeID,
+        uint256 _nonce
+    ) external view returns (address drawnAddress, uint96 fromSubcourtID);
 
     function getJurorBalance(
         address _juror,

contracts/src/arbitration/university/KlerosCoreUniversity.sol

@@ -59,6 +59,7 @@ contract KlerosCoreUniversity is IArbitratorV2, UUPSProxiable, Initializable {
         uint256 repartitions; // A counter of reward repartitions made in this round.
         uint256 pnkPenalties; // The amount of PNKs collected from penalties in this round.
         address[] drawnJurors; // Addresses of the jurors that were drawn in this round.
+        uint96[] drawnJurorFromCourtIDs; // The courtIDs where the juror was drawn from, possibly their stake in a subcourt.
         uint256 sumFeeRewardPaid; // Total sum of arbitration fees paid to coherent jurors as a reward in this round.
         uint256 sumPnkRewardPaid; // Total sum of PNK paid to coherent jurors as a reward in this round.
         IERC20 feeToken; // The token used for paying fees in this round.
@@ -599,13 +600,14 @@ contract KlerosCoreUniversity is IArbitratorV2, UUPSProxiable, Initializable {
         {
             IDisputeKit disputeKit = disputeKits[round.disputeKitID];
             uint256 iteration = round.drawIterations + 1;
-            address drawnAddress = disputeKit.draw(_disputeID, iteration);
+            (address drawnAddress, uint96 fromSubcourtID) = disputeKit.draw(_disputeID, iteration);
             if (drawnAddress == address(0)) {
                 revert NoJurorDrawn();
             }
             sortitionModule.lockStake(drawnAddress, round.pnkAtStakePerJuror);
             emit Draw(drawnAddress, _disputeID, currentRound, round.drawnJurors.length);
             round.drawnJurors.push(drawnAddress);
+            round.drawnJurorFromCourtIDs.push(fromSubcourtID != 0 ? fromSubcourtID : dispute.courtID);
             if (round.drawnJurors.length == round.nbVotes) {
                 sortitionModule.postDrawHook(_disputeID, currentRound);
             }
@@ -774,7 +776,12 @@ contract KlerosCoreUniversity is IArbitratorV2, UUPSProxiable, Initializable {
         sortitionModule.unlockStake(account, penalty);
 
         // Apply the penalty to the staked PNKs.
-        (uint256 pnkBalance, uint256 availablePenalty) = sortitionModule.penalizeStake(account, penalty);
+        uint96 penalizedInCourtID = round.drawnJurorFromCourtIDs[_params.repartition];
+        (uint256 pnkBalance, uint256 availablePenalty) = sortitionModule.setStakePenalty(
+            account,
+            penalizedInCourtID,
+            penalty
+        );
         _params.pnkPenaltiesInRound += availablePenalty;
         emit TokenAndETHShift(
             account,