From c7ab97bd2ddc9f8ed3490e5492751cbbfd634f06 Mon Sep 17 00:00:00 2001
From: Jesse Friedland <jesse.friedland@gmail.com>
Date: Fri, 30 Jun 2023 14:35:45 -0400
Subject: [PATCH] perf(gas): optimize stage per wallet limits gas

- pack into aux on the erc721a ownership storage
- caveat: limits # of stages to 4 (or 8 if max limit per wallet of 255
  is acceptable)
---
 contracts/ERC721M.sol  | 26 ++++++++++++++++++--------
 contracts/IERC721M.sol |  2 +-
 2 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/contracts/ERC721M.sol b/contracts/ERC721M.sol
index 2d5a073..837810b 100644
--- a/contracts/ERC721M.sol
+++ b/contracts/ERC721M.sol
@@ -53,10 +53,6 @@ contract ERC721M is IERC721M, ERC721AQueryable, Ownable, ReentrancyGuard {
     // Mint stage infomation. See MintStageInfo for details.
     MintStageInfo[] private _mintStages;
 
-    // Minted count per stage per wallet.
-    mapping(uint256 => mapping(address => uint32))
-        private _stageMintedCountsPerWallet;
-
     // Minted count per stage.
     mapping(uint256 => uint256) private _stageMintedCounts;
 
@@ -180,6 +176,8 @@ contract ERC721M is IERC721M, ERC721AQueryable, Ownable, ReentrancyGuard {
      * ]
      */
     function setStages(MintStageInfo[] calldata newStages) external onlyOwner {
+        if (newStages.length > 4) revert InvalidStageArgsLength();
+
         uint256 originalSize = _mintStages.length;
         for (uint256 i = 0; i < originalSize; i++) {
             _mintStages.pop();
@@ -309,14 +307,14 @@ contract ERC721M is IERC721M, ERC721AQueryable, Ownable, ReentrancyGuard {
         override
         returns (
             MintStageInfo memory,
-            uint32,
+            uint16,
             uint256
         )
     {
         if (index >= _mintStages.length) {
             revert("InvalidStage");
         }
-        uint32 walletMinted = _stageMintedCountsPerWallet[index][msg.sender];
+        uint16 walletMinted = _numberMintedForStage(msg.sender, index);
         uint256 stageMinted = _stageMintedCounts[index];
         return (_mintStages[index], walletMinted, stageMinted);
     }
@@ -447,7 +445,7 @@ contract ERC721M is IERC721M, ERC721AQueryable, Ownable, ReentrancyGuard {
         // Check wallet limit for stage if applicable, limit == 0 means no limit enforced
         if (stage.walletLimit > 0) {
             if (
-                _stageMintedCountsPerWallet[activeStage][to] + qty >
+                _numberMintedForStage(to, activeStage) + qty >
                 stage.walletLimit
             ) revert WalletStageLimitExceeded();
         }
@@ -462,11 +460,23 @@ contract ERC721M is IERC721M, ERC721AQueryable, Ownable, ReentrancyGuard {
             ) revert InvalidProof();
         }
 
-        _stageMintedCountsPerWallet[activeStage][to] += qty;
+        // Optimizes Stage per wallet mint counts gas
+        // - previously was: _stageMintedCountsPerWallet[to][activeStage] += qty;
+        uint64 packed = _getAux(to);
+        packed += uint64(qty * (1 << (16 * activeStage)));
+        _setAux(to, packed);
+
         _stageMintedCounts[activeStage] += qty;
         _safeMint(to, qty);
     }
 
+    /**
+     * Returns the number of tokens minted by `owner` for a specific stage.
+     */
+    function _numberMintedForStage(address owner, uint256 stageIndex) internal view returns (uint16) {
+        return uint16((_getAux(owner) >> (16 * stageIndex)) & ((1 << 16) - 1));
+    }
+
     /**
      * @dev Mints token(s) by owner.
      *
diff --git a/contracts/IERC721M.sol b/contracts/IERC721M.sol
index 21d2332..fe99466 100644
--- a/contracts/IERC721M.sol
+++ b/contracts/IERC721M.sol
@@ -79,7 +79,7 @@ interface IERC721M is IERC721AQueryable {
         view
         returns (
             MintStageInfo memory,
-            uint32,
+            uint16,
             uint256
         );