From 8cb7245a4cd49064165d0203e1071da9d13eaaea Mon Sep 17 00:00:00 2001
From: jaybuidl <jb@kleros.io>
Date: Sat, 9 Sep 2023 14:48:03 +0800
Subject: [PATCH 01/10] test: reproduction of the staking and locked stakes bug

---
 contracts/test/arbitration/draw.ts | 233 +++++++++++++++++++++++++----
 cspell.json                        |   1 +
 2 files changed, 206 insertions(+), 28 deletions(-)

diff --git a/contracts/test/arbitration/draw.ts b/contracts/test/arbitration/draw.ts
index 40840832b..5e835857b 100644
--- a/contracts/test/arbitration/draw.ts
+++ b/contracts/test/arbitration/draw.ts
@@ -1,17 +1,16 @@
 import { anyValue } from "@nomicfoundation/hardhat-chai-matchers/withArgs";
 import { deployments, ethers, getNamedAccounts, network } from "hardhat";
-import { BigNumber, ContractTransaction, Wallet } from "ethers";
+import { BigNumber, ContractReceipt, ContractTransaction, Wallet } from "ethers";
 import {
   PNK,
   KlerosCore,
   ArbitrableExample,
   HomeGateway,
   DisputeKitClassic,
-  RandomizerRNG,
-  RandomizerMock,
   SortitionModule,
 } from "../../typechain-types";
 import { expect } from "chai";
+import { DrawEvent } from "../../typechain-types/src/kleros-v1/kleros-liquid-xdai/XKlerosLiquidV2";
 
 /* eslint-disable no-unused-vars */
 /* eslint-disable no-unused-expressions */ // https://github.com/standard/standard/issues/690#issuecomment-278533482
