Merge pull request #156 from whetstoneresearch/test/invariant-tests

Test/invariant tests (Draft)

test/invariant/DopplerHandler.sol

@@ -4,8 +4,6 @@ pragma solidity ^0.8.13;
 import {Test} from "forge-std/Test.sol";
 import {AddressSet, LibAddressSet} from "./AddressSet.sol";
 import {DopplerImplementation} from "test/shared/DopplerImplementation.sol";
-import {PoolSwapTest} from "v4-core/src/test/PoolSwapTest.sol";
-import {IPoolManager} from "v4-core/src/interfaces/IPoolManager.sol";
 import {TestERC20} from "v4-core/src/test/TestERC20.sol";
 import {PoolKey} from "v4-core/src/types/PoolKey.sol";
 import {Currency} from "v4-core/src/types/Currency.sol";
@@ -24,15 +22,20 @@ contract DopplerHandler is Test {
     bool public isToken0;
     bool public isUsingEth;
 
+    // Ghost variables are used to mimic the state of the hook contract.
     uint256 public ghost_reserve0;
     uint256 public ghost_reserve1;
+    uint256 public ghost_totalTokensSold;
+    uint256 public ghost_totalProceeds;
 
     mapping(bytes4 => uint256) public calls;
     uint256 public totalCalls;
 
     AddressSet internal actors;
     address internal currentActor;
 
+    mapping(address actor => uint256 balance) public assetBalanceOf;
+
     modifier createActor() {
         currentActor = msg.sender;
         actors.add(msg.sender);
@@ -71,26 +74,123 @@ contract DopplerHandler is Test {
         token1 = TestERC20(Currency.unwrap(poolKey.currency1));
 
         if (isToken0) {
-            numeraire = token0;
-            asset = token1;
-        } else {
-            numeraire = token1;
             asset = token0;
+            numeraire = token1;
+        } else {
+            asset = token1;
+            numeraire = token0;
         }
 
         ghost_reserve0 = token0.balanceOf(address(hook));
         ghost_reserve1 = token1.balanceOf(address(hook));
     }
 
     /// @notice Buys an amount of asset tokens using an exact amount of numeraire tokens
-    function buyExactAmountIn(uint256 amount) public createActor countCall(this.buyExactAmountIn.selector) {
+    function buyExactAmountIn(uint256 amountToSpend) public createActor countCall(this.buyExactAmountIn.selector) {
+        amountToSpend = 1 ether;
+
         if (isUsingEth) {
-            deal(currentActor, amount);
+            deal(currentActor, amountToSpend);
         } else {
-            numeraire.mint(currentActor, amount);
-            numeraire.approve(address(router), amount);
+            numeraire.mint(currentActor, amountToSpend);
+            numeraire.approve(address(router), amountToSpend);
         }
 
-        uint256 bought = router.buyExactIn{value: isUsingEth ? amount : 0}(amount);
+        uint256 bought = router.buyExactIn{value: isUsingEth ? amountToSpend : 0}(amountToSpend);
+        assetBalanceOf[currentActor] += bought;
+        ghost_totalTokensSold += bought;
+        ghost_totalProceeds += amountToSpend;
+
+        if (isToken0) {
+            ghost_reserve0 -= bought;
+            ghost_reserve1 += amountToSpend;
+        } else {
+            ghost_reserve1 -= bought;
+            ghost_reserve0 += amountToSpend;
+        }
+    }
+
+    function buyExactAmountOut(uint256 assetsToBuy) public createActor countCall(this.buyExactAmountOut.selector) {
+        assetsToBuy = 1 ether;
+        uint256 amountInRequired = router.computeBuyExactOut(assetsToBuy);
+
+        if (isUsingEth) {
+            deal(currentActor, amountInRequired);
+        } else {
+            numeraire.mint(currentActor, amountInRequired);
+            numeraire.approve(address(router), amountInRequired);
+        }
+
+        uint256 spent = router.buyExactOut{value: isUsingEth ? amountInRequired : 0}(assetsToBuy);
+        assetBalanceOf[currentActor] += assetsToBuy;
+        ghost_totalTokensSold += assetsToBuy;
+        ghost_totalProceeds += spent;
+
+        if (isToken0) {
+            ghost_reserve0 -= assetsToBuy;
+            ghost_reserve1 += spent;
+        } else {
+            ghost_reserve1 -= assetsToBuy;
+            ghost_reserve0 += spent;
+        }
+    }
+
+    function sellExactIn(uint256 seed)
+        public
+        useActor(uint256(uint160(msg.sender)))
+        countCall(this.sellExactIn.selector)
+    {
+        // If the currentActor is address(0), it means no one has bought any assets yet.
+        if (currentActor == address(0) || assetBalanceOf[currentActor] == 0) return;
+
+        uint256 assetsToSell = seed % assetBalanceOf[currentActor] + 1;
+        TestERC20(asset).approve(address(router), assetsToSell);
+        uint256 received = router.sellExactIn(assetsToSell);
+
+        assetBalanceOf[currentActor] -= assetsToSell;
+        ghost_totalTokensSold -= assetsToSell;
+        ghost_totalProceeds -= received;
+
+        if (isToken0) {
+            ghost_reserve0 += assetsToSell;
+            ghost_reserve1 -= received;
+        } else {
+            ghost_reserve1 += assetsToSell;
+            ghost_reserve0 -= received;
+        }
+    }
+
+    function sellExactOut(uint256 seed)
+        public
+        useActor(uint256(uint160(msg.sender)))
+        countCall(this.sellExactOut.selector)
+    {
+        // If the currentActor is address(0), it means no one has bought any assets yet.
+        if (currentActor == address(0) || assetBalanceOf[currentActor] == 0) return;
+
+        // We compute the maximum amount we can receive from our current balance.
+        uint256 maxAmountToReceive = router.computeSellExactOut(assetBalanceOf[currentActor]);
+
+        // Then we compute a random amount from that maximum.
+        uint256 amountToReceive = seed % maxAmountToReceive + 1;
+
+        TestERC20(asset).approve(address(router), router.computeSellExactOut(amountToReceive));
+        uint256 sold = router.sellExactOut(amountToReceive);
+
+        assetBalanceOf[currentActor] -= sold;
+        ghost_totalTokensSold -= sold;
+        ghost_totalProceeds -= amountToReceive;
+
+        if (isToken0) {
+            ghost_reserve0 += sold;
+            ghost_reserve1 -= amountToReceive;
+        } else {
+            ghost_reserve0 -= amountToReceive;
+            ghost_reserve1 += sold;
+        }
+    }
+
+    function goNextEpoch() public countCall(this.goNextEpoch.selector) {
+        vm.warp(block.timestamp + hook.getEpochLength());
     }
 }

test/invariant/DopplerInvariants.sol

@@ -4,29 +4,130 @@ pragma solidity ^0.8.13;
 import {console} from "forge-std/console.sol";
 import {BaseTest} from "test/shared/BaseTest.sol";
 import {DopplerHandler} from "test/invariant/DopplerHandler.sol";
+import {State} from "src/Doppler.sol";
+import {LiquidityAmounts} from "v4-core/test/utils/LiquidityAmounts.sol";
+import {TickMath} from "v4-core/src/libraries/TickMath.sol";
+
+/*
+
+ ? totalTokensSold can't underflow
+ ? totalProceeds can't underflow
+ ? Computed ticks can't under/overflow
+ X Amount of asset tokens to be supplied to liquidity positions at once <= numTokensToSell - totalTokensSold
+ Likely further have to consider the actual available token balance as dust is lost due to rounding
+ Does not attempt to create liquidity positions with equal upper and lower ticks
+Leads to divide by zero error (even just in computing the liquidity amount)
+Also relevant for just doing relevant math, e.g. retrieving liquidity given an amount can result in a revert
+ Cannot trade the price to below the lower slug range
+Else it allows for price manipulation
+I think this is also a compliance requirement regardless
+ Always places a lower slug if totalTokensSold > 0
+ Selling all tokens back in to the curve must exceed the available liquidity
+ Single tick ranges do not exceed max liquidity per tick
+ Cannot modify the price in any way prior to the start time
+ */
 
 contract DopplerInvariantsTest is BaseTest {
-    // DopplerHandler public handler;
+    DopplerHandler public handler;
+
+    function setUp() public override {
+        super.setUp();
+        handler = new DopplerHandler(key, hook, router, isToken0, usingEth);
+
+        bytes4[] memory selectors = new bytes4[](3);
+        selectors[0] = handler.buyExactAmountIn.selector;
+        selectors[1] = handler.buyExactAmountOut.selector;
+        selectors[2] = handler.sellExactIn.selector;
+        // selectors[3] = handler.sellExactOut.selector;
+
+        targetSelector(FuzzSelector({addr: address(handler), selectors: selectors}));
+        targetContract(address(handler));
+
+        vm.warp(DEFAULT_STARTING_TIME);
+    }
+
+    function afterInvariant() public view {
+        console.log("+-------------------+-----------------------+");
+        console.log("| Function Name     | Calls                 |", handler.totalCalls());
+        console.log("+-------------------+-----------------------+");
+        console.log("| buyExactAmountIn  |", handler.calls(handler.buyExactAmountIn.selector), "                  |");
+        console.log("| buyExactAmountOut |", handler.calls(handler.buyExactAmountOut.selector), "                  |");
+        console.log("| sellExactIn       |", handler.calls(handler.sellExactIn.selector), "                    |");
+        console.log("| sellExactOut      |", handler.calls(handler.sellExactOut.selector), "                    |");
+        console.log("+-------------------+-----------------------+");
+    }
+
+    /// forge-config: default.invariant.fail-on-revert = true
+    function invariant_TracksTotalTokensSoldAndProceeds() public view {
+        (,, uint256 totalTokensSold, uint256 totalProceeds,,) = hook.state();
+        assertEq(totalTokensSold, handler.ghost_totalTokensSold());
+        assertEq(totalProceeds, handler.ghost_totalProceeds());
+    }
+
+    /// forge-config: default.invariant.fail-on-revert = true
+    function invariant_CantSellMoreThanNumTokensToSell() public view {
+        uint256 numTokensToSell = hook.getNumTokensToSell();
+        assertLe(handler.ghost_totalTokensSold(), numTokensToSell);
+    }
+
+    /// forge-config: default.invariant.fail-on-revert = true
+    function invariant_AlwaysProvidesAllAvailableTokens() public {
+        vm.skip(true);
+        uint256 numTokensToSell = hook.getNumTokensToSell();
+        uint256 totalTokensProvided;
+        uint256 slugs = hook.getNumPDSlugs();
+
+        int24 currentTick = hook.getCurrentTick(poolId);
+
+        for (uint256 i = 1; i < 4 + slugs; i++) {
+            (int24 tickLower, int24 tickUpper, uint128 liquidity,) = hook.positions(bytes32(uint256(i)));
+            (uint256 amount0, uint256 amount1) = LiquidityAmounts.getAmountsForLiquidity(
+                TickMath.getSqrtPriceAtTick(currentTick),
+                TickMath.getSqrtPriceAtTick(tickLower),
+                TickMath.getSqrtPriceAtTick(tickUpper),
+                liquidity
+            );
+            totalTokensProvided += isToken0 ? amount0 : amount1;
+        }
+
+        (,, uint256 totalTokensSold,,,) = hook.state();
+        assertEq(totalTokensProvided, numTokensToSell - totalTokensSold);
+    }
+
+    function invariant_LowerSlugWhenTokensSold() public {
+        vm.skip(true);
+        (,, uint256 totalTokensSold,,,) = hook.state();
 
-    // function setUp() public override {
-    //     super.setUp();
-    //     handler = new DopplerHandler(key, hook, router, isToken0, usingEth);
+        if (totalTokensSold > 0) {
+            (,, uint128 liquidity,) = hook.positions(bytes32(uint256(1)));
+            assertTrue(liquidity > 0);
+        }
+    }
 
-    //     bytes4[] memory selectors = new bytes4[](1);
-    //     selectors[0] = handler.buyExactAmountIn.selector;
+    function invariant_CannotTradeUnderLowerSlug() public view {
+        (int24 tickLower,,,) = hook.positions(bytes32(uint256(1)));
+        int24 currentTick = hook.getCurrentTick(poolId);
 
-    //     targetSelector(FuzzSelector({addr: address(handler), selectors: selectors}));
-    //     targetContract(address(handler));
-    // }
+        if (isToken0) {
+            assertTrue(currentTick >= tickLower);
+        } else {
+            assertTrue(currentTick <= tickLower);
+        }
+    }
 
-    // function afterInvariant() public view {
-    //     console.log("Handler address", address(handler));
-    //     console.log("Calls: ", handler.totalCalls());
-    //     console.log("buyExactAmountIn: ", handler.calls(handler.buyExactAmountIn.selector));
-    // }
+    /// forge-config: default.invariant.fail-on-revert = true
+    function invariant_PositionsDifferentTicks() public view {
+        uint256 slugs = hook.getNumPDSlugs();
+        for (uint256 i = 1; i < 4 + slugs; i++) {
+            (int24 tickLower, int24 tickUpper, uint128 liquidity,) = hook.positions(bytes32(uint256(i)));
+            if (liquidity > 0) assertTrue(tickLower != tickUpper);
+        }
+    }
 
-    // /// forge-config: default.invariant.fail-on-revert = true
-    // function invariant_works() public {
-    //     assertTrue(true);
-    // }
+    function invariant_NoPriceChangesBeforeStart() public {
+        vm.skip(true);
+        vm.warp(DEFAULT_STARTING_TIME - 1);
+        // TODO: I think this test is broken because we don't set the tick in the constructor.
+        assertEq(hook.getCurrentTick(poolId), hook.getStartingTick());
+    }
 }

test/shared/BaseTest.sol

@@ -245,6 +245,7 @@ contract BaseTest is Test, Deployers {
 
         // Deploy swapRouter
         swapRouter = new PoolSwapTest(manager);
+        vm.label(address(swapRouter), "SwapRouter");
 
         // Deploy modifyLiquidityRouter
         // Note: Only used to validate that liquidity can't be manually modified
@@ -259,8 +260,10 @@ contract BaseTest is Test, Deployers {
         TestERC20(token1).approve(address(modifyLiquidityRouter), type(uint256).max);
 
         quoter = new Quoter(manager);
+        vm.label(address(quoter), "Quoter");
 
         router = new CustomRouter(swapRouter, quoter, key, isToken0, usingEth);
+        vm.label(address(router), "Router");
     }
 
     function computeBuyExactOut(uint256 amountOut) public returns (uint256) {

test/shared/CustomRouter.sol

@@ -165,6 +165,7 @@ contract CustomRouter is Test {
     /// @return Amount of numeraire tokens received.
     function sell(int256 amount) public returns (uint256, uint256) {
         uint256 approveAmount = amount < 0 ? uint256(-amount) : computeSellExactOut(uint256(amount));
+        TestERC20(asset).transferFrom(msg.sender, address(this), uint256(approveAmount));
         TestERC20(asset).approve(address(swapRouter), uint256(approveAmount));
 
         BalanceDelta delta = swapRouter.swap(