fix: dividend external withdrawal

Author: amarinkovic

src/facets/TokenizedVaultFacet.sol

@@ -184,6 +184,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/libs/LibTokenizedVault.sol

@@ -246,7 +246,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 = LibTokenizedVault._internalTokenSupply(_dividendTokenId);
+            uint256 adjustedDividendAmount = _amount;
+            if (_to == _dividendTokenId) {
+                // withdrawn dividend was adjusted for the previous holder in this case,
+                // therfore 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
@@ -305,7 +313,8 @@ library LibTokenizedVault {
 
         address tokenAddress = LibHelpers._getAddressFromId(_tokenId);
 
-        uint256 depositTotal = s.depositTotal[_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
@@ -318,7 +327,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);
         }
 
@@ -327,7 +336,8 @@ library LibTokenizedVault {
             revert RebasingInterestInsufficient(_tokenId, _amount, accruedAmount);
         }
 
-        s.tokenBalances[_tokenId][_tokenId] += _amount;
-        s.depositTotal[_tokenId] += _amount;
+        // 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,7 +79,7 @@ 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(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

test/T03TokenizedVault.t.sol

@@ -1186,20 +1186,24 @@ contract T03TokenizedVaultTest is D03ProtocolDefaults, MockAccounts {
         changePrank(sm);
         nayms.distributeAccruedInterest(wethId, 1 ether, makeId(LC.OBJECT_TYPE_DIVIDEND, bytes20("0x1")));
 
+        changePrank(david);
+        writeTokenBalance(david.addr, naymsAddress, wethAddress, 4 ether);
+
         changePrank(alice.addr);
-        writeTokenBalance(alice.addr, naymsAddress, wethAddress, depositAmount);
+        writeTokenBalance(alice.addr, naymsAddress, wethAddress, 2 ether);
         nayms.externalDeposit(wethAddress, 1 ether);
         assertEq(nayms.internalBalanceOf(alice.entityId, wethId), 1 ether, "acc0EntityId wethId balance should INCREASE (mint)");
 
         changePrank(bob.addr);
-        writeTokenBalance(bob.addr, naymsAddress, wethAddress, depositAmount);
+        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(4 ether));
+        changePrank(david);
+        IERC20(wethAddress).transfer(naymsAddress, 2 ether);
         assertEq(nayms.accruedInterest(wethId), 2 ether, "Accrued interest should increase");
 
         changePrank(sm);
@@ -1217,6 +1221,31 @@ contract T03TokenizedVaultTest is D03ProtocolDefaults, MockAccounts {
         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,