@@ -49,7 +48,11 @@ describe("Draw Benchmark", async () => {
   let homeGateway;
   let sortitionModule;
   let rng;
-  let randomizer;
+  let parentCourtMinStake: BigNumber;
+  let childCourtMinStake: BigNumber;
+  const RANDOM = BigNumber.from("61688911660239508166491237672720926005752254046266901728404745669596507231249");
+  const PARENT_COURT = 1;
+  const CHILD_COURT = 2;
 
   beforeEach("Setup", async () => {
     ({ deployer, relayer } = await getNamedAccounts());
@@ -63,11 +66,25 @@ describe("Draw Benchmark", async () => {
     core = (await ethers.getContract("KlerosCore")) as KlerosCore;
     homeGateway = (await ethers.getContract("HomeGatewayToEthereum")) as HomeGateway;
     arbitrable = (await ethers.getContract("ArbitrableExample")) as ArbitrableExample;
-    rng = (await ethers.getContract("RandomizerRNG")) as RandomizerRNG;
-    randomizer = (await ethers.getContract("RandomizerMock")) as RandomizerMock;
     sortitionModule = (await ethers.getContract("SortitionModule")) as SortitionModule;
 
-    // CourtId 2
+    parentCourtMinStake = await core
+      .GENERAL_COURT()
+      .then((courtId) => core.courts(courtId))
+      .then((court) => court.minStake);
+
+    childCourtMinStake = BigNumber.from(10).pow(20).mul(3); // 300 PNK
+
+    // Make the tests more deterministic with this dummy RNG
+    rng = await deployments.deploy("IncrementalNG", {
+      from: deployer,
+      args: [RANDOM],
+      log: true,
+    });
+
+    await sortitionModule.changeRandomNumberGenerator(rng.address, 20);
+
+    // CourtId 2 = CHILD_COURT
     const minStake = BigNumber.from(10).pow(20).mul(3); // 300 PNK
     const alpha = 10000;
     const feeForJuror = BigNumber.from(10).pow(17);
@@ -84,16 +101,24 @@ describe("Draw Benchmark", async () => {
     );
   });
 
+  type CountedDraws = { [address: string]: number };
   type SetStake = (wallet: Wallet) => Promise<void>;
   type ExpectFromDraw = (drawTx: Promise<ContractTransaction>) => Promise<void>;
 
-  const draw = async (setStake: SetStake, createDisputeCourtId: string, expectFromDraw: ExpectFromDraw) => {
+  const draw = async (
+    stake: SetStake,
+    createDisputeCourtId: number,
+    expectFromDraw: ExpectFromDraw,
+    unstake: SetStake
+  ) => {
     const arbitrationCost = ONE_TENTH_ETH.mul(3);
     const [bridger] = await ethers.getSigners();
+    const wallets: Wallet[] = [];
 
     // Stake some jurors
     for (let i = 0; i < 16; i++) {
       const wallet = ethers.Wallet.createRandom().connect(ethers.provider);
+      wallets.push(wallet);
 
       await bridger.sendTransaction({
         to: wallet.address,
@@ -106,7 +131,7 @@ describe("Draw Benchmark", async () => {
 
       await pnk.connect(wallet).approve(core.address, ONE_THOUSAND_PNK.mul(10), { gasLimit: 300000 });
 
-      await setStake(wallet);
+      await stake(wallet);
     }
 
     // Create a dispute
@@ -144,75 +169,227 @@ describe("Draw Benchmark", async () => {
       await network.provider.send("evm_mine");
     }
 
-    await randomizer.relay(rng.address, 0, ethers.utils.randomBytes(32));
     await sortitionModule.passPhase(); // Generating -> Drawing
 
     await expectFromDraw(core.draw(0, 1000, { gasLimit: 1000000 }));
+
+    await network.provider.send("evm_increaseTime", [2000]); // Wait for maxDrawingTime
+    await sortitionModule.passPhase(); // Drawing -> Staking
+    expect(await sortitionModule.phase()).to.equal(Phase.staking);
+
+    // Unstake jurors
+    for (const wallet of wallets) {
+      await unstake(wallet);
+    }
+  };
+
+  const countDraws = async (blockNumber: number) => {
+    const draws: Array<DrawEvent> = await core.queryFilter(core.filters.Draw(), blockNumber, blockNumber);
+    return draws.reduce((acc: { [address: string]: number }, draw) => {
+      const address = draw.args._address;
+      acc[address] = acc[address] ? acc[address] + 1 : 1;
+      return acc;
+    }, {});
   };
 
   it("Stakes in parent court and should draw jurors in parent court", async () => {
-    const setStake = async (wallet: Wallet) => {
-      await core.connect(wallet).setStake(1, ONE_THOUSAND_PNK.mul(5), { gasLimit: 5000000 });
+    const stake = async (wallet: Wallet) => {
+      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), // stakedInCourt
+        0, // lockedInCourt
+        PARENT_COURT, // nbOfCourts
+      ]);
     };
-
+    let countedDraws: CountedDraws;
     const expectFromDraw = async (drawTx: Promise<ContractTransaction>) => {
-      await expect(drawTx)
+      const tx = await (await drawTx).wait();
+      expect(tx)
         .to.emit(core, "Draw")
         .withArgs(anyValue, 0, 0, 0)
         .to.emit(core, "Draw")
         .withArgs(anyValue, 0, 0, 1)
         .to.emit(core, "Draw")
         .withArgs(anyValue, 0, 0, 2);
+
+      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), // stakedInCourt
+          parentCourtMinStake.mul(draws), // lockedInCourt
+          1, // nbOfCourts
+        ]);
+        expect(await core.getJurorBalance(address, CHILD_COURT)).to.deep.equal([
+          0, // stakedInCourt
+          0, // lockedInCourt
+          1, // nbOfCourts
+        ]);
+      }
     };
 
-    await draw(setStake, "1", expectFromDraw);
+    const unstake = async (wallet: Wallet) => {
+      await core.connect(wallet).setStake(PARENT_COURT, 0, { gasLimit: 5000000 });
+      const locked = parentCourtMinStake.mul(countedDraws[wallet.address] ?? 0);
+      expect(
+        await core.getJurorBalance(wallet.address, PARENT_COURT),
+        "Drawn jurors have a locked stake in the parent court"
+      ).to.deep.equal([
+        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, // stakedInCourt
+        0, // lockedInCourt
+        0, // nbOfCourts
+      ]);
+    };
+
+    await draw(stake, PARENT_COURT, expectFromDraw, unstake);
   });
 
   it("Stakes in parent court and should draw nobody in subcourt", async () => {
-    const setStake = async (wallet: Wallet) => {
-      await core.connect(wallet).setStake(1, ONE_THOUSAND_PNK.mul(5), { gasLimit: 5000000 });
+    const stake = async (wallet: Wallet) => {
+      await core.connect(wallet).setStake(PARENT_COURT, ONE_THOUSAND_PNK.mul(5), { gasLimit: 5000000 });
     };
 
     const expectFromDraw = async (drawTx: Promise<ContractTransaction>) => {
       await expect(drawTx).to.not.emit(core, "Draw");
     };
 
-    await draw(setStake, "2", expectFromDraw);
+    const unstake = async (wallet: Wallet) => {
+      await core.connect(wallet).setStake(PARENT_COURT, 0, { gasLimit: 5000000 });
+      expect(
+        await core.getJurorBalance(wallet.address, PARENT_COURT),
+        "No locked stake in the parent court"
+      ).to.deep.equal([
+        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, // stakedInCourt
+        0, // lockedInCourt
+        0, // nbOfCourts
+      ]);
+    };
+
+    await draw(stake, CHILD_COURT, expectFromDraw, unstake);
   });
 
   it("Stakes in subcourt and should draw jurors in parent court", async () => {
-    const setStake = async (wallet: Wallet) => {
-      await core.connect(wallet).setStake(2, ONE_THOUSAND_PNK.mul(5), { gasLimit: 5000000 });
+    const stake = async (wallet: Wallet) => {
+      await core.connect(wallet).setStake(CHILD_COURT, ONE_THOUSAND_PNK.mul(5), { gasLimit: 5000000 });
     };
-
+    let countedDraws: CountedDraws;
     const expectFromDraw = async (drawTx: Promise<ContractTransaction>) => {
-      await expect(drawTx)
+      const tx = await (await drawTx).wait();
+      expect(tx)
         .to.emit(core, "Draw")
         .withArgs(anyValue, 0, 0, 0)
         .to.emit(core, "Draw")
         .withArgs(anyValue, 0, 0, 1)
         .to.emit(core, "Draw")
         .withArgs(anyValue, 0, 0, 2);
+
+      countedDraws = await countDraws(tx.blockNumber);
+      for (const [address, draws] of Object.entries(countedDraws)) {
+        expect(await core.getJurorBalance(address, PARENT_COURT)).to.deep.equal([
+          0, // stakedInCourt
+          parentCourtMinStake.mul(draws), // lockedInCourt
+          1, // nbOfCourts
+        ]);
+        expect(await core.getJurorBalance(address, CHILD_COURT)).to.deep.equal([
+          ONE_THOUSAND_PNK.mul(5), // stakedInCourt
+          0, // lockedInCourt
+          1, // nbOfCourts
+        ]);
+      }
     };
 
-    await draw(setStake, "1", expectFromDraw);
+    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, // 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, // stakedInCourt
+        locked, // lockedInCourt
+        0, // nbOfCourts
+      ]);
+    };
+
+    await draw(stake, PARENT_COURT, expectFromDraw, unstake);
   });
 
   it("Stakes in subcourt and should draw jurors in subcourt", async () => {
-    const setStake = async (wallet: Wallet) => {
-      await core.connect(wallet).setStake(2, ONE_THOUSAND_PNK.mul(5), { gasLimit: 5000000 });
+    const stake = async (wallet: Wallet) => {
+      await core.connect(wallet).setStake(CHILD_COURT, ONE_THOUSAND_PNK.mul(5), { gasLimit: 5000000 });
     };
-
+    let countedDraws: CountedDraws;
     const expectFromDraw = async (drawTx: Promise<ContractTransaction>) => {
-      await expect(drawTx)
+      const tx = await (await drawTx).wait();
+      expect(tx)
         .to.emit(core, "Draw")
         .withArgs(anyValue, 0, 0, 0)
         .to.emit(core, "Draw")
         .withArgs(anyValue, 0, 0, 1)
         .to.emit(core, "Draw")
         .withArgs(anyValue, 0, 0, 2);
+
+      countedDraws = await countDraws(tx.blockNumber);
+      for (const [address, draws] of Object.entries(countedDraws)) {
+        expect(await core.getJurorBalance(address, PARENT_COURT)).to.deep.equal([
+          0, // stakedInCourt
+          0, // lockedInCourt
+          1, // nbOfCourts
+        ]);
+        expect(await core.getJurorBalance(address, CHILD_COURT)).to.deep.equal([
+          ONE_THOUSAND_PNK.mul(5), // stakedInCourt
+          childCourtMinStake.mul(draws), // lockedInCourt
+          1, // nbOfCourts
+        ]);
+      }
+    };
+
+    const unstake = async (wallet: Wallet) => {
+      await core.connect(wallet).setStake(CHILD_COURT, 0, { gasLimit: 5000000 });
+      const locked = childCourtMinStake.mul(countedDraws[wallet.address] ?? 0);
+      expect(
+        await core.getJurorBalance(wallet.address, PARENT_COURT),
+        "No locked stake in the parent court"
+      ).to.deep.equal([
+        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, // stakedInCourt
+        locked, // lockedInCourt
+        0, // nbOfCourts
+      ]);
     };
 
-    await draw(setStake, "2", expectFromDraw);
+    await draw(stake, CHILD_COURT, expectFromDraw, unstake);
   });
 });
diff --git a/cspell.json b/cspell.json
index 73009f268..1aa319e95 100644
--- a/cspell.json
+++ b/cspell.json
@@ -36,6 +36,7 @@
     "typechain",
     "uncommify",
     "Unslashed",
+    "unstake",
     "viem",
     "wagmi"
   ],

From 73782d997022e422ec989f61aeb71bb9e79e278f Mon Sep 17 00:00:00 2001
From: unknownunknown1 <time_arrow@bk.ru>
Date: Wed, 16 Aug 2023 19:09:33 +1000
Subject: [PATCH 02/10] fix(KlerosCore): staking logic fix

---
 contracts/src/arbitration/KlerosCore.sol      | 89 ++++++++++---------
 contracts/src/arbitration/SortitionModule.sol | 18 ++--
 .../dispute-kits/DisputeKitClassic.sol        | 10 ++-
 .../dispute-kits/DisputeKitSybilResistant.sol |  5 +-
 .../interfaces/ISortitionModule.sol           |  7 +-
 5 files changed, 65 insertions(+), 64 deletions(-)

diff --git a/contracts/src/arbitration/KlerosCore.sol b/contracts/src/arbitration/KlerosCore.sol
index 819d250d8..cea096c8e 100644
--- a/contracts/src/arbitration/KlerosCore.sol
+++ b/contracts/src/arbitration/KlerosCore.sol
@@ -68,8 +68,9 @@ contract KlerosCore is IArbitratorV2 {
 
     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`.
-        mapping(uint96 => uint256) stakedPnk; // The amount of PNKs the juror has staked in the court in the form `stakedPnk[courtID]`.
-        mapping(uint96 => uint256) lockedPnk; // The amount of PNKs the juror has locked in the court in the form `lockedPnk[courtID]`.
+        uint256 stakedPnk; // The juror's total amount of tokens staked in subcourts. Reflects actual pnk balance.
+        uint256 lockedPnk; // The juror's total amount of tokens locked in disputes. Can reflect actual pnk balance when stakedPnk are fully withdrawn.
+        mapping(uint96 => uint256) stakedPnkByCourt; // The amount of PNKs the juror has staked in the court in the form `stakedPnkByCourt[courtID]`.
     }
 
     struct DisputeKitNode {
@@ -126,7 +127,7 @@ contract KlerosCore is IArbitratorV2 {
     // ************************************* //
 
     event StakeSet(address indexed _address, uint256 _courtID, uint256 _amount);
-    event StakeDelayed(address indexed _address, uint256 _courtID, uint256 _amount, uint256 _penalty);
+    event StakeDelayed(address indexed _address, uint256 _courtID, uint256 _amount);
     event NewPeriod(uint256 indexed _disputeID, Period _period);
     event AppealPossible(uint256 indexed _disputeID, IArbitrableV2 indexed _arbitrable);
     event AppealDecision(uint256 indexed _disputeID, IArbitrableV2 indexed _arbitrable);
@@ -483,12 +484,12 @@ contract KlerosCore is IArbitratorV2 {
     /// @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, 0)) revert StakingFailed();
+        if (!_setStakeForAccount(msg.sender, _courtID, _stake)) revert StakingFailed();
     }
 
-    function setStakeBySortitionModule(address _account, uint96 _courtID, uint256 _stake, uint256 _penalty) external {
+    function setStakeBySortitionModule(address _account, uint96 _courtID, uint256 _stake) external {
         if (msg.sender != address(sortitionModule)) revert WrongCaller();
-        _setStakeForAccount(_account, _courtID, _stake, _penalty);
+        _setStakeForAccount(_account, _courtID, _stake);
     }
 
     /// @inheritdoc IArbitratorV2
@@ -614,7 +615,7 @@ contract KlerosCore is IArbitratorV2 {
         for (uint256 i = startIndex; i < endIndex; i++) {
             address drawnAddress = disputeKit.draw(_disputeID);
             if (drawnAddress != address(0)) {
-                jurors[drawnAddress].lockedPnk[dispute.courtID] += round.pnkAtStakePerJuror;
+                jurors[drawnAddress].lockedPnk += round.pnkAtStakePerJuror;
                 emit Draw(drawnAddress, _disputeID, currentRound, round.drawnJurors.length);
                 round.drawnJurors.push(drawnAddress);
 
@@ -763,16 +764,12 @@ contract KlerosCore is IArbitratorV2 {
 
         // Unlock the PNKs affected by the penalty
         address account = round.drawnJurors[_params.repartition];
-        jurors[account].lockedPnk[dispute.courtID] -= penalty;
-
-        // Apply the penalty to the staked PNKs
-        if (jurors[account].stakedPnk[dispute.courtID] >= courts[dispute.courtID].minStake + penalty) {
-            // The juror still has enough staked PNKs after penalty for this court.
-            uint256 newStake = jurors[account].stakedPnk[dispute.courtID] - penalty;
-            _setStakeForAccount(account, dispute.courtID, newStake, penalty);
-        } else if (jurors[account].stakedPnk[dispute.courtID] != 0) {
-            // The juror does not have enough staked PNKs after penalty for this court, unstake them.
-            _setStakeForAccount(account, dispute.courtID, 0, penalty);
+        jurors[account].lockedPnk -= penalty;
+
+        // Apply the penalty to the staked PNKs if there ara any.
+        // Note that lockedPnk will always cover penalty while stakedPnk can become lower after manual unstaking.
+        if (jurors[account].stakedPnk >= penalty) {
+            jurors[account].stakedPnk -= penalty;
         }
         emit TokenAndETHShift(
             account,
@@ -832,10 +829,10 @@ contract KlerosCore is IArbitratorV2 {
         uint256 pnkLocked = (round.pnkAtStakePerJuror * degreeOfCoherence) / ALPHA_DIVISOR;
 
         // Release the rest of the PNKs of the juror for this round.
-        jurors[account].lockedPnk[dispute.courtID] -= pnkLocked;
+        jurors[account].lockedPnk -= pnkLocked;
 
         // Give back the locked PNKs in case the juror fully unstaked earlier.
-        if (jurors[account].stakedPnk[dispute.courtID] == 0) {
+        if (jurors[account].stakedPnk == 0) {
             pinakion.safeTransfer(account, pnkLocked);
         }
 
@@ -1014,10 +1011,11 @@ contract KlerosCore is IArbitratorV2 {
     function getJurorBalance(
         address _juror,
         uint96 _courtID
-    ) external view returns (uint256 staked, uint256 locked, uint256 nbCourts) {
+    ) external view returns (uint256 totalStaked, uint256 totalLocked, uint256 stakedInCourt, uint256 nbCourts) {
         Juror storage juror = jurors[_juror];
-        staked = juror.stakedPnk[_courtID];
-        locked = juror.lockedPnk[_courtID];
+        totalStaked = juror.stakedPnk;
+        totalLocked = juror.lockedPnk;
+        stakedInCourt = juror.stakedPnkByCourt[_courtID];
         nbCourts = juror.courtIDs.length;
     }
 
@@ -1110,35 +1108,33 @@ contract KlerosCore is IArbitratorV2 {
     /// @param _account The address of the juror.
     /// @param _courtID The ID of the court.
     /// @param _stake The new stake.
-    /// @param _penalty Penalized amount won't be transferred back to juror when the stake is lowered.
     /// @return succeeded True if the call succeeded, false otherwise.
-    function _setStakeForAccount(
-        address _account,
-        uint96 _courtID,
-        uint256 _stake,
-        uint256 _penalty
-    ) internal returns (bool succeeded) {
+    function _setStakeForAccount(address _account, uint96 _courtID, uint256 _stake) internal returns (bool succeeded) {
         if (_courtID == FORKING_COURT || _courtID > courts.length) return false;
 
         Juror storage juror = jurors[_account];
-        uint256 currentStake = juror.stakedPnk[_courtID];
+        uint256 currentStake = juror.stakedPnkByCourt[_courtID];
 
         if (_stake != 0) {
-            // Check against locked PNKs in case the min stake was lowered.
-            if (_stake < courts[_courtID].minStake || _stake < juror.lockedPnk[_courtID]) return false;
+            if (_stake < courts[_courtID].minStake) return false;
         }
 
-        ISortitionModule.preStakeHookResult result = sortitionModule.preStakeHook(_account, _courtID, _stake, _penalty);
+        ISortitionModule.preStakeHookResult result = sortitionModule.preStakeHook(_account, _courtID, _stake);
         if (result == ISortitionModule.preStakeHookResult.failed) {
             return false;
         } else if (result == ISortitionModule.preStakeHookResult.delayed) {
-            emit StakeDelayed(_account, _courtID, _stake, _penalty);
+            emit StakeDelayed(_account, _courtID, _stake);
             return true;
         }
 
         uint256 transferredAmount;
         if (_stake >= currentStake) {
-            transferredAmount = _stake - currentStake;
+            // 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
+                : 0;
             if (transferredAmount > 0) {
                 if (pinakion.safeTransferFrom(_account, address(this), transferredAmount)) {
                     if (currentStake == 0) {
@@ -1150,8 +1146,14 @@ contract KlerosCore is IArbitratorV2 {
             }
         } else {
             if (_stake == 0) {
-                // Keep locked PNKs in the contract and release them after dispute is executed.
-                transferredAmount = currentStake - juror.lockedPnk[_courtID] - _penalty;
+                // 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)) {
                         for (uint256 i = juror.courtIDs.length; i > 0; i--) {
@@ -1166,7 +1168,13 @@ contract KlerosCore is IArbitratorV2 {
                     }
                 }
             } else {
-                transferredAmount = currentStake - _stake - _penalty;
+                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;
@@ -1175,8 +1183,9 @@ contract KlerosCore is IArbitratorV2 {
             }
         }
 
-        // Update juror's records.
-        juror.stakedPnk[_courtID] = _stake;
+        // 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;
 
         sortitionModule.setStake(_account, _courtID, _stake);
         emit StakeSet(_account, _courtID, _stake);
diff --git a/contracts/src/arbitration/SortitionModule.sol b/contracts/src/arbitration/SortitionModule.sol
index f870ed52a..6a21f496c 100644
--- a/contracts/src/arbitration/SortitionModule.sol
+++ b/contracts/src/arbitration/SortitionModule.sol
@@ -36,7 +36,6 @@ contract SortitionModule is ISortitionModule {
         address account; // The address of the juror.
         uint96 courtID; // The ID of the court.
         uint256 stake; // The new stake.
-        uint256 penalty; // Penalty value, in case the stake was set during execution.
     }
 
     // ************************************* //
@@ -185,12 +184,7 @@ contract SortitionModule is ISortitionModule {
 
         for (uint256 i = delayedStakeReadIndex; i < newDelayedStakeReadIndex; i++) {
             DelayedStake storage delayedStake = delayedStakes[i];
-            core.setStakeBySortitionModule(
-                delayedStake.account,
-                delayedStake.courtID,
-                delayedStake.stake,
-                delayedStake.penalty
-            );
+            core.setStakeBySortitionModule(delayedStake.account, delayedStake.courtID, delayedStake.stake);
             delete delayedStakes[i];
         }
         delayedStakeReadIndex = newDelayedStakeReadIndex;
@@ -199,10 +193,9 @@ contract SortitionModule is ISortitionModule {
     function preStakeHook(
         address _account,
         uint96 _courtID,
-        uint256 _stake,
-        uint256 _penalty
+        uint256 _stake
     ) external override onlyByCore returns (preStakeHookResult) {
-        (uint256 currentStake, , uint256 nbCourts) = core.getJurorBalance(_account, _courtID);
+        (, , uint256 currentStake, uint256 nbCourts) = core.getJurorBalance(_account, _courtID);
         if (currentStake == 0 && nbCourts >= MAX_STAKE_PATHS) {
             // Prevent staking beyond MAX_STAKE_PATHS but unstaking is always allowed.
             return preStakeHookResult.failed;
@@ -211,8 +204,7 @@ contract SortitionModule is ISortitionModule {
                 delayedStakes[++delayedStakeWriteIndex] = DelayedStake({
                     account: _account,
                     courtID: _courtID,
-                    stake: _stake,
-                    penalty: _penalty
+                    stake: _stake
                 });
                 return preStakeHookResult.delayed;
             }
@@ -264,7 +256,7 @@ contract SortitionModule is ISortitionModule {
     function setJurorInactive(address _account) external override onlyByCore {
         uint96[] memory courtIDs = core.getJurorCourtIDs(_account);
         for (uint256 j = courtIDs.length; j > 0; j--) {
-            core.setStakeBySortitionModule(_account, courtIDs[j - 1], 0, 0);
+            core.setStakeBySortitionModule(_account, courtIDs[j - 1], 0);
         }
     }
 
diff --git a/contracts/src/arbitration/dispute-kits/DisputeKitClassic.sol b/contracts/src/arbitration/dispute-kits/DisputeKitClassic.sol
index 4eb8fb523..ecdc9a930 100644
--- a/contracts/src/arbitration/dispute-kits/DisputeKitClassic.sol
+++ b/contracts/src/arbitration/dispute-kits/DisputeKitClassic.sol
@@ -558,7 +558,13 @@ contract DisputeKitClassic is BaseDisputeKit, IEvidence {
     // ************************************* //
 
     /// @inheritdoc BaseDisputeKit
-    function _postDrawCheck(uint256 /*_coreDisputeID*/, address /*_juror*/) internal pure override returns (bool) {
-        return true;
+    function _postDrawCheck(uint256 _coreDisputeID, address _juror) internal view override returns (bool) {
+        (uint96 courtID, , , , ) = core.disputes(_coreDisputeID);
+        (, uint256 lockedAmountPerJuror, , , , , , , , ) = core.getRoundInfo(
+            _coreDisputeID,
+            core.getNumberOfRounds(_coreDisputeID) - 1
+        );
+        (uint256 totalStaked, uint256 totalLocked, , ) = core.getJurorBalance(_juror, courtID);
+        return totalStaked >= totalLocked + lockedAmountPerJuror;
     }
 }
diff --git a/contracts/src/arbitration/dispute-kits/DisputeKitSybilResistant.sol b/contracts/src/arbitration/dispute-kits/DisputeKitSybilResistant.sol
index 87528b8b9..9d230de56 100644
--- a/contracts/src/arbitration/dispute-kits/DisputeKitSybilResistant.sol
+++ b/contracts/src/arbitration/dispute-kits/DisputeKitSybilResistant.sol
@@ -587,9 +587,8 @@ contract DisputeKitSybilResistant is BaseDisputeKit, IEvidence {
             _coreDisputeID,
             core.getNumberOfRounds(_coreDisputeID) - 1
         );
-        (uint256 staked, uint256 locked, ) = core.getJurorBalance(_juror, courtID);
-        (, , uint256 minStake, , , , ) = core.courts(courtID);
-        if (staked < locked + lockedAmountPerJuror || staked < minStake) {
+        (uint256 totalStaked, uint256 totalLocked, , ) = core.getJurorBalance(_juror, courtID);
+        if (totalStaked < totalLocked + lockedAmountPerJuror) {
             return false;
         } else {
             return _proofOfHumanity(_juror);
diff --git a/contracts/src/arbitration/interfaces/ISortitionModule.sol b/contracts/src/arbitration/interfaces/ISortitionModule.sol
index 0c5f15c6f..a9ecd56c1 100644
--- a/contracts/src/arbitration/interfaces/ISortitionModule.sol
+++ b/contracts/src/arbitration/interfaces/ISortitionModule.sol
@@ -26,12 +26,7 @@ interface ISortitionModule {
 
     function draw(bytes32 _court, uint256 _coreDisputeID, uint256 _voteID) external view returns (address);
 
-    function preStakeHook(
-        address _account,
-        uint96 _courtID,
-        uint256 _stake,
-        uint256 _penalty
-    ) external returns (preStakeHookResult);
+    function preStakeHook(address _account, uint96 _courtID, uint256 _stake) external returns (preStakeHookResult);
 
     function createDisputeHook(uint256 _disputeID, uint256 _roundID) external;
 

From 17d7a81d3df52c5d5a4e10bcc86887604fd3cb68 Mon Sep 17 00:00:00 2001
From: unknownunknown1 <time_arrow@bk.ru>
Date: Wed, 16 Aug 2023 19:32:59 +1000
Subject: [PATCH 03/10] fix(KlerosCore): drawing iterations don't repeat

---
 contracts/src/arbitration/KlerosCore.sol              | 11 ++++++-----
 contracts/src/arbitration/SortitionModule.sol         |  6 +++---
 .../arbitration/dispute-kits/DisputeKitClassic.sol    |  6 ++++--
 .../dispute-kits/DisputeKitSybilResistant.sol         |  6 ++++--
 contracts/src/arbitration/interfaces/IDisputeKit.sol  |  3 ++-
 .../src/arbitration/interfaces/ISortitionModule.sol   |  2 +-
 6 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/contracts/src/arbitration/KlerosCore.sol b/contracts/src/arbitration/KlerosCore.sol
index cea096c8e..15fa652d4 100644
--- a/contracts/src/arbitration/KlerosCore.sol
+++ b/contracts/src/arbitration/KlerosCore.sol
@@ -64,6 +64,7 @@ contract KlerosCore is IArbitratorV2 {
         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.
+        uint256 drawIterations; // The number of iterations passed drawing the jurors for this round.
     }
 
     struct Juror {
@@ -609,11 +610,10 @@ contract KlerosCore is IArbitratorV2 {
 
         IDisputeKit disputeKit = disputeKitNodes[round.disputeKitID].disputeKit;
 
-        uint256 startIndex = round.drawnJurors.length;
-        uint256 endIndex = startIndex + _iterations <= round.nbVotes ? startIndex + _iterations : round.nbVotes;
-
-        for (uint256 i = startIndex; i < endIndex; i++) {
-            address drawnAddress = disputeKit.draw(_disputeID);
+        uint256 startIndex = round.drawIterations;
+        // TODO: cap the iterations?
+        for (uint256 i = startIndex; i < startIndex + _iterations; i++) {
+            address drawnAddress = disputeKit.draw(_disputeID, i);
             if (drawnAddress != address(0)) {
                 jurors[drawnAddress].lockedPnk += round.pnkAtStakePerJuror;
                 emit Draw(drawnAddress, _disputeID, currentRound, round.drawnJurors.length);
@@ -624,6 +624,7 @@ contract KlerosCore is IArbitratorV2 {
                 }
             }
         }
+        round.drawIterations += _iterations;
     }
 
     /// @dev Appeals the ruling of a specified dispute.
diff --git a/contracts/src/arbitration/SortitionModule.sol b/contracts/src/arbitration/SortitionModule.sol
index 6a21f496c..1cd4e3bf6 100644
--- a/contracts/src/arbitration/SortitionModule.sol
+++ b/contracts/src/arbitration/SortitionModule.sol
@@ -268,7 +268,7 @@ contract SortitionModule is ISortitionModule {
     /// Note that this function reverts if the sum of all values in the tree is 0.
     /// @param _key The key of the tree.
     /// @param _coreDisputeID Index of the dispute in Kleros Core.
-    /// @param _voteID ID of the voter.
+    /// @param _nonce Nonce to hash with random number.
     /// @return drawnAddress The drawn address.
     /// `O(k * log_k(n))` where
     /// `k` is the maximum number of children per node in the tree,
@@ -276,7 +276,7 @@ contract SortitionModule is ISortitionModule {
     function draw(
         bytes32 _key,
         uint256 _coreDisputeID,
-        uint256 _voteID
+        uint256 _nonce
     ) public view override returns (address drawnAddress) {
         require(phase == Phase.drawing, "Wrong phase.");
         SortitionSumTree storage tree = sortitionSumTrees[_key];
@@ -285,7 +285,7 @@ contract SortitionModule is ISortitionModule {
             return address(0); // No jurors staked.
         }
 
-        uint256 currentDrawnNumber = uint256(keccak256(abi.encodePacked(randomNumber, _coreDisputeID, _voteID))) %
+        uint256 currentDrawnNumber = uint256(keccak256(abi.encodePacked(randomNumber, _coreDisputeID, _nonce))) %
             tree.nodes[0];
 
         // While it still has children
diff --git a/contracts/src/arbitration/dispute-kits/DisputeKitClassic.sol b/contracts/src/arbitration/dispute-kits/DisputeKitClassic.sol
index ecdc9a930..12d3639f2 100644
--- a/contracts/src/arbitration/dispute-kits/DisputeKitClassic.sol
+++ b/contracts/src/arbitration/dispute-kits/DisputeKitClassic.sol
@@ -181,9 +181,11 @@ contract DisputeKitClassic is BaseDisputeKit, IEvidence {
     /// @dev Draws the juror from the sortition tree. The drawn address is picked up by Kleros Core.
     /// Note: Access restricted to Kleros Core only.
     /// @param _coreDisputeID The ID of the dispute in Kleros Core.
+    /// @param _nonce Nonce of the drawing iteration.
     /// @return drawnAddress The drawn address.
     function draw(
-        uint256 _coreDisputeID
+        uint256 _coreDisputeID,
+        uint256 _nonce
     ) external override onlyByCore notJumped(_coreDisputeID) returns (address drawnAddress) {
         Dispute storage dispute = disputes[coreDisputeIDToLocal[_coreDisputeID]];
         Round storage round = dispute.rounds[dispute.rounds.length - 1];
@@ -193,7 +195,7 @@ contract DisputeKitClassic is BaseDisputeKit, IEvidence {
         bytes32 key = bytes32(uint256(courtID)); // Get the ID of the tree.
 
         // TODO: Handle the situation when no one has staked yet.
-        drawnAddress = sortitionModule.draw(key, _coreDisputeID, round.votes.length);
+        drawnAddress = sortitionModule.draw(key, _coreDisputeID, _nonce);
 
         if (_postDrawCheck(_coreDisputeID, drawnAddress)) {
             round.votes.push(Vote({account: drawnAddress, commit: bytes32(0), choice: 0, voted: false}));
diff --git a/contracts/src/arbitration/dispute-kits/DisputeKitSybilResistant.sol b/contracts/src/arbitration/dispute-kits/DisputeKitSybilResistant.sol
index 9d230de56..54bc0e8b0 100644
--- a/contracts/src/arbitration/dispute-kits/DisputeKitSybilResistant.sol
+++ b/contracts/src/arbitration/dispute-kits/DisputeKitSybilResistant.sol
@@ -200,9 +200,11 @@ contract DisputeKitSybilResistant is BaseDisputeKit, IEvidence {
     /// @dev Draws the juror from the sortition tree. The drawn address is picked up by Kleros Core.
     /// Note: Access restricted to Kleros Core only.
     /// @param _coreDisputeID The ID of the dispute in Kleros Core.
+    /// @param _nonce Nonce of the drawing iteration.
     /// @return drawnAddress The drawn address.
     function draw(
-        uint256 _coreDisputeID
+        uint256 _coreDisputeID,
+        uint256 _nonce
     ) external override onlyByCore notJumped(_coreDisputeID) returns (address drawnAddress) {
         Dispute storage dispute = disputes[coreDisputeIDToLocal[_coreDisputeID]];
         Round storage round = dispute.rounds[dispute.rounds.length - 1];
@@ -212,7 +214,7 @@ contract DisputeKitSybilResistant is BaseDisputeKit, IEvidence {
         bytes32 key = bytes32(uint256(courtID)); // Get the ID of the tree.
 
         // TODO: Handle the situation when no one has staked yet.
-        drawnAddress = sortitionModule.draw(key, _coreDisputeID, round.votes.length);
+        drawnAddress = sortitionModule.draw(key, _coreDisputeID, _nonce);
 
         if (_postDrawCheck(_coreDisputeID, drawnAddress) && !round.alreadyDrawn[drawnAddress]) {
             round.votes.push(Vote({account: drawnAddress, commit: bytes32(0), choice: 0, voted: false}));
diff --git a/contracts/src/arbitration/interfaces/IDisputeKit.sol b/contracts/src/arbitration/interfaces/IDisputeKit.sol
index bd8f804c0..2905c9cf5 100644
--- a/contracts/src/arbitration/interfaces/IDisputeKit.sol
+++ b/contracts/src/arbitration/interfaces/IDisputeKit.sol
@@ -51,8 +51,9 @@ interface IDisputeKit {
     /// @dev Draws the juror from the sortition tree. The drawn address is picked up by Kleros Core.
     /// Note: Access restricted to Kleros Core only.
     /// @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) external returns (address drawnAddress);
+    function draw(uint256 _coreDisputeID, uint256 _nonce) external returns (address drawnAddress);
 
     // ************************************* //
     // *           Public Views            * //
diff --git a/contracts/src/arbitration/interfaces/ISortitionModule.sol b/contracts/src/arbitration/interfaces/ISortitionModule.sol
index a9ecd56c1..3bee0e270 100644
--- a/contracts/src/arbitration/interfaces/ISortitionModule.sol
+++ b/contracts/src/arbitration/interfaces/ISortitionModule.sol
@@ -24,7 +24,7 @@ interface ISortitionModule {
 
     function notifyRandomNumber(uint256 _drawnNumber) external;
 
-    function draw(bytes32 _court, uint256 _coreDisputeID, uint256 _voteID) external view returns (address);
+    function draw(bytes32 _court, uint256 _coreDisputeID, uint256 _nonce) external view returns (address);
 
     function preStakeHook(address _account, uint96 _courtID, uint256 _stake) external returns (preStakeHookResult);
 

From 62e4e7540faa1da2b90d940b375f8bf611228abf Mon Sep 17 00:00:00 2001
From: unknownunknown1 <time_arrow@bk.ru>
Date: Tue, 29 Aug 2023 16:33:28 +1000
Subject: [PATCH 04/10] fix(KC): small bug fixes

---
 contracts/src/arbitration/KlerosCore.sol | 32 +++++++++++++-----------
 1 file changed, 17 insertions(+), 15 deletions(-)

diff --git a/contracts/src/arbitration/KlerosCore.sol b/contracts/src/arbitration/KlerosCore.sol
index 15fa652d4..a21ce5598 100644
--- a/contracts/src/arbitration/KlerosCore.sol
+++ b/contracts/src/arbitration/KlerosCore.sol
@@ -767,10 +767,12 @@ contract KlerosCore is IArbitratorV2 {
         address account = round.drawnJurors[_params.repartition];
         jurors[account].lockedPnk -= penalty;
 
-        // Apply the penalty to the staked PNKs if there ara any.
+        // Apply the penalty to the staked PNKs.
         // Note that lockedPnk will always cover penalty while stakedPnk can become lower after manual unstaking.
         if (jurors[account].stakedPnk >= penalty) {
             jurors[account].stakedPnk -= penalty;
+        } else {
+            jurors[account].stakedPnk = 0;
         }
         emit TokenAndETHShift(
             account,
@@ -1118,6 +1120,8 @@ contract KlerosCore is IArbitratorV2 {
 
         if (_stake != 0) {
             if (_stake < courts[_courtID].minStake) return false;
+        } else if (currentStake == 0) {
+            return false;
         }
 
         ISortitionModule.preStakeHookResult result = sortitionModule.preStakeHook(_account, _courtID, _stake);
@@ -1137,14 +1141,13 @@ contract KlerosCore is IArbitratorV2 {
                 ? _stake - currentStake - previouslyLocked
                 : 0;
             if (transferredAmount > 0) {
-                if (pinakion.safeTransferFrom(_account, address(this), transferredAmount)) {
-                    if (currentStake == 0) {
-                        juror.courtIDs.push(_courtID);
-                    }
-                } else {
+                if (!pinakion.safeTransferFrom(_account, address(this), transferredAmount)) {
                     return false;
                 }
             }
+            if (currentStake == 0) {
+                juror.courtIDs.push(_courtID);
+            }
         } else {
             if (_stake == 0) {
                 // Make sure locked tokens always stay in the contract. They can only be released during Execution.
@@ -1156,18 +1159,17 @@ contract KlerosCore is IArbitratorV2 {
                     transferredAmount = juror.stakedPnk - juror.lockedPnk;
                 }
                 if (transferredAmount > 0) {
-                    if (pinakion.safeTransfer(_account, transferredAmount)) {
-                        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 (!pinakion.safeTransfer(_account, transferredAmount)) {
                         return false;
                     }
                 }
+                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.

From 55347685378e585233e190b32645992c29f68642 Mon Sep 17 00:00:00 2001
From: jaybuidl <jb@kleros.io>
Date: Wed, 6 Sep 2023 22:44:19 +0800
Subject: [PATCH 05/10] refactor: minor refactor and fixed the tests

---
 contracts/src/arbitration/KlerosCore.sol | 101 +++++++++++------------
 contracts/test/arbitration/draw.ts       |  50 ++++++-----
 contracts/test/integration/index.ts      |  22 +++--
 3 files changed, 90 insertions(+), 83 deletions(-)

diff --git a/contracts/src/arbitration/KlerosCore.sol b/contracts/src/arbitration/KlerosCore.sol
index a21ce5598..81dc736e4 100644
--- 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,20 +1153,20 @@ 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];
@@ -1170,28 +1174,15 @@ contract KlerosCore is IArbitratorV2 {
                         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
index 5e835857b..9979afd09 100644
--- 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
index 86162b3ce..e04d39fa5 100644
--- 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)
+  );
 };

From 66dce00f00d9fdb543eb1dcc22079a2701194a53 Mon Sep 17 00:00:00 2001
From: marino <102478601+kemuru@users.noreply.github.com>
Date: Thu, 7 Sep 2023 10:39:43 +0200
Subject: [PATCH 06/10] fix(subgraph): update subgraph with contract changes

---
 subgraph/subgraph.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/subgraph/subgraph.yaml b/subgraph/subgraph.yaml
index fd19eb740..04487470e 100644
--- a/subgraph/subgraph.yaml
+++ b/subgraph/subgraph.yaml
@@ -48,7 +48,7 @@ dataSources:
           handler: handleDisputeKitEnabled
         - event: StakeSet(indexed address,uint256,uint256)
           handler: handleStakeSet
-        - event: StakeDelayed(indexed address,uint256,uint256,uint256)
+        - event: StakeDelayed(indexed address,uint256,uint256)
           handler: handleStakeDelayed
         - event: TokenAndETHShift(indexed address,indexed uint256,indexed uint256,uint256,int256,int256,address)
           handler: handleTokenAndETHShift

From 7d90db71ed1b2f081d6c405bd1d4f8fe2016fdd0 Mon Sep 17 00:00:00 2001
From: jaybuidl <jb@kleros.io>
Date: Sat, 9 Sep 2023 10:19:16 +0800
Subject: [PATCH 07/10] fix: subgraph update script

---
 subgraph/package.json      |  3 ++-
 subgraph/scripts/update.sh | 35 +++++++++++++++++++++++++----------
 2 files changed, 27 insertions(+), 11 deletions(-)

diff --git a/subgraph/package.json b/subgraph/package.json
index cc9ca7499..7d70d40d7 100644
--- a/subgraph/package.json
+++ b/subgraph/package.json
@@ -5,6 +5,7 @@
     "update:arbitrum-goerli": "./scripts/update.sh arbitrumGoerli arbitrum-goerli",
     "update:arbitrum-goerli-devnet": "./scripts/update.sh arbitrumGoerliDevnet arbitrum-goerli",
     "update:arbitrum": "./scripts/update.sh arbitrum arbitrum",
+    "update:local": "./scripts/update.sh localhost mainnet",
     "codegen": "graph codegen",
     "build": "graph build",
     "clean": "graph clean && rm subgraph.yaml.bak.*",
@@ -14,7 +15,7 @@
     "create-local": "graph create --node http://localhost:8020/ kleros/kleros-v2-core-local",
     "remove-local": "graph remove --node http://localhost:8020/ kleros/kleros-v2-core-local",
     "deploy-local": "graph deploy --node http://localhost:8020/ --ipfs http://localhost:5001 kleros/kleros-v2-core-local --version-label v$(date +%s)",
-    "rebuild-deploy-local": "./scripts/update.sh localhost mainnet && yarn codegen && yarn create-local && yarn deploy-local",
+    "rebuild-deploy-local": "yarn update:local && yarn codegen && yarn create-local && yarn deploy-local",
     "start-local-indexer": "docker compose -f ../services/graph-node/docker-compose.yml up -d && docker compose -f ../services/graph-node/docker-compose.yml logs -f",
     "stop-local-indexer": "docker compose -f ../services/graph-node/docker-compose.yml down && rm -rf ../services/graph-node/data"
   },
diff --git a/subgraph/scripts/update.sh b/subgraph/scripts/update.sh
index 980b70a3c..b36e7ed0b 100755
--- a/subgraph/scripts/update.sh
+++ b/subgraph/scripts/update.sh
@@ -2,18 +2,33 @@
 
 SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
 
-function update() #file #dataSourceIndex #graphNetwork
+function update() #hardhatNetwork #graphNetwork #dataSourceIndex #contract
 {
-    local f="$1"
-    local dataSourceIndex="$2"
-
-    graphNetwork=$3 yq -i  ".dataSources[$dataSourceIndex].network=env(graphNetwork)" "$SCRIPT_DIR"/../subgraph.yaml
+    local hardhatNetwork="$1"
+    local graphNetwork="$2"
+    local dataSourceIndex="$3"
+    local contract="$4"
+    local artifact="$SCRIPT_DIR/../../contracts/deployments/$hardhatNetwork/$contract.json"
+
+    # Set the address
+    address=$(cat "$artifact" | jq '.address')
+    yq -i ".dataSources[$dataSourceIndex].source.address=$address" "$SCRIPT_DIR"/../subgraph.yaml
+    
+    # Set the start block
+    blockNumber="$(cat "$artifact" | jq '.receipt.blockNumber')"
+    yq -i ".dataSources[$dataSourceIndex].source.startBlock=$blockNumber" "$SCRIPT_DIR"/../subgraph.yaml
 
-    address=$(cat "$f" | jq '.address')
-    yq -i  ".dataSources[$dataSourceIndex].source.address=$address" "$SCRIPT_DIR"/../subgraph.yaml
+    # Set the Graph network
+    graphNetwork=$graphNetwork yq -i  ".dataSources[$dataSourceIndex].network=env(graphNetwork)" "$SCRIPT_DIR"/../subgraph.yaml
     
-    blockNumber="$(cat "$f" | jq '.receipt.blockNumber')"
-    yq -i  ".dataSources[$dataSourceIndex].source.startBlock=$blockNumber" "$SCRIPT_DIR"/../subgraph.yaml
+    # Set the ABIs path for this Hardhat network
+    abiIndex=0
+    for f in $(yq e .dataSources[$dataSourceIndex].mapping.abis[].file subgraph.yaml -o json -I 0 | jq -sr '.[]')
+    do
+        f2=$(echo $f | sed "s|\(.*\/deployments\/\).*\/|\1$hardhatNetwork\/|")
+        yq -i ".dataSources[$dataSourceIndex].mapping.abis[$abiIndex].file=\"$f2\"" "$SCRIPT_DIR"/../subgraph.yaml
+        (( ++abiIndex ))
+    done
 }
 
 # as per ../contracts/hardhat.config.js
@@ -28,6 +43,6 @@ cp "$SCRIPT_DIR"/../subgraph.yaml "$SCRIPT_DIR"/../subgraph.yaml.bak.$(date +%s)
 
 for contract in $(yq .dataSources[].name "$SCRIPT_DIR"/../subgraph.yaml)
 do
-    update "$SCRIPT_DIR/../../contracts/deployments/$hardhatNetwork/$contract.json" $i $graphNetwork
+    update $hardhatNetwork $graphNetwork $i $contract
     (( ++i ))
 done

From b714c12e47602f0330a9a6ce0caacef6748eaa73 Mon Sep 17 00:00:00 2001
From: jaybuidl <jb@kleros.io>
Date: Mon, 11 Sep 2023 00:24:59 +0800
Subject: [PATCH 08/10] fix: accounting of round.drawIterations

---
 contracts/src/arbitration/KlerosCore.sol      | 43 +++----------------
 .../dispute-kits/DisputeKitClassic.sol        |  7 ++-
 .../dispute-kits/DisputeKitSybilResistant.sol |  7 ++-
 contracts/test/arbitration/draw.ts            |  9 +++-
 4 files changed, 21 insertions(+), 45 deletions(-)

diff --git a/contracts/src/arbitration/KlerosCore.sol b/contracts/src/arbitration/KlerosCore.sol
index 81dc736e4..5e7616882 100644
--- a/contracts/src/arbitration/KlerosCore.sol
+++ b/contracts/src/arbitration/KlerosCore.sol
@@ -610,9 +610,10 @@ contract KlerosCore is IArbitratorV2 {
 
         IDisputeKit disputeKit = disputeKitNodes[round.disputeKitID].disputeKit;
 
-        uint256 startIndex = round.drawIterations;
-        for (uint256 i = startIndex; i < startIndex + _iterations && round.drawnJurors.length < round.nbVotes; i++) {
-            address drawnAddress = disputeKit.draw(_disputeID, i);
+        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++);
             if (drawnAddress == address(0)) {
                 continue;
             }
@@ -623,7 +624,7 @@ contract KlerosCore is IArbitratorV2 {
                 sortitionModule.postDrawHook(_disputeID, currentRound);
             }
         }
-        round.drawIterations += _iterations;
+        round.drawIterations += i;
     }
 
     /// @dev Appeals the ruling of a specified dispute.
@@ -972,38 +973,8 @@ contract KlerosCore is IArbitratorV2 {
         (ruling, tied, overridden) = disputeKit.currentRuling(_disputeID);
     }
 
-    function getRoundInfo(
-        uint256 _disputeID,
-        uint256 _round
-    )
-        external
-        view
-        returns (
-            uint256 disputeKitID,
-            uint256 pnkAtStakePerJuror,
-            uint256 totalFeesForJurors,
-            uint256 nbVotes,
-            uint256 repartitions,
-            uint256 pnkPenalties,
-            address[] memory drawnJurors,
-            uint256 sumFeeRewardPaid,
-            uint256 sumPnkRewardPaid,
-            IERC20 feeToken
-        )
-    {
-        Round storage round = disputes[_disputeID].rounds[_round];
-        return (
-            round.disputeKitID,
-            round.pnkAtStakePerJuror,
-            round.totalFeesForJurors,
-            round.nbVotes,
-            round.repartitions,
-            round.pnkPenalties,
-            round.drawnJurors,
-            round.sumFeeRewardPaid,
-            round.sumPnkRewardPaid,
-            round.feeToken
-        );
+    function getRoundInfo(uint256 _disputeID, uint256 _round) external view returns (Round memory) {
+        return disputes[_disputeID].rounds[_round];
     }
 
     function getNumberOfRounds(uint256 _disputeID) external view returns (uint256) {
diff --git a/contracts/src/arbitration/dispute-kits/DisputeKitClassic.sol b/contracts/src/arbitration/dispute-kits/DisputeKitClassic.sol
index 12d3639f2..cf4e927b3 100644
--- a/contracts/src/arbitration/dispute-kits/DisputeKitClassic.sol
+++ b/contracts/src/arbitration/dispute-kits/DisputeKitClassic.sol
@@ -562,10 +562,9 @@ contract DisputeKitClassic is BaseDisputeKit, IEvidence {
     /// @inheritdoc BaseDisputeKit
     function _postDrawCheck(uint256 _coreDisputeID, address _juror) internal view override returns (bool) {
         (uint96 courtID, , , , ) = core.disputes(_coreDisputeID);
-        (, uint256 lockedAmountPerJuror, , , , , , , , ) = core.getRoundInfo(
-            _coreDisputeID,
-            core.getNumberOfRounds(_coreDisputeID) - 1
-        );
+        uint256 lockedAmountPerJuror = core
+            .getRoundInfo(_coreDisputeID, core.getNumberOfRounds(_coreDisputeID) - 1)
+            .pnkAtStakePerJuror;
         (uint256 totalStaked, uint256 totalLocked, , ) = core.getJurorBalance(_juror, courtID);
         return totalStaked >= totalLocked + lockedAmountPerJuror;
     }
diff --git a/contracts/src/arbitration/dispute-kits/DisputeKitSybilResistant.sol b/contracts/src/arbitration/dispute-kits/DisputeKitSybilResistant.sol
index 54bc0e8b0..eb907c87e 100644
--- a/contracts/src/arbitration/dispute-kits/DisputeKitSybilResistant.sol
+++ b/contracts/src/arbitration/dispute-kits/DisputeKitSybilResistant.sol
@@ -585,10 +585,9 @@ contract DisputeKitSybilResistant is BaseDisputeKit, IEvidence {
     /// @return Whether the address can be drawn or not.
     function _postDrawCheck(uint256 _coreDisputeID, address _juror) internal view override returns (bool) {
         (uint96 courtID, , , , ) = core.disputes(_coreDisputeID);
-        (, uint256 lockedAmountPerJuror, , , , , , , , ) = core.getRoundInfo(
-            _coreDisputeID,
-            core.getNumberOfRounds(_coreDisputeID) - 1
-        );
+        uint256 lockedAmountPerJuror = core
+            .getRoundInfo(_coreDisputeID, core.getNumberOfRounds(_coreDisputeID) - 1)
+            .pnkAtStakePerJuror;
         (uint256 totalStaked, uint256 totalLocked, , ) = core.getJurorBalance(_juror, courtID);
         if (totalStaked < totalLocked + lockedAmountPerJuror) {
             return false;
diff --git a/contracts/test/arbitration/draw.ts b/contracts/test/arbitration/draw.ts
index 9979afd09..cb3bc2619 100644
--- a/contracts/test/arbitration/draw.ts
+++ b/contracts/test/arbitration/draw.ts
@@ -203,6 +203,8 @@ describe("Draw Benchmark", async () => {
     };
     let countedDraws: CountedDraws;
     const expectFromDraw = async (drawTx: Promise<ContractTransaction>) => {
+      expect(await core.getRoundInfo(0, 0).then((round) => round.drawIterations)).to.equal(3);
+
       const tx = await (await drawTx).wait();
       expect(tx)
         .to.emit(core, "Draw")
@@ -261,7 +263,8 @@ describe("Draw Benchmark", async () => {
     };
 
     const expectFromDraw = async (drawTx: Promise<ContractTransaction>) => {
-      await expect(drawTx).to.not.emit(core, "Draw");
+      expect(await core.getRoundInfo(0, 0).then((round) => round.drawIterations)).to.equal(20);
+      expect(await drawTx).to.not.emit(core, "Draw");
     };
 
     const unstake = async (wallet: Wallet) => {
@@ -295,6 +298,8 @@ describe("Draw Benchmark", async () => {
     };
     let countedDraws: CountedDraws;
     const expectFromDraw = async (drawTx: Promise<ContractTransaction>) => {
+      expect(await core.getRoundInfo(0, 0).then((round) => round.drawIterations)).to.equal(3);
+
       const tx = await (await drawTx).wait();
       expect(tx)
         .to.emit(core, "Draw")
@@ -353,6 +358,8 @@ describe("Draw Benchmark", async () => {
     };
     let countedDraws: CountedDraws;
     const expectFromDraw = async (drawTx: Promise<ContractTransaction>) => {
+      expect(await core.getRoundInfo(0, 0).then((round) => round.drawIterations)).to.equal(3);
+
       const tx = await (await drawTx).wait();
       expect(tx)
         .to.emit(core, "Draw")

From 96245e384695775546b08770c4520767f60703ce Mon Sep 17 00:00:00 2001
From: marino <102478601+kemuru@users.noreply.github.com>
Date: Thu, 7 Sep 2023 15:04:57 +0200
Subject: [PATCH 09/10] feat(web): round info on dispute card

---
 web/src/assets/svgs/icons/round.svg            | 10 ++++++++++
 web/src/components/DisputeCard/DisputeInfo.tsx |  7 +++++--
 web/src/components/DisputeCard/index.tsx       |  4 ++++
 3 files changed, 19 insertions(+), 2 deletions(-)
 create mode 100644 web/src/assets/svgs/icons/round.svg

diff --git a/web/src/assets/svgs/icons/round.svg b/web/src/assets/svgs/icons/round.svg
new file mode 100644
index 000000000..3e36edfdd
--- /dev/null
+++ b/web/src/assets/svgs/icons/round.svg
@@ -0,0 +1,10 @@
+<svg width="15" height="15" viewBox="0 0 15 15" fill="none" xmlns="http://www.w3.org/2000/svg">
+<g clip-path="url(#clip0_16930_50355)">
+<path fill-rule="evenodd" clip-rule="evenodd" d="M1.96382 5.27C1.68215 5.935 1.53215 6.64167 1.51549 7.37C1.51382 7.44167 1.45549 7.5 1.38215 7.5H0.382154C0.364361 7.50001 0.346746 7.49645 0.330348 7.48954C0.31395 7.48263 0.299101 7.47251 0.286675 7.45978C0.274249 7.44704 0.264499 7.43195 0.257997 7.41538C0.251496 7.39882 0.248376 7.38112 0.248821 7.36333C0.320488 3.56 3.42715 0.5 7.24549 0.5C9.48549 0.5 11.4788 1.55167 12.7605 3.19167L13.7005 2.45667C13.7872 2.38667 13.9155 2.45 13.9155 2.56167L13.9022 5.57667C13.9022 5.66167 13.8205 5.725 13.7372 5.705L10.8088 4.98833C10.7844 4.98239 10.7622 4.96968 10.7447 4.95167C10.7273 4.93367 10.7152 4.91109 10.71 4.88654C10.7048 4.86199 10.7066 4.83647 10.7152 4.8129C10.7239 4.78934 10.739 4.76869 10.7588 4.75333L11.7622 3.97C11.6188 3.78667 11.4638 3.61167 11.2988 3.44667C10.7773 2.9222 10.1585 2.50441 9.47715 2.21667C8.77215 1.91833 8.02215 1.76667 7.24715 1.76667C6.47215 1.76667 5.72215 1.91833 5.01549 2.21667C4.33215 2.505 3.71882 2.91833 3.19215 3.44667C2.66549 3.975 2.25215 4.58667 1.96382 5.27ZM13.1155 7.5H14.1155C14.1333 7.49999 14.1509 7.50355 14.1673 7.51046C14.1837 7.51737 14.1985 7.52749 14.211 7.54022C14.2234 7.55296 14.2331 7.56805 14.2396 7.58462C14.2461 7.60118 14.2493 7.61888 14.2488 7.63667C14.1772 11.4383 11.0722 14.5 7.25215 14.5C5.01215 14.5 3.01882 13.4483 1.73715 11.81L0.797154 12.545C0.710488 12.615 0.582154 12.5517 0.582154 12.44L0.593821 9.42333C0.593821 9.33833 0.675488 9.275 0.758821 9.295L3.68715 10.0117C3.71154 10.0176 3.73375 10.0303 3.75123 10.0483C3.76872 10.0663 3.78076 10.0889 3.78599 10.1135C3.79121 10.138 3.7894 10.1635 3.78076 10.1871C3.77212 10.2107 3.75701 10.2313 3.73715 10.2467L2.73382 11.03C2.87715 11.2133 3.03215 11.3883 3.19715 11.5533C3.72843 12.0871 4.36013 12.5103 5.05582 12.7987C5.75151 13.087 6.49741 13.2347 7.25049 13.2333C8.00381 13.2348 8.74998 13.0871 9.44594 12.7987C10.1419 12.5104 10.7739 12.0871 11.3055 11.5533C11.8322 11.0267 12.2455 10.4133 12.5338 9.73C12.8155 9.065 12.9655 8.35833 12.9822 7.63C12.9838 7.55833 13.0422 7.5 13.1155 7.5ZM9.56967 6.68723L9.80169 6.91924C9.89779 7.01535 9.89779 7.17117 9.80169 7.26727L8.52558 8.54337C8.42948 8.63948 8.27367 8.63948 8.17756 8.54337L7.94554 8.31136C7.84943 8.21526 7.84943 8.05944 7.94554 7.96333L8.00355 7.90533L7.59751 7.49929L6.76654 8.33027L6.82455 8.38827C6.95269 8.51641 6.95269 8.72417 6.82455 8.85231L5.64796 10.0289C5.51982 10.157 5.31207 10.157 5.18393 10.0289L4.71989 9.56485C4.59174 9.43671 4.59174 9.22896 4.71989 9.10081L5.89646 7.92422C6.0246 7.79608 6.23235 7.79608 6.3605 7.92422L6.4185 7.98223L7.24948 7.15126L6.84344 6.74522L6.78544 6.80323C6.68934 6.89934 6.53352 6.89934 6.43741 6.80323L6.2054 6.57121C6.10929 6.47511 6.10929 6.3193 6.2054 6.22319L7.4815 4.94708C7.5776 4.85097 7.73342 4.85097 7.82953 4.94708L8.06154 5.17911C8.15765 5.27521 8.15765 5.43102 8.06154 5.52713L8.00354 5.58514L9.16364 6.74523L9.22165 6.68723C9.31775 6.59112 9.47356 6.59112 9.56967 6.68723Z" fill="#B45FFF"/>
+</g>
+<defs>
+<clipPath id="clip0_16930_50355">
+<rect width="14" height="14" fill="white" transform="translate(0.248779 0.5)"/>
+</clipPath>
+</defs>
+</svg>
diff --git a/web/src/components/DisputeCard/DisputeInfo.tsx b/web/src/components/DisputeCard/DisputeInfo.tsx
index 0b3568097..dc36b524f 100644
--- a/web/src/components/DisputeCard/DisputeInfo.tsx
+++ b/web/src/components/DisputeCard/DisputeInfo.tsx
@@ -5,6 +5,7 @@ import BookmarkIcon from "svgs/icons/bookmark.svg";
 import CalendarIcon from "svgs/icons/calendar.svg";
 import LawBalanceIcon from "svgs/icons/law-balance.svg";
 import PileCoinsIcon from "svgs/icons/pile-coins.svg";
+import RoundIcon from "svgs/icons/round.svg";
 import Field from "../Field";
 
 const Container = styled.div`
@@ -33,13 +34,15 @@ export interface IDisputeInfo {
   rewards?: string;
   period?: Periods;
   date?: number;
+  round?: number;
 }
 
-const DisputeInfo: React.FC<IDisputeInfo> = ({ courtId, court, category, rewards, period, date }) => {
+const DisputeInfo: React.FC<IDisputeInfo> = ({ courtId, court, category, rewards, period, date, round }) => {
   return (
     <Container>
-      {category && <Field icon={BookmarkIcon} name="Category" value={category} />}
       {court && courtId && <Field icon={LawBalanceIcon} name="Court" value={court} link={`/courts/${courtId}`} />}
+      {category && <Field icon={BookmarkIcon} name="Category" value={category} />}
+      {round && <Field icon={RoundIcon} name="Round" value={round.toString()} />}
       {rewards && <Field icon={PileCoinsIcon} name="Juror Rewards" value={rewards} />}
       {typeof period !== "undefined" && date && (
         <Field icon={CalendarIcon} name={getPeriodPhrase(period)} value={new Date(date * 1000).toLocaleString()} />
diff --git a/web/src/components/DisputeCard/index.tsx b/web/src/components/DisputeCard/index.tsx
index 2579d0130..6a27b0277 100644
--- a/web/src/components/DisputeCard/index.tsx
+++ b/web/src/components/DisputeCard/index.tsx
@@ -11,6 +11,7 @@ import { useDisputeTemplate } from "queries/useDisputeTemplate";
 import DisputeInfo from "./DisputeInfo";
 import PeriodBanner from "./PeriodBanner";
 import { isUndefined } from "utils/index";
+import { useVotingHistory } from "hooks/queries/useVotingHistory";
 
 const StyledCard = styled(Card)`
   max-width: 380px;
@@ -61,6 +62,8 @@ const DisputeCard: React.FC<CasesPageQuery["disputes"][number]> = ({
   const { data: courtPolicy } = useCourtPolicy(court.id);
   const courtName = courtPolicy?.name;
   const category = disputeTemplate ? disputeTemplate.category : undefined;
+  const { data: votingHistory } = useVotingHistory(id);
+  const localRounds = votingHistory?.dispute?.disputeKitDispute?.localRounds;
   const navigate = useNavigate();
   return (
     <StyledCard hover onClick={() => navigate(`/cases/${id.toString()}`)}>
@@ -71,6 +74,7 @@ const DisputeCard: React.FC<CasesPageQuery["disputes"][number]> = ({
           courtId={court?.id}
           court={courtName}
           period={currentPeriodIndex}
+          round={localRounds?.length}
           {...{ category, rewards, date }}
         />
       </Container>

From f946f6e1cec6dbfbd51f6b17a4963de9ce00cf13 Mon Sep 17 00:00:00 2001
From: marino <102478601+kemuru@users.noreply.github.com>
Date: Sat, 9 Sep 2023 07:23:13 +0200
Subject: [PATCH 10/10] feat(web): adding missing round inside case overview

---
 web/src/pages/Cases/CaseDetails/Overview.tsx | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/web/src/pages/Cases/CaseDetails/Overview.tsx b/web/src/pages/Cases/CaseDetails/Overview.tsx
index 7de07014f..c7e4b0f7f 100644
--- a/web/src/pages/Cases/CaseDetails/Overview.tsx
+++ b/web/src/pages/Cases/CaseDetails/Overview.tsx
@@ -14,6 +14,7 @@ import PolicyIcon from "svgs/icons/policy.svg";
 import { StyledSkeleton } from "components/StyledSkeleton";
 import DisputeInfo from "components/DisputeCard/DisputeInfo";
 import Verdict from "components/Verdict/index";
+import { useVotingHistory } from "hooks/queries/useVotingHistory";
 
 const Container = styled.div`
   width: 100%;
@@ -86,10 +87,13 @@ const Overview: React.FC<IOverview> = ({ arbitrable, courtID, currentPeriodIndex
   const { data: disputeDetails } = useDisputeDetailsQuery(id);
   const { data: courtPolicyURI } = useCourtPolicyURI(courtID);
   const { data: courtPolicy } = useCourtPolicy(courtID);
+  const { data: votingHistory } = useVotingHistory(id);
   const courtName = courtPolicy?.name;
   const court = disputeDetails?.dispute?.court;
   const rewards = court ? `≥ ${formatEther(court.feeForJuror)} ETH` : undefined;
   const category = disputeTemplate ? disputeTemplate.category : undefined;
+  const localRounds = votingHistory?.dispute?.disputeKitDispute?.localRounds;
+
   return (
     <>
       <Container>
@@ -126,7 +130,7 @@ const Overview: React.FC<IOverview> = ({ arbitrable, courtID, currentPeriodIndex
           </>
         )}
 
-        <DisputeInfo courtId={court?.id} court={courtName} {...{ rewards, category }} />
+        <DisputeInfo courtId={court?.id} court={courtName} round={localRounds?.length} {...{ rewards, category }} />
       </Container>
       <ShadeArea>
         <p>Make sure you understand the Policies</p>