add ERC20SnapshotBlockNumber

contracts/mocks/ERC20SnapshotBlockNumberMock.sol

@@ -0,0 +1,29 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.0;
+
+import "../token/ERC20/extensions/ERC20SnapshotBlockNumber.sol";
+
+
+contract ERC20SnapshotBlockNumberMock is ERC20SnapshotBlockNumber {
+    constructor(
+        string memory name,
+        string memory symbol,
+        address initialAccount,
+        uint256 initialBalance
+    ) ERC20(name, symbol) {
+        _mint(initialAccount, initialBalance);
+    }
+
+    function snapshot() public {
+        _snapshot();
+    }
+
+    function mint(address account, uint256 amount) public {
+        _mint(account, amount);
+    }
+
+    function burn(address account, uint256 amount) public {
+        _burn(account, amount);
+    }
+}

contracts/token/ERC20/extensions/ERC20Snapshot.sol

@@ -84,6 +84,20 @@ abstract contract ERC20Snapshot is ERC20 {
         return currentId;
     }
 
+    /**
+     * @dev Create a new snapshot this a specific snapshot id. Usefull to overload {_snapshot} with specific
+     * snapshot id mechanism such as snapshot by block id
+     */
+    function _createSnapshot(uint256 currentId) internal virtual returns (uint256) {
+        if (_currentSnapshotId._value < currentId) {
+            _currentSnapshotId._value = currentId;
+            emit Snapshot(currentId);
+        } else if (_currentSnapshotId._value > currentId) {
+            revert("ERC20Snapshot: snapshot id must increase monotonically");
+        }
+        return currentId;
+    }
+
     /**
      * @dev Retrieves the balance of `account` at the time `snapshotId` was created.
      */

contracts/token/ERC20/extensions/ERC20SnapshotBlockNumber.sol

@@ -0,0 +1,14 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.0;
+
+import "./ERC20Snapshot.sol";
+
+/**
+ * @dev Variant of the ERC20Snapshot extension that uses block number as snapshot ids.
+ */
+abstract contract ERC20SnapshotBlockNumber is ERC20Snapshot {
+    function _snapshot() internal virtual override returns (uint256) {
+        return _createSnapshot(block.number);
+    }
+}

test/token/ERC20/extensions/ERC20SnapshotBlockNumber.test.js

