From 22c5742027d90087615d57369380b93abfd8a383 Mon Sep 17 00:00:00 2001
From: Hans Wang <2709448+cwang25@users.noreply.github.com>
Date: Wed, 6 Dec 2023 15:57:38 -0800
Subject: [PATCH] Address L2/Q4 (#112)

* address l2/q4

* removed seems outdated nextNonce() left in factory, as most use cases are getting nextNonce() directly from stateManager in our scenarios instead of passing through factory

* optimization changes

* address comments

* Update src/QuarkStateManager.sol

Co-authored-by: Kevin Cheng <kevincheng96@hotmail.com>

* limit loop to not have unsafe downcasting uint96 in nonce step

---------

Co-authored-by: Kevin Cheng <kevincheng96@hotmail.com>
---
 src/QuarkStateManager.sol     | 23 ++++++++++++++++++-----
 src/QuarkWalletFactory.sol    | 11 -----------
 test/QuarkWalletFactory.t.sol |  8 ++++----
 3 files changed, 22 insertions(+), 20 deletions(-)

diff --git a/src/QuarkStateManager.sol b/src/QuarkStateManager.sol
index 1aaafd1b..92a1a46e 100644
--- a/src/QuarkStateManager.sol
+++ b/src/QuarkStateManager.sol
@@ -60,18 +60,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 for
      * @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;
+        // Exclude bucket==uint88.max out to prevent from overlfowing, as favor in using unchecked for gas optimization
+        for (uint256 bucket = 0; bucket < type(uint88).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();
     }
 
diff --git a/src/QuarkWalletFactory.sol b/src/QuarkWalletFactory.sol
index 9bd1335d..2753aa57 100644
--- 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/QuarkWalletFactory.t.sol b/test/QuarkWalletFactory.t.sol
index 8b337f19..3cf8f499 100644
--- 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));
         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));
         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));
         QuarkWallet.QuarkOperation memory op = QuarkWallet.QuarkOperation({
             scriptAddress: address(0),
             scriptSource: incrementer,
@@ -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);
         QuarkWallet.QuarkOperation memory op = QuarkWallet.QuarkOperation({
             scriptAddress: address(0),
             scriptSource: getMessageDetails,