From 75dd08377b85563d097f9cb0999d12426f99feab Mon Sep 17 00:00:00 2001
From: Ryan Sauge <71391932+rya-sge@users.noreply.github.com>
Date: Tue, 12 Nov 2024 17:35:26 +0100
Subject: [PATCH] Separate ERC20 calls inside snapshotModule

---
 .../internal/ERC20SnapshotModuleInternal.sol  | 41 +++-------------
 .../internal/base/SnapshotModuleBase.sol      | 48 +++++++++++++++++++
 2 files changed, 55 insertions(+), 34 deletions(-)

diff --git a/contracts/modules/internal/ERC20SnapshotModuleInternal.sol b/contracts/modules/internal/ERC20SnapshotModuleInternal.sol
index 654b7186..e7db352e 100644
--- a/contracts/modules/internal/ERC20SnapshotModuleInternal.sol
+++ b/contracts/modules/internal/ERC20SnapshotModuleInternal.sol
@@ -68,13 +68,7 @@ abstract contract ERC20SnapshotModuleInternal is ICMTATSnapshot, SnapshotModuleB
         uint256 time,
         address owner
     ) public view returns (uint256) {
-        SnapshotModuleBaseStorage storage $ = _getSnapshotModuleBaseStorage();
-        (bool snapshotted, uint256 value) = _valueAt(
-            time,
-            $._accountBalanceSnapshots[owner]
-        );
-
-        return snapshotted ? value : balanceOf(owner);
+        return _snapshotBalanceOf(time, owner, balanceOf(owner));
     }
 
     /**
@@ -83,12 +77,7 @@ abstract contract ERC20SnapshotModuleInternal is ICMTATSnapshot, SnapshotModuleB
     * @return value stored in the snapshot, or the actual totalSupply if no snapshot
     */
     function snapshotTotalSupply(uint256 time) public view returns (uint256) {
-        SnapshotModuleBaseStorage storage $ = _getSnapshotModuleBaseStorage();
-        (bool snapshotted, uint256 value) = _valueAt(
-            time,
-            $._totalSupplySnapshots
-        );
-        return snapshotted ? value : totalSupply();
+        return _snapshotTotalSupply(time, totalSupply());
     }
 
     /*//////////////////////////////////////////////////////////////
@@ -106,34 +95,18 @@ abstract contract ERC20SnapshotModuleInternal is ICMTATSnapshot, SnapshotModuleB
         _setCurrentSnapshot();
         if (from != address(0)) {
             // for both burn and transfer
-            _updateAccountSnapshot(from);
+            _updateAccountSnapshot(from, balanceOf(from));
             if (to != address(0)) {
                 // transfer
-                _updateAccountSnapshot(to);
+                _updateAccountSnapshot(to, balanceOf(to));
             } else {
                 // burn
-                _updateTotalSupplySnapshot();
+                _updateTotalSupplySnapshot(totalSupply());
             }
         } else {
             // mint
-            _updateAccountSnapshot(to);
-            _updateTotalSupplySnapshot();
+            _updateAccountSnapshot(to, balanceOf(to));
+            _updateTotalSupplySnapshot(totalSupply());
         }
     }
-
-    /**
-    * @dev See {OpenZeppelin - ERC20Snapshot}
-    */
-    function _updateAccountSnapshot(address account) private {
-        SnapshotModuleBaseStorage storage $ = _getSnapshotModuleBaseStorage();
-        _updateSnapshot($._accountBalanceSnapshots[account], balanceOf(account));
-    }
-
-    /**
-    * @dev See {OpenZeppelin - ERC20Snapshot}
-    */
-    function _updateTotalSupplySnapshot() private {
-        SnapshotModuleBaseStorage storage $ = _getSnapshotModuleBaseStorage();
-        _updateSnapshot($._totalSupplySnapshots, totalSupply());
-    }
 }
diff --git a/contracts/modules/internal/base/SnapshotModuleBase.sol b/contracts/modules/internal/base/SnapshotModuleBase.sol
index d099e467..09b6b70e 100644
--- a/contracts/modules/internal/base/SnapshotModuleBase.sol
+++ b/contracts/modules/internal/base/SnapshotModuleBase.sol
@@ -397,6 +397,54 @@ abstract contract SnapshotModuleBase is Initializable {
         return (mostRecent, index);
     }
 
+    /* ============ Require balance and total supply ============ */
+     /**
+    * @dev See {OpenZeppelin - ERC20Snapshot}
+    */
+    function _updateAccountSnapshot(address account, uint256 accountBalance) internal {
+        SnapshotModuleBaseStorage storage $ = _getSnapshotModuleBaseStorage();
+        _updateSnapshot($._accountBalanceSnapshots[account], accountBalance);
+    }
+
+    /**
+    * @dev See {OpenZeppelin - ERC20Snapshot}
+    */
+    function _updateTotalSupplySnapshot(uint256 totalSupply) internal {
+        SnapshotModuleBaseStorage storage $ = _getSnapshotModuleBaseStorage();
+        _updateSnapshot($._totalSupplySnapshots, totalSupply);
+    }
+
+       /** 
+    * @notice Return the number of tokens owned by the given owner at the time when the snapshot with the given time was created.
+    * @return value stored in the snapshot, or the actual balance if no snapshot
+    */
+    function _snapshotBalanceOf(
+        uint256 time,
+        address owner,
+        uint256 ownerBalance
+    ) internal view returns (uint256) {
+        SnapshotModuleBaseStorage storage $ = _getSnapshotModuleBaseStorage();
+        (bool snapshotted, uint256 value) = _valueAt(
+            time,
+            $._accountBalanceSnapshots[owner]
+        );
+        return snapshotted ? value :  ownerBalance;
+    }
+
+    /**
+    * @dev See {OpenZeppelin - ERC20Snapshot}
+    * Retrieves the total supply at the specified time.
+    * @return value stored in the snapshot, or the actual totalSupply if no snapshot
+    */
+    function _snapshotTotalSupply(uint256 time, uint256 totalSupply) internal view returns (uint256) {
+        SnapshotModuleBaseStorage storage $ = _getSnapshotModuleBaseStorage();
+        (bool snapshotted, uint256 value) = _valueAt(
+            time,
+            $._totalSupplySnapshots
+        );
+        return snapshotted ? value : totalSupply;
+    }
+
     /* ============ Utility functions ============ */