fix: correctly segment fees on condition that invariant > 0

contracts/Portfolio.sol

@@ -616,8 +616,8 @@ contract Portfolio is ERC1155, IPortfolio {
             // ====== Invariant Check ====== //
 
             bool validInvariant;
-            (validInvariant, inter.segmentFees, inter.nextInvariant) = strategy.validateSwap(
-                poolId, inter.prevInvariant, adjustedVirtualX, adjustedVirtualY, sellAsset, inter.feeAmountUnit
+            (validInvariant, inter.nextInvariant) = strategy.validateSwap(
+                poolId, inter.prevInvariant, adjustedVirtualX, adjustedVirtualY 
             );
 
             if (!validInvariant) {
@@ -632,19 +632,21 @@ contract Portfolio is ERC1155, IPortfolio {
         // Take the protocol fee from the input amount, so that protocol fee is not re-invested into pool.
         // Increases the independent pool reserve by the input amount, including fee and excluding protocol fee.
         // Decrease the dependent pool reserve by the output amount.
-        if (inter.segmentFees) {
+        pool.adjustReserves(
+            args.sellAsset,
+            inter.amountInputUnit - inter.protocolFeeAmountUnit,
+            inter.amountOutputUnit
+        );
+        // returns false if the invariant is > 0
+        bool valid = strategy.checkInvariant(poolId, pool);
+        if (!valid) {
+            // If the invariant is > 0, the pool is segmented and the fees are not reinvested.
             pool.adjustReserves(
-                args.sellAsset,
-                inter.amountInputUnit - inter.protocolFeeAmountUnit - inter.feeAmountUnit,
-                inter.amountOutputUnit
+                !args.sellAsset,
+                0,
+                inter.feeAmountUnit
             );
             inter.feeGrowthGlobal += inter.feeAmountUnit;
-        } else {
-            pool.adjustReserves(
-                args.sellAsset,
-                inter.amountInputUnit - inter.protocolFeeAmountUnit,
-                inter.amountOutputUnit
-            );
         }
 
         // Increasing reserves requires Portfolio's balance of tokens to also increases by the end of `settlement`.

contracts/interfaces/IStrategy.sol

@@ -4,6 +4,7 @@ pragma solidity >=0.8.0;
 import { Order } from "../libraries/SwapLib.sol";
 import { SwapState } from "../libraries/SwapLib.sol";
 import { IPortfolioStrategy } from "./IPortfolio.sol";
+import { PortfolioPool } from "../libraries/PoolLib.sol";
 
 /**
  * @title
@@ -84,8 +85,24 @@ interface IStrategy is IPortfolioStrategy {
         uint64 poolId,
         int256 invariant,
         uint256 reserveX,
-        uint256 reserveY,
-        bool sellAsset,
-        uint256 feeAmountUnit
-    ) external view returns (bool, bool, int256);
+        uint256 reserveY
+    ) external view returns (bool, int256);
+
+    /**
+     * @notice
+     * Checks the validity of the invariant of the in-memory pool during swap execution.
+     *
+     * @dev
+     * Critical function that is responsible for the economic validity of the protocol.
+     * Swaps should not push the invariant over 0 in many pools
+     * This tells portfolio if it should segment fees or not
+     *
+     * @param poolId Id of the pool.
+     * @param pool intermediate pool struct in swap execution
+     * @return success Whether the invariant is positive.
+     */
+    function checkInvariant(
+        uint64 poolId,
+        PortfolioPool memory pool
+    ) external view returns (bool);
 }

contracts/strategies/NormalStrategy.sol

@@ -6,6 +6,7 @@ import "../libraries/BisectionLib.sol";
 import "../libraries/PortfolioLib.sol";
 import "./INormalStrategy.sol";
 import "./NormalStrategyLib.sol";
+import "forge-std/console2.sol";
 
 /// @dev Emitted when a hook is called by a non-portfolio address.
 error NormalStrategy_NotPortfolio();
@@ -109,27 +110,19 @@ contract NormalStrategy is INormalStrategy {
         uint64 poolId,
         int256 invariant,
         uint256 reserveX,
-        uint256 reserveY,
-        bool sellAsset,
-        uint256 feeAmountUnit
-    ) public view returns (bool, bool, int256) {
+        uint256 reserveY
+    ) public view returns (bool, int256) {
         PortfolioPool memory pool = IPortfolioStruct(portfolio).pools(poolId);
-        bool segmentFees = false;
 
         // Update the reserves in memory.
         pool.virtualX = reserveX.safeCastTo128();
         pool.virtualY = reserveY.safeCastTo128();
 
         // Compute the new invariant.
         int256 invariantAfterSwap = pool.getInvariantDown(configs[poolId]);
-        if (invariantAfterSwap > 0) {
-            segmentFees = true;
-            sellAsset ? pool.virtualX -= feeAmountUnit.safeCastTo128() : pool.virtualY -= feeAmountUnit.safeCastTo128();
-            invariantAfterSwap = pool.getInvariantDown(configs[poolId]);
-        }
-        bool valid = _validateSwap(invariant, invariantAfterSwap, segmentFees);
+        bool valid = _validateSwap(invariant, invariantAfterSwap);
 
-        return (valid, segmentFees, invariantAfterSwap);
+        return (valid, invariantAfterSwap);
     }
 
     /**
@@ -152,15 +145,10 @@ contract NormalStrategy is INormalStrategy {
      */
     function _validateSwap(
         int256 invariantBefore,
-        int256 invariantAfter,
-        bool segmentFees
+        int256 invariantAfter
     ) internal pure returns (bool) {
         int256 delta = invariantAfter - invariantBefore;
-        if (segmentFees) {
-            if (invariantBefore < invariantAfter) return false;
-        } else {
-            if (delta < MINIMUM_INVARIANT_DELTA) return false;
-        }
+        if (delta < MINIMUM_INVARIANT_DELTA) return false;
 
         return true;
     }
@@ -307,7 +295,18 @@ contract NormalStrategy is INormalStrategy {
             protocolFee: IPortfolio(portfolio).protocolFee()
         });
 
-        success = _validateSwap(prevInvariant, postInvariant, false);
+        success = _validateSwap(prevInvariant, postInvariant);
+    }
+
+    function checkInvariant(uint64 poolId, PortfolioPool memory pool)
+        public
+        view
+        returns (bool)
+    {
+        int256 invariant = pool.getInvariantDown(configs[poolId]);
+        console2.log("invariant in checkInvariant", invariant);
+        return invariant < 0;
+
     }
 
     /// @inheritdoc IPortfolioStrategy

test/TestPortfolioSwap.t.sol

@@ -121,17 +121,17 @@ contract TestPortfolioSwap is Setup {
     {
         bool sellAsset = true;
         uint128 amtIn = 0.5 ether;
-        int256 initialInvariant = subject().getInvariant(ghost().poolId);
-        console.logInt(initialInvariant);
 
         uint256 i = 0;
-        while (i < 5) {
+        while (i < 365) {
             vm.warp(block.timestamp + 1 days);
             (uint256 reserveAsset, uint256 reserveQuote) =
                 subject().getPoolReserves(ghost().poolId);
+            int256 invariant = subject().getInvariant(ghost().poolId);
+            console2.log("invariant before swap execution", invariant);
             sellAsset = reserveQuote > 1 ether ? true : sellAsset;
             sellAsset = reserveAsset > 1 ether ? false : sellAsset;
-            amtIn = uint128(sellAsset ? reserveQuote - 0.12 ether : reserveAsset - 0.12 ether);
+            amtIn = uint128(sellAsset ? reserveQuote - 0.4 ether : reserveAsset - 0.4 ether);
             // amtIn = 1000;
             console.log("amtIn", amtIn);
             uint128 amtOut = uint128(
@@ -153,8 +153,8 @@ contract TestPortfolioSwap is Setup {
 
             subject().swap(order);
 
-            int256 invariant = subject().getInvariant(ghost().poolId);
-            console.logInt(invariant);
+            invariant = subject().getInvariant(ghost().poolId);
+            console2.log("invariant after swap execution", invariant);
 
             sellAsset = !sellAsset;
             i++;

test/strategies/NormalConfiguration.sol

@@ -20,7 +20,7 @@ import "contracts/strategies/NormalStrategyLib.sol";
 
 uint256 constant NormalConfiguration_DEFAULT_PRICE = 1 ether;
 uint256 constant NormalConfiguration_DEFAULT_STRIKE_WAD = 1 ether;
-uint256 constant NormalConfiguration_DEFAULT_VOLATILITY_BPS = 100 wei; // in bps
+uint256 constant NormalConfiguration_DEFAULT_VOLATILITY_BPS = 500 wei; // in bps
 uint256 constant NormalConfiguration_DEFAULT_DURATION_SEC = 365 days + 100; // in seconds
 uint256 constant NormalConfiguration_DEFAULT_CREATION_TIMESTAMP = 0;
 bool constant NormalConfiguration_DEFAULT_IS_PERPETUAL = false;