@@ -0,0 +1,205 @@
+const { BN, expectEvent, expectRevert } = require('@openzeppelin/test-helpers');
+const ERC20SnapshotBlockNumberMock = artifacts.require('ERC20SnapshotBlockNumberMock');
+
+const { expect } = require('chai');
+
+contract('ERC20SnapshotBlockNumber', function (accounts) {
+  const [ initialHolder, recipient, other ] = accounts;
+
+  const initialSupply = new BN(100);
+
+  const name = 'My Token';
+  const symbol = 'MTKN';
+
+  beforeEach(async function () {
+    this.token = await ERC20SnapshotBlockNumberMock.new(name, symbol, initialHolder, initialSupply);
+  });
+
+  describe('snapshot', function () {
+    it('emits a snapshot event', async function () {
+      const { logs } = await this.token.snapshot();
+      expectEvent.inLogs(logs, 'Snapshot');
+    });
+
+    it('creates increasing snapshots ids, corresponding to the blockNumber', async function () {
+      // eslint-disable-next-line
+      for (const _ of Array(5).fill()) {
+        const { receipt } = await this.token.snapshot();
+        expectEvent(receipt, 'Snapshot', { id: new BN(receipt.blockNumber) });
+      }
+    });
+  });
+
+  describe('totalSupplyAt', function () {
+    it('reverts with a snapshot id of 0', async function () {
+      await expectRevert(this.token.totalSupplyAt(0), 'ERC20Snapshot: id is 0');
+    });
+
+    it('reverts with a not-yet-created snapshot id', async function () {
+      await expectRevert(this.token.totalSupplyAt(1), 'ERC20Snapshot: nonexistent id');
+    });
+
+    context('with initial snapshot', function () {
+      beforeEach(async function () {
+        const { receipt } = await this.token.snapshot();
+        this.initialSnapshotId = new BN(receipt.blockNumber);
+        expectEvent(receipt, 'Snapshot', { id: this.initialSnapshotId });
+      });
+
+      context('with no supply changes after the snapshot', function () {
+        it('returns the current total supply', async function () {
+          expect(await this.token.totalSupplyAt(this.initialSnapshotId)).to.be.bignumber.equal(initialSupply);
+        });
+      });
+
+      context('with supply changes after the snapshot', function () {
+        beforeEach(async function () {
+          await this.token.mint(other, new BN('50'));
+          await this.token.burn(initialHolder, new BN('20'));
+        });
+
+        it('returns the total supply before the changes', async function () {
+          expect(await this.token.totalSupplyAt(this.initialSnapshotId)).to.be.bignumber.equal(initialSupply);
+        });
+
+        context('with a second snapshot after supply changes', function () {
+          beforeEach(async function () {
+            const { receipt } = await this.token.snapshot();
+            this.secondSnapshotId = new BN(receipt.blockNumber);
+            expectEvent(receipt, 'Snapshot', { id: this.secondSnapshotId });
+          });
+
+          it('snapshots return the supply before and after the changes', async function () {
+            expect(await this.token.totalSupplyAt(this.initialSnapshotId)).to.be.bignumber.equal(initialSupply);
+
+            expect(await this.token.totalSupplyAt(this.secondSnapshotId)).to.be.bignumber.equal(
+              await this.token.totalSupply(),
+            );
+          });
+        });
+
+        context('with multiple snapshots after supply changes', function () {
+          beforeEach(async function () {
+            this.secondSnapshotIds = [];
+            // eslint-disable-next-line
+            for (const _ in Array(3).fill()) {
+              const { receipt } = await this.token.snapshot();
+              const id = new BN(receipt.blockNumber);
+              expectEvent(receipt, 'Snapshot', { id });
+              this.secondSnapshotIds.push(id);
+            }
+          });
+
+          it('all posterior snapshots return the supply after the changes', async function () {
+            expect(await this.token.totalSupplyAt(this.initialSnapshotId)).to.be.bignumber.equal(initialSupply);
+
+            const currentSupply = await this.token.totalSupply();
+
+            for (const id of this.secondSnapshotIds) {
+              expect(await this.token.totalSupplyAt(id)).to.be.bignumber.equal(currentSupply);
+            }
+          });
+        });
+      });
+    });
+  });
+
+  describe('balanceOfAt', function () {
+    it('reverts with a snapshot id of 0', async function () {
+      await expectRevert(this.token.balanceOfAt(other, 0), 'ERC20Snapshot: id is 0');
+    });
+
+    it('reverts with a not-yet-created snapshot id', async function () {
+      await expectRevert(this.token.balanceOfAt(other, 1), 'ERC20Snapshot: nonexistent id');
+    });
+
+    context('with initial snapshot', function () {
+      beforeEach(async function () {
+        const { receipt } = await this.token.snapshot();
+        this.initialSnapshotId = new BN(receipt.blockNumber);
+        expectEvent(receipt, 'Snapshot', { id: this.initialSnapshotId });
+      });
+
+      context('with no balance changes after the snapshot', function () {
+        it('returns the current balance for all accounts', async function () {
+          expect(await this.token.balanceOfAt(initialHolder, this.initialSnapshotId))
+            .to.be.bignumber.equal(initialSupply);
+          expect(await this.token.balanceOfAt(recipient, this.initialSnapshotId)).to.be.bignumber.equal('0');
+          expect(await this.token.balanceOfAt(other, this.initialSnapshotId)).to.be.bignumber.equal('0');
+        });
+      });
+
+      context('with balance changes after the snapshot', function () {
+        beforeEach(async function () {
+          await this.token.transfer(recipient, new BN('10'), { from: initialHolder });
+          await this.token.mint(other, new BN('50'));
+          await this.token.burn(initialHolder, new BN('20'));
+        });
+
+        it('returns the balances before the changes', async function () {
+          expect(await this.token.balanceOfAt(initialHolder, this.initialSnapshotId))
+            .to.be.bignumber.equal(initialSupply);
+          expect(await this.token.balanceOfAt(recipient, this.initialSnapshotId)).to.be.bignumber.equal('0');
+          expect(await this.token.balanceOfAt(other, this.initialSnapshotId)).to.be.bignumber.equal('0');
+        });
+
+        context('with a second snapshot after supply changes', function () {
+          beforeEach(async function () {
+            const { receipt } = await this.token.snapshot();
+            this.secondSnapshotId = new BN(receipt.blockNumber);
+            expectEvent(receipt, 'Snapshot', { id: this.secondSnapshotId });
+          });
+
+          it('snapshots return the balances before and after the changes', async function () {
+            expect(await this.token.balanceOfAt(initialHolder, this.initialSnapshotId))
+              .to.be.bignumber.equal(initialSupply);
+            expect(await this.token.balanceOfAt(recipient, this.initialSnapshotId)).to.be.bignumber.equal('0');
+            expect(await this.token.balanceOfAt(other, this.initialSnapshotId)).to.be.bignumber.equal('0');
+
+            expect(await this.token.balanceOfAt(initialHolder, this.secondSnapshotId)).to.be.bignumber.equal(
+              await this.token.balanceOf(initialHolder),
+            );
+            expect(await this.token.balanceOfAt(recipient, this.secondSnapshotId)).to.be.bignumber.equal(
+              await this.token.balanceOf(recipient),
+            );
+            expect(await this.token.balanceOfAt(other, this.secondSnapshotId)).to.be.bignumber.equal(
+              await this.token.balanceOf(other),
+            );
+          });
+        });
+
+        context('with multiple snapshots after supply changes', function () {
+          beforeEach(async function () {
+            this.secondSnapshotIds = [];
+            // eslint-disable-next-line
+            for (const _ in Array(3).fill()) {
+              const { receipt } = await this.token.snapshot();
+              const id = new BN(receipt.blockNumber);
+              expectEvent(receipt, 'Snapshot', { id });
+              this.secondSnapshotIds.push(id);
+            }
+          });
+
+          it('all posterior snapshots return the supply after the changes', async function () {
+            expect(await this.token.balanceOfAt(initialHolder, this.initialSnapshotId))
+              .to.be.bignumber.equal(initialSupply);
+            expect(await this.token.balanceOfAt(recipient, this.initialSnapshotId)).to.be.bignumber.equal('0');
+            expect(await this.token.balanceOfAt(other, this.initialSnapshotId)).to.be.bignumber.equal('0');
+
+            for (const id of this.secondSnapshotIds) {
+              expect(await this.token.balanceOfAt(initialHolder, id)).to.be.bignumber.equal(
+                await this.token.balanceOf(initialHolder),
+              );
+              expect(await this.token.balanceOfAt(recipient, id)).to.be.bignumber.equal(
+                await this.token.balanceOf(recipient),
+              );
+              expect(await this.token.balanceOfAt(other, id)).to.be.bignumber.equal(
+                await this.token.balanceOf(other),
+              );
+            }
+          });
+        });
+      });
+    });
+  });
+});