refactor: merged StakeController._setStake() and _setStakeBySystem(), removed SortitionModule refs

jaybuidl · jaybuidl · commit 9d8737389ad6 · 2025-06-09T12:50:44.000+01:00
diff --git a/contracts/src/arbitration/KlerosCoreXBase.sol b/contracts/src/arbitration/KlerosCoreXBase.sol
@@ -1131,7 +1131,7 @@ abstract contract KlerosCoreXBase is IArbitratorV2, Initializable, UUPSProxiable
         uint96 _courtID,
         uint256 _newStake
     ) internal virtual returns (bool success) {
-        (uint256 pnkDeposit, uint256 pnkWithdrawal, StakingResult stakingResult) = stakeController.setStakeBySystem(
+        (uint256 pnkDeposit, uint256 pnkWithdrawal, StakingResult stakingResult) = stakeController.setStake(
             _account,
             _courtID,
             _newStake
@@ -1213,7 +1213,6 @@ abstract contract KlerosCoreXBase is IArbitratorV2, Initializable, UUPSProxiable
     error GovernorOnly();
     error GuardianOrGovernorOnly();
     error DisputeKitOnly();
-    error SortitionModuleOnly();
     error StakeControllerOnly();
     error UnsuccessfulCall();
     error InvalidDisputKitParent();
diff --git a/contracts/src/arbitration/StakeController.sol b/contracts/src/arbitration/StakeController.sol
@@ -10,7 +10,7 @@ import {RNG} from "../rng/RNG.sol";
 
 /// @title StakeController
 /// @notice Basic implementation of the Stake Controller
-/// @dev Coordinates between Vault and SortitionModule for the new architecture
+/// @dev Coordinates between Vault and SortitionSumTree for the new architecture
 contract StakeController is StakeControllerBase {
     string public constant override version = "1.0.0";
 
@@ -27,7 +27,7 @@ contract StakeController is StakeControllerBase {
     /// @param _governor The governor's address.
     /// @param _core The KlerosCore contract.
     /// @param _vault The Vault contract.
-    /// @param _sortitionModule The SortitionModule contract.
+    /// @param _sortition The SortitionSumTree contract.
     /// @param _minStakingTime The minimum staking time.
     /// @param _maxDrawingTime The maximum drawing time.
     /// @param _rng The random number generator.
@@ -36,7 +36,7 @@ contract StakeController is StakeControllerBase {
         address _governor,
         KlerosCoreXBase _core,
         IVault _vault,
-        ISortitionSumTree _sortitionModule,
+        ISortitionSumTree _sortition,
         uint256 _minStakingTime,
         uint256 _maxDrawingTime,
         RNG _rng,
@@ -46,7 +46,7 @@ contract StakeController is StakeControllerBase {
             _governor,
             _core,
             _vault,
-            _sortitionModule,
+            _sortition,
             _minStakingTime,
             _maxDrawingTime,
             _rng,
diff --git a/contracts/src/arbitration/StakeControllerBase.sol b/contracts/src/arbitration/StakeControllerBase.sol
@@ -13,7 +13,7 @@ import {RNG} from "../rng/RNG.sol";
 import "../libraries/Constants.sol";
 
 /// @title StakeControllerBase
-/// @notice Abstract base contract for coordinating between Vault and SortitionModule
+/// @notice Abstract base contract for coordinating between Vault and SortitionSumTree
 /// @dev Manages phases, delayed stakes, and coordination logic
 abstract contract StakeControllerBase is IStakeController, Initializable, UUPSProxiable {
     // ************************************* //
@@ -39,7 +39,7 @@ abstract contract StakeControllerBase is IStakeController, Initializable, UUPSPr
     address public governor; // The governor of the contract.
     KlerosCoreXBase public core; // The core arbitrator contract.
     IVault public vault; // The PNK vault for token management.
-    ISortitionSumTree public sortitionModule; // The sortition module for drawing logic.
+    ISortitionSumTree public sortition; // The sortition sum tree for drawing.
 
     // Phase management
     Phase public override phase; // The current phase. Uses Phase from IStakeController.
@@ -82,7 +82,7 @@ abstract contract StakeControllerBase is IStakeController, Initializable, UUPSPr
         address _governor,
         KlerosCoreXBase _core,
         IVault _vault,
-        ISortitionSumTree _sortitionModule,
+        ISortitionSumTree _sortition,
         uint256 _minStakingTime,
         uint256 _maxDrawingTime,
         RNG _rng,
@@ -91,7 +91,7 @@ abstract contract StakeControllerBase is IStakeController, Initializable, UUPSPr
         governor = _governor;
         core = _core;
         vault = _vault;
-        sortitionModule = _sortitionModule;
+        sortition = _sortition;
         minStakingTime = _minStakingTime;
         maxDrawingTime = _maxDrawingTime;
         lastPhaseChange = block.timestamp;
@@ -116,10 +116,10 @@ abstract contract StakeControllerBase is IStakeController, Initializable, UUPSPr
         vault = _vault;
     }
 
-    /// @dev Changes the `sortitionModule` storage variable.
-    /// @param _sortitionModule The new sortition module address.
-    function changeSortitionModule(ISortitionSumTree _sortitionModule) external onlyByGovernor {
-        sortitionModule = _sortitionModule;
+    /// @dev Changes the `sortition` storage variable.
+    /// @param _sortition The new sortition module address.
+    function changeSortitionSumTree(ISortitionSumTree _sortition) external onlyByGovernor {
+        sortition = _sortition;
     }
 
     /// @dev Changes the `minStakingTime` storage variable.
@@ -198,15 +198,6 @@ abstract contract StakeControllerBase is IStakeController, Initializable, UUPSPr
     // *         Stake Management          * //
     // ************************************* //
 
-    /// @inheritdoc IStakeController
-    function setStakeBySystem(
-        address _account,
-        uint96 _courtID,
-        uint256 _newStake
-    ) external override onlyByCore returns (uint256 pnkDeposit, uint256 pnkWithdrawal, StakingResult stakingResult) {
-        return _setStakeBySystem(_account, _courtID, _newStake);
-    }
-
     /// @inheritdoc IStakeController
     function setStake(
         address _account,
@@ -243,7 +234,7 @@ abstract contract StakeControllerBase is IStakeController, Initializable, UUPSPr
         uint96[] storage courtIDsForJuror = jurorStakes[_account].stakedCourtIDs;
         while (courtIDsForJuror.length > 0) {
             uint96 courtID = courtIDsForJuror[0];
-            _setStakeBySystem(_account, courtID, 0);
+            _setStake(_account, courtID, 0);
         }
         jurorStakes[_account].totalStake = 0;
         pnkToWithdraw = vault.getAvailableBalance(_account);
@@ -256,14 +247,14 @@ abstract contract StakeControllerBase is IStakeController, Initializable, UUPSPr
 
     /// @inheritdoc IStakeController
     function createTree(bytes32 _key, bytes memory _extraData) external override onlyByCore {
-        sortitionModule.createTree(_key, _extraData);
+        sortition.createTree(_key, _extraData);
     }
 
     /// @inheritdoc IStakeController
     function draw(bytes32 _court, uint256 _coreDisputeID, uint256 _nonce) external view override returns (address) {
         if (phase != Phase.drawing) revert NotDrawingPhase();
         if (randomNumber == 0) revert RandomNumberNotReadyYet();
-        return sortitionModule.draw(_court, _coreDisputeID, _nonce, randomNumber);
+        return sortition.draw(_court, _coreDisputeID, _nonce, randomNumber);
     }
 
     /// @inheritdoc IStakeController
@@ -329,71 +320,54 @@ abstract contract StakeControllerBase is IStakeController, Initializable, UUPSPr
     // *            Internal               * //
     // ************************************* //
 
-    /// @dev Internal implementation of setStakeBySystem
+    /// @dev Internal implementation of setStake with phase-aware delayed stake logic
     /// @param _account The account to set the stake for.
     /// @param _courtID The ID of the court to set the stake for.
     /// @param _newStake The new stake.
     /// @return pnkDeposit The amount of PNK to deposit.
     /// @return pnkWithdrawal The amount of PNK to withdraw.
     /// @return stakingResult The result of the staking operation.
-    function _setStakeBySystem(
+    function _setStake(
         address _account,
         uint96 _courtID,
         uint256 _newStake
     ) internal virtual returns (uint256 pnkDeposit, uint256 pnkWithdrawal, StakingResult stakingResult) {
         JurorStake storage currentJurorStake = jurorStakes[_account];
         uint256 currentStakeInCourt = currentJurorStake.stakes[_courtID];
 
-        if (currentStakeInCourt == 0) {
-            if (_newStake == 0)
-                // No change needed
-                return (0, 0, StakingResult.CannotStakeZeroWhenNoStake);
-            else if (_newStake > 0 && currentJurorStake.stakedCourtIDs.length >= MAX_STAKE_PATHS)
-                // Cannot stake in more courts
-                return (0, 0, StakingResult.CannotStakeInMoreCourts);
-        }
-
-        currentJurorStake.stakes[_courtID] = _newStake;
+        if (phase == Phase.staking) {
+            if (currentStakeInCourt == 0) {
+                if (_newStake == 0)
+                    // No change needed
+                    return (0, 0, StakingResult.CannotStakeZeroWhenNoStake);
+                else if (_newStake > 0 && currentJurorStake.stakedCourtIDs.length >= MAX_STAKE_PATHS)
+                    // Cannot stake in more courts
+                    return (0, 0, StakingResult.CannotStakeInMoreCourts);
+            }
 
-        if (_newStake > currentStakeInCourt) {
-            pnkDeposit = _newStake - currentStakeInCourt;
-            currentJurorStake.totalStake += pnkDeposit;
-        } else if (_newStake < currentStakeInCourt) {
-            pnkWithdrawal = currentStakeInCourt - _newStake;
-            currentJurorStake.totalStake -= pnkWithdrawal;
-        }
+            currentJurorStake.stakes[_courtID] = _newStake;
 
-        // Manage stakedCourtIDs
-        if (currentStakeInCourt == 0 && _newStake > 0) {
-            currentJurorStake.stakedCourtIDs.push(_courtID);
-        } else if (currentStakeInCourt > 0 && _newStake == 0) {
-            _removeCourt(currentJurorStake.stakedCourtIDs, _courtID);
-        }
+            if (_newStake > currentStakeInCourt) {
+                pnkDeposit = _newStake - currentStakeInCourt;
+                currentJurorStake.totalStake += pnkDeposit;
+            } else if (_newStake < currentStakeInCourt) {
+                pnkWithdrawal = currentStakeInCourt - _newStake;
+                currentJurorStake.totalStake -= pnkWithdrawal;
+            }
 
-        sortitionModule.setStake(_account, _courtID, _newStake);
+            // Manage stakedCourtIDs
+            if (currentStakeInCourt == 0 && _newStake > 0) {
+                currentJurorStake.stakedCourtIDs.push(_courtID);
+            } else if (currentStakeInCourt > 0 && _newStake == 0) {
+                _removeCourt(currentJurorStake.stakedCourtIDs, _courtID);
+            }
 
-        emit StakeSet(_account, _courtID, _newStake, currentJurorStake.totalStake);
-    }
+            sortition.setStake(_account, _courtID, _newStake);
 
-    /// @dev Internal implementation of setStake with phase-aware delayed stake logic
-    /// @param _account The account to set the stake for.
-    /// @param _courtID The ID of the court to set the stake for.
-    /// @param _newStake The new stake.
-    /// @return pnkDeposit The amount of PNK to deposit.
-    /// @return pnkWithdrawal The amount of PNK to withdraw.
-    /// @return stakingResult The result of the staking operation.
-    function _setStake(
-        address _account,
-        uint96 _courtID,
-        uint256 _newStake
-    ) internal virtual returns (uint256 pnkDeposit, uint256 pnkWithdrawal, StakingResult stakingResult) {
-        if (phase == Phase.staking) {
-            return _setStakeBySystem(_account, _courtID, _newStake);
+            emit StakeSet(_account, _courtID, _newStake, currentJurorStake.totalStake);
+            return (pnkDeposit, pnkWithdrawal, StakingResult.Successful);
         }
 
-        JurorStake storage currentJurorStake = jurorStakes[_account];
-        uint256 currentStakeInCourt = currentJurorStake.stakes[_courtID];
-
         if (_newStake > currentStakeInCourt) {
             pnkDeposit = _newStake - currentStakeInCourt;
         } else if (_newStake < currentStakeInCourt) {
@@ -470,9 +444,9 @@ abstract contract StakeControllerBase is IStakeController, Initializable, UUPSPr
                 }
 
                 // _setStakeBySystem will update local juror stake mappings (jurorStakes)
-                // AND call sortitionModule.setStake.
+                // AND call sortition.setStake.
                 // The pnkDeposit/pnkWithdrawal are calculated but not used by this import function.
-                (, , StakingResult stakingResult) = _setStakeBySystem(account, courtID, stakeToImport);
+                (, , StakingResult stakingResult) = _setStake(account, courtID, stakeToImport);
 
                 if (stakingResult == StakingResult.Successful) {
                     totalImportedSuccess++;
diff --git a/contracts/src/arbitration/dispute-kits/DisputeKitClassicBase.sol b/contracts/src/arbitration/dispute-kits/DisputeKitClassicBase.sol
@@ -207,7 +207,7 @@ abstract contract DisputeKitClassicBase is IDisputeKit, Initializable, UUPSProxi
         emit DisputeCreation(_coreDisputeID, _numberOfChoices, _extraData);
     }
 
-    /// @dev Draws the juror from the sortition tree. The drawn address is picked up by Kleros Core.
+    /// @dev Draws the juror from the Stake Controller.
     /// Note: Access restricted to Kleros Core only.
     /// @param _coreDisputeID The ID of the dispute in Kleros Core.
     /// @param _nonce Nonce of the drawing iteration.
@@ -221,11 +221,11 @@ abstract contract DisputeKitClassicBase is IDisputeKit, Initializable, UUPSProxi
         uint256 localRoundID = dispute.rounds.length - 1;
         Round storage round = dispute.rounds[localRoundID];
 
-        IStakeController sortitionModule = core.stakeController();
+        IStakeController stakeController = core.stakeController();
         (uint96 courtID, , , , ) = core.disputes(_coreDisputeID);
         bytes32 key = bytes32(uint256(courtID)); // Get the ID of the tree.
 
-        drawnAddress = sortitionModule.draw(key, _coreDisputeID, _nonce);
+        drawnAddress = stakeController.draw(key, _coreDisputeID, _nonce);
 
         if (_postDrawCheck(round, _coreDisputeID, drawnAddress)) {
             round.votes.push(Vote({account: drawnAddress, commit: bytes32(0), choice: 0, voted: false}));
diff --git a/contracts/src/arbitration/interfaces/IStakeController.sol b/contracts/src/arbitration/interfaces/IStakeController.sol
@@ -6,7 +6,7 @@ import {KlerosCoreXBase} from "../KlerosCoreXBase.sol";
 import "../../libraries/Constants.sol";
 
 /// @title IStakeController
-/// @notice Interface for the Stake Controller that coordinates between Vault and SortitionModule
+/// @notice Interface for the Stake Controller that coordinates between Vault and SortitionSumTree
 /// @dev Combines phase management, delayed stakes, and coordination between vault and sortition
 interface IStakeController {
     // ************************************* //
@@ -62,20 +62,6 @@ interface IStakeController {
     // *         Stake Management          * //
     // ************************************* //
 
-    /// @notice System-level update to a juror's stake directly in the SortitionModule.
-    /// @dev Called by KlerosCoreXBase for executing delayed stakes. Skips regular phase checks and delayed stake creation.
-    /// @param _account The juror's account
-    /// @param _courtID The ID of the court
-    /// @param _newStake The new stake amount
-    /// @return pnkDeposit The amount of PNK to deposit
-    /// @return pnkWithdrawal The amount of PNK to withdraw
-    /// @return stakingResult The result of the staking operation
-    function setStakeBySystem(
-        address _account,
-        uint96 _courtID,
-        uint256 _newStake
-    ) external returns (uint256 pnkDeposit, uint256 pnkWithdrawal, StakingResult stakingResult);
-
     /// @notice Set stake for a juror with vault coordination
     /// @param _account The juror's account
     /// @param _courtID The ID of the court
@@ -128,12 +114,12 @@ interface IStakeController {
     // *         Sortition Delegation      * //
     // ************************************* //
 
-    /// @notice Create a sortition tree (delegated to SortitionModule)
+    /// @notice Create a sortition tree (delegated to SortitionSumTree)
     /// @param _key The key of the tree
     /// @param _extraData Extra data for tree configuration
     function createTree(bytes32 _key, bytes memory _extraData) external;
 
-    /// @notice Draw a juror for dispute (delegated to SortitionModule)
+    /// @notice Draw a juror for dispute (delegated to SortitionSumTree)
     /// @param _court The court identifier
     /// @param _coreDisputeID The core dispute ID
     /// @param _nonce The drawing nonce
