Remove claimRewards double call

morpho-org · Dec 15, 2022 · 2304cf4 · 2304cf4
1 parent 36b26a4
commit 2304cf4
Showing 1 changed file with 22 additions and 10 deletions.
diff --git a/src/compound/SupplyVault.sol b/src/compound/SupplyVault.sol
@@ -116,28 +116,40 @@ contract SupplyVault is ISupplyVault, SupplyVaultBase {
         address to,
         uint256 amount
     ) internal override {
-        _accrueUnclaimedRewards(from);
-        _accrueUnclaimedRewards(to);
+        uint256 newRewardsIndex = _claimRewards();
+        _accrueUnclaimedRewardsFromRewardsIndex(from, newRewardsIndex);
+        _accrueUnclaimedRewardsFromRewardsIndex(to, newRewardsIndex);
+
         super._beforeTokenTransfer(from, to, amount);
     }
 
-    function _accrueUnclaimedRewards(address _user) internal returns (uint256 unclaimed) {
-        uint256 supply = totalSupply();
-        uint256 rewardsIndexMem = rewardsIndex;
+    function _claimRewards() internal returns (uint256 newRewardsIndex) {
+        newRewardsIndex = rewardsIndex;
 
+        uint256 supply = totalSupply();
         if (supply > 0) {
             address[] memory poolTokens = new address[](1);
             poolTokens[0] = poolToken;
-            rewardsIndexMem += morpho.claimRewards(poolTokens, false).divWadDown(supply);
-            rewardsIndex = rewardsIndexMem;
+
+            newRewardsIndex += morpho.claimRewards(poolTokens, false).divWadDown(supply);
+            rewardsIndex = newRewardsIndex;
         }
+    }
+
+    function _accrueUnclaimedRewards(address _user) internal returns (uint256 unclaimed) {
+        return _accrueUnclaimedRewardsFromRewardsIndex(_user, _claimRewards());
+    }
 
+    function _accrueUnclaimedRewardsFromRewardsIndex(address _user, uint256 _newRewardsIndex)
+        internal
+        returns (uint256 unclaimed)
+    {
         UserRewardsData storage userRewardsData = userRewards[_user];
         uint256 rewardsIndexDiff;
 
         // Safe because we always have `rewardsIndex` >= `userRewardsData.index`.
         unchecked {
-            rewardsIndexDiff = rewardsIndexMem - userRewardsData.index;
+            rewardsIndexDiff = _newRewardsIndex - userRewardsData.index;
         }
 
         unclaimed = userRewardsData.unclaimed;
@@ -146,8 +158,8 @@ contract SupplyVault is ISupplyVault, SupplyVaultBase {
             userRewardsData.unclaimed = unclaimed.safeCastTo128();
         }
 
-        userRewardsData.index = rewardsIndexMem.safeCastTo128();
+        userRewardsData.index = _newRewardsIndex.safeCastTo128();
 
-        emit Accrued(_user, rewardsIndexMem, unclaimed);
+        emit Accrued(_user, _newRewardsIndex, unclaimed);
     }
 }