fluxswap/test/FluxClInvariant.t.sol

// SPDX-License-Identifier: MIT
pragma solidity 0.8.36;

import {Test} from "forge-std/Test.sol";
import {StdInvariant} from "forge-std/StdInvariant.sol";
import {FluxClFactory} from "../src/FluxClFactory.sol";
import {FluxClPool} from "../src/FluxClPool.sol";
import {TestErc20} from "../src/TestErc20.sol";
import {IErc20} from "../src/interfaces/IErc20.sol";
import {IFluxMintCallback} from "../src/interfaces/IFluxMintCallback.sol";
import {IFluxSwapCallback} from "../src/interfaces/IFluxSwapCallback.sol";
import {TickMath} from "../src/libraries/TickMath.sol";

contract ClHandler is IFluxMintCallback, IFluxSwapCallback {
    FluxClPool public pool;
    TestErc20 public token0;
    TestErc20 public token1;
    int24 private constant spacing = 60;

    int256 public ghost0;
    int256 public ghost1;

    struct Range {
        int24 lower;
        int24 upper;
    }

    Range[] public ranges;
    mapping(bytes32 => uint128) public liquidityOf;

    constructor(FluxClPool poolAddress, TestErc20 first, TestErc20 second) {
        pool = poolAddress;
        token0 = first;
        token1 = second;
        token0.mint(address(this), 1e33);
        token1.mint(address(this), 1e33);
    }

    function fluxMintCallback(uint256 amount0Owed, uint256 amount1Owed, bytes calldata) external {
        if (amount0Owed > 0) token0.transfer(msg.sender, amount0Owed);
        if (amount1Owed > 0) token1.transfer(msg.sender, amount1Owed);
        ghost0 += int256(amount0Owed);
        ghost1 += int256(amount1Owed);
    }

    function fluxSwapCallback(int256 amount0Delta, int256 amount1Delta, bytes calldata) external {
        if (amount0Delta > 0) token0.transfer(msg.sender, uint256(amount0Delta));
        if (amount1Delta > 0) token1.transfer(msg.sender, uint256(amount1Delta));
        ghost0 += amount0Delta;
        ghost1 += amount1Delta;
    }

    function doMint(uint128 amountSeed, int24 lowerSeed, int24 widthSeed) external {
        uint128 amount = uint128(bound(uint256(amountSeed), 1e12, 1e24));
        int24 lower = int24(bound(int256(lowerSeed), -30000, 30000)) / spacing * spacing;
        int24 width = int24(bound(int256(widthSeed), 1, 100)) * spacing;
        int24 upper = lower + width;
        if (upper > TickMath.maxTick) return;

        try pool.mint(address(this), lower, upper, amount, "") {
            bytes32 key = keccak256(abi.encodePacked(lower, upper));
            if (liquidityOf[key] == 0) ranges.push(Range({lower: lower, upper: upper}));
            liquidityOf[key] += amount;
        } catch {}
    }

    function doSwap(uint256 amountSeed, bool zeroForOne) external {
        int256 amountSpecified = int256(bound(amountSeed, 1e6, 1e26));
        uint160 limit = zeroForOne ? TickMath.minSqrtRatio + 1 : TickMath.maxSqrtRatio - 1;
        try pool.swap(address(this), zeroForOne, amountSpecified, limit, "") {} catch {}
    }

    function doBurn(uint256 rangeSeed, uint128 amountSeed) external {
        if (ranges.length == 0) return;
        Range memory r = ranges[rangeSeed % ranges.length];
        bytes32 key = keccak256(abi.encodePacked(r.lower, r.upper));
        uint128 have = liquidityOf[key];
        if (have == 0) return;
        uint128 amount = uint128(bound(uint256(amountSeed), 1, have));

        try pool.burn(r.lower, r.upper, amount) {
            liquidityOf[key] -= amount;
            (uint128 c0, uint128 c1) =
                pool.collect(address(this), r.lower, r.upper, type(uint128).max, type(uint128).max);
            ghost0 -= int256(uint256(c0));
            ghost1 -= int256(uint256(c1));
        } catch {}
    }

    function bound(uint256 x, uint256 min, uint256 max) private pure returns (uint256) {
        if (max <= min) return min;
        return min + (x % (max - min + 1));
    }

    function bound(int256 x, int256 min, int256 max) private pure returns (int256) {
        if (max <= min) return min;
        uint256 range = uint256(max - min) + 1;
        uint256 offset = uint256(x < 0 ? -x : x) % range;
        return min + int256(offset);
    }

    function rangesLength() external view returns (uint256) {
        return ranges.length;
    }
}

contract FluxClInvariantTest is StdInvariant, Test {
    uint160 private constant priceOne = 79228162514264337593543950336;

    FluxClFactory private factory;
    FluxClPool private pool;
    TestErc20 private token0;
    TestErc20 private token1;
    ClHandler private handler;

    function setUp() public {
        factory = new FluxClFactory();
        TestErc20 tokenA = new TestErc20("Alpha", "ALPHA", 0);
        TestErc20 tokenB = new TestErc20("Beta", "BETA", 0);
        (address a, address b) =
            address(tokenA) < address(tokenB) ? (address(tokenA), address(tokenB)) : (address(tokenB), address(tokenA));
        token0 = TestErc20(a);
        token1 = TestErc20(b);

        pool = FluxClPool(factory.createPool(a, b, 3000));
        pool.initialize(priceOne);

        handler = new ClHandler(pool, token0, token1);

        bytes4[] memory selectors = new bytes4[](3);
        selectors[0] = ClHandler.doMint.selector;
        selectors[1] = ClHandler.doSwap.selector;
        selectors[2] = ClHandler.doBurn.selector;
        targetSelector(FuzzSelector({addr: address(handler), selectors: selectors}));
        targetContract(address(handler));
    }

    function invariantPoolBalanceMatchesAccounting() public view {
        assertEq(int256(token0.balanceOf(address(pool))), handler.ghost0());
        assertEq(int256(token1.balanceOf(address(pool))), handler.ghost1());
    }

    function invariantPriceStaysWithinTickRange() public view {
        (uint160 sqrtP, int24 tick,,,,) = pool.slot0();
        assertGe(sqrtP, TickMath.minSqrtRatio);
        assertLt(sqrtP, TickMath.maxSqrtRatio);
        assertGe(tick, TickMath.minTick);
        assertLe(tick, TickMath.maxTick);
    }

    function invariantPoolNeverOwesMoreThanItHolds() public view {
        assertGe(handler.ghost0(), 0);
        assertGe(handler.ghost1(), 0);
    }
}