feat: optimize total checkpoints

Co-authored-by: 0xRaccoon <[email protected]>

solidity/contracts/governance/utils/WonderVotes.sol

@@ -40,7 +40,7 @@ abstract contract WonderVotes is Context, EIP712, Nonces, IERC6372, IWonderVotes
 
   mapping(address delegatee => mapping(uint8 proposalType => Checkpoints.Trace208)) private _delegateCheckpoints;
 
-  mapping(uint8 proposalType => Checkpoints.Trace208) private _totalCheckpoints;
+  Checkpoints.Trace208 private _totalCheckpoints;
 
   mapping(address account => bool) private _nonDelegableAddresses;
 
@@ -109,19 +109,19 @@ abstract contract WonderVotes is Context, EIP712, Nonces, IERC6372, IWonderVotes
    *
    * - `timepoint` must be in the past. If operating using block numbers, the block must be already mined.
    */
-  function getPastTotalSupply(uint8 proposalType, uint256 timepoint) public view virtual returns (uint256) {
+  function getPastTotalSupply(uint256 timepoint) public view virtual returns (uint256) {
     uint48 currentTimepoint = clock();
     if (timepoint >= currentTimepoint) {
       revert ERC5805FutureLookup(timepoint, currentTimepoint);
     }
-    return _totalCheckpoints[proposalType].upperLookupRecent(SafeCast.toUint48(timepoint));
+    return _totalCheckpoints.upperLookupRecent(SafeCast.toUint48(timepoint));
   }
 
   /**
    * @dev Returns the current total supply of votes for a given `proposalType`.
    */
-  function _getTotalSupply(uint8 proposalType) internal view virtual returns (uint256) {
-    return _totalCheckpoints[proposalType].latest();
+  function _getTotalSupply() internal view virtual returns (uint256) {
+    return _totalCheckpoints.latest();
   }
 
   /**
@@ -285,29 +285,25 @@ abstract contract WonderVotes is Context, EIP712, Nonces, IERC6372, IWonderVotes
     _moveDelegateVotes(proposalType, _oldDelegates, delegatees, _getVotingUnits(account));
   }
 
-  /**
-   * @dev Loops the proposalTypes implemented and calls the `_transferVotingUnits` helper method.
-   */
-  function _transferVotingUnits(address from, address to, uint256 amount) internal virtual {
-    uint8[] memory _proposalTypes = _getProposalTypes();
-
-    for (uint256 i = 0; i < _proposalTypes.length; i++) {
-      _transferVotingUnits(_proposalTypes[i], from, to, amount);
-    }
-  }
-
   /**
    * @dev Transfers, mints, or burns voting units. To register a mint, `from` should be zero. To register a burn, `to`
    * should be zero. Total supply of voting units will be adjusted with mints and burns.
+   * Loops the proposalTypes implemented and calls the `_moveDelegateVotes` helper method.
    */
-  function _transferVotingUnits(uint8 proposalType, address from, address to, uint256 amount) private {
+  function _transferVotingUnits(address from, address to, uint256 amount) internal virtual {
     if (from == address(0)) {
-      _push(_totalCheckpoints[proposalType], _add, SafeCast.toUint208(amount));
+      _push(_totalCheckpoints, _add, SafeCast.toUint208(amount));
     }
     if (to == address(0)) {
-      _push(_totalCheckpoints[proposalType], _subtract, SafeCast.toUint208(amount));
+      _push(_totalCheckpoints, _subtract, SafeCast.toUint208(amount));
+    }
+
+    uint8[] memory _proposalTypes = _getProposalTypes();
+
+    for (uint256 _i = 0; _i < _proposalTypes.length; _i++) {
+      uint8 _proposalType = _proposalTypes[_i];
+      _moveDelegateVotes(_proposalType, delegates(from, _proposalType), delegates(to, _proposalType), amount);
     }
-    _moveDelegateVotes(proposalType, delegates(from, proposalType), delegates(to, proposalType), amount);
   }
 
   /**

solidity/interfaces/governance/utils/IWonderVotes.sol

@@ -71,14 +71,14 @@ interface IWonderVotes {
   function getPastVotes(address account, uint8 proposalType, uint256 timepoint) external view returns (uint256);
 
   /**
-   * @dev Returns for a given `proposalType` the total supply of votes available at a specific moment in the past. If the `clock()` is
+   * @dev Returns the total supply of votes available at a specific moment in the past. If the `clock()` is
    * configured to use block numbers, this will return the value at the end of the corresponding block.
    *
-   * NOTE: This value is the sum of all available votes for a `proposalType`, which is not necessarily the sum of all delegated votes.
+   * NOTE: This value is the sum of all available votes, which is not necessarily the sum of all delegated votes.
    * Votes that have not been delegated are still part of total supply, even though they would not participate in a
    * vote.
    */
-  function getPastTotalSupply(uint8 proposalType, uint256 timepoint) external view returns (uint256);
+  function getPastTotalSupply(uint256 timepoint) external view returns (uint256);
 
   /**
    * @dev Returns the delegates that `account` has chosen for a given `proposalType`.

solidity/test/unit/WonderVotes.t.sol

@@ -1085,3 +1085,61 @@ contract Unit_SuspendDelegation is BaseTest {
     rabbitToken.delegate(_delegate);
   }
 }
+
+contract Unit_PastTotalSupply is BaseTest {
+  function test_TotalSupplyIsZeroIfNothingMinted() public {
+    assertEq(rabbitToken.totalSupply(), 0);
+  }
+
+  function test_Revert_IfGettingPastTotalSupplyforCurrentBlock() public {
+    uint256 _blockNumber = block.number;
+
+    vm.expectRevert(abi.encodeWithSelector(WonderVotes.ERC5805FutureLookup.selector, _blockNumber, _blockNumber));
+
+    rabbitToken.getPastTotalSupply(_blockNumber);
+  }
+
+  function test_Revert_IfGettingPastTotalSupplyforFutureBlock() public {
+    uint256 _blockNumber = block.number;
+
+    vm.expectRevert(abi.encodeWithSelector(WonderVotes.ERC5805FutureLookup.selector, _blockNumber + 1, _blockNumber));
+
+    rabbitToken.getPastTotalSupply(_blockNumber + 1);
+  }
+
+  function test_PastTotalSupplyIsZeroIfNothingMinted() public {
+    uint256 _blockNumber = block.number;
+    vm.roll(_blockNumber + 1); // To avoid ERC5805FutureLookup error
+    assertEq(rabbitToken.getPastTotalSupply(_blockNumber), 0);
+  }
+
+  function test_PastTotalSupplyIncreaseIfTokensMinted(address _account, uint256 _amount) public {
+    vm.assume(_amount > 0 && _amount <= type(uint208).max);
+    vm.assume(_account != address(0));
+    uint256 _blockNumber = block.number;
+
+    // Mint and advance block
+    WonderVotesForTest(address(rabbitToken)).mint(_account, _amount);
+    vm.roll(_blockNumber + 1);
+
+    assertEq(rabbitToken.getPastTotalSupply(_blockNumber), _amount);
+  }
+
+  function test_PastTotalSupplyDecreaseIfTokensBurned(address _account, uint256 _amount) public {
+    vm.assume(_amount > 0 && _amount <= type(uint208).max);
+    vm.assume(_account != address(0));
+    uint256 _blockNumber = block.number;
+
+    // Mint and advance block
+    WonderVotesForTest(address(rabbitToken)).mint(_account, _amount);
+    vm.roll(_blockNumber + 1);
+
+    // Burn and advance block
+    vm.prank(_account);
+    WonderVotesForTest(address(rabbitToken)).burn(_amount);
+    vm.roll(_blockNumber + 2);
+
+    assertEq(rabbitToken.getPastTotalSupply(_blockNumber), _amount);
+    assertEq(rabbitToken.getPastTotalSupply(_blockNumber + 1), 0);
+  }
+}