Account for yieldFeeBalance in twabSupplyLimit

src/PrizeVault.sol

@@ -14,6 +14,9 @@ import { ILiquidationSource } from "pt-v5-liquidator-interfaces/ILiquidationSour
 import { PrizePool } from "pt-v5-prize-pool/PrizePool.sol";
 import { TwabController, SPONSORSHIP_ADDRESS } from "pt-v5-twab-controller/TwabController.sol";
 
+/// @dev The TWAB supply limit is the max number of shares that can be minted in the TWAB controller.
+uint256 constant TWAB_SUPPLY_LIMIT = type(uint96).max;
+
 /// @title  PoolTogether V5 Prize Vault
 /// @author G9 Software Inc.
 /// @notice The prize vault takes deposits of an asset and earns yield with the deposits through an underlying yield
@@ -249,6 +252,10 @@ contract PrizeVault is TwabERC20, Claimable, IERC4626, ILiquidationSource, Ownab
     /// @param totalSupply The total shares minted and internally accounted for by the vault
     error LossyDeposit(uint256 totalAssets, uint256 totalSupply);
 
+    /// @notice Thrown when the mint limit is exceeded after increasing an external or internal share balance.
+    /// @param excess The amount in excess over the limit
+    error MintLimitExceeded(uint256 excess);
+
     ////////////////////////////////////////////////////////////////////////////////
     // Modifiers
     ////////////////////////////////////////////////////////////////////////////////
@@ -363,28 +370,27 @@ contract PrizeVault is TwabERC20, Claimable, IERC4626, ILiquidationSource, Ownab
     }
 
     /// @inheritdoc IERC4626
-    /// @dev Considers the uint96 limit on total share supply in the TwabController
+    /// @dev Considers the TWAB mint limit
     /// @dev Returns zero if any deposit would result in a loss of assets
     /// @dev Any latent balance of assets in the prize vault will be swept in with the deposit as a part of
     /// the "dust collection strategy". This means that the max deposit must account for the latent balance
     /// by subtracting it from the max deposit available otherwise.
     function maxDeposit(address /* receiver */) public view returns (uint256) {
-        uint256 _totalSupply = totalSupply();
-        uint256 totalDebt_ = _totalDebt(_totalSupply);
-        if (totalAssets() < totalDebt_) return 0;
+        uint256 _totalDebt = totalSupply() + yieldFeeBalance;
+        if (totalAssets() < _totalDebt) return 0;
 
-        // the vault will never mint more than 1 share per asset, so no need to convert supply limit to assets
-        uint256 twabSupplyLimit_ = _twabSupplyLimit(_totalSupply);
-        uint256 _maxDeposit;
         uint256 _latentBalance = _asset.balanceOf(address(this));
         uint256 _maxYieldVaultDeposit = yieldVault.maxDeposit(address(this));
         if (_latentBalance >= _maxYieldVaultDeposit) {
             return 0;
         } else {
+            // the vault will never mint more than 1 share per asset, so no need to convert mint limit to assets
+            uint256 _depositLimit = _mintLimit(_totalDebt);
+            uint256 _maxDeposit;
             unchecked {
                 _maxDeposit = _maxYieldVaultDeposit - _latentBalance;
             }
-            return twabSupplyLimit_ < _maxDeposit ? twabSupplyLimit_ : _maxDeposit;
+            return _depositLimit < _maxDeposit ? _depositLimit : _maxDeposit;
         }
     }
 
@@ -568,7 +574,7 @@ contract PrizeVault is TwabERC20, Claimable, IERC4626, ILiquidationSource, Ownab
     /// @notice Returns the total assets that are owed to share holders and any other internal balances.
     /// @return The total asset debt of the vault
     function totalDebt() public view returns (uint256) {
-        return _totalDebt(totalSupply());
+        return totalSupply() + yieldFeeBalance;
     }
 
     ////////////////////////////////////////////////////////////////////////////////
