Sortition module: externalize delayed stakes to a library

contracts/src/arbitration/SortitionModuleBase.sol

@@ -8,11 +8,14 @@ import {IDisputeKit} from "./interfaces/IDisputeKit.sol";
 import {Initializable} from "../proxy/Initializable.sol";
 import {UUPSProxiable} from "../proxy/UUPSProxiable.sol";
 import {IRNG} from "../rng/IRNG.sol";
+import {DelayedStakes} from "../libraries/DelayedStakes.sol";
 import "../libraries/Constants.sol";
 
 /// @title SortitionModuleBase
 /// @dev A factory of trees that keeps track of staked values for sortition.
 abstract contract SortitionModuleBase is ISortitionModule, Initializable, UUPSProxiable {
+    using DelayedStakes for DelayedStakes.Queue;
+
     // ************************************* //
     // *         Enums / Structs           * //
     // ************************************* //
@@ -26,13 +29,6 @@ abstract contract SortitionModuleBase is ISortitionModule, Initializable, UUPSPr
         mapping(uint256 nodeIndex => bytes32 stakePathID) nodeIndexesToIDs;
     }
 
-    struct DelayedStake {
-        address account; // The address of the juror.
-        uint96 courtID; // The ID of the court.
-        uint256 stake; // The new stake.
-        bool alreadyTransferred; // DEPRECATED. True if tokens were already transferred before delayed stake's execution.
-    }
-
     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. PNK balance including locked PNK and penalty deductions.
@@ -54,12 +50,9 @@ abstract contract SortitionModuleBase is ISortitionModule, Initializable, UUPSPr
     IRNG public rng; // The random number generator.
     uint256 public randomNumber; // Random number returned by RNG.
     uint256 public rngLookahead; // DEPRECATED: to be removed in the next redeploy
-    uint256 public delayedStakeWriteIndex; // The index of the last `delayedStake` item that was written to the array. 0 index is skipped.
-    uint256 public delayedStakeReadIndex; // The index of the next `delayedStake` item that should be processed. Starts at 1 because 0 index is skipped.
+    DelayedStakes.Queue public delayedStakesQueue; // Queue for managing delayed stakes during Drawing phase.
     mapping(bytes32 treeHash => SortitionSumTree) sortitionSumTrees; // The mapping trees by keys.
     mapping(address account => Juror) public jurors; // The jurors.
-    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 jurorAccount => mapping(uint96 courtId => uint256)) public latestDelayedStakeIndex; // DEPRECATED. 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].
 
     // ************************************* //
     // *              Events               * //
@@ -111,7 +104,7 @@ abstract contract SortitionModuleBase is ISortitionModule, Initializable, UUPSPr
         maxDrawingTime = _maxDrawingTime;
         lastPhaseChange = block.timestamp;
         rng = _rng;
-        delayedStakeReadIndex = 1;
+        delayedStakesQueue.initialize();
     }
 
     // ************************************* //
@@ -200,19 +193,9 @@ abstract contract SortitionModuleBase is ISortitionModule, Initializable, UUPSPr
     /// @param _iterations The number of delayed stakes to execute.
     function executeDelayedStakes(uint256 _iterations) external {
         if (phase != Phase.staking) revert NotStakingPhase();
-        if (delayedStakeWriteIndex < delayedStakeReadIndex) revert NoDelayedStakeToExecute();
-
-        uint256 actualIterations = (delayedStakeReadIndex + _iterations) - 1 > delayedStakeWriteIndex
-            ? (delayedStakeWriteIndex - delayedStakeReadIndex) + 1
-            : _iterations;
-        uint256 newDelayedStakeReadIndex = delayedStakeReadIndex + actualIterations;
-
-        for (uint256 i = delayedStakeReadIndex; i < newDelayedStakeReadIndex; i++) {
-            DelayedStake storage delayedStake = delayedStakes[i];
-            core.setStakeBySortitionModule(delayedStake.account, delayedStake.courtID, delayedStake.stake);
-            delete delayedStakes[i];
+        if (!delayedStakesQueue.execute(core, _iterations)) {
+            revert NoDelayedStakeToExecute();
         }
-        delayedStakeReadIndex = newDelayedStakeReadIndex;
     }
 
     function createDisputeHook(uint256 /*_disputeID*/, uint256 /*_roundID*/) external override onlyByCore {
@@ -262,10 +245,7 @@ abstract contract SortitionModuleBase is ISortitionModule, Initializable, UUPSPr
 
         if (phase != Phase.staking) {
             // Store the stake change as delayed, to be applied when the phase switches back to Staking.
-            DelayedStake storage delayedStake = delayedStakes[++delayedStakeWriteIndex];
-            delayedStake.account = _account;
-            delayedStake.courtID = _courtID;
-            delayedStake.stake = _newStake;
+            delayedStakesQueue.add(_account, _courtID, _newStake);
             emit StakeDelayed(_account, _courtID, _newStake);
             return (pnkDeposit, pnkWithdrawal, StakingResult.Delayed);
         }
@@ -547,6 +527,16 @@ abstract contract SortitionModuleBase is ISortitionModule, Initializable, UUPSPr
         }
     }
 
