Adjust withdrawn amount when same coin is used to pay out a dividends

package.json

@@ -48,7 +48,7 @@
     "eslint-plugin-prettier": "^4.2.1",
     "eslint-plugin-promise": "^6.0.0",
     "ethers": "^5.6.9",
-    "gemforge": "^2.9.1",
+    "gemforge": "^2.14.0",
     "glob": "^8.0.3",
     "prettier": "^2.7.1",
     "prettier-plugin-solidity": "^1.0.0-beta.19",

src/facets/TokenizedVaultFacet.sol

@@ -174,6 +174,6 @@ contract TokenizedVaultFacet is Modifiers, ReentrancyGuard {
         // The _claimRebasingInterest method verifies the token is valid, and that there is available interest.
         // No need to do it again.
         LibTokenizedVault._claimRebasingInterest(_tokenId, _amount);
-        LibTokenizedVault._payDividend(_guid, _tokenId, _tokenId, _tokenId, _amount);
+        LibTokenizedVault._payDividend(_guid, LibAdmin._getSystemId(), _tokenId, _tokenId, _amount);
     }
 }

src/init/DividendPatchInitDiamond.sol

@@ -0,0 +1,18 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.20;
+
+import { AppStorage, LibAppStorage } from "../shared/AppStorage.sol";
+
+contract DividendPatchInitDiamond {
+    function init() external {
+        AppStorage storage s = LibAppStorage.diamondStorage();
+
+        bytes32 _tokenId = 0x59d9356e565ab3a36dd77763fc0d87feaf85508c000000000000000000000000; // USDM
+        bytes32 _dividendTokenId = 0x59d9356e565ab3a36dd77763fc0d87feaf85508c000000000000000000000000; // USDM
+        bytes32 _ownerId = 0x454e54495459000000000000653eeaf2a1e51089765e3c6291e780a640085943; // ILW
+
+        s.tokenSupply[_tokenId] = 533337609745764857001482;
+        s.totalDividends[_tokenId][_dividendTokenId] = 16406802958540022191510;
+        s.withdrawnDividendPerOwner[_tokenId][_dividendTokenId][_ownerId] = 16171508827360588570817;
+    }
+}

src/libs/LibTokenizedVault.sol