@@ -624,28 +630,27 @@ contract PrizeVault is TwabERC20, Claimable, IERC4626, ILiquidationSource, Ownab
     /// @dev Returns the liquid amount of `_tokenOut` minus any yield fees.
     /// @dev Supports the liquidation of either assets or prize vault shares.
     function liquidatableBalanceOf(address _tokenOut) public view returns (uint256) {
-        uint256 _totalSupply = totalSupply();
+        uint256 _totalDebt = totalSupply() + yieldFeeBalance;
         uint256 _maxAmountOut;
         if (_tokenOut == address(this)) {
-            // Liquidation of vault shares is capped to the TWAB supply limit.
-            _maxAmountOut = _twabSupplyLimit(_totalSupply);
+            // Liquidation of vault shares is capped to the mint limit.
+            _maxAmountOut = _mintLimit(_totalDebt);
         } else if (_tokenOut == address(_asset)) {
             // Liquidation of yield assets is capped at the max yield vault withdraw plus any latent balance.
             _maxAmountOut = _maxYieldVaultWithdraw() + _asset.balanceOf(address(this));
         } else {
             return 0;
         }
 
-        // The liquid yield is computed by taking the available yield balance and multiplying it
-        // by (1 - yieldFeePercentage), rounding down, to ensure that enough yield is left for the
-        // yield fee.
-        uint256 _liquidYield = 
-            _availableYieldBalance(totalAssets(), _totalDebt(_totalSupply))
-            .mulDiv(FEE_PRECISION - yieldFeePercentage, FEE_PRECISION);
-
         // The liquid yield is limited by the max that can be minted or withdrawn, depending on
         // `_tokenOut`.
-        return _liquidYield >= _maxAmountOut ? _maxAmountOut : _liquidYield;
+        uint256 _availableYield = _availableYieldBalance(totalAssets(), _totalDebt);
+        uint256 _liquidYield = _availableYield >= _maxAmountOut ? _maxAmountOut : _availableYield;
+
+        // The final balance is computed by taking the liquid yield and multiplying it by
+        // (1 - yieldFeePercentage), rounding down, to ensure that enough yield is left for
+        // the yield fee.
+        return _liquidYield.mulDiv(FEE_PRECISION - yieldFeePercentage, FEE_PRECISION);
     }
 
     /// @inheritdoc ILiquidationSource
@@ -676,8 +681,9 @@ contract PrizeVault is TwabERC20, Claimable, IERC4626, ILiquidationSource, Ownab
         }
 
         // Increase yield fee balance:
+        uint256 _newYieldFeeBalance = yieldFeeBalance + _yieldFee;
         if (_yieldFee > 0) {
-            yieldFeeBalance += _yieldFee;
+            yieldFeeBalance = _newYieldFeeBalance;
         }
 
         // Mint or withdraw amountOut to `_receiver`:
@@ -689,6 +695,8 @@ contract PrizeVault is TwabERC20, Claimable, IERC4626, ILiquidationSource, Ownab
             revert LiquidationTokenOutNotSupported(_tokenOut);
         }
 
+        _enforceMintLimit(totalSupply(), _newYieldFeeBalance);
+
         emit TransferYieldOut(msg.sender, _tokenOut, _receiver, _amountOut, _yieldFee);
 
         return "";
@@ -777,22 +785,25 @@ contract PrizeVault is TwabERC20, Claimable, IERC4626, ILiquidationSource, Ownab
         return (false, 0);
     }
 
