Browse Source

Use beforeTokenTransfer hook in ERC20Snapshot (#2312)

Julian M. Rodriguez 5 years ago
parent
commit
9700e6b4bd
1 changed files with 12 additions and 19 deletions
  1. 12 19
      contracts/token/ERC20/ERC20Snapshot.sol

+ 12 - 19
contracts/token/ERC20/ERC20Snapshot.sol

@@ -104,28 +104,21 @@ abstract contract ERC20Snapshot is ERC20 {
         return snapshotted ? value : totalSupply();
     }
 
-    // _transfer, _mint and _burn are the only functions where the balances are modified, so it is there that the
-    // snapshots are updated. Note that the update happens _before_ the balance change, with the pre-modified value.
-    // The same is true for the total supply and _mint and _burn.
-    function _transfer(address from, address to, uint256 value) internal virtual override {
-        _updateAccountSnapshot(from);
-        _updateAccountSnapshot(to);
-
-        super._transfer(from, to, value);
-    }
 
-    function _mint(address account, uint256 value) internal virtual override {
-        _updateAccountSnapshot(account);
-        _updateTotalSupplySnapshot();
+    // Update balance and/or total supply snapshots before the values are modified. This is implemented
+    // in the _beforeTokenTransfer hook, which is executed for _mint, _burn, and _transfer operations.
+    function _beforeTokenTransfer(address from, address to, uint256 amount) internal virtual override {
+      super._beforeTokenTransfer(from, to, amount);
 
-        super._mint(account, value);
-    }
-
-    function _burn(address account, uint256 value) internal virtual override {
-        _updateAccountSnapshot(account);
+      if (from == address(0) || to == address(0)) {
+        // mint or burn
+        _updateAccountSnapshot(from);
         _updateTotalSupplySnapshot();
-
-        super._burn(account, value);
+      } else {
+        // transfer
+        _updateAccountSnapshot(from);
+        _updateAccountSnapshot(to);
+      }
     }
 
     function _valueAt(uint256 snapshotId, Snapshots storage snapshots)