harb/onchain/script/backtesting/EventReplayer.sol

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

import { MockToken } from "./MockToken.sol";
import { StrategyExecutor } from "./StrategyExecutor.sol";
import { TickMath } from "@aperture/uni-v3-lib/TickMath.sol";
import { IUniswapV3Pool } from "@uniswap-v3-core/interfaces/IUniswapV3Pool.sol";
import { IUniswapV3MintCallback } from "@uniswap-v3-core/interfaces/callback/IUniswapV3MintCallback.sol";
import { IUniswapV3SwapCallback } from "@uniswap-v3-core/interfaces/callback/IUniswapV3SwapCallback.sol";

import { Vm } from "forge-std/Vm.sol";
import { console2 } from "forge-std/console2.sol";

/**
 * @title EventReplayer
 * @notice Replays historical Swap/Mint/Burn events from a JSON Lines cache
 *         against a shadow Uniswap V3 pool.
 *
 * Architecture note:
 *   EventReplayer is the owner of all shadow-pool positions because
 *   IUniswapV3Pool.mint() and pool.swap() call back to msg.sender to collect
 *   tokens.  vm.prank() cannot be combined with pool callbacks (the pranked
 *   address has no code).  Historical owner/sender addresses are therefore
 *   ignored for position accounting; only tick ranges and amounts matter.
 *
 *   fetch-events.ts serialises BigInt Solidity values (int256, uint160, int24,
 *   uint128, …) as **decimal strings** in JSON.  Use vm.parseJsonString() +
 *   vm.parseInt()/vm.parseUint() for those fields.  The `block` and `logIndex`
 *   fields are JS numbers and may be read directly with vm.parseJsonUint().
 */
