From 4a17de5504a46687b2baeffc37872ccf6a1b592a Mon Sep 17 00:00:00 2001
From: Joaquin Gonzalez <jpgonzalezra@gmail.com>
Date: Mon, 27 May 2024 18:05:18 -0300
Subject: [PATCH] refactor: add function to calculate total accrued rewards

---
 src/CompotaToken.sol    | 35 +++++++++++++++++++++++++----------
 test/CompotaToken.t.sol | 14 +++++++++-----
 2 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/src/CompotaToken.sol b/src/CompotaToken.sol
index 21bdd30..b7e4fdf 100644
--- a/src/CompotaToken.sol
+++ b/src/CompotaToken.sol
@@ -27,8 +27,8 @@ contract CompotaToken is ICompotaToken, ERC20Extended, Owned {
 
     uint16 public yearlyRate;
 
+    uint256 public latestUpdateTimestamp;
     uint256 internal _totalSupply;
-    address[] internal _earners;
 
     mapping(address => uint256) internal _balances;
     mapping(address => uint256) internal _lastUpdateTimestamp;
@@ -97,12 +97,7 @@ contract CompotaToken is ICompotaToken, ERC20Extended, Owned {
      * @inheritdoc IERC20
      */
     function totalSupply() external view returns (uint256 totalSupply_) {
-        totalSupply_ = _totalSupply;
-        uint256 length = _earners.length;
-        for (uint256 i = 0; i < length; i++) {
-            totalSupply_ += _calculateCurrentRewards(_earners[i]);
-        }
-        return totalSupply_;
+        return _totalSupply + _calculateTotalCurrentRewards();
     }
 
     /**
@@ -179,9 +174,9 @@ contract CompotaToken is ICompotaToken, ERC20Extended, Owned {
      */
     function _updateRewards(address account_) internal {
         uint256 timestamp = block.timestamp;
+        latestUpdateTimestamp = timestamp;
         if (_lastUpdateTimestamp[account_] == 0) {
             _lastUpdateTimestamp[account_] = timestamp;
-            _earners.push(account_);
             emit StartedEarningRewards(account_);
             return;
         }
@@ -200,12 +195,32 @@ contract CompotaToken is ICompotaToken, ERC20Extended, Owned {
      */
     function _calculateCurrentRewards(address account_) internal view returns (uint256) {
         if (_lastUpdateTimestamp[account_] == 0) return 0;
+        return _calculateRewards(_balances[account_], _lastUpdateTimestamp[account_]);
+    }
+
+    /**
+     * @notice Calculates the total current accrued rewards for the entire supply since the last update.
+     * @return The amount of rewards accrued since the last update.
+     */
+    function _calculateTotalCurrentRewards() internal view returns (uint256) {
+        if (latestUpdateTimestamp == 0) return 0;
+        return _calculateRewards(_totalSupply, latestUpdateTimestamp);
+    }
+
+    /**
+     * @notice Generalized function to calculate rewards based on an amount and a timestamp.
+     * @param amount_ The amount of tokens to calculate rewards for.
+     * @param timestamp_ The timestamp to calculate rewards from.
+     * @return The amount of rewards accrued since the last update.
+     */
+    function _calculateRewards(uint256 amount_, uint256 timestamp_) internal view returns (uint256) {
+        if (timestamp_ == 0) return 0;
         uint256 timeElapsed;
         // Safe to use unchecked here, since `block.timestamp` is always greater than `_lastUpdateTimestamp[account_]`.
         unchecked {
-            timeElapsed = block.timestamp - _lastUpdateTimestamp[account_];
+            timeElapsed = block.timestamp - timestamp_;
         }
-        return (_balances[account_] * timeElapsed * yearlyRate) / (SCALE_FACTOR * uint256(SECONDS_PER_YEAR));
+        return (amount_ * timeElapsed * yearlyRate) / (SCALE_FACTOR * uint256(SECONDS_PER_YEAR));
     }
 
     /**
diff --git a/test/CompotaToken.t.sol b/test/CompotaToken.t.sol
index 9e8d524..aafc58b 100644
--- a/test/CompotaToken.t.sol
+++ b/test/CompotaToken.t.sol
@@ -279,18 +279,22 @@ contract CompotaTokenTest is Test {
     }
 
     function testTotalSupplyWithUnclaimedRewards() external {
-        uint256 initialMint = 1000 * 10e6; // Initial mint amount
+        uint256 aliceInitialMint = 1000 * 10e6;
+        uint256 bobInitialMint = aliceInitialMint * 2;
 
-        _mint(owner, alice, initialMint);
-        assertEq(token.totalSupply(), initialMint);
+        _mint(owner, alice, aliceInitialMint);
+        _mint(owner, bob, bobInitialMint);
+
+        uint256 totalSupply = aliceInitialMint + bobInitialMint;
+        assertEq(token.totalSupply(), totalSupply);
 
         vm.warp(block.timestamp + 180 days);
 
         // Calculate expected rewards for 180 days with the initial rate
-        uint256 expectedRewards = (initialMint * INTEREST_RATE * 180 days) / (10_000 * 365 days);
+        uint256 expectedRewards = (totalSupply * INTEREST_RATE * 180 days) / (10_000 * 365 days);
 
         // Verify total supply includes unclaimed rewards
-        assertEq(token.totalSupply(), initialMint + expectedRewards);
+        assertEq(token.totalSupply(), totalSupply + expectedRewards);
     }
 
     function testClaimRewardsWithNoRewards() public {