feat: Add invariant hooks, refactor to handle multiple suites (SC-481)

foundry.toml

@@ -12,6 +12,7 @@ runs = 1000
 [invariant]
 runs = 20
 depth = 1000
+shrink_run_limit = 0
 
 [profile.pr.invariant]
 runs = 200

test/invariant/Invariants.t.sol

@@ -5,36 +5,33 @@ import "forge-std/Test.sol";
 
 import { PSMTestBase } from "test/PSMTestBase.sol";
 
-import { LpHandler }      from "test/invariant/handlers/LpHandler.sol";
-import { SwapperHandler } from "test/invariant/handlers/SwapperHandler.sol";
+import { LpHandler }       from "test/invariant/handlers/LpHandler.sol";
+import { SwapperHandler }  from "test/invariant/handlers/SwapperHandler.sol";
+import { TransferHandler } from "test/invariant/handlers/TransferHandler.sol";
 
-contract PSMInvariantTests is PSMTestBase {
+abstract contract PSMInvariantTestBase is PSMTestBase {
 
-    LpHandler      public lpHandler;
-    SwapperHandler public swapperHandler;
+    LpHandler       public lpHandler;
+    SwapperHandler  public swapperHandler;
+    TransferHandler public transferHandler;
 
     address BURN_ADDRESS = makeAddr("burn-address");
 
     // NOTE [CRITICAL]: All invariant tests are operating under the assumption that the initial seed
     //                  deposit of 1e18 shares has been made. This is a key requirement and
     //                  assumption for all invariant tests.
-    function setUp() public override {
+    function setUp() public virtual override {
         super.setUp();
 
         // Seed the pool with 1e18 shares (1e18 of value)
         _deposit(address(dai), BURN_ADDRESS, 1e18);
-
-        lpHandler      = new LpHandler(psm, dai, usdc, sDai, 3);
-        swapperHandler = new SwapperHandler(psm, dai, usdc, sDai, 3);
-
-        // TODO: Add rate updates
-        rateProvider.__setConversionRate(1.25e27);
-
-        targetContract(address(lpHandler));
-        targetContract(address(swapperHandler));
     }
 
-    function invariant_A() public view {
+    /**********************************************************************************************/
+    /*** Invariant assertion functions                                                          ***/
+    /**********************************************************************************************/
+
+    function _checkInvariant_A() public view {
         assertEq(
             psm.shares(address(lpHandler.lps(0))) +
             psm.shares(address(lpHandler.lps(1))) +
@@ -44,15 +41,15 @@ contract PSMInvariantTests is PSMTestBase {
         );
     }
 
-    function invariant_B() public view {
+    function _checkInvariant_B() public view {
         assertApproxEqAbs(
             psm.getPsmTotalValue(),
             psm.convertToAssetValue(psm.totalShares()),
             2
         );
     }
 
-    function invariant_C() public view {
+    function _checkInvariant_C() public view {
         assertApproxEqAbs(
             psm.convertToAssetValue(psm.shares(address(lpHandler.lps(0)))) +
             psm.convertToAssetValue(psm.shares(address(lpHandler.lps(1)))) +
@@ -63,7 +60,11 @@ contract PSMInvariantTests is PSMTestBase {
         );
     }
 
-    function invariant_logs() public view {
+    /**********************************************************************************************/
+    /*** Helper functions                                                                       ***/
+    /**********************************************************************************************/
+
+    function _logHandlerCallCounts() public view {
         console.log("depositCount    ", lpHandler.depositCount());
         console.log("withdrawCount   ", lpHandler.withdrawCount());
         console.log("swapCount       ", swapperHandler.swapCount());
@@ -77,4 +78,178 @@ contract PSMInvariantTests is PSMTestBase {
         );
     }
 
+    function _getLpTokenValue(address lp) internal view returns (uint256) {
+        uint256 daiValue  = dai.balanceOf(lp);
+        uint256 usdcValue = usdc.balanceOf(lp) * 1e12;
+        uint256 sDaiValue = sDai.balanceOf(lp) * rateProvider.getConversionRate() / 1e27;
+
+        return daiValue + usdcValue + sDaiValue;
+    }
+
+    /**********************************************************************************************/
+    /*** After invariant hook functions                                                         ***/
+    /**********************************************************************************************/
+
+    function _withdrawAllPositions() public {
+        address lp0 = lpHandler.lps(0);
+        address lp1 = lpHandler.lps(1);
+        address lp2 = lpHandler.lps(2);
+
+        // Get value of each LPs current deposits.
+        uint256 lp0DepositsValue = psm.convertToAssetValue(psm.shares(lp0));
+        uint256 lp1DepositsValue = psm.convertToAssetValue(psm.shares(lp1));
+        uint256 lp2DepositsValue = psm.convertToAssetValue(psm.shares(lp2));
+
+        // Get value of each LPs token holdings from previous withdrawals.
+        uint256 lp0WithdrawsValue = _getLpTokenValue(lp0);
+        uint256 lp1WithdrawsValue = _getLpTokenValue(lp1);
+        uint256 lp2WithdrawsValue = _getLpTokenValue(lp2);
+
+        uint256 psmTotalValue = psm.getPsmTotalValue();
+
+        uint256 startingSeedValue = psm.convertToAssetValue(1e18);
+
+        // Liquidity is unknown so withdraw all assets for all users to empty PSM.
+        _withdraw(address(dai),  lp0, type(uint256).max);
+        _withdraw(address(usdc), lp0, type(uint256).max);
+        _withdraw(address(sDai), lp0, type(uint256).max);
+
+        _withdraw(address(dai),  lp1, type(uint256).max);
+        _withdraw(address(usdc), lp1, type(uint256).max);
+        _withdraw(address(sDai), lp1, type(uint256).max);
+
+        _withdraw(address(dai),  lp2, type(uint256).max);
+        _withdraw(address(usdc), lp2, type(uint256).max);
+        _withdraw(address(sDai), lp2, type(uint256).max);
+
+        // All funds are completely withdrawn.
+        assertEq(psm.shares(lp0), 0);
+        assertEq(psm.shares(lp1), 0);
+        assertEq(psm.shares(lp2), 0);
+
+        uint256 seedValue = psm.convertToAssetValue(1e18);
+
+        // PSM is empty (besides seed amount).
+        assertEq(psm.totalShares(),      1e18);
+        assertEq(psm.getPsmTotalValue(), seedValue);
+
+        // Tokens held by LPs are equal to the sum of their previous balance
+        // plus the amount of value originally represented in the PSM's shares.
+        // There can be rounding here because of share burning up to 1e12 when withdrawing USDC.
+        // It should be noted that LP2 here has a rounding error of 2e12 since both LP0 and LP1
+        // could have rounding errors that accumulate to LP2.
+        assertApproxEqAbs(_getLpTokenValue(lp0), lp0DepositsValue + lp0WithdrawsValue, 1e12);
+        assertApproxEqAbs(_getLpTokenValue(lp1), lp1DepositsValue + lp1WithdrawsValue, 1e12);
+        assertApproxEqAbs(_getLpTokenValue(lp2), lp2DepositsValue + lp2WithdrawsValue, 2e12);
+
+        // All rounding errors from LPs can accrue to the burn address after withdrawals are made.
+        assertApproxEqAbs(seedValue, startingSeedValue, 3e12);
+
+        // Current value of all LPs' token holdings.
+        uint256 sumLpValue = _getLpTokenValue(lp0) + _getLpTokenValue(lp1) + _getLpTokenValue(lp2);
+
+        // Total amount just withdrawn from the PSM.
+        uint256 totalWithdrawals
+            = sumLpValue - (lp0WithdrawsValue + lp1WithdrawsValue + lp2WithdrawsValue);
+
+        // Assert that all funds were withdrawn equals the original value of the PSM minus the
+        // 1e18 share seed deposit.
+        assertApproxEqAbs(totalWithdrawals, psmTotalValue - seedValue, 2);
+
+        // Get the starting sum of all LPs' deposits and withdrawals.
+        uint256 sumStartingValue =
+            (lp0DepositsValue  + lp1DepositsValue  + lp2DepositsValue) +
+            (lp0WithdrawsValue + lp1WithdrawsValue + lp2WithdrawsValue);
+
+        // Assert that the sum of all LPs' deposits and withdrawals equals
+        // the sum of all LPs' resulting token holdings. Rounding errors are accumulated to the
+        // burn address.
+        assertApproxEqAbs(sumLpValue, sumStartingValue, seedValue - startingSeedValue + 2);
+
+        // NOTE: Below logic is not realistic, shown to demonstrate precision.
+
+        _withdraw(address(dai),  BURN_ADDRESS, type(uint256).max);
+        _withdraw(address(usdc), BURN_ADDRESS, type(uint256).max);
+        _withdraw(address(sDai), BURN_ADDRESS, type(uint256).max);
+
+        // When all funds are completely withdrawn, the sum of all funds withdrawn is equal to the
+        // sum of value of all LPs including the burn address. All rounding errors get reduced to
+        // a few wei.
+        assertApproxEqAbs(
+            sumLpValue + _getLpTokenValue(BURN_ADDRESS),
+            sumStartingValue + startingSeedValue,
+            5
+        );
+
+        // All funds can always be withdrawn completely.
+        assertEq(psm.totalShares(),      0);
+        assertEq(psm.getPsmTotalValue(), 0);
+    }
+
+}
+
+contract PSMInvariants_ConstantRate_NoTransfer is PSMInvariantTestBase {
+
+    function setUp() public override {
+        super.setUp();
+
+        lpHandler      = new LpHandler(psm, dai, usdc, sDai, 3);
+        swapperHandler = new SwapperHandler(psm, dai, usdc, sDai, 3);
+
+        rateProvider.__setConversionRate(1.25e27);
+
+        targetContract(address(lpHandler));
+        targetContract(address(swapperHandler));
+    }
+
+    function invariant_A() public view {
+        _checkInvariant_A();
+    }
+
+    function invariant_B() public view {
+        _checkInvariant_B();
+    }
+
+    function invariant_C() public view {
+        _checkInvariant_C();
+    }
+
+    function afterInvariant() public {
+        _withdrawAllPositions();
+    }
+
+}
+
+contract PSMInvariants_ConstantRate_WithTransfers is PSMInvariantTestBase {
+
+    function setUp() public override {
+        super.setUp();
+
+        lpHandler       = new LpHandler(psm, dai, usdc, sDai, 3);
+        swapperHandler  = new SwapperHandler(psm, dai, usdc, sDai, 3);
+        transferHandler = new TransferHandler(psm, dai, usdc, sDai);
+
+        rateProvider.__setConversionRate(1.25e27);
+
+        targetContract(address(lpHandler));
+        targetContract(address(swapperHandler));
+        targetContract(address(transferHandler));
+    }
+
+    function invariant_A() public view {
+        _checkInvariant_A();
+    }
+
+    function invariant_B() public view {
+        _checkInvariant_B();
+    }
+
+    function invariant_C() public view {
+        _checkInvariant_C();
+    }
+
+    function afterInvariant() public {
+        _withdrawAllPositions();
+    }
+
 }

test/invariant/handlers/LpHandler.sol

@@ -14,8 +14,6 @@ contract LpHandler is HandlerBase {
     uint256 public depositCount;
     uint256 public withdrawCount;
 
-    uint256 public constant TRILLION = 1e12;
-
     constructor(
         PSM3      psm_,
         MockERC20 asset0,
@@ -24,7 +22,7 @@ contract LpHandler is HandlerBase {
         uint256   lpCount
     ) HandlerBase(psm_, asset0, asset1, asset2) {
         for (uint256 i = 0; i < lpCount; i++) {
-            lps.push(makeAddr(string(abi.encodePacked("lp-", i))));
+            lps.push(makeAddr(string(abi.encodePacked("lp-", vm.toString(i)))));
         }
     }
 
@@ -36,7 +34,7 @@ contract LpHandler is HandlerBase {
         MockERC20 asset = _getAsset(assetSeed);
         address   lp    = _getLP(lpSeed);
 
-        amount = _bound(amount, 1, TRILLION * 10 ** asset.decimals());
+        amount = _bound(amount, 1, 1e12 * 10 ** asset.decimals());
 
         vm.startPrank(lp);
         asset.mint(lp, amount);
@@ -51,7 +49,7 @@ contract LpHandler is HandlerBase {
         MockERC20 asset = _getAsset(assetSeed);
         address   lp    = _getLP(lpSeed);
 
-        amount = _bound(amount, 1, TRILLION * 10 ** asset.decimals());
+        amount = _bound(amount, 1, 1e12 * 10 ** asset.decimals());
 
         vm.prank(lp);
         psm.withdraw(address(asset), lp, amount);

test/invariant/handlers/SwapperHandler.sol

@@ -22,7 +22,7 @@ contract SwapperHandler is HandlerBase {
         uint256   lpCount
     ) HandlerBase(psm_, asset0, asset1, asset2) {
         for (uint256 i = 0; i < lpCount; i++) {
-            swappers.push(makeAddr(string(abi.encodePacked("swapper-", i))));
+            swappers.push(makeAddr(string(abi.encodePacked("swapper-", vm.toString(i)))));
         }
     }
 

test/invariant/handlers/TransferHandler.sol

@@ -0,0 +1,38 @@
+// SPDX-License-Identifier: AGPL-3.0-or-later
+pragma solidity ^0.8.13;
+
+import { MockERC20 } from "erc20-helpers/MockERC20.sol";
+
+import { HandlerBase } from "test/invariant/handlers/HandlerBase.sol";
+
+import { PSM3 } from "src/PSM3.sol";
+
+contract TransferHandler is HandlerBase {
+
+    uint256 public transferCount;
+
+    constructor(
+        PSM3      psm_,
+        MockERC20 asset0,
+        MockERC20 asset1,
+        MockERC20 asset2
+    ) HandlerBase(psm_, asset0, asset1, asset2) {}
+
+    function transfer(uint256 assetSeed, string memory senderSeed, uint256 amount) external {
+        MockERC20 asset  = _getAsset(assetSeed);
+        address   sender = makeAddr(senderSeed);
+
+        // Bounding to 10 million here because 1 trillion introduces unrealistic conditions with
+        // large rounding errors. Would rather keep tolerances smaller with a lower upper bound
+        // on transfer amounts.
+        amount = _bound(amount, 1, 10_000_000 * 10 ** asset.decimals());
+
+        asset.mint(sender, amount);
+
+        vm.prank(sender);
+        asset.transfer(address(psm), amount);
+
+        transferCount += 1;
+    }
+
+}