From 9700e6b4bdcd3dace1405c17cd4f048d455d4014 Mon Sep 17 00:00:00 2001
From: "Julian M. Rodriguez" <56316686+julianmrodri@users.noreply.github.com>
Date: Tue, 11 Aug 2020 16:51:58 -0300
Subject: [PATCH] Use beforeTokenTransfer hook in ERC20Snapshot (#2312)

---
 contracts/token/ERC20/ERC20Snapshot.sol | 31 ++++++++++---------------
 1 file changed, 12 insertions(+), 19 deletions(-)

diff --git a/contracts/token/ERC20/ERC20Snapshot.sol b/contracts/token/ERC20/ERC20Snapshot.sol
index 54d91c0573a..cc25d618e97 100644
--- a/contracts/token/ERC20/ERC20Snapshot.sol
+++ b/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)