harb/onchain/test/LiquidityManager.t.sol

// SPDX-License-Identifier: GPL-3.0-or-later
pragma solidity ^0.8.19;

import "forge-std/Test.sol";
import "@aperture/uni-v3-lib/TickMath.sol";
import {LiquidityAmounts} from "@aperture/uni-v3-lib/LiquidityAmounts.sol";
import {WETH} from "solmate/tokens/WETH.sol";
import {PoolAddress, PoolKey} from "@aperture/uni-v3-lib/PoolAddress.sol";
import "@uniswap-v3-core/interfaces/IUniswapV3Factory.sol";
import "@uniswap-v3-core/interfaces/IUniswapV3Pool.sol";
import "../src/interfaces/IWETH9.sol";
import {Harberg} from "../src/Harberg.sol";

import {Stake, ExceededAvailableStake} from "../src/Stake.sol";
import {LiquidityManager} from "../src/LiquidityManager.sol";
import "../src/helpers/UniswapHelpers.sol";
import {CSVHelper} from "./helpers/CSVHelper.sol";
import {CSVManager} from "./helpers/CSVManager.sol";
import {UniswapTestBase} from "./helpers/UniswapTestBase.sol";
import "../src/Sentimenter.sol";
import "../test/mocks/MockSentimenter.sol";

address constant TAX_POOL = address(2);
// default fee of 1%
uint24 constant FEE = uint24(10_000);
int24 constant TICK_SPACING = 200;
int24 constant ANCHOR_SPACING = 5 * TICK_SPACING;

// Dummy.sol
contract Dummy {
    // This contract can be empty as it is only used to affect the nonce
}

