Staking refactor including bug fixes

contracts/deploy/00-home-chain-arbitration.ts

@@ -2,11 +2,11 @@ import { HardhatRuntimeEnvironment } from "hardhat/types";
 import { DeployFunction } from "hardhat-deploy/types";
 import { getContractAddress } from "./utils/getContractAddress";
 import { deployUpgradable } from "./utils/deployUpgradable";
-import { changeCurrencyRate } from "./utils/klerosCoreHelper";
 import { HomeChains, isSkipped, isDevnet, PNK, ETH } from "./utils";
 import { getContractOrDeploy, getContractOrDeployUpgradable } from "./utils/getContractOrDeploy";
 import { deployERC20AndFaucet } from "./utils/deployTokens";
 import { ChainlinkRNG, DisputeKitClassic, KlerosCore } from "../typechain-types";
+import { changeCurrencyRate } from "./utils/klerosCoreHelper";
 
 const deployArbitration: DeployFunction = async (hre: HardhatRuntimeEnvironment) => {
   const { ethers, deployments, getNamedAccounts, getChainId } = hre;
@@ -29,65 +29,104 @@ const deployArbitration: DeployFunction = async (hre: HardhatRuntimeEnvironment)
 
   await getContractOrDeployUpgradable(hre, "EvidenceModule", { from: deployer, args: [deployer], log: true });
 
-  const disputeKit = await deployUpgradable(deployments, "DisputeKitClassic", {
+  const disputeKit = await deployUpgradable(deployments, "DisputeKitClassicV2", {
     from: deployer,
-    args: [deployer, ZeroAddress],
+    contract: "DisputeKitClassic",
+    args: [
+      deployer,
+      ZeroAddress, // Placeholder for KlerosCore address, configured later
+    ],
     log: true,
   });
 
-  let klerosCoreAddress = await deployments.getOrNull("KlerosCore").then((deployment) => deployment?.address);
-  if (!klerosCoreAddress) {
-    const nonce = await ethers.provider.getTransactionCount(deployer);
-    klerosCoreAddress = getContractAddress(deployer, nonce + 3); // deployed on the 4th tx (nonce+3): SortitionModule Impl tx, SortitionModule Proxy tx, KlerosCore Impl tx, KlerosCore Proxy tx
-    console.log("calculated future KlerosCore address for nonce %d: %s", nonce + 3, klerosCoreAddress);
-  }
+  // Calculate future addresses for circular dependencies
+  const nonce = await ethers.provider.getTransactionCount(deployer);
+
+  const vaultAddress = getContractAddress(deployer, nonce + 1); // deployed on the 2nd tx (nonce+1): Vault Impl tx, Vault Proxy tx
+  console.log("calculated future Vault address for nonce %d: %s", nonce + 1, vaultAddress);
+
+  const stakeControllerAddress = getContractAddress(deployer, nonce + 5); // deployed on the 6th tx (nonce+5): Vault Impl tx, Vault Proxy tx, SortitionModule Impl tx, SortitionModule Proxy tx,, StakeController Impl tx, StakeController Proxy tx
+  console.log("calculated future StakeController address for nonce %d: %s", nonce + 5, stakeControllerAddress);
+
+  const klerosCoreAddress = getContractAddress(deployer, nonce + 7); // deployed on the 8th tx (nonce+7): Vault Impl tx, Vault Proxy tx, SortitionModule Impl tx, SortitionModule Proxy tx, StakeController Impl tx, StakeController Proxy tx, KlerosCore Impl tx, KlerosCore Proxy tx
+  console.log("calculated future KlerosCore address for nonce %d: %s", nonce + 7, klerosCoreAddress);
+
+  const vault = await deployUpgradable(deployments, "Vault", {
+    from: deployer,
+    args: [deployer, pnk.target, stakeControllerAddress, klerosCoreAddress],
+    log: true,
+  }); // nonce (implementation), nonce + 1 (proxy)
+
+  // Deploy SortitionSumTree
+  const sortitionModuleV2 = await deployUpgradable(deployments, "SortitionSumTree", {
+    from: deployer,
+    args: [deployer, stakeControllerAddress],
+    log: true,
+  }); // nonce + 2 (implementation), nonce + 3 (proxy)
+
+  // Deploy StakeController
   const devnet = isDevnet(hre.network);
   const minStakingTime = devnet ? 180 : 1800;
-  const maxFreezingTime = devnet ? 600 : 1800;
-  const rng = (await ethers.getContract("ChainlinkRNG")) as ChainlinkRNG;
-  const sortitionModule = await deployUpgradable(deployments, "SortitionModule", {
+  const maxDrawingTime = devnet ? 600 : 1800;
+  const rng = await ethers.getContract<ChainlinkRNG>("ChainlinkRNG");
+  const stakeController = await deployUpgradable(deployments, "StakeController", {
     from: deployer,
-    args: [deployer, klerosCoreAddress, minStakingTime, maxFreezingTime, rng.target, RNG_LOOKAHEAD],
+    args: [
+      deployer,
+      klerosCoreAddress,
+      vault.address,
+      sortitionModuleV2.address,
+      rng.target,
+      minStakingTime,
+      maxDrawingTime,
+      RNG_LOOKAHEAD,
+    ],
     log: true,
-  }); // nonce (implementation), nonce+1 (proxy)
+  }); // nonce + 4 (implementation), nonce + 5 (proxy)
 
   const minStake = PNK(200);
   const alpha = 10000;
   const feeForJuror = ETH(0.1);
   const jurorsForCourtJump = 256;
-  const klerosCore = await deployUpgradable(deployments, "KlerosCore", {
+
+  // Deploy KlerosCore
+  const klerosCoreV2 = await deployUpgradable(deployments, "KlerosCore", {
     from: deployer,
     args: [
       deployer,
       deployer,
-      pnk.target,
-      ZeroAddress, // KlerosCore is configured later
+      ZeroAddress, // JurorProsecutionModule, not implemented yet
       disputeKit.address,
       false,
       [minStake, alpha, feeForJuror, jurorsForCourtJump],
       [0, 0, 0, 10], // evidencePeriod, commitPeriod, votePeriod, appealPeriod
       ethers.toBeHex(5), // Extra data for sortition module will return the default value of K
-      sortitionModule.address,
+      stakeController.address,
+      vault.address,
     ],
     log: true,
-  }); // nonce+2 (implementation), nonce+3 (proxy)
+  });
+
+  // Configure cross-dependencies
+  console.log("Configuring cross-dependencies...");
 
   // disputeKit.changeCore() only if necessary
-  const disputeKitContract = (await ethers.getContract("DisputeKitClassic")) as DisputeKitClassic;
+  const disputeKitContract = await ethers.getContract<DisputeKitClassic>("DisputeKitClassicV2");
   const currentCore = await disputeKitContract.core();
-  if (currentCore !== klerosCore.address) {
-    console.log(`disputeKit.changeCore(${klerosCore.address})`);
-    await disputeKitContract.changeCore(klerosCore.address);
+  if (currentCore !== klerosCoreV2.address) {
+    console.log(`disputeKit.changeCore(${klerosCoreV2.address})`);
+    await disputeKitContract.changeCore(klerosCoreV2.address);
   }
 
   // rng.changeSortitionModule() only if necessary
+  // Note: the RNG's `sortitionModule` variable is misleading, it's only for access control and should be renamed to `consumer`.
   const rngSortitionModule = await rng.sortitionModule();
-  if (rngSortitionModule !== sortitionModule.address) {
-    console.log(`rng.changeSortitionModule(${sortitionModule.address})`);
-    await rng.changeSortitionModule(sortitionModule.address);
+  if (rngSortitionModule !== stakeController.address) {
+    console.log(`rng.changeSortitionModule(${stakeController.address})`);
+    await rng.changeSortitionModule(stakeController.address);
   }
 
-  const core = (await hre.ethers.getContract("KlerosCore")) as KlerosCore;
+  const core = await hre.ethers.getContract<KlerosCore>("KlerosCore");
   try {
     await changeCurrencyRate(core, await pnk.getAddress(), true, 12225583, 12);
     await changeCurrencyRate(core, await dai.getAddress(), true, 60327783, 11);
@@ -98,9 +137,16 @@ const deployArbitration: DeployFunction = async (hre: HardhatRuntimeEnvironment)
 
   await deploy("KlerosCoreSnapshotProxy", {
     from: deployer,
+    contract: "KlerosCoreSnapshotProxy",
     args: [deployer, core.target],
     log: true,
   });
+
+  console.log("✅ V2 Architecture deployment completed successfully!");
+  console.log(`📦 Vault: ${vault.address}`);
+  console.log(`🎯 SortitionSumTree: ${sortitionModuleV2.address}`);
+  console.log(`🎮 StakeController: ${stakeController.address}`);
+  console.log(`⚖️ KlerosCore: ${klerosCoreV2.address}`);
 };
 
 deployArbitration.tags = ["Arbitration"];

contracts/deploy/utils/klerosCoreHelper.ts

@@ -1,8 +1,8 @@
-import { KlerosCore, KlerosCoreNeo, KlerosCoreRuler, KlerosCoreUniversity } from "../../typechain-types";
+import { KlerosCore, KlerosCoreNeo, KlerosCoreRuler, KlerosCoreUniversity, KlerosCore } from "../../typechain-types";
 import { BigNumberish, toBigInt } from "ethers";
 
 export const changeCurrencyRate = async (
-  core: KlerosCore | KlerosCoreNeo | KlerosCoreRuler | KlerosCoreUniversity,
+  core: KlerosCore | KlerosCoreNeo | KlerosCoreRuler | KlerosCoreUniversity | KlerosCore,
   erc20: string,
   accepted: boolean,
   rateInEth: BigNumberish,

contracts/src/arbitration/KlerosCore.sol

@@ -2,71 +2,67 @@
 
 pragma solidity 0.8.24;
 
-import {KlerosCoreBase, IDisputeKit, ISortitionModule, IERC20} from "./KlerosCoreBase.sol";
+import "./KlerosCoreBase.sol";
 
 /// @title KlerosCore
-/// Core arbitrator contract for Kleros v2.
-/// Note that this contract trusts the PNK token, the dispute kit and the sortition module contracts.
+/// @notice KlerosCore implementation with new StakeController architecture for testing environments
 contract KlerosCore is KlerosCoreBase {
-    string public constant override version = "0.9.4";
+    /// @notice Version of the implementation contract
+    string public constant override version = "0.10.0";
 
     // ************************************* //
     // *            Constructor            * //
     // ************************************* //
 
-    /// @custom:oz-upgrades-unsafe-allow constructor
+    /// @notice Constructor, initializing the implementation to reduce attack surface.
     constructor() {
         _disableInitializers();
     }
 
-    /// @dev Initializer (constructor equivalent for upgradable contracts).
-    /// @param _governor The governor's address.
-    /// @param _guardian The guardian's address.
-    /// @param _pinakion The address of the token contract.
-    /// @param _jurorProsecutionModule The address of the juror prosecution module.
-    /// @param _disputeKit The address of the default dispute kit.
-    /// @param _hiddenVotes The `hiddenVotes` property value of the general court.
-    /// @param _courtParameters Numeric parameters of General court (minStake, alpha, feeForJuror and jurorsForCourtJump respectively).
-    /// @param _timesPerPeriod The `timesPerPeriod` property value of the general court.
-    /// @param _sortitionExtraData The extra data for sortition module.
-    /// @param _sortitionModuleAddress The sortition module responsible for sortition of the jurors.
+    /// @notice Initialization function for UUPS proxy
+    /// @param _governor The governor of the contract.
+    /// @param _guardian The guardian able to pause asset withdrawals.
+    /// @param _jurorProsecutionModule The module for juror's prosecution.
+    /// @param _disputeKit The dispute kit responsible for the dispute logic.
+    /// @param _hiddenVotes Whether to use commit and reveal or not.
+    /// @param _courtParameters [0]: minStake, [1]: alpha, [2]: feeForJuror, [3]: jurorsForCourtJump
+    /// @param _timesPerPeriod The timesPerPeriod array for courts
+    /// @param _sortitionExtraData Extra data for sortition module setup
+    /// @param _stakeController The stake controller for coordination
+    /// @param _vault The vault for coordination
     function initialize(
         address _governor,
         address _guardian,
-        IERC20 _pinakion,
         address _jurorProsecutionModule,
         IDisputeKit _disputeKit,
         bool _hiddenVotes,
         uint256[4] memory _courtParameters,
         uint256[4] memory _timesPerPeriod,
         bytes memory _sortitionExtraData,
-        ISortitionModule _sortitionModuleAddress
-    ) external reinitializer(1) {
+        IStakeController _stakeController,
+        IVault _vault
+    ) external initializer {
         __KlerosCoreBase_initialize(
             _governor,
             _guardian,
-            _pinakion,
             _jurorProsecutionModule,
             _disputeKit,
             _hiddenVotes,
             _courtParameters,
             _timesPerPeriod,
             _sortitionExtraData,
-            _sortitionModuleAddress
+            _stakeController,
+            _vault
         );
     }
 
-    function initialize5() external reinitializer(5) {
-        // NOP
-    }
-
     // ************************************* //
     // *             Governance            * //
     // ************************************* //
 
-    /// @dev Access Control to perform implementation upgrades (UUPS Proxiable)
-    ///      Only the governor can perform upgrades (`onlyByGovernor`)
+    /// @notice Access Control to perform implementation upgrades (UUPS Proxiable)
+    ///         Only the governor can perform upgrades (`onlyByGovernor`)
     function _authorizeUpgrade(address) internal view override onlyByGovernor {
-        // NOP
+        // Empty block: access control implemented by `onlyByGovernor` modifier
     }
 }