contract EventReplayer is IUniswapV3MintCallback, IUniswapV3SwapCallback {
    // -------------------------------------------------------------------------
    // Constants
    // -------------------------------------------------------------------------

    /// @dev Foundry cheatcode address (same constant as in Vm.sol).
    Vm internal constant vm = Vm(address(uint160(uint256(keccak256("hevm cheat code")))));

    /// @dev Log progress and check drift every N events.
    uint256 internal constant LOG_INTERVAL = 100;

    /// @dev Approximate seconds per block on Base mainnet (~2 s).
    uint256 internal constant BLOCK_DURATION = 2;

    // -------------------------------------------------------------------------
    // Immutables
    // -------------------------------------------------------------------------

    IUniswapV3Pool public immutable pool;
    MockToken public immutable token0;
    MockToken public immutable token1;
    /// @dev Pool tick spacing — cached once at construction (immutable for the pool lifetime).
    int24 public immutable tickSpacing;

    // -------------------------------------------------------------------------
    // Stats (accumulated across the replay)
    // -------------------------------------------------------------------------

    /// @dev Count of LOG_INTERVAL checkpoints (and the final sample) where tick drift > 0.
    uint256 public checkpointsWithDrift;
    uint256 public totalAbsDrift;
    uint256 public maxDrift;
    uint256 public skippedCount;

    // -------------------------------------------------------------------------
    // Internal chain-state tracking
    // -------------------------------------------------------------------------

    uint256 internal _lastBlock;
    uint256 internal _lastTimestamp;

    // -------------------------------------------------------------------------
    // Constructor
    // -------------------------------------------------------------------------

    constructor(IUniswapV3Pool _pool, MockToken _token0, MockToken _token1) {
        pool = _pool;
        token0 = _token0;
        token1 = _token1;
        tickSpacing = _pool.tickSpacing();
        _lastBlock = block.number;
        _lastTimestamp = block.timestamp;
    }

    // -------------------------------------------------------------------------
    // Public entry point
    // -------------------------------------------------------------------------

    /**
     * @notice Read every line of the JSON Lines events file and replay it against
     *         the shadow pool.
     * @param eventsFile  Absolute (or project-relative) path to the .jsonl cache.
     * @param totalEvents Total number of events in the file (for progress display).
     *                    Pass 0 to omit the denominator from progress logs.
     */
    function replay(string memory eventsFile, uint256 totalEvents) external {
        _replayWithStrategy(eventsFile, totalEvents, StrategyExecutor(address(0)));
    }

    /**
     * @notice Replay events and trigger KrAIken recenter logic after each block.
     * @param eventsFile       Path to the .jsonl events cache.
     * @param totalEvents      Total event count (for progress logs).
     * @param strategyExecutor StrategyExecutor to call after every block advancement.
     *                         Pass address(0) to disable strategy integration.
     */
    function replay(string memory eventsFile, uint256 totalEvents, StrategyExecutor strategyExecutor) external {
        _replayWithStrategy(eventsFile, totalEvents, strategyExecutor);
    }

    function _replayWithStrategy(string memory eventsFile, uint256 totalEvents, StrategyExecutor strategyExecutor) internal {
        uint256 idx = 0;

        // Track the last Swap event's expected state for drift measurement.
        int24 lastExpectedTick;
        uint160 lastExpectedSqrtPrice;
        bool hasSwapRef;

        while (true) {
            string memory line;
            try vm.readLine(eventsFile) returns (string memory l) {
                line = l;
            } catch {
                break; // EOF or read error — done
            }
            if (bytes(line).length == 0) break; // empty line signals EOF

            string memory eventName = vm.parseJsonString(line, ".event");
            uint256 blockNum = vm.parseJsonUint(line, ".block");

            _advanceChain(blockNum);

            // After advancing chain state, give the strategy a chance to recenter.
            // The call is best-effort: any revert is caught inside StrategyExecutor.
            if (address(strategyExecutor) != address(0)) {
                strategyExecutor.maybeRecenter(blockNum);
            }

            if (_streq(eventName, "Swap")) {
                (int24 expTick, uint160 expSqrtPrice) = _replaySwap(line);
                // Update reference state only when the swap was not skipped.
                if (expSqrtPrice != 0) {
                    lastExpectedTick = expTick;
                    lastExpectedSqrtPrice = expSqrtPrice;
                    hasSwapRef = true;
                }
            } else if (_streq(eventName, "Mint")) {
                _replayMint(line);
            } else if (_streq(eventName, "Burn")) {
                _replayBurn(line);
            }

            idx++;

            // Progress + drift validation every LOG_INTERVAL events.
            if (idx % LOG_INTERVAL == 0 && hasSwapRef) {
                _logCheckpoint(idx, totalEvents, lastExpectedTick, lastExpectedSqrtPrice);
            }
        }

        // Fetch final pool state once — used both for the trailing drift sample and _logSummary.
        (uint160 finalSqrtPrice, int24 finalTick,,,,,) = pool.slot0();

        // Final drift sample: captures trailing events after the last checkpoint.
        // Guard: when idx is an exact multiple of LOG_INTERVAL, _logCheckpoint already fired for
        // this identical pool state inside the loop — accumulating stats again would double-count
        // that measurement in totalAbsDrift and checkpointsWithDrift.
        if (hasSwapRef && idx % LOG_INTERVAL != 0) {
            int256 diff = int256(finalTick) - int256(lastExpectedTick);
            uint256 absDrift = diff >= 0 ? uint256(diff) : uint256(-diff);
            totalAbsDrift += absDrift;
            if (absDrift > 0) checkpointsWithDrift++;
            if (absDrift > maxDrift) maxDrift = absDrift;
            // Log sqrtPrice deviation when it exceeds ~0.01% (filters rounding noise).
            if (finalSqrtPrice != lastExpectedSqrtPrice) {
                uint256 priceDelta =
                    finalSqrtPrice > lastExpectedSqrtPrice ? uint256(finalSqrtPrice - lastExpectedSqrtPrice) : uint256(lastExpectedSqrtPrice - finalSqrtPrice);
                if (lastExpectedSqrtPrice > 0 && priceDelta * 10_000 > uint256(lastExpectedSqrtPrice)) {
                    console2.log("  final sqrtPrice divergence:", priceDelta);
                }
            }
        }

        _logSummary(idx, finalSqrtPrice, finalTick);
    }

    // -------------------------------------------------------------------------
    // Swap
    // -------------------------------------------------------------------------

    /**
     * @notice Replay a Swap event.
     * @return expectedTick       Cached tick from the event (0 if skipped).
     * @return expectedSqrtPrice  Cached sqrtPriceX96 from the event (0 if skipped).
     */
    function _replaySwap(string memory line) internal returns (int24 expectedTick, uint160 expectedSqrtPrice) {
        // BigInt fields are serialised as decimal strings by fetch-events.ts.
        int256 amount0 = vm.parseInt(vm.parseJsonString(line, ".amount0"));
        int256 amount1 = vm.parseInt(vm.parseJsonString(line, ".amount1"));
        uint160 targetSqrtPrice = uint160(vm.parseUint(vm.parseJsonString(line, ".sqrtPriceX96")));
        int24 evtTick = int24(int256(vm.parseInt(vm.parseJsonString(line, ".tick"))));

        // Skip degenerate zero-amount events.
        if (amount0 == 0 && amount1 == 0) {
            skippedCount++;
            return (0, 0);
        }

        // amount0 > 0  →  caller paid token0 into pool (zeroForOne = true).
        // amount1 > 0  →  caller paid token1 into pool (zeroForOne = false).
        bool zeroForOne = amount0 > 0;

        // amountSpecified: exact-input using the historical token amount.
        // Guard against degenerate events where the expected input side is non-positive,
        // which would unintentionally switch pool.swap into exact-output mode.
        int256 amountSpecified = zeroForOne ? amount0 : amount1;
        if (amountSpecified <= 0) {
            skippedCount++;
            return (0, 0);
        }

        // Skip if the pool is already at or past the ORIGINAL target price.
        // Perform this check against targetSqrtPrice — before any clamping — so that
        // events targeting a price exactly at a hard limit are not incorrectly skipped.
        (uint160 currentSqrtPrice,,,,,,) = pool.slot0();
        if (zeroForOne && currentSqrtPrice <= targetSqrtPrice) {
            skippedCount++;
            return (0, 0);
        }
        if (!zeroForOne && currentSqrtPrice >= targetSqrtPrice) {
            skippedCount++;
            return (0, 0);
        }

        // Clamp the price LIMIT for pool.swap() to the exclusive valid range.
        // Applied after the skip check so extreme-price events are not silently dropped;
        // only the pool.swap parameter is adjusted.
        uint160 sqrtPriceLimitX96 = targetSqrtPrice;
        if (zeroForOne && sqrtPriceLimitX96 <= TickMath.MIN_SQRT_RATIO) {
            sqrtPriceLimitX96 = TickMath.MIN_SQRT_RATIO + 1;
        }
        if (!zeroForOne && sqrtPriceLimitX96 >= TickMath.MAX_SQRT_RATIO) {
            sqrtPriceLimitX96 = TickMath.MAX_SQRT_RATIO - 1;
        }

        // Pre-fund so the swap callback can pay the pool without needing to mint
        // inside the callback (which avoids a re-entrant MockToken.mint call).
        // amountSpecified > 0 is guaranteed by the guard above, so the cast is safe.
        uint256 fundAmount = uint256(amountSpecified) + 1;
        if (zeroForOne) {
            token0.mint(address(this), fundAmount);
        } else {
            token1.mint(address(this), fundAmount);
        }

        try pool.swap(address(this), zeroForOne, amountSpecified, sqrtPriceLimitX96, abi.encode(zeroForOne)) {
            // success
        } catch {
            skippedCount++;
            return (0, 0);
        }

        return (evtTick, targetSqrtPrice);
    }

    // -------------------------------------------------------------------------
    // Mint
    // -------------------------------------------------------------------------

    function _replayMint(string memory line) internal {
        uint128 amount = uint128(vm.parseUint(vm.parseJsonString(line, ".amount")));
        int24 tickLower = int24(int256(vm.parseInt(vm.parseJsonString(line, ".tickLower"))));
        int24 tickUpper = int24(int256(vm.parseInt(vm.parseJsonString(line, ".tickUpper"))));

        // Skip zero-liquidity and degenerate tick ranges.
        if (amount == 0 || tickLower >= tickUpper) {
            skippedCount++;
            return;
        }

        // Skip ticks outside Uniswap's hard limits.
        if (tickLower < TickMath.MIN_TICK || tickUpper > TickMath.MAX_TICK) {
            skippedCount++;
            return;
        }

        // Skip ticks not aligned to the pool's tick spacing (uses cached immutable).
        if (tickLower % tickSpacing != 0 || tickUpper % tickSpacing != 0) {
            skippedCount++;
            return;
        }

        // The mint callback (uniswapV3MintCallback) will mint tokens as needed.
        try pool.mint(address(this), tickLower, tickUpper, amount, "") {
            // success — position recorded under address(this)
        } catch {
            skippedCount++;
        }
    }

    // -------------------------------------------------------------------------
    // Burn
    // -------------------------------------------------------------------------

    function _replayBurn(string memory line) internal {
        uint128 amount = uint128(vm.parseUint(vm.parseJsonString(line, ".amount")));
        int24 tickLower = int24(int256(vm.parseInt(vm.parseJsonString(line, ".tickLower"))));
        int24 tickUpper = int24(int256(vm.parseInt(vm.parseJsonString(line, ".tickUpper"))));

        if (amount == 0) {
            skippedCount++;
            return;
        }

        // Look up the shadow position owned by this contract.
        bytes32 posKey = keccak256(abi.encodePacked(address(this), tickLower, tickUpper));
        (uint128 posLiquidity,,,,) = pool.positions(posKey);

        // No shadow position to burn (historical mint predates our event window).
        if (posLiquidity == 0) {
            skippedCount++;
            return;
        }

        // Clamp to available shadow liquidity (partial-burn edge case).
        if (amount > posLiquidity) amount = posLiquidity;

        try pool.burn(tickLower, tickUpper, amount) {
            // success
        } catch {
            skippedCount++;
        }
    }

    // -------------------------------------------------------------------------
    // Uniswap V3 callbacks
    // -------------------------------------------------------------------------

    /// @inheritdoc IUniswapV3MintCallback
    function uniswapV3MintCallback(uint256 amount0Owed, uint256 amount1Owed, bytes calldata) external override {
        require(msg.sender == address(pool), "EventReplayer: bad mint callback");
        // Mint exactly what the pool needs and transfer it immediately.
        if (amount0Owed > 0) {
            token0.mint(address(this), amount0Owed);
            token0.transfer(msg.sender, amount0Owed);
        }
        if (amount1Owed > 0) {
            token1.mint(address(this), amount1Owed);
            token1.transfer(msg.sender, amount1Owed);
        }
    }

    /// @inheritdoc IUniswapV3SwapCallback
    function uniswapV3SwapCallback(int256 amount0Delta, int256 amount1Delta, bytes calldata data) external override {
        require(msg.sender == address(pool), "EventReplayer: bad swap callback");
        bool zeroForOne = abi.decode(data, (bool));
        // Pay the positive delta (the input token the pool is owed).
        if (zeroForOne && amount0Delta > 0) {
            token0.transfer(msg.sender, uint256(amount0Delta));
        } else if (!zeroForOne && amount1Delta > 0) {
            token1.transfer(msg.sender, uint256(amount1Delta));
        }
    }

    // -------------------------------------------------------------------------
    // Chain-state helpers
    // -------------------------------------------------------------------------

    /**
     * @notice Advance block number and timestamp to match the historical event's
     *         block.  Only moves forward; events in the same block are no-ops.
     */
    function _advanceChain(uint256 blockNum) internal {
        if (blockNum <= _lastBlock) return;
        uint256 delta = blockNum - _lastBlock;
        _lastTimestamp += delta * BLOCK_DURATION;
        _lastBlock = blockNum;
        vm.roll(blockNum);
        vm.warp(_lastTimestamp);
    }

    // -------------------------------------------------------------------------
    // Logging helpers
    // -------------------------------------------------------------------------

    /**
     * @notice Emit a progress line and accumulate drift statistics for one checkpoint.
     */
    function _logCheckpoint(uint256 idx, uint256 totalEvents, int24 expectedTick, uint160 expectedSqrtPrice) internal {
        (uint160 currentSqrtPrice, int24 currentTick,,,,,) = pool.slot0();
        int256 diff = int256(currentTick) - int256(expectedTick);
        uint256 absDrift = diff >= 0 ? uint256(diff) : uint256(-diff);

        totalAbsDrift += absDrift;
        if (absDrift > 0) checkpointsWithDrift++;
        if (absDrift > maxDrift) maxDrift = absDrift;

        string memory denominator = totalEvents > 0 ? string.concat("/", vm.toString(totalEvents)) : "";
        console2.log(
            string.concat(
                "[",
                vm.toString(idx),
                denominator,
                "] tick=",
                vm.toString(int256(currentTick)),
                " expected=",
                vm.toString(int256(expectedTick)),
                " drift=",
                vm.toString(absDrift)
            )
        );

        // Log sqrtPrice deviation when it exceeds ~0.01% (filters rounding noise).
        if (currentSqrtPrice != expectedSqrtPrice) {
            uint256 priceDelta =
                currentSqrtPrice > expectedSqrtPrice ? uint256(currentSqrtPrice - expectedSqrtPrice) : uint256(expectedSqrtPrice - currentSqrtPrice);
            if (expectedSqrtPrice > 0 && priceDelta * 10_000 > uint256(expectedSqrtPrice)) {
                console2.log("  sqrtPrice divergence:", priceDelta);
            }
        }
    }

    // -------------------------------------------------------------------------
    // Utilities
    // -------------------------------------------------------------------------

    function _streq(string memory a, string memory b) internal pure returns (bool) {
        return keccak256(bytes(a)) == keccak256(bytes(b));
    }

    function _logSummary(uint256 totalReplayed, uint160 finalSp, int24 finalTick) internal view {
        console2.log("=== Replay Complete ===");
        console2.log("Total events:    ", totalReplayed);
        console2.log("Skipped:         ", skippedCount);
        console2.log("Drift checkpoints:", checkpointsWithDrift);
        console2.log("Total abs drift: ", totalAbsDrift);
        console2.log("Max drift:       ", maxDrift);
        if (checkpointsWithDrift > 0) {
            console2.log("Avg tick drift:  ", totalAbsDrift / checkpointsWithDrift);
        }
        console2.log("Final tick:      ", int256(finalTick));
        console2.log("Final sqrtPriceX96:", uint256(finalSp));
    }
}