-    /// @notice Returns the total assets that are owed to share holders and any other internal balances.
-    /// @dev The yield fee balance is included since it's cheaper to keep track of those shares
-    ///      internally instead of doing an additional TWAB mint on every liquidation.
-    /// @param _totalSupply The total share supply of the vault
-    /// @return The total asset debt of the vault
-    function _totalDebt(uint256 _totalSupply) internal view returns (uint256) {
-        return _totalSupply + yieldFeeBalance;
+    /// @notice Returns the shares that can be minted without exceeding the TwabController supply limit.
+    /// @dev The TwabController limits the total supply for each vault.
+    /// @param _existingShares The current allocated prize vault shares (internal and external)
+    /// @return The remaining shares that can be minted without exceeding TWAB limits
+    function _mintLimit(uint256 _existingShares) internal pure returns (uint256) {
+        return TWAB_SUPPLY_LIMIT - _existingShares;
     }
 
-    /// @notice Returns the remaining supply that can be minted without exceeding the TwabController limits.
-    /// @dev The TwabController limits the total supply for each vault to uint96
-    /// @param _totalSupply The total share supply of the vault
-    /// @return The remaining supply that can be minted without exceeding TWAB limits
-    function _twabSupplyLimit(uint256 _totalSupply) internal pure returns (uint256) {
-        unchecked {
-            return type(uint96).max - _totalSupply;
+    /// @notice Verifies that the mint limit can support the new share balance.
+    /// @dev Reverts if the mint limit is exceeded.
+    /// @dev This MUST be called anytime there is a positive increase in the net total shares.
+    /// @param _existingShares The total existing prize vault shares (internal and external)
+    /// @param _newShares The new shares
+    function _enforceMintLimit(uint256 _existingShares, uint256 _newShares) internal pure {
+        uint256 _limit = _mintLimit(_existingShares);
+        if (_newShares > _limit) {
+            unchecked {
+                revert MintLimitExceeded(_newShares - _limit);
+            }
         }
     }
 
@@ -860,9 +871,14 @@ contract PrizeVault is TwabERC20, Claimable, IERC4626, ILiquidationSource, Ownab
         uint256 _yieldVaultShares = yieldVault.previewDeposit(_assetsWithDust);
         uint256 _assetsUsed = yieldVault.mint(_yieldVaultShares, address(this));
 
-        _mint(_receiver, _shares);
+        // Enforce the mint limit and protect against lossy deposits.
+        uint256 _totalDebtBeforeMint = totalSupply() + yieldFeeBalance;
+        _enforceMintLimit(_totalDebtBeforeMint, _shares);
+        if (totalAssets() < _totalDebtBeforeMint + _shares) {
+            revert LossyDeposit(totalAssets(), _totalDebtBeforeMint + _shares);
+        }
 
-        if (totalAssets() < totalDebt()) revert LossyDeposit(totalAssets(), totalDebt());
+        _mint(_receiver, _shares);
 
         emit Deposit(_caller, _receiver, _assets, _shares);
     }

test/invariant/PrizeVault/LossyPrizeVaultInvariant.t.sol

@@ -35,6 +35,12 @@ contract LossyPrizeVaultInvariant is Test {
         }
     }
 
+    function invariantYieldFeeBalanceAlwaysClaimable() external {
+        uint256 supplyLimit = type(uint96).max - lossyVaultHarness.vault().totalSupply();
+        uint256 yieldFeeBalance = lossyVaultHarness.vault().yieldFeeBalance();
+        assertLe(yieldFeeBalance, supplyLimit);
+    }
+
     function invariantAllAssetsAccountedFor() external {
         uint256 totalAssets = lossyVaultHarness.vault().totalAssets();
         uint256 totalDebt = lossyVaultHarness.vault().totalDebt();

test/invariant/PrizeVault/PrizeVaultInvariant.t.sol

@@ -47,6 +47,12 @@ contract PrizeVaultInvariant is Test {
         assertLe(liquidBalance, availableYieldBalance);
     }
 
+    function invariantYieldFeeBalanceAlwaysClaimable() external useCurrentTime {
+        uint256 supplyLimit = type(uint96).max - vaultHarness.vault().totalSupply();
+        uint256 yieldFeeBalance = vaultHarness.vault().yieldFeeBalance();
+        assertLe(yieldFeeBalance, supplyLimit);
+    }
+
     function invariantAllAssetsAccountedFor() external useCurrentTime {
         PrizeVault vault = vaultHarness.vault();
         uint256 totalAssets = vault.totalAssets();

test/unit/PrizeVault/Liquidate.t.sol

@@ -92,6 +92,37 @@ contract PrizeVaultLiquidationTest is UnitBaseSetup {
         assertEq(vault.liquidatableBalanceOf(address(vault)), supplyCapLeft); // less than available yield since shares are capped at uint96 max
     }
 
+    function testLiquidatableBalanceOf_respectsMaxShareMintWithFee() public {
+        vault.setYieldFeePercentage(1e8); // 10%
+        vault.setYieldFeeRecipient(address(this));
+        vault.setLiquidationPair(address(this));
+
+        uint256 supplyCapLeft = 100;
+
+        // make a large deposit to use most of the shares:
+        underlyingAsset.mint(address(alice), type(uint96).max);
+        vm.startPrank(alice);
+        underlyingAsset.approve(address(vault), type(uint96).max - supplyCapLeft);
+        vault.deposit(type(uint96).max - supplyCapLeft, alice);
+        vm.stopPrank();
+
+        underlyingAsset.mint(address(vault), 1e18);
+        uint256 availableYield = vault.availableYieldBalance();
+        assertApproxEqAbs(availableYield, 1e18 - vault.yieldBuffer(), 1);
+
+        assertLt(supplyCapLeft, availableYield);
+
+        uint256 amountOut = (supplyCapLeft * 9) / 10;
+        assertEq(vault.liquidatableBalanceOf(address(vault)), amountOut);
+        vault.transferTokensOut(address(0), address(this), address(vault), amountOut);
+
+        assertEq(vault.liquidatableBalanceOf(address(vault)), 0);
+        assertEq(vault.yieldFeeBalance(), supplyCapLeft - amountOut);
+
+        // ensure the yield fee can be minted
+        vault.claimYieldFeeShares(supplyCapLeft - amountOut);
+    }
+
     /* ============ transferTokensOut ============ */
 
     function testTransferTokensOut_noFee() public {
@@ -244,6 +275,34 @@ contract PrizeVaultLiquidationTest is UnitBaseSetup {
         vault.transferTokensOut(address(0), bob, address(vault), amountOut + 1);
     }
 
+    function testTransferTokensOut_YieldFeeExceedsSupplyCap() public {
+        vault.setYieldFeePercentage(1e8); // 10%
+        vault.setYieldFeeRecipient(bob);
+        vault.setLiquidationPair(address(this));
+
+        uint256 supplyCapLeft = 100;
+
+        // make a large deposit to use most of the shares:
+        underlyingAsset.mint(address(alice), type(uint96).max);
+        vm.startPrank(alice);
+        underlyingAsset.approve(address(vault), type(uint96).max - supplyCapLeft);
+        vault.deposit(type(uint96).max - supplyCapLeft, alice);
+        vm.stopPrank();
+
+        underlyingAsset.mint(address(vault), 1e18);
+        uint256 availableYield = vault.availableYieldBalance();
+        assertApproxEqAbs(availableYield, 1e18 - vault.yieldBuffer(), 1);
+
+        assertLt(supplyCapLeft, availableYield);
+        assertEq((supplyCapLeft * 9) / 10, vault.liquidatableBalanceOf(address(vault)));
+
+        uint256 amountOut = supplyCapLeft; // 10 assets too much
+        // (even though there is available yield, the supply cap will be exceeded by the yield fee)
+
+        vm.expectRevert(abi.encodeWithSelector(PrizeVault.MintLimitExceeded.selector, 11)); // yield fee is 11
+        vault.transferTokensOut(address(0), address(this), address(vault), amountOut);
+    }
+
     /* ============ verifyTokensIn ============ */
 
     function testVerifyTokensIn() public {