+    /// @dev Get a specific delayed stake by index
+    function delayedStakes(uint256 index) external view returns (DelayedStakes.Stake memory) {
+        return delayedStakesQueue.stakes[index];
+    }
+
+    /// @dev Get the number of pending delayed stakes
+    function pendingDelayedStakesCount() external view returns (uint256) {
+        return delayedStakesQueue.pendingStakesCount();
+    }
+
     // ************************************* //
     // *            Internal               * //
     // ************************************* //

contracts/src/libraries/DelayedStakes.sol

@@ -0,0 +1,92 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.24;
+
+import {KlerosCore} from "../arbitration/KlerosCore.sol";
+
+/// @title DelayedStakes
+/// @notice Library for managing delayed stakes in the sortition module
+library DelayedStakes {
+    struct Stake {
+        address account; // The address of the juror.
+        uint96 courtID; // The ID of the court.
+        uint256 stake; // The new stake.
+        bool alreadyTransferred; // DEPRECATED. True if tokens were already transferred before delayed stake's execution.
+    }
+
+    struct Queue {
+        uint256 writeIndex; // The index of the last `delayedStake` item that was written to the array. 0 index is skipped.
+        uint256 readIndex; // The index of the next `delayedStake` item that should be processed. Starts at 1 because 0 index is skipped.
+        mapping(uint256 => Stake) stakes; // The delayed stakes storage
+        mapping(address jurorAccount => mapping(uint96 courtId => uint256)) latestDelayedStakeIndex; // DEPRECATED. Maps the juror to its latest delayed stake.
+    }
+
+    /// @notice Initialize the queue
+    /// @param queue The queue to initialize
+    function initialize(Queue storage queue) internal {
+        queue.readIndex = 1; // Skip index 0
+    }
+
+    /// @notice Add a new delayed stake to the queue
+    /// @param queue The delayed stakes queue
+    /// @param account The address of the juror
+    /// @param courtID The ID of the court
+    /// @param stake The new stake amount
+    function add(Queue storage queue, address account, uint96 courtID, uint256 stake) internal {
+        uint256 newWriteIndex = ++queue.writeIndex;
+        Stake storage delayedStake = queue.stakes[newWriteIndex];
+        delayedStake.account = account;
+        delayedStake.courtID = courtID;
+        delayedStake.stake = stake;
+    }
+
+    /// @notice Execute pending delayed stakes
+    /// @param queue The delayed stakes queue
+    /// @param core The KlerosCore instance
+    /// @param iterations The maximum number of stakes to execute
+    /// @return success True if there were stakes to execute, false otherwise
+    function execute(Queue storage queue, KlerosCore core, uint256 iterations) internal returns (bool success) {
+        if (queue.writeIndex < queue.readIndex) {
+            return false; // Nothing to execute.
+        }
+
+        uint256 actualIterations = calculateIterations(queue.writeIndex, queue.readIndex, iterations);
+        uint256 newReadIndex = queue.readIndex + actualIterations;
+
+        for (uint256 i = queue.readIndex; i < newReadIndex; i++) {
+            Stake storage delayedStake = queue.stakes[i];
+            core.setStakeBySortitionModule(delayedStake.account, delayedStake.courtID, delayedStake.stake);
+            delete queue.stakes[i];
+        }
+
+        queue.readIndex = newReadIndex;
+        return true;
+    }
+
+    /// @notice Calculate the actual number of iterations to process
+    /// @param writeIndex The current write index
+    /// @param readIndex The current read index
+    /// @param requestedIterations The requested number of iterations
+    /// @return The actual number of iterations that can be processed
+    function calculateIterations(
+        uint256 writeIndex,
+        uint256 readIndex,
+        uint256 requestedIterations
+    ) internal pure returns (uint256) {
+        if (readIndex + requestedIterations - 1 > writeIndex) {
+            return writeIndex - readIndex + 1;
+        } else {
+            return requestedIterations;
+        }
+    }
+
+    /// @notice Get the number of pending stakes
+    /// @param queue The delayed stakes queue
+    /// @return The number of stakes waiting to be executed
+    function pendingStakesCount(Queue storage queue) internal view returns (uint256) {
+        if (queue.writeIndex < queue.readIndex) {
+            return 0;
+        }
+        return queue.writeIndex - queue.readIndex + 1;
+    }
+}