contract LiquidityManagerTest is UniswapTestBase, CSVManager {
    using UniswapHelpers for IUniswapV3Pool;
    using CSVHelper for *;

    IUniswapV3Factory factory;
    Stake stake;
    LiquidityManager lm;
    address feeDestination = makeAddr("fees");
    
    struct Response {
	uint256 ethFloor;
	uint256 ethAnchor;
	uint256 ethDiscovery;
	uint256 harbergFloor;
	uint256 harbergAnchor;
	uint256 harbergDiscovery;
    }

    // Utility to deploy dummy contracts
    function deployDummies(uint count) internal {
        for (uint i = 0; i < count; i++) {
            new Dummy(); // Just increment the nonce
        }
    }

    function setUpCustomToken0(bool token0shouldBeWeth) public {
        factory = UniswapHelpers.deployUniswapFactory();

        bool setupComplete = false;
        uint retryCount = 0;
        while (!setupComplete && retryCount < 5) {
            // Clean slate if retrying
            if (retryCount > 0) {
                deployDummies(1); // Deploy a dummy contract to shift addresses
            }

            weth = IWETH9(address(new WETH()));
            harberg = new Harberg("HARB", "HARB");

            // Check if the setup meets the required condition
            if (token0shouldBeWeth == address(weth) < address(harberg)) {
                setupComplete = true;
            } else {
                // Clear current instances for re-deployment
                delete weth;
                delete harberg;
                retryCount++;
            }
        }
        require(setupComplete, "Setup failed to meet the condition after several retries");

        pool = IUniswapV3Pool(factory.createPool(address(weth), address(harberg), FEE));

        token0isWeth = address(weth) < address(harberg);
        pool.initializePoolFor1Cent(token0isWeth);

        stake = new Stake(address(harberg), feeDestination);
        harberg.setStakingPool(address(stake));
	Sentimenter senti = Sentimenter(address(new MockSentimenter()));
	senti.initialize(address(harberg), address(stake));
        lm = new LiquidityManager(address(factory), address(weth), address(harberg), address(senti));
        lm.setFeeDestination(feeDestination);
        vm.prank(feeDestination);
        harberg.setLiquidityManager(address(lm));
        vm.deal(address(lm), 10 ether);
        initializePositionsCSV(); // Set up the CSV header
    }

    function recenter(bool last) internal {
        // have some time pass to record prices in uni oracle
        uint256 timeBefore = block.timestamp;
        vm.warp(timeBefore + (60 * 60 * 5));

	try lm.recenter() returns (bool isUp, uint256 sentiment) {

		// Check liquidity positions after slide
		Response memory rsp;
		rsp = checkLiquidity(isUp ? "shift" : "slide", sentiment);
		assertGt(rsp.ethFloor, rsp.ethAnchor, "slide - Floor should hold more ETH than Anchor");
		assertGt(rsp.harbergDiscovery, rsp.harbergAnchor * 5, "slide - Discovery should hold more HARB than Anchor");
		assertEq(rsp.harbergFloor, 0, "slide - Floor should have no HARB");
		assertEq(rsp.ethDiscovery, 0, "slide - Discovery should have no ETH");

	} catch Error(string memory reason) {
		if (keccak256(abi.encodePacked(reason)) == keccak256(abi.encodePacked("amplitude not reached."))) {
			console.log("slide failed on amplitude");
		} else {
			if (!last) {
				revert(reason); // Rethrow the error if it's not the expected message
			}
		}
	}
    }

    function getBalancesPool(LiquidityManager.Stage s) internal view returns (int24 currentTick, int24 tickLower, int24 tickUpper, uint256 ethAmount, uint256 harbergAmount) {
	    (,tickLower, tickUpper) = lm.positions(s);
	    (uint128 liquidity, , , ,) = pool.positions(keccak256(abi.encodePacked(address(lm), tickLower, tickUpper)));

	    // Fetch the current price from the pool
	    uint160 sqrtPriceX96;
	    (sqrtPriceX96, currentTick, , , , , ) = pool.slot0();
	    uint160 sqrtPriceAX96 = TickMath.getSqrtRatioAtTick(tickLower);
	    uint160 sqrtPriceBX96 = TickMath.getSqrtRatioAtTick(tickUpper);

	    // Calculate amounts based on the current tick position relative to provided ticks
	    if (token0isWeth) {
		    if (currentTick < tickLower) {
			    // Current price is below the lower bound of the liquidity position
			    ethAmount = LiquidityAmounts.getAmount0ForLiquidity(sqrtPriceAX96, sqrtPriceBX96, liquidity);
			    harbergAmount = 0;  // All liquidity is in token0 (ETH)
		    } else if (currentTick > tickUpper) {
			    // Current price is above the upper bound of the liquidity position
			    ethAmount = 0;  // All liquidity is in token1 (HARB)
			    harbergAmount = LiquidityAmounts.getAmount1ForLiquidity(sqrtPriceAX96, sqrtPriceBX96, liquidity);
		    } else {
			    // Current price is within the bounds of the liquidity position
			    ethAmount = LiquidityAmounts.getAmount0ForLiquidity(sqrtPriceX96, sqrtPriceBX96, liquidity);
			    harbergAmount = LiquidityAmounts.getAmount1ForLiquidity(sqrtPriceAX96, sqrtPriceX96, liquidity);
		    }
	    } else {
		    if (currentTick < tickLower) {
			    // Current price is below the lower bound of the liquidity position
			    harbergAmount = LiquidityAmounts.getAmount0ForLiquidity(sqrtPriceAX96, sqrtPriceBX96, liquidity);
			    ethAmount = 0;  // All liquidity is in token1 (ETH)
		    } else if (currentTick > tickUpper) {
			    // Current price is above the upper bound of the liquidity position
			    harbergAmount = 0;  // All liquidity is in token0 (HARB)
			    ethAmount = LiquidityAmounts.getAmount1ForLiquidity(sqrtPriceAX96, sqrtPriceBX96, liquidity);
		    } else {
			    // Current price is within the bounds of the liquidity position
			    harbergAmount = LiquidityAmounts.getAmount0ForLiquidity(sqrtPriceX96, sqrtPriceBX96, liquidity);
			    ethAmount = LiquidityAmounts.getAmount1ForLiquidity(sqrtPriceAX96, sqrtPriceX96, liquidity);
		    }
	    }
    }

    function checkLiquidity(string memory eventName, uint256 sentiment) internal returns (Response memory) {
	Response memory rsp;
	int24 currentTick;
	string memory floorData; 
	string memory anchorData; 
	string memory discoveryData; 
	{
	    int24 tickLower;
	    int24 tickUpper;
	    uint256 eth;
	    uint256 harb;
	    {
		(currentTick, tickLower, tickUpper, eth, harb) = getBalancesPool(LiquidityManager.Stage.FLOOR);
		floorData = string(abi.encodePacked(CSVHelper.intToStr(tickLower), ",", CSVHelper.intToStr(tickUpper), ",", CSVHelper.uintToStr(eth), ",", CSVHelper.uintToStr(harb), ","));
		rsp.ethFloor = eth;
		rsp.harbergFloor = harb;
	    }
	    {
		(,tickLower, tickUpper, eth, harb) = getBalancesPool(LiquidityManager.Stage.ANCHOR);
		anchorData = string(abi.encodePacked(CSVHelper.intToStr(tickLower), ",", CSVHelper.intToStr(tickUpper), ",", CSVHelper.uintToStr(eth), ",", CSVHelper.uintToStr(harb), ","));
		rsp.ethAnchor = eth;
		rsp.harbergAnchor = harb;
	    }
	    {
		(,tickLower, tickUpper, eth, harb) = getBalancesPool(LiquidityManager.Stage.DISCOVERY);
		discoveryData = string(abi.encodePacked(CSVHelper.intToStr(tickLower), ",", CSVHelper.intToStr(tickUpper), ",", CSVHelper.uintToStr(eth), ",", CSVHelper.uintToStr(harb), ","));
		rsp.ethDiscovery = eth;
		rsp.harbergDiscovery = harb;
	    }
	}

	string memory newRow = string(abi.encodePacked(eventName, ",", CSVHelper.intToStr(currentTick), ",", CSVHelper.uintToStr(sentiment / 1e12), ",", floorData, anchorData, discoveryData));
        appendCSVRow(newRow); // Append the new row to the CSV
	return rsp;
    }

    function buy(uint256 amountEth) internal {
        performSwap(amountEth, true);
        checkLiquidity(string.concat("buy ", CSVHelper.uintToStr(amountEth)), 0);
    }

    function sell(uint256 amountHarb) internal {
        performSwap(amountHarb, false);
        checkLiquidity(string.concat("sell ", CSVHelper.uintToStr(amountHarb)), 0);
    }

    receive() external payable {}


    function writeCsv() public {
        writeCSVToFile("./out/positions.csv"); // Write CSV to file
    }

    function testHandleCumulativeOverflow() public {
        setUpCustomToken0(false);
        vm.deal(account, 201 ether);
        vm.prank(account);
        weth.deposit{value: 201 ether}();

        // Setup initial liquidity
        recenter(false);

        vm.store(
            address(lm),
            bytes32(uint256(0)),
            bytes32(uint256(type(uint256).max - 10))
        );

        vm.store(
            address(lm),
            bytes32(uint256(1)),
            bytes32(uint256((type(uint256).max - 10) / (3000 * 10**20)))
        );


        uint256 cumulativeVolumeWeightedPriceX96 = lm.cumulativeVolumeWeightedPriceX96();
        uint256 beforeCumulativeVolume = lm.cumulativeVolume();

        assertGt(cumulativeVolumeWeightedPriceX96, type(uint256).max / 2, "Initial cumulativeVolumeWeightedPrice is not near max uint256");

        buy(25 ether);

        recenter(false);

        cumulativeVolumeWeightedPriceX96 = lm.cumulativeVolumeWeightedPriceX96();
        uint256 cumulativeVolume = lm.cumulativeVolume();

        // Assert that the values after wrap-around are valid and smaller than max uint256
        assertGt(beforeCumulativeVolume, cumulativeVolume, "cumulativeVolume after wrap-around is smaller than before");

        // Assert that the price is reasonable
        uint256 calculatedPrice = cumulativeVolumeWeightedPriceX96 / cumulativeVolume;
        assertTrue(calculatedPrice > 0 && calculatedPrice < 10**40, "Calculated price after wrap-around is not within a reasonable range");
    }

    // function testScenarioBuyAll() public {
    //     setUpCustomToken0(false);
    //     vm.deal(account, 300 ether);
    //     vm.prank(account);
    //     weth.deposit{value: 300 ether}();

    //     uint256 traderBalanceBefore = weth.balanceOf(account);

    //     // Setup initial liquidity
    //     recenter(false);

    //     buy(200 ether);

    //     recenter(false);

    //     //revert();

    //     sell(harberg.balanceOf(account));

    //     recenter(true);

    //     writeCsv();

    //     uint256 traderBalanceAfter = weth.balanceOf(account);

    //     assertGt(traderBalanceBefore, traderBalanceAfter, "trader should not have made profit");
    // }

    // function testScenarioB() public {
    //     setUpCustomToken0(false);
    //     vm.deal(account, 501 ether);
    //     vm.prank(account);
    //     weth.deposit{value: 501 ether}();

    //     uint256 traderBalanceBefore = weth.balanceOf(account);

    //     // Setup initial liquidity
    //     recenter(false);

    //     buy(25 ether);

    //     recenter(false);

    //     buy(45 ether);

    //     recenter(false);

    //     buy(80 ether);

    //     recenter(false);

    //     buy(120 ether);

    //     recenter(false);

    //     sell(harberg.balanceOf(account) / 4);

    //     recenter(true);

    //     sell(harberg.balanceOf(account) / 4);

    //     recenter(true);

    //     sell(harberg.balanceOf(account) / 4);

    //     recenter(true);

    //     sell(harberg.balanceOf(account));

    //     recenter(true);

    //     writeCsv();
    //     uint256 traderBalanceAfter = weth.balanceOf(account);
    //     console.log(traderBalanceBefore);
    //     console.log(traderBalanceAfter);
    //     assertGt(traderBalanceBefore, traderBalanceAfter, "trader should not have made profit");
    //     revert();
    // }

    function testScenarioFuzz(uint8 numActions, uint8 frequency, uint8[] calldata amounts) public {
        vm.assume(numActions > 5);
        vm.assume(frequency > 0);
        vm.assume(frequency < 20);
        vm.assume(amounts.length >= numActions);

        setUpCustomToken0(numActions % 2 == 0 ? true : false);
        vm.deal(account, 20 ether);
        vm.prank(account);
        weth.deposit{value: 20 ether}();


        // Setup initial liquidity
        recenter(false);

        uint256 traderBalanceBefore = weth.balanceOf(account);
        uint8 f = 0;
        for (uint i = 0; i < numActions; i++) {
            uint256 amount = (uint256(amounts[i]) * 1 ether) + 1 ether;
            uint256 harbergBal = harberg.balanceOf(account);
            if (harbergBal == 0) {
                amount = amount % (weth.balanceOf(account) / 2);
                amount = amount == 0 ? weth.balanceOf(account) : amount;
                buy(amount);
            } else if (weth.balanceOf(account) == 0) {
                sell(amount % harbergBal);
            } else {
                if (amount % 2 == 0) {
                    amount = amount % (weth.balanceOf(account) / 2);
                    amount = amount == 0 ? weth.balanceOf(account) : amount;
                    buy(amount);
                } else {
                    sell(amount % harbergBal);
                }
            }

            (, int24 currentTick, , , , , ) = pool.slot0();
            if (currentTick < -887270) {
               // buy(1000000000000000);
                sell(100000000000000);
            }
            if (currentTick > 887270) {
               buy(1000000000000000);
               // sell(100000000000000);
            }
            if (f >= frequency) {
		recenter(false);
                f = 0;
            } else {
                f++;
            }
        }

        // Simulate large sell to push price down to floor
        sell(harberg.balanceOf(account));

        recenter(true);

        uint256 traderBalanceAfter = weth.balanceOf(account);

        if (traderBalanceAfter > traderBalanceBefore){
            writeCsv();
        }
        // TODO: take 1% fee into account
        assertGt(traderBalanceBefore, traderBalanceAfter, "trader should not have made profit");
    }

}