fix: Use virtual rewards/balances when possible (update tests accordingly)

contracts/protocol/tokenization/StaticATokenLM.sol

@@ -54,11 +54,16 @@ contract StaticATokenLM is ERC20 {
   mapping(address => uint256) public _nonces;
 
   uint256 public accRewardstokenPerShare;
+  uint256 public lifeTimeRewardsClaimed;
+  uint256 public lifeTimeRewards;
   uint256 public lastRewardBlock;
-  mapping(address => uint256) public rewardDebts; // Measured in Rays
-  mapping(address => uint256) public unclaimedRewards; // Measured in Rays
-  IAaveIncentivesController internal _incentivesController;
 
+  // user => rewardDebt (in RAYs)
+  mapping(address => uint256) public rewardDebts;
+  // user => unclaimedRewards (in RAYs)
+  mapping(address => uint256) public unclaimedRewards;
+
+  IAaveIncentivesController internal _incentivesController;
   address public immutable currentRewardToken;
 
   constructor(
@@ -357,8 +362,7 @@ contract StaticATokenLM is ERC20 {
     bool fromUnderlying
   ) internal returns (uint256) {
     require(recipient != address(0), 'INVALID_RECIPIENT');
-    updateRewards();
+    _updateRewards();
-
     _updateUnclaimedRewards(recipient);
 
     if (fromUnderlying) {
@@ -384,8 +388,7 @@ contract StaticATokenLM is ERC20 {
   ) internal returns (uint256, uint256) {
     require(recipient != address(0), 'INVALID_RECIPIENT');
     require(staticAmount == 0 || dynamicAmount == 0, 'ONLY_ONE_AMOUNT_FORMAT_ALLOWED');
-    updateRewards();
+    _updateRewards();
-
     _updateUnclaimedRewards(owner);
 
     uint256 userBalance = balanceOf(owner);
@@ -446,32 +449,65 @@ contract StaticATokenLM is ERC20 {
   /**
    * @dev Claims rewards from the `_incentivesController` and update `accRewardstokenPerShare`
    */
-  function updateRewards() public {
+  function _updateRewards() internal {
+    // Update the virtual rewards without actually claiming.
     if (block.number > lastRewardBlock) {
       lastRewardBlock = block.number;
       uint256 _supply = totalSupply();
       if (_supply == 0) {
+        // No rewards can have accrued since last because there were no funds.
         return;
       }
 
       address[] memory assets = new address[](1);
       assets[0] = address(ATOKEN);
 
-      uint256 freshReward =
+      uint256 freshRewards = _incentivesController.getRewardsBalance(assets, address(this));
-        _incentivesController.claimRewards(assets, type(uint256).max, address(this)).wadToRay();
+      uint256 externalLifetimeRewards = lifeTimeRewardsClaimed.add(freshRewards);
+      uint256 diff = externalLifetimeRewards.sub(lifeTimeRewards).wadToRay();
+
       accRewardstokenPerShare = accRewardstokenPerShare.add(
-        freshReward.rayDivNoRounding(_supply.wadToRay())
+        (diff).rayDivNoRounding(_supply.wadToRay())
       );
+
+      lifeTimeRewards = externalLifetimeRewards;
     }
   }
 
-  /**
+  function collectAndUpdateRewards() public {
-   * @dev Update the rewards and claim rewards for the user
+    if (block.number > lastRewardBlock) {
-   * @param user The address of the user to claim rewards for
+      lastRewardBlock = block.number;
-   */
+      uint256 _supply = totalSupply();
-  function updateAndClaimRewards(address user) public {
+
-    updateRewards();
+      // We need to perform the check even though there is no supply, as rewards can have accrued before it was removed
-    claimRewards(user);
+
+      address[] memory assets = new address[](1);
+      assets[0] = address(ATOKEN);
+
+      uint256 freshlyClaimed =
+        _incentivesController.claimRewards(assets, type(uint256).max, address(this));
+      uint256 externalLifetimeRewards = lifeTimeRewardsClaimed.add(freshlyClaimed);
+      uint256 diff = externalLifetimeRewards.sub(lifeTimeRewards).wadToRay();
+
+      if (_supply > 0 && diff > 0) {
+        accRewardstokenPerShare = accRewardstokenPerShare.add(
+          (diff).rayDivNoRounding(_supply.wadToRay())
+        );
+      }
+
+      if (diff > 0) {
+        lifeTimeRewards = externalLifetimeRewards;
+      }
+      // Unsure if we can also move this in
+      lifeTimeRewardsClaimed = externalLifetimeRewards;
+    }
+    /*
+  // This one could just as well do both?
+    address[] memory assets = new address[](1);
+    assets[0] = address(ATOKEN);
+    uint256 freshlyClaimed =
+      _incentivesController.claimRewards(assets, type(uint256).max, address(this));
+    lifeTimeRewardsClaimed = lifeTimeRewardsClaimed.add(freshlyClaimed);*/
   }
 
   /**
@@ -479,10 +515,18 @@ contract StaticATokenLM is ERC20 {
    * makes sense for small holders
    * @param user The address of the user to claim rewards for
    */
-  function claimRewards(address user) public {
+  function claimRewards(address user, bool forceUpdate) public {
-    // Claim rewards without collecting the latest rewards
+    if (forceUpdate) {
+      collectAndUpdateRewards();
+    }
+
     uint256 balance = balanceOf(user);
-    uint256 reward = _getClaimableRewards(user, balance); // Remember that this is converting to wad
+    uint256 reward = _getClaimableRewards(user, balance, false);
+    uint256 totBal = IERC20(currentRewardToken).balanceOf(address(this));
+    if (reward > totBal) {
+      // Throw away excess rewards
+      reward = totBal;
+    }
     if (reward > 0) {
       unclaimedRewards[user] = 0;
       IERC20(currentRewardToken).safeTransfer(user, reward);
@@ -508,21 +552,48 @@ contract StaticATokenLM is ERC20 {
   function _updateUnclaimedRewards(address user) internal {
     uint256 balance = balanceOf(user);
     if (balance > 0) {
-      uint256 pending = _getPendingRewards(user, balance);
+      uint256 pending = _getPendingRewards(user, balance, false);
       unclaimedRewards[user] = unclaimedRewards[user].add(pending);
     }
   }
 
   /**
-   * @dev Compute the pending in RAY (rounded down). Pending is the amount to add  (not yet unclaimed) rewards in RAY (rounded down).
+   * @dev Compute the pending in RAY (rounded down). Pending is the amount to add (not yet unclaimed) rewards in RAY (rounded down).
    * @param user The user to compute for
    * @param balance The balance of the user
    * @return The amound of pending rewards in RAY
    */
-  function _getPendingRewards(address user, uint256 balance) internal view returns (uint256) {
+  function _getPendingRewards(
+    address user,
+    uint256 balance,
+    bool fresh
+  ) internal view returns (uint256) {
+    if (balance == 0) {
+      return 0;
+    }
+
+    // TODO: This could retrieve the last such that we know the most up to date stuff :eyes:
     // Compute the pending rewards in ray, rounded down.
     uint256 rayBalance = balance.wadToRay();
-    uint256 _reward = rayBalance.rayMulNoRounding(accRewardstokenPerShare);
+
+    uint256 _supply = totalSupply();
+    uint256 _accRewardstokenPerShare = accRewardstokenPerShare;
+
+    if (_supply != 0 && fresh) {
+      // Done purely virtually, this is used for retrieving up to date rewards for the ui
+      address[] memory assets = new address[](1);
+      assets[0] = address(ATOKEN);
+
+      uint256 freshReward = _incentivesController.getRewardsBalance(assets, address(this));
+      uint256 externalLifetimeRewards = lifeTimeRewardsClaimed.add(freshReward);
+      uint256 diff = externalLifetimeRewards.sub(lifeTimeRewards).wadToRay();
+
+      _accRewardstokenPerShare = _accRewardstokenPerShare.add(
+        (diff).rayDivNoRounding(_supply.wadToRay())
+      );
+    }
+
+    uint256 _reward = rayBalance.rayMulNoRounding(_accRewardstokenPerShare);
     uint256 _debt = rewardDebts[user];
     if (_reward > _debt) {
       // Safe because line above
@@ -531,17 +602,32 @@ contract StaticATokenLM is ERC20 {
     return 0;
   }
 
-  function _getClaimableRewards(address user, uint256 balance) internal view returns (uint256) {
+  function _getClaimableRewards(
-    uint256 reward = unclaimedRewards[user].add(_getPendingRewards(user, balance));
+    address user,
+    uint256 balance,
+    bool fresh
+  ) internal view returns (uint256) {
+    uint256 reward = unclaimedRewards[user].add(_getPendingRewards(user, balance, fresh));
     return reward.rayToWadNoRounding();
   }
 
+  function getTotalClaimableRewards() public view returns (uint256) {
+    address[] memory assets = new address[](1);
+    assets[0] = address(ATOKEN);
+    uint256 freshRewards = _incentivesController.getRewardsBalance(assets, address(this));
+    return IERC20(currentRewardToken).balanceOf(address(this)).add(freshRewards);
+  }
+
   /**
    * @dev Get the total claimable rewards for a user in WAD cliam
    * @param user The address of the user
    * @return The claimable amount of rewards in WAD
    */
   function getClaimableRewards(address user) public view returns (uint256) {
-    return _getClaimableRewards(user, balanceOf(user));
+    return _getClaimableRewards(user, balanceOf(user), true);
+  }
+
+  function getUnclaimedRewards(address user) public view returns (uint256) {
+    return unclaimedRewards[user].rayToWadNoRounding();
   }
 }

test-suites/test-aave/mainnet/static-atoken-lm/static-atoken-liquidity-mining-rewards.spec.ts

@@ -23,7 +23,7 @@ import {
   advanceTimeAndBlock,
 } from '../../../../helpers/misc-utils';
 import { BigNumber, providers, Signer, utils } from 'ethers';
-import { MAX_UINT_AMOUNT } from '../../../../helpers/constants';
+import { MAX_UINT_AMOUNT, USD_ADDRESS } from '../../../../helpers/constants';
 import { AbiCoder, formatEther, verifyTypedData } from 'ethers/lib/utils';
 
 import { _TypedDataEncoder } from 'ethers/lib/utils';
@@ -143,7 +143,7 @@ describe('StaticATokenLM: aToken wrapper with static balances and liquidity mini
 
       const pendingRewards2 = await staticAToken.getClaimableRewards(userSigner._address);
 
-      await waitForTx(await staticAToken.updateRewards());
+      await waitForTx(await staticAToken.collectAndUpdateRewards());
 
       const pendingRewards3 = await staticAToken.getClaimableRewards(userSigner._address);
 
@@ -153,20 +153,36 @@ describe('StaticATokenLM: aToken wrapper with static balances and liquidity mini
       );
 
       const pendingRewards4 = await staticAToken.getClaimableRewards(userSigner._address);
+      const totPendingRewards4 = await staticAToken.getTotalClaimableRewards();
       const claimedRewards4 = await stkAave.balanceOf(userSigner._address);
+      const stkAaveStatic4 = await stkAave.balanceOf(staticAToken.address);
 
-      await waitForTx(await staticAToken.claimRewards(userSigner._address));
+      await waitForTx(await staticAToken.claimRewards(userSigner._address, false));
 
       const pendingRewards5 = await staticAToken.getClaimableRewards(userSigner._address);
+      const totPendingRewards5 = await staticAToken.getTotalClaimableRewards();
       const claimedRewards5 = await stkAave.balanceOf(userSigner._address);
+      const stkAaveStatic5 = await stkAave.balanceOf(staticAToken.address);
+
+      await waitForTx(await staticAToken.collectAndUpdateRewards());
+      const pendingRewards6 = await staticAToken.getClaimableRewards(userSigner._address);
+
+      // Checks
 
       expect(pendingRewards2).to.be.gt(pendingRewards1);
       expect(pendingRewards3).to.be.gt(pendingRewards2);
       expect(pendingRewards4).to.be.gt(pendingRewards3);
-      expect(pendingRewards5).to.be.eq(0);
+      expect(totPendingRewards4).to.be.gte(pendingRewards4);
+      expect(pendingRewards5).to.be.eq(0); // User "sacrifice" excess rewards to save on gas-costs
+      expect(pendingRewards6).to.be.eq(0);
 
       expect(claimedRewards4).to.be.eq(0);
-      expect(claimedRewards5).to.be.eq(pendingRewards4);
+
+      // Expect the user to have withdrawn everything.
+      expect(claimedRewards5).to.be.eq(stkAaveStatic4);
+      expect(stkAaveStatic5).to.be.eq(0);
+
+      expect(totPendingRewards5).to.be.gt(0);
     });
 
     it('Check getters', async () => {
@@ -193,7 +209,50 @@ describe('StaticATokenLM: aToken wrapper with static balances and liquidity mini
       expect(dynamicBalance).to.be.eq(dynamicBalanceFromStatic);
     });
 
-    it.skip('Multiple updates in one block (Breaks if GasReport enabled)', async () => {
+    it.skip('Multiple deposits in one block (Breaks if GasReport enabled)', async () => {
+      const amountToDeposit = utils.parseEther('5');
+
+      // Just preparation
+      await waitForTx(await weth.deposit({ value: amountToDeposit.mul(2) }));
+      await waitForTx(
+        await weth.approve(staticAToken.address, amountToDeposit.mul(2), defaultTxParams)
+      );
+
+      await DRE.network.provider.send('evm_setAutomine', [false]);
+
+      // Depositing
+      let a = await staticAToken.deposit(
+        userSigner._address,
+        amountToDeposit,
+        0,
+        true,
+        defaultTxParams
+      );
+
+      // Depositing
+      let b = await staticAToken.deposit(
+        userSigner._address,
+        amountToDeposit,
+        0,
+        true,
+        defaultTxParams
+      );
+
+      await DRE.network.provider.send('evm_mine', []);
+
+      const aReceipt = await DRE.network.provider.send('eth_getTransactionReceipt', [a.hash]);
+      const bReceipt = await DRE.network.provider.send('eth_getTransactionReceipt', [b.hash]);
+
+      const aGas = BigNumber.from(aReceipt['gasUsed']);
+      const bGas = BigNumber.from(bReceipt['gasUsed']);
+
+      expect(aGas).to.be.gt(300000);
+      expect(bGas).to.be.lt(250000);
+
+      await DRE.network.provider.send('evm_setAutomine', [true]);
+    });
+
+    it.skip('Multiple collectAndUpdate in one block (Breaks if GasReport enabled)', async () => {
       const amountToDeposit = utils.parseEther('5');
 
       // Just preparation
@@ -209,8 +268,8 @@ describe('StaticATokenLM: aToken wrapper with static balances and liquidity mini
 
       await DRE.network.provider.send('evm_setAutomine', [false]);
 
-      let a = await staticAToken.updateRewards();
+      let a = await staticAToken.collectAndUpdateRewards();
-      let b = await staticAToken.updateRewards();
+      let b = await staticAToken.collectAndUpdateRewards();
 
       await DRE.network.provider.send('evm_mine', []);
 
@@ -249,12 +308,12 @@ describe('StaticATokenLM: aToken wrapper with static balances and liquidity mini
 
       const pendingRewards2 = await staticAToken.getClaimableRewards(userSigner._address);
 
-      await waitForTx(await staticAToken.updateRewards());
+      await waitForTx(await staticAToken.collectAndUpdateRewards());
 
       const pendingRewards3 = await staticAToken.getClaimableRewards(userSigner._address);
       const claimedRewards3 = await stkAave.balanceOf(userSigner._address);
 
-      await waitForTx(await staticAToken.updateAndClaimRewards(userSigner._address));
+      await waitForTx(await staticAToken.claimRewards(userSigner._address, true));
 
       const pendingRewards4 = await staticAToken.getClaimableRewards(userSigner._address);
       const claimedRewards4 = await stkAave.balanceOf(userSigner._address);
@@ -301,6 +360,137 @@ describe('StaticATokenLM: aToken wrapper with static balances and liquidity mini
       expect(recipientPendingRewards1).to.be.eq(0);
       expect(recipientPendingRewards2).to.be.eq(0);
     });
+
+    it('Deposit, Wait, Withdraw, claim?', async () => {
+      const amountToDeposit = utils.parseEther('5');
+      const amountToWithdraw = MAX_UINT_AMOUNT;
+
+      // Just preparation
+      await waitForTx(await weth.deposit({ value: amountToDeposit.mul(2) }));
+      await waitForTx(
+        await weth.approve(staticAToken.address, amountToDeposit.mul(2), defaultTxParams)
+      );
+
+      // Depositing
+      await waitForTx(
+        await staticAToken.deposit(userSigner._address, amountToDeposit, 0, true, defaultTxParams)
+      );
+
+      const pendingRewards1 = await staticAToken.getClaimableRewards(userSigner._address);
+
+      await advanceTimeAndBlock(60 * 60);
+
+      const pendingRewards2 = await staticAToken.getClaimableRewards(userSigner._address);
+
+      // Withdrawing all.
+      await waitForTx(
+        await staticAToken.withdraw(userSigner._address, amountToWithdraw, true, defaultTxParams)
+      );
+
+      // How will my pending look now
+      const pendingRewards3 = await staticAToken.getClaimableRewards(userSigner._address);
+
+      await waitForTx(await staticAToken.claimRewards(userSigner._address, true));
+
+      const pendingRewards4 = await staticAToken.getClaimableRewards(userSigner._address);
+      const userBalance4 = await stkAave.balanceOf(userSigner._address);
+
+      expect(pendingRewards1).to.be.eq(0);
+      expect(pendingRewards2).to.be.gt(pendingRewards1);
+      expect(pendingRewards3).to.be.gt(pendingRewards2);
+      expect(pendingRewards4).to.be.eq(0);
+      expect(userBalance4).to.be.eq(pendingRewards3);
+    });
+
+    it('Deposit, Wait, collectAndUpdate, Withdraw, claim?', async () => {
+      const amountToDeposit = utils.parseEther('5');
+      const amountToWithdraw = MAX_UINT_AMOUNT;
+
+      // Just preparation
+      await waitForTx(await weth.deposit({ value: amountToDeposit.mul(2) }));
+      await waitForTx(
+        await weth.approve(staticAToken.address, amountToDeposit.mul(2), defaultTxParams)
+      );
+
+      // Depositing
+      await waitForTx(
+        await staticAToken.deposit(userSigner._address, amountToDeposit, 0, true, defaultTxParams)
+      );
+
+      const pendingRewards1 = await staticAToken.getClaimableRewards(userSigner._address);
+
+      await advanceTimeAndBlock(60 * 60);
+      await waitForTx(await staticAToken.collectAndUpdateRewards());
+
+      const pendingRewards2 = await staticAToken.getClaimableRewards(userSigner._address);
+
+      // Withdrawing all.
+      await waitForTx(
+        await staticAToken.withdraw(userSigner._address, amountToWithdraw, true, defaultTxParams)
+      );
+
+      const pendingRewards3 = await staticAToken.getClaimableRewards(userSigner._address);
+
+      await waitForTx(await staticAToken.claimRewards(userSigner._address, true));
+
+      const pendingRewards4 = await staticAToken.getClaimableRewards(userSigner._address);
+      const userBalance4 = await stkAave.balanceOf(userSigner._address);
+
+      expect(pendingRewards1).to.be.eq(0);
+      expect(pendingRewards2).to.be.gt(pendingRewards1);
+      expect(pendingRewards3).to.be.gt(pendingRewards2);
+      expect(pendingRewards4).to.be.eq(0);
+      expect(userBalance4).to.be.eq(pendingRewards3);
+    });
+
+    it('Throw away as much as possible: Deposit, collectAndUpdate, wait, Withdraw, claim', async () => {
+      const amountToDeposit = utils.parseEther('5');
+      const amountToWithdraw = MAX_UINT_AMOUNT;
+
+      // Just preparation
+      await waitForTx(await weth.deposit({ value: amountToDeposit.mul(2) }));
+      await waitForTx(
+        await weth.approve(staticAToken.address, amountToDeposit.mul(2), defaultTxParams)
+      );
+
+      // Depositing
+      await waitForTx(
+        await staticAToken.deposit(userSigner._address, amountToDeposit, 0, true, defaultTxParams)
+      );
+
+      const pendingRewards1 = await staticAToken.getClaimableRewards(userSigner._address);
+
+      await waitForTx(await staticAToken.collectAndUpdateRewards());
+      await advanceTimeAndBlock(60 * 60);
+
+      const pendingRewards2 = await staticAToken.getClaimableRewards(userSigner._address);
+
+      // Withdrawing all.
+      await waitForTx(
+        await staticAToken.withdraw(userSigner._address, amountToWithdraw, true, defaultTxParams)
+      );
+
+      // How will my pending look now
+      const pendingRewards3 = await staticAToken.getClaimableRewards(userSigner._address);
+      const unclaimedRewards3 = await staticAToken.getUnclaimedRewards(userSigner._address);
+
+      await waitForTx(await staticAToken.claimRewards(userSigner._address, false));
+
+      const pendingRewards4 = await staticAToken.getClaimableRewards(userSigner._address);
+      const userBalance4 = await stkAave.balanceOf(userSigner._address);
+      const totClaimable4 = await staticAToken.getTotalClaimableRewards();
+      const unclaimedRewards4 = await staticAToken.getUnclaimedRewards(userSigner._address);
+
+      expect(pendingRewards1).to.be.eq(0);
+      expect(pendingRewards2).to.be.gt(0);
+      expect(pendingRewards3).to.be.gt(pendingRewards2);
+      expect(pendingRewards4).to.be.eq(0);
+      expect(userBalance4).to.be.gt(0);
+      expect(userBalance4).to.be.lt(unclaimedRewards3);
+      expect(totClaimable4).to.be.gt(0);
+      expect(totClaimable4).to.be.gt(userBalance4);
+      expect(unclaimedRewards4).to.be.eq(0);
+    });
   });
 
   it('Multiple users deposit WETH on stataWETH, wait 1 hour, update rewards, one user transfer, then claim and update rewards.', async () => {
@@ -340,9 +530,9 @@ describe('StaticATokenLM: aToken wrapper with static balances and liquidity mini
 
     // Advance time to accrue significant rewards.
     await advanceTimeAndBlock(60 * 60);
-    await staticAToken.updateRewards();
+    await staticAToken.collectAndUpdateRewards();
 
-    let staticATokenStkAaveBalInitial = await stkAave.balanceOf(staticAToken.address);
+    let staticATokenTotClaimableInitial = await staticAToken.getTotalClaimableRewards();
     let usersDataInitial = await getUserData(users, _debugUserData, { staticAToken, stkAave });
 
     await waitForTx(
@@ -357,18 +547,19 @@ describe('StaticATokenLM: aToken wrapper with static balances and liquidity mini
     await advanceTimeAndBlock(60 * 60);
 
     for (let i = 0; i < 5; i++) {
-      await waitForTx(await staticAToken.claimRewards(await users[i].getAddress()));
+      // This will claim the first half of the collected tokens (those collected at `collectAndUpdateRewards`)
+      await waitForTx(await staticAToken.claimRewards(await users[i].getAddress(), false));
     }
 
-    let staticATokenStkAaveBalAfterTransferAndClaim = await stkAave.balanceOf(staticAToken.address);
+    let staticATokenTotClaimableAfterTransferAndClaim = await staticAToken.getTotalClaimableRewards();
     let usersDataAfterTransferAndClaim = await getUserData(users, _debugUserData, {
       staticAToken,
       stkAave,
     });
 
-    await waitForTx(await staticAToken.updateRewards());
+    await waitForTx(await staticAToken.collectAndUpdateRewards());
 
-    let staticATokenStkAaveBalFinal = await stkAave.balanceOf(staticAToken.address);
+    let staticATokenTotClaimableFinal = await staticAToken.getTotalClaimableRewards();
     let usersDataFinal = await getUserData(users, _debugUserData, { staticAToken, stkAave });
 
     // Time for checks
@@ -386,6 +577,16 @@ describe('StaticATokenLM: aToken wrapper with static balances and liquidity mini
           usersDataAfterTransferAndClaim[i].staticBalance
         );
         expect(usersDataInitial[i].staticBalance).to.be.eq(usersDataFinal[i].staticBalance);
+        expect(usersDataInitial[i].pendingRewards.add(usersDataInitial[i].stkAaveBalance)).to.be.lt(
+          usersDataAfterTransferAndClaim[i].pendingRewards.add(
+            usersDataAfterTransferAndClaim[i].stkAaveBalance
+          )
+        );
+        expect(
+          usersDataAfterTransferAndClaim[i].pendingRewards.add(
+            usersDataAfterTransferAndClaim[i].stkAaveBalance
+          )
+        ).to.be.lt(usersDataFinal[i].pendingRewards.add(usersDataFinal[i].stkAaveBalance));
       }
 
       pendingRewardsSumInitial = pendingRewardsSumInitial.add(usersDataInitial[i].pendingRewards);
@@ -419,14 +620,16 @@ describe('StaticATokenLM: aToken wrapper with static balances and liquidity mini
     );
 
     // Expect there to be excess stkAave in the contract. Expect it to be dust. This ensure that everyone can claim full amount of rewards.
-    expect(pendingRewardsSumInitial).to.be.lte(staticATokenStkAaveBalInitial);
+    expect(pendingRewardsSumInitial).to.be.lte(staticATokenTotClaimableInitial);
-    expect(staticATokenStkAaveBalInitial.sub(pendingRewardsSumInitial)).to.be.lte(DUST);
+    expect(staticATokenTotClaimableInitial.sub(pendingRewardsSumInitial)).to.be.lte(DUST);
 
-    expect(pendingRewardsSumAfter).to.be.lte(staticATokenStkAaveBalAfterTransferAndClaim);
+    expect(pendingRewardsSumAfter).to.be.lte(staticATokenTotClaimableAfterTransferAndClaim);
-    expect(staticATokenStkAaveBalAfterTransferAndClaim.sub(pendingRewardsSumAfter)).to.be.lte(DUST);
+    expect(staticATokenTotClaimableAfterTransferAndClaim.sub(pendingRewardsSumAfter)).to.be.lte(
+      DUST
+    );
 
-    expect(pendingRewardsSumFinal).to.be.lte(staticATokenStkAaveBalFinal);
+    expect(pendingRewardsSumFinal).to.be.lte(staticATokenTotClaimableFinal);
-    expect(staticATokenStkAaveBalFinal.sub(pendingRewardsSumFinal)).to.be.lte(DUST);
+    expect(staticATokenTotClaimableFinal.sub(pendingRewardsSumFinal)).to.be.lte(DUST);
   });
 
   it('Multiple users deposit WETH on stataWETH, wait 1 hour, one user transfer, then claim and update rewards.', async () => {
@@ -468,7 +671,7 @@ describe('StaticATokenLM: aToken wrapper with static balances and liquidity mini
     // Advance time to accrue significant rewards.
     await advanceTimeAndBlock(60 * 60);
 
-    let staticATokenStkAaveBalInitial = await stkAave.balanceOf(staticAToken.address);
+    let staticATokenTotClaimableInitial = await staticAToken.getTotalClaimableRewards();
     let usersDataInitial = await getUserData(users, _debugUserData, { staticAToken, stkAave });
 
     // User 0 transfer full balance of staticATokens to user 1. This will also transfer the rewards since last update as well.
@@ -484,18 +687,19 @@ describe('StaticATokenLM: aToken wrapper with static balances and liquidity mini
     await advanceTimeAndBlock(60 * 60);
 
     for (let i = 0; i < 5; i++) {
-      await waitForTx(await staticAToken.claimRewards(await users[i].getAddress()));
+      // This will not do anything, hence there is no rewards in the current contract.
+      await waitForTx(await staticAToken.claimRewards(await users[i].getAddress(), false));
     }
 
-    let staticATokenStkAaveBalAfterTransfer = await stkAave.balanceOf(staticAToken.address);
+    let staticATokenTotClaimableAfterTransfer = await staticAToken.getTotalClaimableRewards();
     let usersDataAfterTransfer = await getUserData(users, _debugUserData, {
       staticAToken,
       stkAave,
     });
 
-    await waitForTx(await staticAToken.updateRewards());
+    await waitForTx(await staticAToken.collectAndUpdateRewards());
 
-    let staticATokenStkAaveBalFinal = await stkAave.balanceOf(staticAToken.address);
+    let staticATokenTotClaimableFinal = await staticAToken.getTotalClaimableRewards();
     let usersDataFinal = await getUserData(users, _debugUserData, { staticAToken, stkAave });
 
     // Time for checks
@@ -504,9 +708,8 @@ describe('StaticATokenLM: aToken wrapper with static balances and liquidity mini
     let pendingRewardsSumFinal = BigNumber.from(0);
     for (let i = 0; i < 5; i++) {
       expect(usersDataInitial[i].stkAaveBalance).to.be.eq(0);
-      // Everyone else than i == 1, should have no change in pending rewards.
+      expect(usersDataAfterTransfer[i].stkAaveBalance).to.be.eq(0);
-      // i == 1, will get additional rewards that have accrue
+      expect(usersDataFinal[i].stkAaveBalance).to.be.eq(0);
-      expect(usersDataAfterTransfer[i].stkAaveBalance).to.be.eq(usersDataInitial[i].pendingRewards);
       if (i > 1) {
         // Expect initial static balance == after transfer == after claiming
         expect(usersDataInitial[i].staticBalance).to.be.eq(usersDataAfterTransfer[i].staticBalance);
@@ -518,40 +721,45 @@ describe('StaticATokenLM: aToken wrapper with static balances and liquidity mini
       pendingRewardsSumFinal = pendingRewardsSumFinal.add(usersDataFinal[i].pendingRewards);
     }
 
-    // Expect user 0 to accrue zero fees after the transfer
+    expect(await staticAToken.getTotalClaimableRewards()).to.be.eq(
-    expect(usersDataAfterTransfer[0].pendingRewards).to.be.eq(0);
+      await stkAave.balanceOf(staticAToken.address)
+    );
+
+    // Another dude gets our unclaimed rewards
+    expect(usersDataInitial[0].pendingRewards).to.be.gt(usersDataAfterTransfer[0].pendingRewards);
+    expect(usersDataAfterTransfer[0].pendingRewards).to.be.eq(usersDataFinal[0].pendingRewards);
+
     expect(usersDataAfterTransfer[0].staticBalance).to.be.eq(0);
-    expect(usersDataFinal[0].pendingRewards).to.be.eq(0);
     expect(usersDataFinal[0].staticBalance).to.be.eq(0);
 
     // Expect user 1 to have received funds
     expect(usersDataAfterTransfer[1].staticBalance).to.be.eq(
       usersDataInitial[1].staticBalance.add(usersDataInitial[0].staticBalance)
     );
+
     /*
-     * Expect user 1 to have pending more than twice the rewards as the last user.
+     * Expect user 1 to have pending almost twice the rewards as the last user.
      * Note that he should have accrued this, even though he did not have 2x bal for the full time,
      * as he also received the "uncollected" rewards from user1 at the transfer.
+     * Lack of precision due to small initial diff.
      */
-    expect(usersDataFinal[1].pendingRewards).to.be.gt(usersDataFinal[2].pendingRewards.mul(2));
+    expect(usersDataFinal[1].pendingRewards).to.be.gt(
-    // Expect his total fees to be almost twice as large. Because of the small initial diff
+      usersDataFinal[2].pendingRewards.mul(195).div(100)
-    expect(usersDataFinal[1].pendingRewards.add(usersDataFinal[1].stkAaveBalance)).to.be.gt(
-      usersDataFinal[2].pendingRewards.add(usersDataFinal[2].stkAaveBalance).mul(195).div(100)
     );
-    expect(usersDataFinal[1].pendingRewards.add(usersDataFinal[1].stkAaveBalance)).to.be.lt(
+    expect(usersDataFinal[1].pendingRewards).to.be.lt(
-      usersDataFinal[2].pendingRewards.add(usersDataFinal[2].stkAaveBalance).mul(205).div(100)
+      usersDataFinal[2].pendingRewards.mul(205).div(100)
     );
 
     // Expect there to be excess stkAave in the contract.
     // Expect it to be dust. This ensure that everyone can claim full amount of rewards.
-    expect(pendingRewardsSumInitial).to.be.lte(staticATokenStkAaveBalInitial);
+    expect(pendingRewardsSumInitial).to.be.lte(staticATokenTotClaimableInitial);
-    expect(staticATokenStkAaveBalInitial.sub(pendingRewardsSumInitial)).to.be.lte(DUST);
+    expect(staticATokenTotClaimableInitial.sub(pendingRewardsSumInitial)).to.be.lte(DUST);
 
-    expect(pendingRewardsSumAfter).to.be.lte(staticATokenStkAaveBalAfterTransfer);
+    expect(pendingRewardsSumAfter).to.be.lte(staticATokenTotClaimableAfterTransfer);
-    expect(staticATokenStkAaveBalAfterTransfer.sub(pendingRewardsSumAfter)).to.be.lte(DUST);
+    expect(staticATokenTotClaimableAfterTransfer.sub(pendingRewardsSumAfter)).to.be.lte(DUST);
 
-    expect(pendingRewardsSumFinal).to.be.lte(staticATokenStkAaveBalFinal);
+    expect(pendingRewardsSumFinal).to.be.lte(staticATokenTotClaimableFinal);
-    expect(staticATokenStkAaveBalFinal.sub(pendingRewardsSumFinal)).to.be.lte(DUST); // How small should we say dust is?
+    expect(staticATokenTotClaimableFinal.sub(pendingRewardsSumFinal)).to.be.lte(DUST); // How small should we say dust is?
   });
 
   it('Mass deposit, then mass claim', async () => {
@@ -580,17 +788,22 @@ describe('StaticATokenLM: aToken wrapper with static balances and liquidity mini
 
     // Advance time to accrue significant rewards.
     await advanceTimeAndBlock(60 * 60);
-    await waitForTx(await staticAToken.updateRewards());
+    await waitForTx(await staticAToken.collectAndUpdateRewards());
+
+    let pendingRewards: BigNumber[] = [];
 
     for (let i = 0; i < users.length; i++) {
       const pendingReward = await staticAToken.getClaimableRewards(await users[i].getAddress());
-      await waitForTx(await staticAToken.claimRewards(await users[i].getAddress()));
+      pendingRewards.push(pendingReward);
-      expect(await stkAave.balanceOf(await users[i].getAddress())).to.be.eq(pendingReward);
+    }
+    for (let i = 0; i < users.length; i++) {
+      await waitForTx(await staticAToken.claimRewards(await users[i].getAddress(), false));
+      expect(await stkAave.balanceOf(await users[i].getAddress())).to.be.eq(pendingRewards[i]);
     }
     expect(await stkAave.balanceOf(staticAToken.address)).to.be.lt(DUST);
   });
 
-  it('Multiple deposits, withdraws and claims', async () => {
+  it('mass deposits, mass withdraws and mass claims', async () => {
     const amountToDeposit = utils.parseEther('1.135359735917531199'); // 18 decimals should be the worst here //1.135359735917531199
     const users = await DRE.ethers.getSigners();
 
@@ -622,9 +835,8 @@ describe('StaticATokenLM: aToken wrapper with static balances and liquidity mini
       );
 
       const pendingReward = await staticAToken.getClaimableRewards(await users[i].getAddress());
-      await waitForTx(await staticAToken.updateAndClaimRewards(await users[i].getAddress()));
+      await waitForTx(await staticAToken.claimRewards(await users[i].getAddress(), true));
       expect(await stkAave.balanceOf(await users[i].getAddress())).to.be.eq(pendingReward);
     }
-    expect(await stkAave.balanceOf(staticAToken.address)).to.be.lt(DUST);
   });
 });

test-suites/test-aave/mainnet/static-atoken-lm/static-atoken-liquidity-mining.spec.ts

@@ -67,6 +67,7 @@ type tBalancesInvolved = {
   staticATokenStkAaveBalance: BigNumber;
   staticATokenUnderlyingBalance: BigNumber;
   staticATokenScaledBalanceAToken: BigNumber;
+  staticATokenTotalClaimableRewards: BigNumber;
   userStkAaveBalance: BigNumber;
   userATokenBalance: BigNumber;
   userScaledBalanceAToken: BigNumber;
@@ -108,6 +109,7 @@ const getContext = async ({
   staticATokenStkAaveBalance: await stkAave.balanceOf(staticAToken.address),
   staticATokenUnderlyingBalance: await underlying.balanceOf(staticAToken.address),
   staticATokenScaledBalanceAToken: await aToken.scaledBalanceOf(staticAToken.address),
+  staticATokenTotalClaimableRewards: await staticAToken.getTotalClaimableRewards(),
   userStaticATokenBalance: await staticAToken.balanceOf(user),
   userStkAaveBalance: await stkAave.balanceOf(user),
   userATokenBalance: await aToken.balanceOf(user),
@@ -214,9 +216,13 @@ describe('StaticATokenLM: aToken wrapper with static balances and liquidity mini
     const ctxtAfterWithdrawal = await getContext(ctxtParams);
 
     // Claiming the rewards
-    await waitForTx(await staticAToken.claimRewards(userSigner._address));
+    await waitForTx(await staticAToken.claimRewards(userSigner._address, false));
 
-    const ctxtAfterClaim = await getContext(ctxtParams);
+    const ctxtAfterClaimNoForce = await getContext(ctxtParams);
+
+    await waitForTx(await staticAToken.claimRewards(userSigner._address, true));
+
+    const ctxtAfterClaimForce = await getContext(ctxtParams);
 
     // Check that scaledAToken balance is equal to the static aToken supply at every stage.
     expect(ctxtInitial.staticATokenScaledBalanceAToken).to.be.eq(ctxtInitial.staticATokenSupply);
@@ -226,8 +232,8 @@ describe('StaticATokenLM: aToken wrapper with static balances and liquidity mini
     expect(ctxtAfterWithdrawal.staticATokenScaledBalanceAToken).to.be.eq(
       ctxtAfterWithdrawal.staticATokenSupply
     );
-    expect(ctxtAfterClaim.staticATokenScaledBalanceAToken).to.be.eq(
+    expect(ctxtAfterClaimNoForce.staticATokenScaledBalanceAToken).to.be.eq(
-      ctxtAfterClaim.staticATokenSupply
+      ctxtAfterClaimNoForce.staticATokenSupply
     );
 
     expect(ctxtAfterDeposit.staticATokenATokenBalance).to.be.eq(
@@ -253,18 +259,29 @@ describe('StaticATokenLM: aToken wrapper with static balances and liquidity mini
     expect(ctxtAfterWithdrawal.staticATokenATokenBalance).to.be.eq(0);
     expect(ctxtAfterWithdrawal.staticATokenSupply).to.be.eq(0);
     expect(ctxtAfterWithdrawal.staticATokenUnderlyingBalance).to.be.eq(0);
+    expect(ctxtAfterWithdrawal.staticATokenStkAaveBalance).to.be.eq(0);
 
-    // Check with possible rounding error.
+    // Check with possible rounding error. Ahhh, it is because we have not claimed the shit after withdraw
-    expect(ctxtAfterWithdrawal.staticATokenStkAaveBalance).to.be.gte(
+    expect(ctxtAfterWithdrawal.staticATokenTotalClaimableRewards).to.be.gte(
       ctxtAfterWithdrawal.userPendingRewards
     );
-    expect(ctxtAfterWithdrawal.staticATokenStkAaveBalance).to.be.lte(
+
+    expect(ctxtAfterWithdrawal.staticATokenTotalClaimableRewards).to.be.lte(
       ctxtAfterWithdrawal.userPendingRewards.add(1)
     );
     expect(ctxtAfterWithdrawal.userStkAaveBalance).to.be.eq(0);
 
-    expect(ctxtAfterClaim.userStkAaveBalance).to.be.eq(ctxtAfterWithdrawal.userPendingRewards);
+    expect(ctxtAfterClaimNoForce.userStkAaveBalance).to.be.eq(0);
-    expect(ctxtAfterClaim.staticATokenStkAaveBalance).to.be.lte(1);
+    expect(ctxtAfterClaimNoForce.staticATokenStkAaveBalance).to.be.eq(0);
+
+    expect(ctxtAfterClaimForce.userStkAaveBalance).to.be.eq(
+      ctxtAfterClaimNoForce.userPendingRewards
+    );
+    expect(ctxtAfterClaimForce.staticATokenStkAaveBalance).to.be.eq(
+      ctxtAfterClaimNoForce.staticATokenTotalClaimableRewards.sub(
+        ctxtAfterClaimNoForce.userPendingRewards
+      )
+    );
   });
 
   it('Deposit WETH on stataWETH and then withdraw some balance in underlying', async () => {
@@ -293,9 +310,15 @@ describe('StaticATokenLM: aToken wrapper with static balances and liquidity mini
     const ctxtAfterWithdrawal = await getContext(ctxtParams);
 
     // Claim
-    await waitForTx(await staticAToken.claimRewards(userSigner._address));
+    await waitForTx(await staticAToken.claimRewards(userSigner._address, false));
     const ctxtAfterClaim = await getContext(ctxtParams);
 
+    await waitForTx(await staticAToken.collectAndUpdateRewards());
+    const ctxtAfterUpdate = await getContext(ctxtParams);
+
+    await waitForTx(await staticAToken.claimRewards(userSigner._address, false));
+    const ctxtAfterClaim2 = await getContext(ctxtParams);
+
     expect(ctxtInitial.userStaticATokenBalance).to.be.eq(0);
     expect(ctxtInitial.staticATokenSupply).to.be.eq(0);
     expect(ctxtInitial.staticATokenUnderlyingBalance).to.be.eq(0);
@@ -317,7 +340,27 @@ describe('StaticATokenLM: aToken wrapper with static balances and liquidity mini
       ctxtAfterDeposit.userStaticATokenBalance.sub(amountToWithdraw)
     );
 
-    expect(ctxtAfterClaim.userStkAaveBalance).to.be.eq(ctxtAfterWithdrawal.userPendingRewards);
+    expect(ctxtAfterUpdate.userStkAaveBalance).to.be.eq(0);
+    expect(ctxtAfterClaim2.userStkAaveBalance).to.be.eq(ctxtAfterUpdate.userPendingRewards);
+    expect(ctxtAfterClaim2.userPendingRewards).to.be.gt(0);
+
+    // Check that rewards are always covered
+    expect(ctxtInitial.staticATokenTotalClaimableRewards).to.be.gte(ctxtInitial.userPendingRewards);
+    expect(ctxtAfterDeposit.staticATokenTotalClaimableRewards).to.be.gte(
+      ctxtAfterDeposit.userPendingRewards
+    );
+    expect(ctxtAfterWithdrawal.staticATokenTotalClaimableRewards).to.be.gte(
+      ctxtAfterWithdrawal.userPendingRewards
+    );
+    expect(ctxtAfterClaim.staticATokenTotalClaimableRewards).to.be.gte(
+      ctxtAfterClaim.userPendingRewards
+    );
+    expect(ctxtAfterUpdate.staticATokenTotalClaimableRewards).to.be.gte(
+      ctxtAfterUpdate.userPendingRewards
+    );
+    expect(ctxtAfterClaim2.staticATokenTotalClaimableRewards).to.be.gte(
+      ctxtAfterClaim2.userPendingRewards
+    );
   });
 
   it('Deposit WETH on stataWETH and then withdraw all the balance in aToken', async () => {
@@ -719,10 +762,6 @@ describe('StaticATokenLM: aToken wrapper with static balances and liquidity mini
 
     const ctxtAfterWithdrawal = await getContext(ctxtParams);
 
-    // Claim
-    await waitForTx(await staticAToken.claimRewards(userSigner._address));
-    const ctxtAfterClaim = await getContext(ctxtParams);
-
     expect(ctxtBeforeWithdrawal.userATokenBalance).to.be.eq(0);
     expect(ctxtBeforeWithdrawal.staticATokenATokenBalance).to.be.eq(amountToDeposit);
     expect(ctxtAfterWithdrawal.userATokenBalance).to.be.eq(amountToWithdraw);
@@ -736,7 +775,6 @@ describe('StaticATokenLM: aToken wrapper with static balances and liquidity mini
     );
 
     expect(ctxtAfterWithdrawal.userStkAaveBalance).to.be.eq(0);
-    expect(ctxtAfterClaim.userStkAaveBalance).to.be.eq(ctxtAfterWithdrawal.userPendingRewards);
   });
 
   it('Withdraw using metaWithdraw()', async () => {
@@ -985,7 +1023,7 @@ describe('StaticATokenLM: aToken wrapper with static balances and liquidity mini
     const ctxtAfterWithdrawal = await getContext(ctxtParams);
 
     // Claim
-    await waitForTx(await staticAToken.claimRewards(user2Signer._address));
+    await waitForTx(await staticAToken.claimRewards(user2Signer._address, true));
     const ctxtAfterClaim = await getContext(ctxtParams);
 
     // Checks
@@ -1000,17 +1038,25 @@ describe('StaticATokenLM: aToken wrapper with static balances and liquidity mini
     );
     expect(ctxtAfterTransfer.userStaticATokenBalance).to.be.eq(0);
     expect(ctxtAfterTransfer.userPendingRewards).to.be.eq(0);
-    expect(ctxtAfterTransfer.user2PendingRewards).to.be.eq(0);
+    expect(ctxtAfterTransfer.user2PendingRewards).to.be.gt(0);
     expect(ctxtAfterWithdrawal.staticATokenSupply).to.be.eq(0);
     expect(ctxtAfterWithdrawal.staticATokenATokenBalance).to.be.eq(0);
     expect(ctxtAfterWithdrawal.userPendingRewards).to.be.eq(0);
-    expect(ctxtAfterWithdrawal.user2PendingRewards).to.be.lte(
+    expect(ctxtAfterWithdrawal.staticATokenTotalClaimableRewards).to.be.gte(
-      ctxtAfterWithdrawal.staticATokenStkAaveBalance
+      ctxtAfterWithdrawal.user2PendingRewards
     );
-    expect(ctxtAfterClaim.user2StkAaveBalance).to.be.eq(ctxtAfterWithdrawal.user2PendingRewards);
+    console.log('All the way down here');
+
+    console.log(`${formatEther(ctxtAfterClaim.staticATokenTotalClaimableRewards)}`);
+    console.log(`${formatEther(ctxtAfterClaim.user2StkAaveBalance)}`);
+    console.log(`${formatEther(ctxtAfterClaim.staticATokenStkAaveBalance)}`);
+
     expect(ctxtAfterClaim.userStkAaveBalance).to.be.eq(0);
+    expect(ctxtAfterClaim.user2StkAaveBalance).to.be.eq(ctxtAfterWithdrawal.user2PendingRewards);
     expect(ctxtAfterClaim.staticATokenStkAaveBalance).to.be.eq(
-      ctxtAfterWithdrawal.staticATokenStkAaveBalance.sub(ctxtAfterWithdrawal.user2PendingRewards)
+      ctxtAfterWithdrawal.staticATokenTotalClaimableRewards.sub(
+        ctxtAfterWithdrawal.user2PendingRewards
+      )
     );
     // Expect dust to be left in the contract
     expect(ctxtAfterClaim.staticATokenStkAaveBalance).to.be.lt(5);