refactor: minor refactor and fixed the tests

jaybuidl · jaybuidl · commit 55347685378e · 2023-09-09T14:51:27.000+08:00
diff --git a/contracts/src/arbitration/KlerosCore.sol b/contracts/src/arbitration/KlerosCore.sol
@@ -483,14 +483,14 @@ contract KlerosCore is IArbitratorV2 {
 
     /// @dev Sets the caller's stake in a court.
     /// @param _courtID The ID of the court.
-    /// @param _stake The new stake.
-    function setStake(uint96 _courtID, uint256 _stake) external {
-        if (!_setStakeForAccount(msg.sender, _courtID, _stake)) revert StakingFailed();
+    /// @param _newStake The new stake.
+    function setStake(uint96 _courtID, uint256 _newStake) external {
+        if (!_setStakeForAccount(msg.sender, _courtID, _newStake)) revert StakingFailed();
     }
 
-    function setStakeBySortitionModule(address _account, uint96 _courtID, uint256 _stake) external {
+    function setStakeBySortitionModule(address _account, uint96 _courtID, uint256 _newStake) external {
         if (msg.sender != address(sortitionModule)) revert WrongCaller();
-        _setStakeForAccount(_account, _courtID, _stake);
+        _setStakeForAccount(_account, _courtID, _newStake);
     }
 
     /// @inheritdoc IArbitratorV2
@@ -611,17 +611,16 @@ contract KlerosCore is IArbitratorV2 {
         IDisputeKit disputeKit = disputeKitNodes[round.disputeKitID].disputeKit;
 
         uint256 startIndex = round.drawIterations;
-        // TODO: cap the iterations?
-        for (uint256 i = startIndex; i < startIndex + _iterations; i++) {
+        for (uint256 i = startIndex; i < startIndex + _iterations && round.drawnJurors.length < round.nbVotes; i++) {
             address drawnAddress = disputeKit.draw(_disputeID, i);
-            if (drawnAddress != address(0)) {
-                jurors[drawnAddress].lockedPnk += round.pnkAtStakePerJuror;
-                emit Draw(drawnAddress, _disputeID, currentRound, round.drawnJurors.length);
-                round.drawnJurors.push(drawnAddress);
-
-                if (round.drawnJurors.length == round.nbVotes) {
-                    sortitionModule.postDrawHook(_disputeID, currentRound);
-                }
+            if (drawnAddress == address(0)) {
+                continue;
+            }
+            jurors[drawnAddress].lockedPnk += round.pnkAtStakePerJuror;
+            emit Draw(drawnAddress, _disputeID, currentRound, round.drawnJurors.length);
+            round.drawnJurors.push(drawnAddress);
+            if (round.drawnJurors.length == round.nbVotes) {
+                sortitionModule.postDrawHook(_disputeID, currentRound);
             }
         }
         round.drawIterations += _iterations;
@@ -1110,35 +1109,40 @@ contract KlerosCore is IArbitratorV2 {
     /// and `j` is the maximum number of jurors that ever staked in one of these courts simultaneously.
     /// @param _account The address of the juror.
     /// @param _courtID The ID of the court.
-    /// @param _stake The new stake.
+    /// @param _newStake The new stake.
     /// @return succeeded True if the call succeeded, false otherwise.
-    function _setStakeForAccount(address _account, uint96 _courtID, uint256 _stake) internal returns (bool succeeded) {
+    function _setStakeForAccount(
+        address _account,
+        uint96 _courtID,
+        uint256 _newStake
+    ) internal returns (bool succeeded) {
         if (_courtID == FORKING_COURT || _courtID > courts.length) return false;
 
         Juror storage juror = jurors[_account];
         uint256 currentStake = juror.stakedPnkByCourt[_courtID];
 
-        if (_stake != 0) {
-            if (_stake < courts[_courtID].minStake) return false;
+        if (_newStake != 0) {
+            if (_newStake < courts[_courtID].minStake) return false;
         } else if (currentStake == 0) {
             return false;
         }
 
-        ISortitionModule.preStakeHookResult result = sortitionModule.preStakeHook(_account, _courtID, _stake);
+        ISortitionModule.preStakeHookResult result = sortitionModule.preStakeHook(_account, _courtID, _newStake);
         if (result == ISortitionModule.preStakeHookResult.failed) {
             return false;
         } else if (result == ISortitionModule.preStakeHookResult.delayed) {
-            emit StakeDelayed(_account, _courtID, _stake);
+            emit StakeDelayed(_account, _courtID, _newStake);
             return true;
         }
 
         uint256 transferredAmount;
-        if (_stake >= currentStake) {
+        if (_newStake >= currentStake) {
+            // Stake increase
             // When stakedPnk becomes lower than lockedPnk count the locked tokens in when transferring tokens from juror.
             // (E.g. stakedPnk = 0, lockedPnk = 150) which can happen if the juror unstaked fully while having some tokens locked.
-            uint256 previouslyLocked = (juror.lockedPnk >= juror.stakedPnk) ? juror.lockedPnk - juror.stakedPnk : 0;
-            transferredAmount = (_stake >= currentStake + previouslyLocked)
-                ? _stake - currentStake - previouslyLocked
+            uint256 previouslyLocked = (juror.lockedPnk >= juror.stakedPnk) ? juror.lockedPnk - juror.stakedPnk : 0; // underflow guard
+            transferredAmount = (_newStake >= currentStake + previouslyLocked) // underflow guard
+                ? _newStake - currentStake - previouslyLocked
                 : 0;
             if (transferredAmount > 0) {
                 if (!pinakion.safeTransferFrom(_account, address(this), transferredAmount)) {
@@ -1149,49 +1153,36 @@ contract KlerosCore is IArbitratorV2 {
                 juror.courtIDs.push(_courtID);
             }
         } else {
-            if (_stake == 0) {
-                // Make sure locked tokens always stay in the contract. They can only be released during Execution.
-                if (juror.stakedPnk >= currentStake + juror.lockedPnk) {
-                    // We have enough pnk staked to afford withdrawal of the current stake while keeping locked tokens.
-                    transferredAmount = currentStake;
-                } else if (juror.stakedPnk >= juror.lockedPnk) {
-                    // Can't afford withdrawing the current stake fully. Take whatever is available while keeping locked tokens.
-                    transferredAmount = juror.stakedPnk - juror.lockedPnk;
-                }
-                if (transferredAmount > 0) {
-                    if (!pinakion.safeTransfer(_account, transferredAmount)) {
-                        return false;
-                    }
+            // Stake decrease: make sure locked tokens always stay in the contract. They can only be released during Execution.
+            if (juror.stakedPnk >= currentStake - _newStake + juror.lockedPnk) {
+                // We have enough pnk staked to afford withdrawal while keeping locked tokens.
+                transferredAmount = currentStake - _newStake;
+            } else if (juror.stakedPnk >= juror.lockedPnk) {
+                // Can't afford withdrawing the current stake fully. Take whatever is available while keeping locked tokens.
+                transferredAmount = juror.stakedPnk - juror.lockedPnk;
+            }
+            if (transferredAmount > 0) {
+                if (!pinakion.safeTransfer(_account, transferredAmount)) {
+                    return false;
                 }
+            }
+            if (_newStake == 0) {
                 for (uint256 i = juror.courtIDs.length; i > 0; i--) {
                     if (juror.courtIDs[i - 1] == _courtID) {
                         juror.courtIDs[i - 1] = juror.courtIDs[juror.courtIDs.length - 1];
                         juror.courtIDs.pop();
                         break;
                     }
                 }
-            } else {
-                if (juror.stakedPnk >= currentStake - _stake + juror.lockedPnk) {
-                    // We have enough pnk staked to afford withdrawal while keeping locked tokens.
-                    transferredAmount = currentStake - _stake;
-                } else if (juror.stakedPnk >= juror.lockedPnk) {
-                    // Can't afford withdrawing the current stake fully. Take whatever is available while keeping locked tokens.
-                    transferredAmount = juror.stakedPnk - juror.lockedPnk;
-                }
-                if (transferredAmount > 0) {
-                    if (!pinakion.safeTransfer(_account, transferredAmount)) {
-                        return false;
-                    }
-                }
             }
         }
 
         // Note that stakedPnk can become async with currentStake (e.g. after penalty).
-        juror.stakedPnk = (juror.stakedPnk >= currentStake) ? juror.stakedPnk - currentStake + _stake : _stake;
-        juror.stakedPnkByCourt[_courtID] = _stake;
+        juror.stakedPnk = (juror.stakedPnk >= currentStake) ? juror.stakedPnk - currentStake + _newStake : _newStake;
+        juror.stakedPnkByCourt[_courtID] = _newStake;
 
-        sortitionModule.setStake(_account, _courtID, _stake);
-        emit StakeSet(_account, _courtID, _stake);
+        sortitionModule.setStake(_account, _courtID, _newStake);
+        emit StakeSet(_account, _courtID, _newStake);
         return true;
     }
 
diff --git a/contracts/test/arbitration/draw.ts b/contracts/test/arbitration/draw.ts
@@ -15,7 +15,6 @@ import { DrawEvent } from "../../typechain-types/src/kleros-v1/kleros-liquid-xda
 /* eslint-disable no-unused-vars */
 /* eslint-disable no-unused-expressions */ // https://github.com/standard/standard/issues/690#issuecomment-278533482
 
-// FIXME: This test fails on Github actions, cannot figure why, skipping for now.
 describe("Draw Benchmark", async () => {
   const ONE_TENTH_ETH = BigNumber.from(10).pow(17);
   const ONE_THOUSAND_PNK = BigNumber.from(10).pow(21);
@@ -170,8 +169,7 @@ describe("Draw Benchmark", async () => {
     }
 
     await sortitionModule.passPhase(); // Generating -> Drawing
-
-    await expectFromDraw(core.draw(0, 1000, { gasLimit: 1000000 }));
+    await expectFromDraw(core.draw(0, 20, { gasLimit: 1000000 }));
 
     await network.provider.send("evm_increaseTime", [2000]); // Wait for maxDrawingTime
     await sortitionModule.passPhase(); // Drawing -> Staking
@@ -197,8 +195,9 @@ describe("Draw Benchmark", async () => {
       await core.connect(wallet).setStake(PARENT_COURT, ONE_THOUSAND_PNK.mul(5), { gasLimit: 5000000 });
 
       expect(await core.getJurorBalance(wallet.address, 1)).to.deep.equal([
+        ONE_THOUSAND_PNK.mul(5), // totalStaked
+        0, // totalLocked
         ONE_THOUSAND_PNK.mul(5), // stakedInCourt
-        0, // lockedInCourt
         PARENT_COURT, // nbOfCourts
       ]);
     };
@@ -216,13 +215,15 @@ describe("Draw Benchmark", async () => {
       countedDraws = await countDraws(tx.blockNumber);
       for (const [address, draws] of Object.entries(countedDraws)) {
         expect(await core.getJurorBalance(address, PARENT_COURT)).to.deep.equal([
+          ONE_THOUSAND_PNK.mul(5), // totalStaked
+          parentCourtMinStake.mul(draws), // totalLocked
           ONE_THOUSAND_PNK.mul(5), // stakedInCourt
-          parentCourtMinStake.mul(draws), // lockedInCourt
           1, // nbOfCourts
         ]);
         expect(await core.getJurorBalance(address, CHILD_COURT)).to.deep.equal([
+          ONE_THOUSAND_PNK.mul(5), // totalStaked
+          parentCourtMinStake.mul(draws), // totalLocked
           0, // stakedInCourt
-          0, // lockedInCourt
           1, // nbOfCourts
         ]);
       }
@@ -235,16 +236,18 @@ describe("Draw Benchmark", async () => {
         await core.getJurorBalance(wallet.address, PARENT_COURT),
         "Drawn jurors have a locked stake in the parent court"
       ).to.deep.equal([
+        0, // totalStaked
+        locked, // totalLocked
         0, // stakedInCourt
-        locked, // lockedInCourt
         0, // nbOfCourts
       ]);
       expect(
         await core.getJurorBalance(wallet.address, CHILD_COURT),
         "No locked stake in the child court"
       ).to.deep.equal([
+        0, // totalStaked
+        locked, // totalLocked
         0, // stakedInCourt
-        0, // lockedInCourt
         0, // nbOfCourts
       ]);
     };
@@ -267,16 +270,18 @@ describe("Draw Benchmark", async () => {
         await core.getJurorBalance(wallet.address, PARENT_COURT),
         "No locked stake in the parent court"
       ).to.deep.equal([
+        0, // totalStaked
+        0, // totalLocked
         0, // stakedInCourt
-        0, // lockedInCourt
         0, // nbOfCourts
       ]);
       expect(
         await core.getJurorBalance(wallet.address, CHILD_COURT),
         "No locked stake in the child court"
       ).to.deep.equal([
+        0, // totalStaked
+        0, // totalLocked
         0, // stakedInCourt
-        0, // lockedInCourt
         0, // nbOfCourts
       ]);
     };
@@ -302,13 +307,15 @@ describe("Draw Benchmark", async () => {
       countedDraws = await countDraws(tx.blockNumber);
       for (const [address, draws] of Object.entries(countedDraws)) {
         expect(await core.getJurorBalance(address, PARENT_COURT)).to.deep.equal([
+          ONE_THOUSAND_PNK.mul(5), // totalStaked
+          parentCourtMinStake.mul(draws), // totalLocked
           0, // stakedInCourt
-          parentCourtMinStake.mul(draws), // lockedInCourt
           1, // nbOfCourts
         ]);
         expect(await core.getJurorBalance(address, CHILD_COURT)).to.deep.equal([
+          ONE_THOUSAND_PNK.mul(5), // totalStaked
+          parentCourtMinStake.mul(draws), // totalLocked
           ONE_THOUSAND_PNK.mul(5), // stakedInCourt
-          0, // lockedInCourt
           1, // nbOfCourts
         ]);
       }
@@ -317,21 +324,22 @@ describe("Draw Benchmark", async () => {
     const unstake = async (wallet: Wallet) => {
       await core.connect(wallet).setStake(CHILD_COURT, 0, { gasLimit: 5000000 });
       const locked = parentCourtMinStake.mul(countedDraws[wallet.address] ?? 0);
-      console.log(`draws for ${wallet.address}: ${countedDraws[wallet.address] ?? 0}, locked: ${locked}`);
       expect(
         await core.getJurorBalance(wallet.address, PARENT_COURT),
         "No locked stake in the parent court"
       ).to.deep.equal([
+        0, // totalStaked
+        locked, // totalLocked
         0, // stakedInCourt
-        0, // lockedInCourt
         0, // nbOfCourts
       ]);
       expect(
         await core.getJurorBalance(wallet.address, CHILD_COURT),
         "Drawn jurors have a locked stake in the child court"
       ).to.deep.equal([
+        0, // totalStaked
+        locked, // totalLocked
         0, // stakedInCourt
-        locked, // lockedInCourt
         0, // nbOfCourts
       ]);
     };
@@ -357,13 +365,15 @@ describe("Draw Benchmark", async () => {
       countedDraws = await countDraws(tx.blockNumber);
       for (const [address, draws] of Object.entries(countedDraws)) {
         expect(await core.getJurorBalance(address, PARENT_COURT)).to.deep.equal([
+          ONE_THOUSAND_PNK.mul(5), // totalStaked
+          childCourtMinStake.mul(draws), // totalLocked
           0, // stakedInCourt
-          0, // lockedInCourt
           1, // nbOfCourts
         ]);
         expect(await core.getJurorBalance(address, CHILD_COURT)).to.deep.equal([
+          ONE_THOUSAND_PNK.mul(5), // totalStaked
+          childCourtMinStake.mul(draws), // totalLocked
           ONE_THOUSAND_PNK.mul(5), // stakedInCourt
-          childCourtMinStake.mul(draws), // lockedInCourt
           1, // nbOfCourts
         ]);
       }
@@ -376,16 +386,18 @@ describe("Draw Benchmark", async () => {
         await core.getJurorBalance(wallet.address, PARENT_COURT),
         "No locked stake in the parent court"
       ).to.deep.equal([
+        0, // totalStaked
+        locked, // totalLocked
         0, // stakedInCourt
-        0, // lockedInCourt
         0, // nbOfCourts
       ]);
       expect(
         await core.getJurorBalance(wallet.address, CHILD_COURT),
         "Drawn jurors have a locked stake in the child court"
       ).to.deep.equal([
+        0, // totalStaked
+        locked, // totalLocked
         0, // stakedInCourt
-        locked, // lockedInCourt
         0, // nbOfCourts
       ]);
     };
diff --git a/contracts/test/integration/index.ts b/contracts/test/integration/index.ts
@@ -67,29 +67,29 @@ describe("Integration tests", async () => {
 
     await core.setStake(1, ONE_THOUSAND_PNK);
     await core.getJurorBalance(deployer, 1).then((result) => {
-      expect(result.staked).to.equal(ONE_THOUSAND_PNK);
-      expect(result.locked).to.equal(0);
+      expect(result.totalStaked).to.equal(ONE_THOUSAND_PNK);
+      expect(result.totalLocked).to.equal(0);
       logJurorBalance(result);
     });
 
     await core.setStake(1, ONE_HUNDRED_PNK.mul(5));
     await core.getJurorBalance(deployer, 1).then((result) => {
-      expect(result.staked).to.equal(ONE_HUNDRED_PNK.mul(5));
-      expect(result.locked).to.equal(0);
+      expect(result.totalStaked).to.equal(ONE_HUNDRED_PNK.mul(5));
+      expect(result.totalLocked).to.equal(0);
       logJurorBalance(result);
     });
 
     await core.setStake(1, 0);
     await core.getJurorBalance(deployer, 1).then((result) => {
-      expect(result.staked).to.equal(0);
-      expect(result.locked).to.equal(0);
+      expect(result.totalStaked).to.equal(0);
+      expect(result.totalLocked).to.equal(0);
       logJurorBalance(result);
     });
 
     await core.setStake(1, ONE_THOUSAND_PNK.mul(4));
     await core.getJurorBalance(deployer, 1).then((result) => {
-      expect(result.staked).to.equal(ONE_THOUSAND_PNK.mul(4));
-      expect(result.locked).to.equal(0);
+      expect(result.totalStaked).to.equal(ONE_THOUSAND_PNK.mul(4));
+      expect(result.totalLocked).to.equal(0);
       logJurorBalance(result);
     });
     const tx = await arbitrable.functions["createDispute(string)"]("future of france", {
@@ -183,5 +183,9 @@ describe("Integration tests", async () => {
 });
 
 const logJurorBalance = async (result) => {
-  console.log("staked=%s, locked=%s", ethers.utils.formatUnits(result.staked), ethers.utils.formatUnits(result.locked));
+  console.log(
+    "staked=%s, locked=%s",
+    ethers.utils.formatUnits(result.totalStaked),
+    ethers.utils.formatUnits(result.totalLocked)
+  );
 };
