refactor: moved more logic to _postDrawHook() to avoid extra push/pop

Author: jaybuidl

contracts/src/arbitration/dispute-kits/DisputeKitClassicBase.sol

@@ -252,7 +252,7 @@ abstract contract DisputeKitClassicBase is IDisputeKit, Initializable, UUPSProxi
     function draw(
         uint256 _coreDisputeID,
         uint256 _nonce,
-        uint256 /*_roundNbVotes*/
+        uint256 _roundNbVotes
     )
         public
         virtual
@@ -274,7 +274,7 @@ abstract contract DisputeKitClassicBase is IDisputeKit, Initializable, UUPSProxi
             return (drawnAddress, fromSubcourtID);
         }
 
-        if (_postDrawCheck(round, _coreDisputeID, drawnAddress)) {
+        if (_postDrawCheck(round, _coreDisputeID, drawnAddress, _roundNbVotes)) {
             Vote storage vote = round.votes.push();
             vote.account = drawnAddress;
             round.alreadyDrawn[drawnAddress] = true;
@@ -786,19 +786,20 @@ abstract contract DisputeKitClassicBase is IDisputeKit, Initializable, UUPSProxi
     ///
     /// @param _coreDisputeID ID of the dispute in the core contract.
     /// @param _juror Chosen address.
-    /// @return result Whether the address passes the check or not.
+    /// @return Whether the address passes the check or not.
     function _postDrawCheck(
         Round storage /*_round*/,
         uint256 _coreDisputeID,
-        address _juror
-    ) internal view virtual returns (bool result) {
+        address _juror,
+        uint256 /*_roundNbVotes*/
+    ) internal view virtual returns (bool) {
         if (singleDrawPerJuror) {
             uint256 localDisputeID = coreDisputeIDToLocal[_coreDisputeID];
             Dispute storage dispute = disputes[localDisputeID];
             Round storage round = dispute.rounds[dispute.rounds.length - 1];
-            result = !round.alreadyDrawn[_juror];
+            return !round.alreadyDrawn[_juror];
         } else {
-            result = true;
+            return true;
         }
     }
 

contracts/src/arbitration/dispute-kits/DisputeKitGated.sol

@@ -129,9 +129,10 @@ contract DisputeKitGated is DisputeKitClassicBase {
     function _postDrawCheck(
         Round storage _round,
         uint256 _coreDisputeID,
-        address _juror
+        address _juror,
+        uint256 _roundNbVotes
     ) internal view override returns (bool) {
-        if (!super._postDrawCheck(_round, _coreDisputeID, _juror)) return false;
+        if (!super._postDrawCheck(_round, _coreDisputeID, _juror, _roundNbVotes)) return false;
 
         // Get the local dispute and extract token info from extraData
         uint256 localDisputeID = coreDisputeIDToLocal[_coreDisputeID];

contracts/src/arbitration/dispute-kits/DisputeKitGatedArgentinaConsumerProtection.sol

@@ -96,25 +96,11 @@ contract DisputeKitGatedArgentinaConsumerProtection is DisputeKitClassicBase {
 
         if (drawnAddress == address(0)) return (drawnAddress, fromSubcourtID);
 
-        uint256 localDisputeID = coreDisputeIDToLocal[_coreDisputeID];
-        Dispute storage dispute = disputes[localDisputeID];
-        uint256 localRoundID = dispute.rounds.length - 1;
         if (IBalanceHolder(accreditedConsumerProtectionLawyerToken).balanceOf(drawnAddress) > 0) {
             // The drawnAddress is a consumer protection lawyer.
+            uint256 localDisputeID = coreDisputeIDToLocal[_coreDisputeID];
+            uint256 localRoundID = disputes[localDisputeID].rounds.length - 1;
             drawnConsumerProtectionLawyer[localDisputeID][localRoundID] = true;
-        } else {
-            // The drawnAddress is not a consumer protection lawyer.
-            if (
-                dispute.rounds[localRoundID].votes.length == _roundNbVotes &&
-                !drawnConsumerProtectionLawyer[localDisputeID][localRoundID]
-            ) {
-                // This is the last draw iteration and we still have not drawn a consumer protection lawyer.
-                // Drop the last round.votes pushed by super.draw(), so that another iteration can try again.
-                drawnAddress = address(0);
-                dispute.rounds[localRoundID].votes.pop();
-                // Note that round.alreadyDrawn[drawnAddress] is not cleared because we don't know if it has been drawn more than once.
-                // It's fine because this DisputeKit does not enable singleDrawPerJuror.
-            }
         }
         return (drawnAddress, fromSubcourtID);
     }
@@ -127,12 +113,28 @@ contract DisputeKitGatedArgentinaConsumerProtection is DisputeKitClassicBase {
     function _postDrawCheck(
         Round storage _round,
         uint256 _coreDisputeID,
-        address _juror
+        address _juror,
+        uint256 _roundNbVotes
     ) internal view override returns (bool) {
-        if (!super._postDrawCheck(_round, _coreDisputeID, _juror)) return false;
-        return
-            (IBalanceHolder(accreditedLawyerToken).balanceOf(_juror) > 0) ||
-            (IBalanceHolder(accreditedConsumerProtectionLawyerToken).balanceOf(_juror) > 0);
+        if (IBalanceHolder(accreditedConsumerProtectionLawyerToken).balanceOf(_juror) == 0) {
+            // The juror is not a consumer protection lawyer.
+            uint256 localDisputeID = coreDisputeIDToLocal[_coreDisputeID];
+            Dispute storage dispute = disputes[localDisputeID];
+            uint256 localRoundID = dispute.rounds.length - 1;
+            if (
+                dispute.rounds[localRoundID].votes.length == _roundNbVotes - 1 &&
+                !drawnConsumerProtectionLawyer[localDisputeID][localRoundID]
+            ) {
+                // This is the last draw iteration and we still have not drawn a consumer protection lawyer.
+                // Reject this draw so that another iteration can try again later.
+                return false;
+            }
+            if (IBalanceHolder(accreditedLawyerToken).balanceOf(_juror) == 0) {
+                // The juror does not hold either of the tokens.
+                return false;
+            }
+        }
+        return super._postDrawCheck(_round, _coreDisputeID, _juror, _roundNbVotes);
     }
 
     // ************************************* //

contracts/src/arbitration/dispute-kits/DisputeKitGatedShutter.sol

@@ -256,9 +256,10 @@ contract DisputeKitGatedShutter is DisputeKitClassicBase {
     function _postDrawCheck(
         Round storage _round,
         uint256 _coreDisputeID,
-        address _juror
+        address _juror,
+        uint256 _roundNbVotes
     ) internal view override returns (bool) {
-        if (!super._postDrawCheck(_round, _coreDisputeID, _juror)) return false;
+        if (!super._postDrawCheck(_round, _coreDisputeID, _juror, _roundNbVotes)) return false;
 
         // Get the local dispute and extract token info from extraData
         uint256 localDisputeID = coreDisputeIDToLocal[_coreDisputeID];

contracts/src/arbitration/dispute-kits/DisputeKitSybilResistant.sol

@@ -70,8 +70,9 @@ contract DisputeKitSybilResistant is DisputeKitClassicBase {
     function _postDrawCheck(
         Round storage _round,
         uint256 _coreDisputeID,
-        address _juror
+        address _juror,
+        uint256 _roundNbVotes
     ) internal view override returns (bool) {
-        return super._postDrawCheck(_round, _coreDisputeID, _juror) && poh.isRegistered(_juror);
+        return super._postDrawCheck(_round, _coreDisputeID, _juror, _roundNbVotes) && poh.isRegistered(_juror);
     }
 }