Add totalSupply checkpoints to ER20Votes (#2695)

Co-authored-by: Francisco Giordano <frangio.1@gmail.com>

contracts/mocks/ERC20VotesMock.sol

@@ -6,13 +6,17 @@ pragma solidity ^0.8.0;
 import "../token/ERC20/extensions/ERC20Votes.sol";
 
 contract ERC20VotesMock is ERC20Votes {
-    constructor (
-        string memory name,
-        string memory symbol,
-        address initialAccount,
-        uint256 initialBalance
-    ) payable ERC20(name, symbol) ERC20Permit(name) {
-        _mint(initialAccount, initialBalance);
+    constructor (string memory name, string memory symbol)
+    ERC20(name, symbol)
+    ERC20Permit(name)
+    {}
+
+    function mint(address account, uint256 amount) public {
+        _mint(account, amount);
+    }
+
+    function burn(address account, uint256 amount) public {
+        _burn(account, amount);
     }
 
     function getChainId() external view returns (uint256) {

contracts/token/ERC20/extensions/ERC20Votes.sol

@@ -27,6 +27,7 @@ abstract contract ERC20Votes is IERC20Votes, ERC20Permit {
 
     mapping (address => address) private _delegates;
     mapping (address => Checkpoint[]) private _checkpoints;
+    Checkpoint[] private _totalSupplyCheckpoints;
 
     /**
      * @dev Get the `pos`-th checkpoint for `account`.
@@ -62,9 +63,22 @@ abstract contract ERC20Votes is IERC20Votes, ERC20Permit {
      */
     function getPriorVotes(address account, uint256 blockNumber) external view override returns (uint256) {
         require(blockNumber < block.number, "ERC20Votes::getPriorVotes: not yet determined");
+        return _checkpointsLookup(_checkpoints[account], blockNumber);
+    }
 
-        Checkpoint[] storage ckpts = _checkpoints[account];
+    /**
+     * @dev Determine the totalSupply at the begining of `blockNumber`. Note, this value is the sum of all balances.
+     * It is but NOT the sum of all the delegated votes!
+     */
+    function getPriorTotalSupply(uint256 blockNumber) external view override returns(uint256) {
+        require(blockNumber < block.number, "ERC20Votes::getPriorTotalSupply: not yet determined");
+        return _checkpointsLookup(_totalSupplyCheckpoints, blockNumber);
+    }
 
+    /**
+     * @dev Lookup a value in a list of (sorted) checkpoints.
+     */
+    function _checkpointsLookup(Checkpoint[] storage ckpts, uint256 blockNumber) private view returns (uint256) {
         // We run a binary search to look for the earliest checkpoint taken after `blockNumber`.
         //
         // During the loop, the index of the wanted checkpoint remains in the range [low, high).
@@ -117,6 +131,32 @@ abstract contract ERC20Votes is IERC20Votes, ERC20Permit {
         return _delegate(signer, delegatee);
     }
 
+    /**
+     * @dev snapshot the totalSupply after it has been increassed.
+     */
+    function _mint(address account, uint256 amount) internal virtual override {
+        super._mint(account, amount);
+        require(totalSupply() <= type(uint224).max, "ERC20Votes: total supply exceeds 2**224");
+
+        _writeCheckpoint(_totalSupplyCheckpoints, add, amount);
+    }
+
+    /**
+     * @dev snapshot the totalSupply after it has been decreased.
+     */
+    function _burn(address account, uint256 amount) internal virtual override {
+        super._burn(account, amount);
+
+        _writeCheckpoint(_totalSupplyCheckpoints, subtract, amount);
+    }
+
+    /**
+     * @dev move voting power when tokens are transferred.
+     */
+    function _beforeTokenTransfer(address from, address to, uint256 amount) internal virtual override {
+        _moveVotingPower(delegates(from), delegates(to), amount);
+    }
+
     /**
      * @dev Change delegation for `delegator` to `delegatee`.
      */
@@ -133,40 +173,43 @@ abstract contract ERC20Votes is IERC20Votes, ERC20Permit {
     function _moveVotingPower(address src, address dst, uint256 amount) private {
         if (src != dst && amount > 0) {
             if (src != address(0)) {
-                uint256 srcCkptLen = _checkpoints[src].length;
-                uint256 srcCkptOld = srcCkptLen == 0 ? 0 : _checkpoints[src][srcCkptLen - 1].votes;
-                uint256 srcCkptNew = srcCkptOld - amount;
-                _writeCheckpoint(src, srcCkptLen, srcCkptOld, srcCkptNew);
+                (uint256 oldWeight, uint256 newWeight) = _writeCheckpoint(_checkpoints[src], subtract, amount);
+                emit DelegateVotesChanged(src, oldWeight, newWeight);
             }
 
             if (dst != address(0)) {
-                uint256 dstCkptLen = _checkpoints[dst].length;
-                uint256 dstCkptOld = dstCkptLen == 0 ? 0 : _checkpoints[dst][dstCkptLen - 1].votes;
-                uint256 dstCkptNew = dstCkptOld + amount;
-                _writeCheckpoint(dst, dstCkptLen, dstCkptOld, dstCkptNew);
+                (uint256 oldWeight, uint256 newWeight) = _writeCheckpoint(_checkpoints[dst], add, amount);
+                emit DelegateVotesChanged(dst, oldWeight, newWeight);
             }
         }
     }
 
-    function _writeCheckpoint(address delegatee, uint256 pos, uint256 oldWeight, uint256 newWeight) private {
-      if (pos > 0 && _checkpoints[delegatee][pos - 1].fromBlock == block.number) {
-          _checkpoints[delegatee][pos - 1].votes = SafeCast.toUint224(newWeight);
-      } else {
-          _checkpoints[delegatee].push(Checkpoint({
-              fromBlock: SafeCast.toUint32(block.number),
-              votes: SafeCast.toUint224(newWeight)
-          }));
-      }
-
-      emit DelegateVotesChanged(delegatee, oldWeight, newWeight);
+    function _writeCheckpoint(
+        Checkpoint[] storage ckpts,
+        function (uint256, uint256) view returns (uint256) op,
+        uint256 delta
+    )
+        private returns (uint256 oldWeight, uint256 newWeight)
+    {
+        uint256 pos = ckpts.length;
+        oldWeight = pos == 0 ? 0 : ckpts[pos - 1].votes;
+        newWeight = op(oldWeight, delta);
+
+        if (pos > 0 && ckpts[pos - 1].fromBlock == block.number) {
+            ckpts[pos - 1].votes = SafeCast.toUint224(newWeight);
+        } else {
+            ckpts.push(Checkpoint({
+                fromBlock: SafeCast.toUint32(block.number),
+                votes: SafeCast.toUint224(newWeight)
+            }));
+        }
     }
 
-    function _mint(address account, uint256 amount) internal virtual override {
-        super._mint(account, amount);
-        require(totalSupply() <= type(uint224).max, "ERC20Votes: total supply exceeds 2**224");
+    function add(uint256 a, uint256 b) private pure returns (uint256) {
+        return a + b;
     }
 
-    function _beforeTokenTransfer(address from, address to, uint256 amount) internal virtual override {
-        _moveVotingPower(delegates(from), delegates(to), amount);
+    function subtract(uint256 a, uint256 b) private pure returns (uint256) {
+        return a - b;
     }
 }

contracts/token/ERC20/extensions/IERC20Votes.sol

@@ -18,6 +18,7 @@ interface IERC20Votes is IERC20 {
     function numCheckpoints(address account) external view returns (uint32);
     function getCurrentVotes(address account) external view returns (uint256);
     function getPriorVotes(address account, uint256 blockNumber) external view returns (uint256);
+    function getPriorTotalSupply(uint256 blockNumber) external view returns(uint256);
     function delegate(address delegatee) external;
     function delegateBySig(address delegatee, uint nonce, uint expiry, uint8 v, bytes32 r, bytes32 s) external;
 }

test/token/ERC20/extensions/ERC20Votes.test.js

@@ -58,11 +58,10 @@ contract('ERC20Votes', function (accounts) {
   const name = 'My Token';
   const symbol = 'MTKN';
   const version = '1';
-
   const supply = new BN('10000000000000000000000000');
 
   beforeEach(async function () {
-    this.token = await ERC20VotesMock.new(name, symbol, holder, supply);
+    this.token = await ERC20VotesMock.new(name, symbol);
 
     // We get the chain id from the contract because Ganache (used for coverage) does not return the same chain id
     // from within the EVM as from the JSON RPC interface.
@@ -85,7 +84,7 @@ contract('ERC20Votes', function (accounts) {
   it('minting restriction', async function () {
     const amount = new BN('2').pow(new BN('224'));
     await expectRevert(
-      ERC20VotesMock.new(name, symbol, holder, amount),
+      this.token.mint(holder, amount),
       'ERC20Votes: total supply exceeds 2**224',
     );
   });
@@ -93,6 +92,7 @@ contract('ERC20Votes', function (accounts) {
   describe('set delegation', function () {
     describe('call', function () {
       it('delegation with balance', async function () {
+        await this.token.mint(holder, supply);
         expect(await this.token.delegates(holder)).to.be.equal(ZERO_ADDRESS);
 
         const { receipt } = await this.token.delegate(holder, { from: holder });
@@ -116,17 +116,17 @@ contract('ERC20Votes', function (accounts) {
       });
 
       it('delegation without balance', async function () {
-        expect(await this.token.delegates(recipient)).to.be.equal(ZERO_ADDRESS);
+        expect(await this.token.delegates(holder)).to.be.equal(ZERO_ADDRESS);
 
-        const { receipt } = await this.token.delegate(recipient, { from: recipient });
+        const { receipt } = await this.token.delegate(holder, { from: holder });
         expectEvent(receipt, 'DelegateChanged', {
-          delegator: recipient,
+          delegator: holder,
           fromDelegate: ZERO_ADDRESS,
-          toDelegate: recipient,
+          toDelegate: holder,
         });
         expectEvent.notEmitted(receipt, 'DelegateVotesChanged');
 
-        expect(await this.token.delegates(recipient)).to.be.equal(recipient);
+        expect(await this.token.delegates(holder)).to.be.equal(holder);
       });
     });
 
@@ -143,7 +143,7 @@ contract('ERC20Votes', function (accounts) {
       }});
 
       beforeEach(async function () {
-        await this.token.transfer(delegatorAddress, supply, { from: holder });
+        await this.token.mint(delegatorAddress, supply);
       });
 
       it('accept signed delegation', async function () {
@@ -249,6 +249,7 @@ contract('ERC20Votes', function (accounts) {
 
   describe('change delegation', function () {
     beforeEach(async function () {
+      await this.token.mint(holder, supply);
       await this.token.delegate(holder, { from: holder });
     });
 
@@ -285,6 +286,10 @@ contract('ERC20Votes', function (accounts) {
   });
 
   describe('transfers', function () {
+    beforeEach(async function () {
+      await this.token.mint(holder, supply);
+    });
+
     it('no delegation', async function () {
       const { receipt } = await this.token.transfer(recipient, 1, { from: holder });
       expectEvent(receipt, 'Transfer', { from: holder, to: recipient, value: '1' });
@@ -343,6 +348,10 @@ contract('ERC20Votes', function (accounts) {
 
   // The following tests are a adaptation of https://github.com/compound-finance/compound-protocol/blob/master/tests/Governance/CompTest.js.
   describe('Compound test suite', function () {
+    beforeEach(async function () {
+      await this.token.mint(holder, supply);
+    });
+
     describe('balanceOf', function () {
       it('grants to initial account', async function () {
         expect(await this.token.balanceOf(holder)).to.be.bignumber.equal('10000000000000000000000000');
@@ -455,4 +464,66 @@ contract('ERC20Votes', function (accounts) {
       });
     });
   });
+
+  describe('getPriorTotalSupply', function () {
+    beforeEach(async function () {
+      await this.token.delegate(holder, { from: holder });
+    });
+
+    it('reverts if block number >= current block', async function () {
+      await expectRevert(
+        this.token.getPriorTotalSupply(5e10),
+        'ERC20Votes::getPriorTotalSupply: not yet determined',
+      );
+    });
+
+    it('returns 0 if there are no checkpoints', async function () {
+      expect(await this.token.getPriorTotalSupply(0)).to.be.bignumber.equal('0');
+    });
+
+    it('returns the latest block if >= last checkpoint block', async function () {
+      t1 = await this.token.mint(holder, supply);
+
+      await time.advanceBlock();
+      await time.advanceBlock();
+
+      expect(await this.token.getPriorTotalSupply(t1.receipt.blockNumber)).to.be.bignumber.equal(supply);
+      expect(await this.token.getPriorTotalSupply(t1.receipt.blockNumber + 1)).to.be.bignumber.equal(supply);
+    });
+
+    it('returns zero if < first checkpoint block', async function () {
+      await time.advanceBlock();
+      const t1 = await this.token.mint(holder, supply);
+      await time.advanceBlock();
+      await time.advanceBlock();
+
+      expect(await this.token.getPriorTotalSupply(t1.receipt.blockNumber - 1)).to.be.bignumber.equal('0');
+      expect(await this.token.getPriorTotalSupply(t1.receipt.blockNumber + 1)).to.be.bignumber.equal('10000000000000000000000000');
+    });
+
+    it('generally returns the voting balance at the appropriate checkpoint', async function () {
+      const t1 = await this.token.mint(holder, supply);
+      await time.advanceBlock();
+      await time.advanceBlock();
+      const t2 = await this.token.burn(holder, 10);
+      await time.advanceBlock();
+      await time.advanceBlock();
+      const t3 = await this.token.burn(holder, 10);
+      await time.advanceBlock();
+      await time.advanceBlock();
+      const t4 = await this.token.mint(holder, 20);
+      await time.advanceBlock();
+      await time.advanceBlock();
+
+      expect(await this.token.getPriorTotalSupply(t1.receipt.blockNumber - 1)).to.be.bignumber.equal('0');
+      expect(await this.token.getPriorTotalSupply(t1.receipt.blockNumber)).to.be.bignumber.equal('10000000000000000000000000');
+      expect(await this.token.getPriorTotalSupply(t1.receipt.blockNumber + 1)).to.be.bignumber.equal('10000000000000000000000000');
+      expect(await this.token.getPriorTotalSupply(t2.receipt.blockNumber)).to.be.bignumber.equal('9999999999999999999999990');
+      expect(await this.token.getPriorTotalSupply(t2.receipt.blockNumber + 1)).to.be.bignumber.equal('9999999999999999999999990');
+      expect(await this.token.getPriorTotalSupply(t3.receipt.blockNumber)).to.be.bignumber.equal('9999999999999999999999980');
+      expect(await this.token.getPriorTotalSupply(t3.receipt.blockNumber + 1)).to.be.bignumber.equal('9999999999999999999999980');
+      expect(await this.token.getPriorTotalSupply(t4.receipt.blockNumber)).to.be.bignumber.equal('10000000000000000000000000');
+      expect(await this.token.getPriorTotalSupply(t4.receipt.blockNumber + 1)).to.be.bignumber.equal('10000000000000000000000000');
+    });
+  });
 });