@@ -147,6 +147,7 @@ library LibTokenizedVault {
 
         _withdrawAllDividends(_from, _tokenId);
         _normalizeDividendsBurn(_tokenId, _amount);
+
         s.tokenSupply[_tokenId] -= _amount;
         s.tokenBalances[_tokenId][_from] -= _amount;
 
@@ -184,9 +185,19 @@ library LibTokenizedVault {
         uint256 withdrawnSoFar = s.withdrawnDividendPerOwner[_tokenId][_dividendTokenId][_ownerId];
 
         uint256 withdrawableDividend = _getWithdrawableDividendAndDeductionMath(amountOwned, supply, totalDividend, withdrawnSoFar);
+
         if (withdrawableDividend > 0) {
             // Bump the withdrawn dividends for the owner
-            s.withdrawnDividendPerOwner[_tokenId][_dividendTokenId][_ownerId] += withdrawableDividend;
+            /// Special Case: (_tokenId == _dividendTokenId), i.e distributing accrued interest for rebasing coins like USDM
+            /// withdrawnDividendPerOwner should be adjusted before tha update, so that the user cannot claim additional dividend based on the amount he just received as dividend
+            /// dividend is calculated based on a ratio between users balance and the total, but in this case claiming the dividend his balance increases and
+            /// thus his share of the total increases as well, which entitles him to claim more of the dividend, potentially draining out the entirety of it if repeated infinitely
+            if (_tokenId == _dividendTokenId) {
+                uint256 withdrawableDividendAdjusted = _getWithdrawableDividendAndDeductionMath(amountOwned + withdrawableDividend, supply, totalDividend, withdrawnSoFar);
+                s.withdrawnDividendPerOwner[_tokenId][_dividendTokenId][_ownerId] += withdrawableDividendAdjusted;
+            } else {
+                s.withdrawnDividendPerOwner[_tokenId][_dividendTokenId][_ownerId] += withdrawableDividend;
+            }
 
             // Move the dividend
             s.tokenBalances[_dividendTokenId][dividendBankId] -= withdrawableDividend;
@@ -237,7 +248,15 @@ library LibTokenizedVault {
             // issue dividend. if you are owed dividends on the _dividendTokenId, they will be collected
             // Check for possible infinite loop, but probably not
             _internalTransfer(_from, dividendBankId, _dividendTokenId, _amount);
-            s.totalDividends[_to][_dividendTokenId] += _amount;
+            uint256 tokenSupply = _internalTokenSupply(_dividendTokenId);
+            uint256 adjustedDividendAmount = _amount;
+            if (_to == _dividendTokenId) {
+                // withdrawn dividend was adjusted for the previous holder in this case,
+                // therefore dividend amount should be increased in order to give existing token holders the correct amount
+                adjustedDividendAmount = (_amount * (tokenSupply)) / (tokenSupply - _amount);
+            }
+
+            s.totalDividends[_to][_dividendTokenId] += adjustedDividendAmount;
 
             // keep track of the dividend denominations
             // if dividend has not yet been issued in this token, add it to the list and update mappings
@@ -289,7 +308,7 @@ library LibTokenizedVault {
 
         address tokenAddress = LibHelpers._getAddressFromId(_tokenId);
 
-        uint256 depositTotal = s.depositTotal[_tokenId];
+        uint256 depositTotal = s.tokenSupply[_tokenId];
         uint256 total = LibERC20.balanceOf(tokenAddress, address(this));
 
         // If the Nayms balance of the rebasing token has decreased and is lower than the deposit total, revert
@@ -302,7 +321,7 @@ library LibTokenizedVault {
     function _claimRebasingInterest(bytes32 _tokenId, uint256 _amount) internal {
         AppStorage storage s = LibAppStorage.diamondStorage();
 
-        if (s.depositTotal[_tokenId] == 0) {
+        if (s.tokenSupply[_tokenId] == 0) {
             revert RebasingInterestNotInitialized(_tokenId);
         }
 
@@ -311,7 +330,6 @@ library LibTokenizedVault {
             revert RebasingInterestInsufficient(_tokenId, _amount, accruedAmount);
         }
 
-        s.tokenBalances[_tokenId][_tokenId] += _amount;
-        s.depositTotal[_tokenId] += _amount;
+        _internalMint(LibAdmin._getSystemId(), _tokenId, _amount);
     }
 }

src/libs/LibTokenizedVaultIO.sol

@@ -34,9 +34,6 @@ library LibTokenizedVaultIO {
         // Only mint what has been collected.
         LibTokenizedVault._internalMint(_receiverId, internalTokenId, mintAmount);
 
-        AppStorage storage s = LibAppStorage.diamondStorage();
-        s.depositTotal[internalTokenId] += mintAmount;
-
         // emit event
         emit ExternalDeposit(_receiverId, _externalTokenAddress, mintAmount);
     }
@@ -55,9 +52,6 @@ library LibTokenizedVaultIO {
         // transfer AFTER burn
         LibERC20.transfer(_externalTokenAddress, _receiver, _amount);
 
-        AppStorage storage s = LibAppStorage.diamondStorage();
-        s.depositTotal[internalTokenId] -= _amount;
-
         // emit event
         emit ExternalWithdraw(_entityId, _receiver, _externalTokenAddress, _amount);
     }

src/shared/AppStorage.sol

@@ -79,8 +79,8 @@ struct AppStorage {
     mapping(string tokenSymbol => bytes32 objectId) tokenSymbolObjectId; // reverse mapping token symbol => object ID, to ensure symbol uniqueness
     mapping(bytes32 entityId => mapping(uint256 feeScheduleTypeId => FeeSchedule)) feeSchedules; // map entity ID to a fee schedule type and then to array of FeeReceivers (feeScheduleType (1-premium, 2-trading, n-others))
     mapping(bytes32 objectId => uint256 minimumSell) objectMinimumSell; // map object ID to minimum sell amount
-    mapping(bytes32 objectId => uint256) depositTotal; // total amount deposited into contract, for rebasing tokens support
-    mapping(address userAddress => EntityApproval) selfOnboarding; // map address => { entityId, roleId }
+    mapping(bytes32 objectId => uint256) depositTotal; // note: DEPRECATED: total amount deposited into contract, for rebasing tokens support
+    mapping(address userAddress => EntityApproval) selfOnboarding; // note: DEPRECATED
     /// Staking
     mapping(bytes32 entityId => StakingConfig) stakingConfigs; // StakingConfig for an entity
     mapping(bytes32 vTokenId => mapping(bytes32 stakerId => uint256 balance)) stakeBalance; // [vTokenId][ownerId] balance at interval
@@ -90,6 +90,7 @@ struct AppStorage {
     mapping(bytes32 vTokenId => bytes32 denomination) stakingDistributionDenomination; // [vTokenId] Reward currency
     mapping(bytes32 entityId => mapping(bytes32 stakerId => uint64 interval)) stakingSynced; // last interval when data was synced into storage for staker
     mapping(bytes32 vTokenId => mapping(bytes32 stakerId => uint256 balance)) stakeBalanceAdded; // raw balance staked at an interval, withouth any boost included, only for reading future intervals (to calculate the total boosted balance)
+    // mapping(uint256 => bool) initComplete; // think about adding this in the future
 }
 
 /// Staking-Related Mappings

test/T03TokenizedVault.t.sol

@@ -1179,33 +1179,70 @@ contract T03TokenizedVaultTest is D03ProtocolDefaults, MockAccounts {
     }
 
     function testRebasingTokenInterest() public {
-        bytes32 acc0EntityId = nayms.getEntity(account0Id);
-        nayms.assignRole(em.id, acc0EntityId, LC.ROLE_ENTITY_MANAGER);
-
-        assertEq(nayms.internalBalanceOf(account0Id, wethId), 0, "acc0EntityId wethId balance should start at 0");
-
         vm.expectRevert(abi.encodeWithSelector(RebasingInterestNotInitialized.selector, wethId));
         changePrank(sm);
         nayms.distributeAccruedInterest(wethId, 1 ether, makeId(LC.OBJECT_TYPE_DIVIDEND, bytes20("0x1")));
 
-        changePrank(account0);
-        writeTokenBalance(account0, naymsAddress, wethAddress, depositAmount);
+        changePrank(david);
+        writeTokenBalance(david.addr, naymsAddress, wethAddress, 4 ether);
 
+        changePrank(alice.addr);
+        writeTokenBalance(alice.addr, naymsAddress, wethAddress, 2 ether);
         nayms.externalDeposit(wethAddress, 1 ether);
-        assertEq(nayms.internalBalanceOf(acc0EntityId, wethId), 1 ether, "acc0EntityId wethId balance should INCREASE (mint)");
+        assertEq(nayms.internalBalanceOf(alice.entityId, wethId), 1 ether, "acc0EntityId wethId balance should INCREASE (mint)");
+
+        changePrank(bob.addr);
+        writeTokenBalance(bob.addr, naymsAddress, wethAddress, 2 ether);
+        nayms.externalDeposit(wethAddress, 1 ether);
+        assertEq(nayms.internalBalanceOf(bob.entityId, wethId), 1 ether, "bobEntityId wethId balance should INCREASE (mint)");
+
+        vm.warp(2 weeks);
 
         assertEq(nayms.accruedInterest(wethId), 0, "Accrued interest should be zero");
-        vm.mockCall(wethAddress, abi.encodeWithSelector(IERC20.balanceOf.selector), abi.encode(2 ether));
-        assertEq(nayms.accruedInterest(wethId), 1 ether, "Accrued interest should increase");
+        changePrank(david);
+        IERC20(wethAddress).transfer(naymsAddress, 2 ether);
+        assertEq(nayms.accruedInterest(wethId), 2 ether, "Accrued interest should increase");
 
         changePrank(sm);
-        vm.expectRevert(abi.encodeWithSelector(RebasingInterestInsufficient.selector, wethId, 5 ether, 1 ether));
+        vm.expectRevert(abi.encodeWithSelector(RebasingInterestInsufficient.selector, wethId, 5 ether, 2 ether));
         nayms.distributeAccruedInterest(wethId, 5 ether, makeId(LC.OBJECT_TYPE_DIVIDEND, bytes20("0x1")));
 
-        nayms.distributeAccruedInterest(wethId, 1 ether, makeId(LC.OBJECT_TYPE_DIVIDEND, bytes20("0x1")));
+        nayms.distributeAccruedInterest(wethId, 2 ether, makeId(LC.OBJECT_TYPE_DIVIDEND, bytes20("0x1")));
 
-        nayms.withdrawDividend(acc0EntityId, wethId, wethId);
-        assertEq(nayms.internalBalanceOf(acc0EntityId, wethId), 2 ether, "acc0EntityId wethId balance should INCREASE (mint)");
+        changePrank(alice.addr);
+        nayms.withdrawDividend(alice.entityId, wethId, wethId);
+        assertEq(nayms.internalBalanceOf(alice.entityId, wethId), 2 ether, "Alice's wethId balance should INCREASE (mint)".red());
+        assertEq(nayms.getWithdrawableDividend(alice.entityId, wethId, wethId), 0, "alice's withdrawable divident should be zero".red());
+
+        changePrank(bob.addr);
+        nayms.withdrawDividend(bob.entityId, wethId, wethId);
+        assertEq(nayms.internalBalanceOf(bob.entityId, wethId), 2 ether, "Bob's wethId balance should INCREASE (mint)".red());
+        assertEq(nayms.getWithdrawableDividend(bob.entityId, wethId, wethId), 0, "bob's withdrawable divident should be zero".red());
+
+        vm.warp(4 weeks);
+
+        changePrank(david);
+        IERC20(wethAddress).transfer(naymsAddress, 2 ether);
+        assertEq(nayms.accruedInterest(wethId), 2 ether, "Accrued interest should increase");
+
+        changePrank(sm);
+        nayms.distributeAccruedInterest(wethId, 2 ether, makeId(LC.OBJECT_TYPE_DIVIDEND, bytes20("0x2")));
+
+        changePrank(alice.addr);
+        nayms.withdrawDividend(alice.entityId, wethId, wethId);
+        assertEq(nayms.internalBalanceOf(alice.entityId, wethId), 3 ether, "Alice's wethId balance should INCREASE (mint)".red());
+        assertEq(nayms.getWithdrawableDividend(alice.entityId, wethId, wethId), 0, "alice's withdrawable divident should be zero".red());
+
+        nayms.externalWithdrawFromEntity(alice.entityId, alice.addr, wethAddress, 3 ether);
+        assertEq(IERC20(wethAddress).balanceOf(alice.addr), 4 ether, "Alice's private balance should be increased".red());
+
+        changePrank(bob.addr);
+        nayms.withdrawDividend(bob.entityId, wethId, wethId);
+        assertEq(nayms.internalBalanceOf(bob.entityId, wethId), 3 ether, "Bob's wethId balance should INCREASE (mint)".red());
+        assertEq(nayms.getWithdrawableDividend(bob.entityId, wethId, wethId), 0, "Bob's withdrawable divident should be zero".red());
+
+        nayms.externalWithdrawFromEntity(bob.entityId, bob.addr, wethAddress, 3 ether);
+        assertEq(IERC20(wethAddress).balanceOf(bob.addr), 4 ether, "Bob's private balance should be increased".red());
     }
 
     // note withdrawAllDividends() will still succeed even if there are 0 dividends to be paid out,

yarn.lock

@@ -2678,10 +2678,10 @@ functions-have-names@^1.2.3:
   resolved "https://registry.yarnpkg.com/functions-have-names/-/functions-have-names-1.2.3.tgz#0404fe4ee2ba2f607f0e0ec3c80bae994133b834"
   integrity sha512-xckBUXyTIqT97tq2x2AMb+g163b5JFysYk0x4qxNFwbfQkmNZoiRHb6sPzI9/QV33WeuvVYBUIiD4NzNIyqaRQ==
 
-gemforge@^2.9.1:
-  version "2.9.1"
-  resolved "https://registry.yarnpkg.com/gemforge/-/gemforge-2.9.1.tgz#a4c73378764db4c6ad3ba09e566755137c6604cb"
-  integrity sha512-F6GcJm660xaDmRlYKI7yfI6AOwAfpXqvJrzI44fOF89rbQYv5YAOVHe+CslQZiM6aGdQ4KWf7aqbdxkmoUwjOg==
+gemforge@^2.14.0:
+  version "2.14.0"
+  resolved "https://registry.yarnpkg.com/gemforge/-/gemforge-2.14.0.tgz#36127a21e1a95523001def1da7427ef66188f16b"
+  integrity sha512-b9Ji69WJXXLaCvcGSsTaT+1T1eby4luZT8yS/HOaRRhtJrq8nFUbA2ZN2RqTOOo9eke+7LjWQlDlPeQLY4ANiQ==
   dependencies:
     "@solidity-parser/parser" "^0.16.1"
     bigval "^1.7.0"