compound-finance · cwang25 · Dec 6, 2023 · Dec 5, 2023 · Dec 5, 2023 · Dec 5, 2023
diff --git a/src/QuarkStateManager.sol b/src/QuarkStateManager.sol
@@ -16,6 +16,7 @@ interface IExecutor {
 contract QuarkStateManager {
     event ClearNonce(address indexed wallet, uint96 nonce);
 
+    error InvalidNonce();
     error NoActiveNonce();
     error NoUnusedNonces();
     error NonceAlreadySet();
@@ -46,6 +47,10 @@ contract QuarkStateManager {
      * @return Whether the nonce has been exhausted
      */
     function isNonceSet(address wallet, uint96 nonce) public view returns (bool) {
+        if (nonce == type(uint96).max) {
+            revert InvalidNonce();
+        }
+
         (uint256 bucket, uint256 mask) = getBucket(nonce);
         return isNonceSetInternal(wallet, bucket, mask);
     }
@@ -60,18 +65,31 @@ contract QuarkStateManager {
      * @dev Any unset nonce is valid to use, but using this method
      * increases the likelihood that the nonce you use will be in a bucket that
      * has already been written to, which costs less gas
+     * @param wallet Address of the wallet to find the next nonce
+     * @param offset Bucket offset to start searching for the next nonce
      * @return The next unused nonce
      */
-    function nextNonce(address wallet) external view returns (uint96) {
-        for (uint96 i = 0; i <= type(uint96).max;) {
-            if (!isNonceSet(wallet, i) && (nonceScriptAddress[wallet][i] == address(0))) {
-                return i;
+    function nextNonce(address wallet, uint256 offset) external view returns (uint96) {
+        for (uint256 bucket = offset; bucket < type(uint256).max;) {
+            uint256 bucketNonces = nonces[wallet][bucket];
+            uint96 bucketValue = uint96(bucket << 8);
+            for (uint256 maskOffset = 0; maskOffset < 256;) {
+                uint256 mask = 1 << maskOffset;
+                if ((bucketNonces & mask) == 0) {
+                    uint96 nonce = uint96(bucketValue + maskOffset);
+                    if (nonceScriptAddress[wallet][nonce] == address(0)) {
+                        return nonce;
+                    }
+                }
+                unchecked {
+                    ++maskOffset;
+                }
             }
-
             unchecked {
-                ++i;
+                ++bucket;
             }
         }
+
         revert NoUnusedNonces();
     }
 
@@ -111,6 +129,10 @@ contract QuarkStateManager {
      * @param nonce Nonce to set for the calling wallet
      */
     function setNonce(uint96 nonce) external {
+        if (nonce == type(uint96).max) {
+            revert InvalidNonce();
+        }
+
         // TODO: should we check whether there exists a nonceScriptAddress?
         (uint256 bucket, uint256 setMask) = getBucket(nonce);
         setNonceInternal(bucket, setMask);
@@ -132,6 +154,10 @@ contract QuarkStateManager {
         external
         returns (bytes memory)
     {
+        if (nonce == type(uint96).max) {
+            revert InvalidNonce();
+        }
+
         // retrieve the (bucket, mask) pair that addresses the nonce in memory
         (uint256 bucket, uint256 setMask) = getBucket(nonce);
 

diff --git a/src/QuarkWalletFactory.sol b/src/QuarkWalletFactory.sol
@@ -120,17 +120,6 @@ contract QuarkWalletFactory {
         );
     }
 
-    /**
-     * @notice Returns the next unset nonce for the wallet corresponding to the given signer and salt
-     * @dev Any unset nonce is valid to use, but using this method increases
-     * the likelihood that the nonce you use will be on a bucket that has
-     * already been written to, which costs less gas
-     * @return The next unused nonce
-     */
-    function nextNonce(address signer, bytes32 salt) external view returns (uint96) {
-        return stateManager.nextNonce(walletAddressForSignerWithSalt(signer, salt));
-    }
-
     /**
      * @notice Returns the EIP-712 domain separator used for signing operations for the given salted wallet
      * @dev Only use for wallets deployed by this factory, or counterfactual wallets that will be deployed;

diff --git a/test/EIP712.t.sol b/test/EIP712.t.sol
@@ -137,7 +137,7 @@ contract EIP712Test is Test {
         assertEq(counter.number(), 0);
 
         // alice's nonce is not incremented
-        assertEq(stateManager.nextNonce(address(wallet)), op.nonce);
+        assertEq(stateManager.nextNonce(address(wallet), 0), op.nonce);
     }
 
     function testRevertsOnReusedNonce() public {
@@ -154,7 +154,7 @@ contract EIP712Test is Test {
         wallet.executeQuarkOperation(op, v, r, s);
 
         assertEq(counter.number(), 3);
-        assertEq(stateManager.nextNonce(address(wallet)), op.nonce + 1);
+        assertEq(stateManager.nextNonce(address(wallet), 0), op.nonce + 1);
 
         // submitter tries to reuse the same signature twice, for a non-replayable operation
         vm.expectRevert(QuarkStateManager.NonceAlreadySet.selector);
@@ -180,7 +180,7 @@ contract EIP712Test is Test {
         wallet.executeQuarkOperation(op, v, r, s);
 
         assertEq(counter.number(), 0);
-        assertEq(stateManager.nextNonce(address(wallet)), op.nonce);
+        assertEq(stateManager.nextNonce(address(wallet), 0), op.nonce);
     }
 
     function testRevertsInvalidS() public {
@@ -202,7 +202,7 @@ contract EIP712Test is Test {
         wallet.executeQuarkOperation(op, v, r, invalidS);
 
         assertEq(counter.number(), 0);
-        assertEq(stateManager.nextNonce(address(wallet)), op.nonce);
+        assertEq(stateManager.nextNonce(address(wallet), 0), op.nonce);
     }
 
     function testNonceIsNotSetForReplayableOperation() public {
@@ -278,7 +278,7 @@ contract EIP712Test is Test {
         wallet.executeQuarkOperation(op, v, r, s);
 
         assertEq(counter.number(), 0);
-        assertEq(stateManager.nextNonce(address(wallet)), op.nonce);
+        assertEq(stateManager.nextNonce(address(wallet), 0), op.nonce);
     }
 
     function testRequirements() public {

diff --git a/test/Executor.t.sol b/test/Executor.t.sol
@@ -60,12 +60,12 @@ contract ExecutorTest is Test {
 
         // execute counter.increment(5) as bob from alice's wallet (that is, from bob's wallet's executor)
         aliceWallet.executeScript(
-            stateManager.nextNonce(address(aliceWallet)),
+            stateManager.nextNonce(address(aliceWallet), 0),
             executeOnBehalfAddress,
             abi.encodeWithSignature(
                 "run(address,uint96,address,bytes)",
                 address(bobWallet),
-                stateManager.nextNonce(address(bobWallet)),
+                stateManager.nextNonce(address(bobWallet), 0),
                 address(ethcallAddress),
                 abi.encodeWithSignature(
                     "run(address,bytes,uint256)", address(counter), abi.encodeWithSignature("increment(uint256)", 5), 0
@@ -91,7 +91,7 @@ contract ExecutorTest is Test {
             abi.encodeWithSignature(
                 "run(address,uint96,address,bytes)",
                 address(bobWallet),
-                stateManager.nextNonce(address(bobWallet)),
+                stateManager.nextNonce(address(bobWallet), 0),
                 address(ethcallAddress),
                 abi.encodeWithSignature(
                     "run(address,bytes,uint256)", address(counter), abi.encodeWithSignature("increment(uint256)", 3), 0

diff --git a/test/Nonce.t.sol b/test/Nonce.t.sol
@@ -36,12 +36,12 @@ contract NonceTest is Test {
     }
 
     function testNextUnusedNonce() public {
-        uint96 nonce1 = stateManager.nextNonce(address(this));
+        uint96 nonce1 = stateManager.nextNonce(address(this), 0);
 
         stateManager.setNonce(nonce1);
-        assertEq(stateManager.nextNonce(address(this)), nonce1 + 1);
+        assertEq(stateManager.nextNonce(address(this), 0), nonce1 + 1);
 
         stateManager.setNonce(nonce1 + 1);
-        assertEq(stateManager.nextNonce(address(this)), nonce1 + 2);
+        assertEq(stateManager.nextNonce(address(this), 0), nonce1 + 2);
     }
 }
diff --git a/test/QuarkWallet.t.sol b/test/QuarkWallet.t.sol
@@ -89,7 +89,7 @@ contract QuarkWalletTest is Test {
         vm.pauseGasMetering();
         QuarkWallet aliceWalletExecutable = new QuarkWallet(aliceAccount, aliceAccount, codeJar, stateManager);
         bytes memory getMessageDetails = new YulHelper().getDeployed("GetMessageDetails.sol/GetMessageDetails.json");
-        uint96 nonce = stateManager.nextNonce(address(aliceWalletExecutable));
+        uint96 nonce = stateManager.nextNonce(address(aliceWalletExecutable), 0);
         address scriptAddress = codeJar.saveCode(getMessageDetails);
         bytes memory call = abi.encodeWithSignature("getMsgSenderAndValue()");
 
@@ -113,7 +113,7 @@ contract QuarkWalletTest is Test {
         vm.pauseGasMetering();
 
         QuarkWallet.QuarkOperation memory op = QuarkWallet.QuarkOperation({
-            nonce: stateManager.nextNonce(address(aliceWallet)),
+            nonce: stateManager.nextNonce(address(aliceWallet), 0),
             scriptAddress: address(0),
             scriptSource: bytes(""),
             scriptCalldata: bytes(""),
@@ -129,7 +129,7 @@ contract QuarkWalletTest is Test {
         assertEq(stateManager.isNonceSet(address(aliceWallet), op.nonce), true);
 
         // direct execution of the null script with no calldata is allowed
-        uint96 nonce = stateManager.nextNonce(address(aliceWallet));
+        uint96 nonce = stateManager.nextNonce(address(aliceWallet), 0);
         vm.prank(aliceWallet.executor());
         aliceWallet.executeScript(nonce, address(0), bytes(""));
         assertEq(stateManager.isNonceSet(address(aliceWallet), nonce), true);
@@ -140,7 +140,7 @@ contract QuarkWalletTest is Test {
         vm.pauseGasMetering();
 
         QuarkWallet.QuarkOperation memory op = QuarkWallet.QuarkOperation({
-            nonce: stateManager.nextNonce(address(aliceWallet)),
+            nonce: stateManager.nextNonce(address(aliceWallet), 0),
             scriptAddress: address(0xc0c0),
             scriptSource: bytes("f00f00"),
             scriptCalldata: bytes("feefee"),
@@ -282,7 +282,7 @@ contract QuarkWalletTest is Test {
         vm.startPrank(aliceWallet.signer());
 
         // pre-compute execution parameters so that the revert is expected from the right call
-        uint96 nonce = stateManager.nextNonce(address(aliceWallet));
+        uint96 nonce = stateManager.nextNonce(address(aliceWallet), 0);
         address target = codeJar.saveCode(incrementer);
         bytes memory call = abi.encodeWithSignature("incrementCounter(address)", counter);
 
@@ -304,7 +304,7 @@ contract QuarkWalletTest is Test {
         assertEq(counter.number(), 0);
 
         // pre-compute execution parameters so that the revert is expected from the right call
-        uint96 nonce = stateManager.nextNonce(address(aliceWallet));
+        uint96 nonce = stateManager.nextNonce(address(aliceWallet), 0);
         address target = codeJar.saveCode(incrementer);
         bytes memory call = abi.encodeWithSignature("incrementCounter(address)", counter);
 
@@ -515,7 +515,7 @@ contract QuarkWalletTest is Test {
             scriptAddress: address(0x1),
             scriptSource: "",
             scriptCalldata: abi.encode(testHash, vt, rt, st),
-            nonce: aliceWallet.stateManager().nextNonce(address(aliceWallet)),
+            nonce: aliceWallet.stateManager().nextNonce(address(aliceWallet), 0),
             expiry: block.timestamp + 1000
         });
 
@@ -549,7 +549,7 @@ contract QuarkWalletTest is Test {
             scriptAddress: address(0x2),
             scriptSource: "",
             scriptCalldata: abi.encode(numberToHash),
-            nonce: aliceWallet.stateManager().nextNonce(address(aliceWallet)),
+            nonce: aliceWallet.stateManager().nextNonce(address(aliceWallet), 0),
             expiry: block.timestamp + 1000
         });
 
@@ -582,7 +582,7 @@ contract QuarkWalletTest is Test {
             scriptAddress: address(0x3),
             scriptSource: "",
             scriptCalldata: testBytes,
-            nonce: aliceWallet.stateManager().nextNonce(address(aliceWallet)),
+            nonce: aliceWallet.stateManager().nextNonce(address(aliceWallet), 0),
             expiry: block.timestamp + 1000
         });
         (uint8 v, bytes32 r, bytes32 s) = new SignatureHelper().signOp(alicePrivateKey, aliceWallet, op);
@@ -614,7 +614,7 @@ contract QuarkWalletTest is Test {
             scriptAddress: address(0x4),
             scriptSource: "",
             scriptCalldata: testBytes,
-            nonce: aliceWallet.stateManager().nextNonce(address(aliceWallet)),
+            nonce: aliceWallet.stateManager().nextNonce(address(aliceWallet), 0),
             expiry: block.timestamp + 1000
         });
         (uint8 v, bytes32 r, bytes32 s) = new SignatureHelper().signOp(alicePrivateKey, aliceWallet, op);
@@ -654,7 +654,7 @@ contract QuarkWalletTest is Test {
             scriptAddress: address(0x5),
             scriptSource: "",
             scriptCalldata: abi.encode(uint256(0x20), uint256(0x20), uint256(0x20), base, exponent, modulus),
-            nonce: aliceWallet.stateManager().nextNonce(address(aliceWallet)),
+            nonce: aliceWallet.stateManager().nextNonce(address(aliceWallet), 0),
             expiry: block.timestamp + 1000
         });
         (uint8 v, bytes32 r, bytes32 s) = new SignatureHelper().signOp(alicePrivateKey, aliceWallet, op);
@@ -691,7 +691,7 @@ contract QuarkWalletTest is Test {
             scriptAddress: address(0x6),
             scriptSource: "",
             scriptCalldata: abi.encode(input),
-            nonce: aliceWallet.stateManager().nextNonce(address(aliceWallet)),
+            nonce: aliceWallet.stateManager().nextNonce(address(aliceWallet), 0),
             expiry: block.timestamp + 1000
         });
         (uint8 v, bytes32 r, bytes32 s) = new SignatureHelper().signOp(alicePrivateKey, aliceWallet, op);
@@ -729,7 +729,7 @@ contract QuarkWalletTest is Test {
             scriptAddress: address(0x7),
             scriptSource: "",
             scriptCalldata: abi.encode(input),
-            nonce: aliceWallet.stateManager().nextNonce(address(aliceWallet)),
+            nonce: aliceWallet.stateManager().nextNonce(address(aliceWallet), 0),
             expiry: block.timestamp + 1000
         });
         (uint8 v, bytes32 r, bytes32 s) = new SignatureHelper().signOp(alicePrivateKey, aliceWallet, op);
@@ -807,7 +807,7 @@ contract QuarkWalletTest is Test {
             scriptAddress: address(0x9),
             scriptSource: "",
             scriptCalldata: abi.encodePacked(rounds, h[0], h[1], m[0], m[1], m[2], m[3], t[0], t[1], f),
-            nonce: aliceWallet.stateManager().nextNonce(address(aliceWallet)),
+            nonce: aliceWallet.stateManager().nextNonce(address(aliceWallet), 0),
             expiry: block.timestamp + 1000
         });
         (uint8 v, bytes32 r, bytes32 s) = new SignatureHelper().signOp(alicePrivateKey, aliceWallet, op);

diff --git a/test/QuarkWalletFactory.t.sol b/test/QuarkWalletFactory.t.sol
@@ -96,7 +96,7 @@ contract QuarkWalletFactoryTest is Test {
         vm.pauseGasMetering();
         bytes memory incrementer = new YulHelper().getDeployed("Incrementer.sol/Incrementer.json");
 
-        uint96 nonce = factory.nextNonce(alice, factory.DEFAULT_SALT());
+        uint96 nonce = factory.stateManager().nextNonce(factory.walletAddressForSigner(alice), 0);
         QuarkWallet.QuarkOperation memory op = QuarkWallet.QuarkOperation({
             scriptAddress: address(0),
             scriptSource: incrementer,
@@ -132,7 +132,7 @@ contract QuarkWalletFactoryTest is Test {
         vm.pauseGasMetering();
         bytes memory incrementer = new YulHelper().getDeployed("Incrementer.sol/Incrementer.json");
 
-        uint96 nonce = factory.nextNonce(alice, factory.DEFAULT_SALT());
+        uint96 nonce = factory.stateManager().nextNonce(factory.walletAddressForSigner(alice), 0);
         QuarkWallet.QuarkOperation memory op = QuarkWallet.QuarkOperation({
             scriptAddress: address(0),
             scriptSource: incrementer,
@@ -172,7 +172,7 @@ contract QuarkWalletFactoryTest is Test {
         vm.pauseGasMetering();
         bytes memory incrementer = new YulHelper().getDeployed("Incrementer.sol/Incrementer.json");
 
-        uint96 nonce = factory.nextNonce(alice, factory.DEFAULT_SALT());
+        uint96 nonce = factory.stateManager().nextNonce(factory.walletAddressForSigner(alice), 0);
         QuarkWallet.QuarkOperation memory op = QuarkWallet.QuarkOperation({
             scriptAddress: address(0),
             scriptSource: incrementer,
@@ -212,7 +212,7 @@ contract QuarkWalletFactoryTest is Test {
         vm.pauseGasMetering();
         bytes memory getMessageDetails = new YulHelper().getDeployed("GetMessageDetails.sol/GetMessageDetails.json");
         address aliceWallet = factory.walletAddressForSigner(alice);
-        uint96 nonce = factory.stateManager().nextNonce(aliceWallet);
+        uint96 nonce = factory.stateManager().nextNonce(aliceWallet, 0);
         QuarkWallet.QuarkOperation memory op = QuarkWallet.QuarkOperation({
             scriptAddress: address(0),
             scriptSource: getMessageDetails,
@@ -244,8 +244,8 @@ contract QuarkWalletFactoryTest is Test {
         vm.pauseGasMetering();
         bytes memory getMessageDetails = new YulHelper().getDeployed("GetMessageDetails.sol/GetMessageDetails.json");
         bytes32 salt = bytes32("salty salt salt");
-        uint96 nonce = factory.nextNonce(alice, salt);
         address aliceWallet = factory.walletAddressForSignerWithSalt(alice, salt);
+        uint96 nonce = factory.stateManager().nextNonce(aliceWallet, 0);
         QuarkWallet.QuarkOperation memory op = QuarkWallet.QuarkOperation({
             scriptAddress: address(0),
             scriptSource: getMessageDetails,
@@ -298,13 +298,13 @@ contract QuarkWalletFactoryTest is Test {
             scriptCalldata: abi.encodeWithSignature(
                 "run(address,uint96,address,bytes)",
                 address(aliceWalletSecondary),
-                factory.stateManager().nextNonce(address(aliceWalletSecondary)),
+                factory.stateManager().nextNonce(address(aliceWalletSecondary), 0),
                 ethcallAddress,
                 abi.encodeWithSignature(
                     "run(address,bytes,uint256)", address(counter), abi.encodeWithSignature("increment(uint256)", 7), 0
                 )
                 ),
-            nonce: factory.stateManager().nextNonce(address(aliceWalletPrimary)),
+            nonce: factory.stateManager().nextNonce(address(aliceWalletPrimary), 0),
             expiry: block.timestamp + 1000
         });
         (uint8 v, bytes32 r, bytes32 s) = new SignatureHelper().signOp(alicePrivateKey, aliceWalletPrimary, op);

diff --git a/test/core_scripts/Multicall.t.sol b/test/core_scripts/Multicall.t.sol
@@ -331,7 +331,7 @@ contract MulticallTest is Test {
         wallets[0] = address(walletA);
         walletCalls[0] = abi.encodeWithSignature(
             "executeScript(uint96,address,bytes)",
-            factory.stateManager().nextNonce(address(walletA)),
+            factory.stateManager().nextNonce(address(walletA), 0),
             ethcallAddress,
             abi.encodeWithSelector(
                 Ethcall.run.selector,
@@ -342,7 +342,7 @@ contract MulticallTest is Test {
         );
 
         // 2. approve Comet cUSDCv3 to receive 0.5 WETH from wallet B
-        uint96 walletBNextNonce = factory.stateManager().nextNonce(address(walletB));
+        uint96 walletBNextNonce = factory.stateManager().nextNonce(address(walletB), 0);
         wallets[1] = address(walletB);
         walletCalls[1] = abi.encodeWithSignature(
             "executeScript(uint96,address,bytes)",
@@ -428,7 +428,7 @@ contract MulticallTest is Test {
         deal(WETH, address(wallet), 100 ether);
 
         address subWallet1 = factory.walletAddressForSignerWithSalt(alice, bytes32("1"));
-        uint96 nonce = factory.stateManager().nextNonce(subWallet1);
+        uint96 nonce = factory.stateManager().nextNonce(subWallet1, 0);
         // Steps: Wallet#1: Supply WETH to Comet -> Borrow USDC from Comet(USDC) to subwallet -> Create subwallet
         // -> Swap USDC to WETH on Uniswap -> Supply WETH to Comet(WETH)
         address[] memory callContracts = new address[](5);