lmsr-amm/src/PartyPool.sol

// SPDX-License-Identifier: UNLICENSED
pragma solidity ^0.8.30;

import "forge-std/console2.sol";
import "@abdk/ABDKMath64x64.sol";
import "@openzeppelin/contracts/token/ERC20/ERC20.sol";
import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
import "@openzeppelin/contracts/utils/ReentrancyGuard.sol";
import "./LMSRStabilized.sol";
import "./LMSRStabilizedBalancedPair.sol";
import "./IPartyPool.sol";
import "./IPartyFlashCallback.sol";

/// @title PartyPool - LMSR-backed multi-asset pool with LP ERC20 token
/// @notice A multi-asset liquidity pool backed by the LMSRStabilized pricing model.
/// The pool issues an ERC20 LP token representing proportional ownership.
/// It supports:
/// - Proportional minting and burning of LP tokens,
/// - Single-token mint (swapMint) and single-asset withdrawal (burnSwap),
/// - Exact-input swaps and swaps-to-price-limits,
/// - Flash loans via a callback interface.
///
/// @dev The contract stores per-token uint "bases" used to scale token units into the internal Q64.64
/// representation used by the LMSR library. Cached on-chain uint balances are kept to reduce balanceOf calls.
/// The contract uses ceiling/floor rules described in function comments to bias rounding in favor of the pool
/// (i.e., floor outputs to users, ceil inputs/fees where appropriate).
contract PartyPool is IPartyPool, ERC20, ReentrancyGuard {
    using ABDKMath64x64 for int128;
    using LMSRStabilized for LMSRStabilized.State;
    using SafeERC20 for IERC20;


    //
    // Immutable pool configuration
    //

    /// @notice Token addresses comprising the pool. Effectively immutable after construction.
    /// @dev tokens[i] corresponds to the i-th asset and maps to index i in the internal LMSR arrays.
    IERC20[] public tokens; // effectively immutable since there is no interface to change the tokens

    /// @inheritdoc IPartyPool
    function numTokens() external view returns (uint256) { return tokens.length; }

    /// @inheritdoc IPartyPool
    function allTokens() external view returns (IERC20[] memory) { return tokens; }

    // NOTE that the slippage target is only exactly achieved in completely balanced pools where all assets are
    // priced the same. This target is actually a minimum slippage that the pool imposes on traders, and the actual
    // slippage cost can be multiples bigger in practice due to pool inventory imbalances.

    /// @notice Trade fraction (Q64.64) representing a reference trade size as fraction of one asset's inventory.
    /// @dev Used by the LMSR stabilization logic to compute target slippage.
    int128 public immutable tradeFrac; // slippage target trade size as a fraction of one asset's inventory

    /// @notice Target slippage (Q64.64) applied for the reference trade size specified by tradeFrac.
    int128 public immutable targetSlippage; // target slippage applied to that trade size

    /// @notice Per-swap fee in parts-per-million (ppm). Fee is taken from input amounts before LMSR computations.
    uint256 public immutable swapFeePpm;

    /// @notice Flash-loan fee in parts-per-million (ppm) applied to flash borrow amounts.
    uint256 public immutable flashFeePpm;

    //
    // Internal state
    //

    LMSRStabilized.State internal lmsr;

    /// @notice If true and there are exactly two assets, an optimized 2-asset stable-pair path is used for some computations.
    bool immutable private _stablePair; // if true, the optimized LMSRStabilizedBalancedPair optimization path is enabled

    // Cached on-chain balances (uint) and internal 64.64 representation
    // balance / base = internal
    uint256[] internal cachedUintBalances;

    /// @notice Per-token uint base denominators used to convert uint token amounts <-> internal Q64.64 representation.
    /// @dev denominators()[i] is the base for tokens[i]. These bases are chosen by deployer and must match token decimals.
    uint256[] internal bases; // per-token uint base used to scale token amounts <-> internal

    /// @inheritdoc IPartyPool
    function denominators() external view returns (uint256[] memory) { return bases; }

    /// @notice Mapping from token address => (index+1). A zero value indicates the token is not in the pool.
    /// @dev Use index = tokenAddressToIndexPlusOne[token] - 1 when non-zero.
    mapping(IERC20=>uint) public tokenAddressToIndexPlusOne; // Uses index+1 so a result of 0 indicates a failed lookup

    /// @notice Scale factor used when converting LMSR Q64.64 totals to LP token units (uint).
    /// @dev LP tokens are minted in units equal to ABDK.mulu(lastTotalQ64x64, LP_SCALE).
    uint256 public constant LP_SCALE = 1e18; // Scale used to convert LMSR lastTotal (Q64.64) into LP token units (uint)

    /// @param name_ LP token name
    /// @param symbol_ LP token symbol
    /// @param _tokens token addresses (n)
    /// @param _bases scaling bases for each token (n) - used when converting to/from internal 64.64 amounts
    /// @param _tradeFrac trade fraction in 64.64 fixed-point (as used by LMSR)
    /// @param _targetSlippage target slippage in 64.64 fixed-point (as used by LMSR)
    /// @param _swapFeePpm fee in parts-per-million, taken from swap input amounts before LMSR calculations
    /// @param _flashFeePpm fee in parts-per-million, taken for flash loans
    /// @param _stable if true and assets.length==2, then the optimization for 2-asset stablecoin pools is activated.
    constructor(
        string memory name_,
        string memory symbol_,
        IERC20[] memory _tokens,
        uint256[] memory _bases,
        int128 _tradeFrac,
        int128 _targetSlippage,
        uint256 _swapFeePpm,
        uint256 _flashFeePpm,
        bool _stable
    ) ERC20(name_, symbol_) {
        require(_tokens.length > 1, "Pool: need >1 asset");
        require(_tokens.length == _bases.length, "Pool: lengths mismatch");
        tokens = _tokens;
        bases = _bases;
        tradeFrac = _tradeFrac;
        targetSlippage = _targetSlippage;
        require(_swapFeePpm < 1_000_000, "Pool: fee >= ppm");
        swapFeePpm = _swapFeePpm;
        require(_flashFeePpm < 1_000_000, "Pool: flash fee >= ppm");
        flashFeePpm = _flashFeePpm;
        _stablePair = _stable && _tokens.length == 2;

        uint256 n = _tokens.length;

        // Initialize LMSR state nAssets; full init occurs on first mint when quantities are known.
        lmsr.nAssets = n;

        // Initialize token address to index mapping
        for (uint i = 0; i < n;) {
            tokenAddressToIndexPlusOne[_tokens[i]] = i + 1;
            unchecked {i++;}
        }

        // Initialize caches to zero
        cachedUintBalances = new uint256[](n);
    }


    /* ----------------------
       Initialization / Mint / Burn (LP token managed)
       ---------------------- */

    /// @inheritdoc IPartyPool
    function mintDepositAmounts(uint256 lpTokenAmount) public view returns (uint256[] memory depositAmounts) {
        uint256 n = tokens.length;
        depositAmounts = new uint256[](n);

        // If this is the first mint or pool is empty, return zeros
        // For first mint, tokens should already be transferred to the pool
        if (totalSupply() == 0 || lmsr.nAssets == 0) {
            return depositAmounts; // Return zeros, initial deposit handled differently
        }

        // Calculate deposit based on current proportions
        uint256 totalLpSupply = totalSupply();

        // lpTokenAmount / totalLpSupply = depositAmount / currentBalance
        // Therefore: depositAmount = (lpTokenAmount * currentBalance) / totalLpSupply
        // We round up to protect the pool
        for (uint i = 0; i < n; i++) {
            uint256 currentBalance = cachedUintBalances[i];
            // Calculate with rounding up: (a * b + c - 1) / c
            depositAmounts[i] = (lpTokenAmount * currentBalance + totalLpSupply - 1) / totalLpSupply;
        }

        return depositAmounts;
    }

    /// @notice Initial mint to set up pool for the first time.
    /// @dev Assumes tokens have already been transferred to the pool prior to calling.
    ///      Can only be called when the pool is uninitialized (totalSupply() == 0 or lmsr.nAssets == 0).
    /// @param receiver address that receives the LP tokens
    /// @param lpTokens The number of LP tokens to issue for this mint. If 0, then the number of tokens returned will equal the LMSR internal q total
    function initialMint(address receiver, uint256 lpTokens) external nonReentrant
    returns (uint256 lpMinted) {
        uint256 n = tokens.length;

        // Check if this is initial deposit - revert if not
        bool isInitialDeposit = totalSupply() == 0 || lmsr.nAssets == 0;
        require(isInitialDeposit, "initialMint: pool already initialized");

        // Update cached balances for all assets
        int128[] memory newQInternal = new int128[](n);
        uint256[] memory depositAmounts = new uint256[](n);
        for (uint i = 0; i < n; ) {
            uint256 bal = IERC20(tokens[i]).balanceOf(address(this));
            cachedUintBalances[i] = bal;
            newQInternal[i] = _uintToInternalFloor(bal, bases[i]);
            depositAmounts[i] = bal;
            unchecked { i++; }
        }

        // Initialize the stabilized LMSR state
        lmsr.init(newQInternal, tradeFrac, targetSlippage);

        // Compute actual LP tokens to mint based on size metric (scaled)
        if( lpTokens != 0 )
            lpMinted = lpTokens;
        else {
            int128 newTotal = _computeSizeMetric(newQInternal);
            lpMinted = ABDKMath64x64.mulu(newTotal, LP_SCALE);
        }

        require(lpMinted > 0, "initialMint: zero LP amount");
        _mint(receiver, lpMinted);
        emit Mint(address(0), receiver, depositAmounts, lpMinted);
    }

    /// @notice Proportional mint for existing pool.
    /// @dev Payer must approve the required token amounts before calling.
    ///      Can only be called when pool is already initialized (totalSupply() > 0 and lmsr.nAssets > 0).
    ///      Rounds follow the pool-favorable conventions documented in helpers (ceil inputs, floor outputs).
    /// @param payer address that provides the input tokens
    /// @param receiver address that receives the LP tokens
    /// @param lpTokenAmount desired amount of LP tokens to mint
    /// @param deadline timestamp after which the transaction will revert. Pass 0 to ignore.
    function mint(address payer, address receiver, uint256 lpTokenAmount, uint256 deadline) external nonReentrant 
    returns (uint256 lpMinted) {
        require(deadline == 0 || block.timestamp <= deadline, "mint: deadline exceeded");
        uint256 n = tokens.length;

        // Check if this is NOT initial deposit - revert if it is
        bool isInitialDeposit = totalSupply() == 0 || lmsr.nAssets == 0;
        require(!isInitialDeposit, "mint: use initialMint for pool initialization");
        require(lpTokenAmount > 0, "mint: zero LP amount");

        // Capture old pool size metric (scaled) by computing from current balances
        int128 oldTotal = _computeSizeMetric(lmsr.qInternal);
        uint256 oldScaled = ABDKMath64x64.mulu(oldTotal, LP_SCALE);

        // Calculate required deposit amounts for the desired LP tokens
        uint256[] memory depositAmounts = mintDepositAmounts(lpTokenAmount);

        // Transfer in all token amounts
        for (uint i = 0; i < n; ) {
            if (depositAmounts[i] > 0) {
                tokens[i].safeTransferFrom(payer, address(this), depositAmounts[i]);
            }
            unchecked { i++; }
        }

        // Update cached balances for all assets
        int128[] memory newQInternal = new int128[](n);
        for (uint i = 0; i < n; ) {
            uint256 bal = IERC20(tokens[i]).balanceOf(address(this));
            cachedUintBalances[i] = bal;
            newQInternal[i] = _uintToInternalFloor(bal, bases[i]);
            unchecked { i++; }
        }

        // Update for proportional change
        lmsr.updateForProportionalChange(newQInternal);

        // Compute actual LP tokens to mint based on change in size metric (scaled)
        // floor truncation rounds in favor of the pool
        int128 newTotal = _computeSizeMetric(newQInternal);
        uint256 newScaled = ABDKMath64x64.mulu(newTotal, LP_SCALE);
        uint256 actualLpToMint;

        require(oldScaled > 0, "mint: oldScaled zero");
        uint256 delta = (newScaled > oldScaled) ? (newScaled - oldScaled) : 0;
        // Proportional issuance: totalSupply * delta / oldScaled
        if (delta > 0) {
            // floor truncation rounds in favor of the pool
            actualLpToMint = (totalSupply() * delta) / oldScaled;
        } else {
            actualLpToMint = 0;
        }

        // Ensure the calculated LP amount is not too different from requested
        require(actualLpToMint > 0, "mint: zero LP minted");

        // Allow actual amount to be at most 0.00001% less than requested
        // This accounts for rounding in deposit calculations
        uint256 minAcceptable = lpTokenAmount * 99_999 / 100_000;
        require(actualLpToMint >= minAcceptable, "mint: insufficient LP minted");

        _mint(receiver, actualLpToMint);
        emit Mint(payer, receiver, depositAmounts, actualLpToMint);
        return actualLpToMint;
    }

    /// @inheritdoc IPartyPool
    function burnReceiveAmounts(uint256 lpTokenAmount) external view returns (uint256[] memory withdrawAmounts) {
        return _burnReceiveAmounts(lpTokenAmount);
    }

    function _burnReceiveAmounts(uint256 lpTokenAmount) internal view returns (uint256[] memory withdrawAmounts) {
        uint256 n = tokens.length;
        withdrawAmounts = new uint256[](n);

        // If supply is zero or pool uninitialized, return zeros
        if (totalSupply() == 0 || lmsr.nAssets == 0) {
            return withdrawAmounts; // Return zeros, nothing to withdraw
        }

        // Calculate withdrawal amounts based on current proportions
        uint256 totalLpSupply = totalSupply();

        // withdrawAmount = floor(lpTokenAmount * currentBalance / totalLpSupply)
        for (uint i = 0; i < n; i++) {
            uint256 currentBalance = cachedUintBalances[i];
            withdrawAmounts[i] = (lpTokenAmount * currentBalance) / totalLpSupply;
        }

        return withdrawAmounts;
    }

    /// @notice Burn LP tokens and withdraw the proportional basket to receiver.
    /// @dev Payer must own or approve the LP tokens being burned. The function updates LMSR state
    ///      proportionally to reflect the reduced pool size after the withdrawal.
    /// @param payer address that provides the LP tokens to burn
    /// @param receiver address that receives the withdrawn tokens
    /// @param lpAmount amount of LP tokens to burn (proportional withdrawal)
    /// @param deadline timestamp after which the transaction will revert. Pass 0 to ignore.
    function burn(address payer, address receiver, uint256 lpAmount, uint256 deadline) external nonReentrant {
        require(deadline == 0 || block.timestamp <= deadline, "burn: deadline exceeded");
        uint256 n = tokens.length;
        require(lpAmount > 0, "burn: zero lp");

        uint256 supply = totalSupply();
        require(supply > 0, "burn: empty supply");
        require(lmsr.nAssets > 0, "burn: uninit pool");
        require(balanceOf(payer) >= lpAmount, "burn: insufficient LP");

        // Refresh cached balances to reflect current on-chain balances before computing withdrawal amounts
        for (uint i = 0; i < n; ) {
            uint256 bal = IERC20(tokens[i]).balanceOf(address(this));
            cachedUintBalances[i] = bal;
            unchecked { i++; }
        }

        // Compute proportional withdrawal amounts for the requested LP amount (rounded down)
        uint256[] memory withdrawAmounts = _burnReceiveAmounts(lpAmount);

        // Transfer underlying tokens out to receiver according to computed proportions
        for (uint i = 0; i < n; ) {
            if (withdrawAmounts[i] > 0) {
                tokens[i].safeTransfer(receiver, withdrawAmounts[i]);
            }
            unchecked { i++; }
        }

        // Update cached balances and internal q for all assets
        int128[] memory newQInternal = new int128[](n);
        for (uint i = 0; i < n; ) {
            uint256 bal = IERC20(tokens[i]).balanceOf(address(this));
            cachedUintBalances[i] = bal;
            newQInternal[i] = _uintToInternalFloor(bal, bases[i]);
            unchecked { i++; }
        }

        // Apply proportional update or deinitialize if drained
        bool allZero = true;
        for (uint i = 0; i < n; ) {
            if (newQInternal[i] != int128(0)) {
                allZero = false;
                break;
            }
            unchecked { i++; }
        }

        if (allZero) {
            lmsr.deinit();
        } else {
            lmsr.updateForProportionalChange(newQInternal);
        }

        // Burn exactly the requested LP amount from payer (authorization via allowance)
        if (msg.sender != payer) {
            uint256 allowed = allowance(payer, msg.sender);
            require(allowed >= lpAmount, "burn: allowance insufficient");
            _approve(payer, msg.sender, allowed - lpAmount);
        }
        _burn(payer, lpAmount);

        emit Burn(payer, receiver, withdrawAmounts, lpAmount);
    }

    /* ----------------------
       Swaps
       ---------------------- */

    /// @inheritdoc IPartyPool
    function swapAmounts(
        uint256 inputTokenIndex,
        uint256 outputTokenIndex,
        uint256 maxAmountIn,
        int128 limitPrice
    ) external view returns (uint256 amountIn, uint256 amountOut, uint256 fee) {
        (uint256 grossIn, uint256 outUint,,,, uint256 feeUint) = _quoteSwapExactIn(inputTokenIndex, outputTokenIndex, maxAmountIn, limitPrice);
        return (grossIn, outUint, feeUint);
    }

    /// @inheritdoc IPartyPool
    function swapToLimitAmounts(
        uint256 inputTokenIndex,
        uint256 outputTokenIndex,
        int128 limitPrice
    ) external view returns (uint256 amountIn, uint256 amountOut, uint256 fee) {
        (uint256 grossIn, uint256 outUint,,,, uint256 feeUint) = _quoteSwapToLimit(inputTokenIndex, outputTokenIndex, limitPrice);
        return (grossIn, outUint, feeUint);
    }


    /// @notice Swap input token i -> token j. Payer must approve token i.
    /// @dev This function transfers the exact gross input (including fee) from payer and sends the computed output to receiver.
    ///      Non-standard tokens (fee-on-transfer, rebasers) are rejected via balance checks.
    /// @param payer address of the account that pays for the swap
    /// @param receiver address that will receive the output tokens
    /// @param inputTokenIndex index of input asset
    /// @param outputTokenIndex index of output asset
    /// @param maxAmountIn maximum amount of token i (uint256) to transfer in (inclusive of fees)
    /// @param limitPrice maximum acceptable marginal price (64.64 fixed point). Pass 0 to ignore.
    /// @param deadline timestamp after which the transaction will revert. Pass 0 to ignore.
    /// @return amountIn actual input used (uint256), amountOut actual output sent (uint256), fee fee taken from the input (uint256)
    function swap(
        address payer,
        address receiver,
        uint256 inputTokenIndex,
        uint256 outputTokenIndex,
        uint256 maxAmountIn,
        int128 limitPrice,
        uint256 deadline
    ) external nonReentrant returns (uint256 amountIn, uint256 amountOut, uint256 fee) {
        uint256 n = tokens.length;
        require(inputTokenIndex < n && outputTokenIndex < n, "swap: idx");
        require(maxAmountIn > 0, "swap: input zero");
        require(deadline == 0 || block.timestamp <= deadline, "swap: deadline exceeded");

        // Read previous balances for affected assets
        uint256 prevBalI = IERC20(tokens[inputTokenIndex]).balanceOf(address(this));
        uint256 prevBalJ = IERC20(tokens[outputTokenIndex]).balanceOf(address(this));

        // Compute amounts using the same path as views
        (uint256 totalTransferAmount, uint256 amountOutUint, int128 amountInInternalUsed, int128 amountOutInternal, , uint256 feeUint) =
            _quoteSwapExactIn(inputTokenIndex, outputTokenIndex, maxAmountIn, limitPrice);

        // Transfer the exact amount from payer and require exact receipt (revert on fee-on-transfer)
        tokens[inputTokenIndex].safeTransferFrom(payer, address(this), totalTransferAmount);
        uint256 balIAfter = IERC20(tokens[inputTokenIndex]).balanceOf(address(this));
        require(balIAfter == prevBalI + totalTransferAmount, "swap: non-standard tokenIn");

        // Transfer output to receiver and verify exact decrease
        tokens[outputTokenIndex].safeTransfer(receiver, amountOutUint);
        uint256 balJAfter = IERC20(tokens[outputTokenIndex]).balanceOf(address(this));
        require(balJAfter == prevBalJ - amountOutUint, "swap: non-standard tokenOut");

        // Update cached uint balances for i and j using actual balances
        cachedUintBalances[inputTokenIndex] = balIAfter;
        cachedUintBalances[outputTokenIndex] = balJAfter;

        // Apply swap to LMSR state with the internal amounts actually used
        lmsr.applySwap(inputTokenIndex, outputTokenIndex, amountInInternalUsed, amountOutInternal);

        emit Swap(payer, receiver, tokens[inputTokenIndex], tokens[outputTokenIndex], totalTransferAmount, amountOutUint);

        return (totalTransferAmount, amountOutUint, feeUint);
    }

    /// @notice Swap up to the price limit; computes max input to reach limit then performs swap.
    /// @dev If balances prevent fully reaching the limit, the function caps and returns actuals.
    ///      The payer must transfer the exact gross input computed by the view.
    /// @param deadline timestamp after which the transaction will revert. Pass 0 to ignore.
    function swapToLimit(
        address payer,
        address receiver,
        uint256 inputTokenIndex,
        uint256 outputTokenIndex,
        int128 limitPrice,
        uint256 deadline
    ) external returns (uint256 amountInUsed, uint256 amountOut, uint256 fee) {
        uint256 n = tokens.length;
        require(inputTokenIndex < n && outputTokenIndex < n, "swapToLimit: idx");
        require(limitPrice > int128(0), "swapToLimit: limit <= 0");
        require(deadline == 0 || block.timestamp <= deadline, "swapToLimit: deadline exceeded");

        // Read previous balances for affected assets
        uint256 prevBalI = IERC20(tokens[inputTokenIndex]).balanceOf(address(this));
        uint256 prevBalJ = IERC20(tokens[outputTokenIndex]).balanceOf(address(this));

        // Compute amounts using the same path as views
        (uint256 totalTransferAmount, uint256 amountOutUint, int128 amountInInternalMax, int128 amountOutInternal, uint256 amountInUsedUint, uint256 feeUint) =
            _quoteSwapToLimit(inputTokenIndex, outputTokenIndex, limitPrice);

        // Transfer the exact amount needed from payer and require exact receipt (revert on fee-on-transfer)
        tokens[inputTokenIndex].safeTransferFrom(payer, address(this), totalTransferAmount);
        uint256 balIAfter = IERC20(tokens[inputTokenIndex]).balanceOf(address(this));
        require(balIAfter == prevBalI + totalTransferAmount, "swapToLimit: non-standard tokenIn");

        // Transfer output to receiver and verify exact decrease
        tokens[outputTokenIndex].safeTransfer(receiver, amountOutUint);
        uint256 balJAfter = IERC20(tokens[outputTokenIndex]).balanceOf(address(this));
        require(balJAfter == prevBalJ - amountOutUint, "swapToLimit: non-standard tokenOut");

        // Update caches to actual balances
        cachedUintBalances[inputTokenIndex] = balIAfter;
        cachedUintBalances[outputTokenIndex] = balJAfter;

        // Apply swap to LMSR state with the internal amounts
        lmsr.applySwap(inputTokenIndex, outputTokenIndex, amountInInternalMax, amountOutInternal);

        // Maintain original event semantics (logs input without fee)
        emit Swap(payer, receiver, tokens[inputTokenIndex], tokens[outputTokenIndex], amountInUsedUint, amountOutUint);

        return (amountInUsedUint, amountOutUint, feeUint);
    }

    /// @notice Ceiling fee helper: computes ceil(x * feePpm / 1_000_000)
    /// @dev Internal helper; public-facing functions use this to ensure fees round up in favor of pool.
    function _ceilFee(uint256 x, uint256 feePpm) internal pure returns (uint256) {
        if (feePpm == 0) return 0;
        // ceil division: (num + denom - 1) / denom
        return (x * feePpm + 1_000_000 - 1) / 1_000_000;
    }

    /// @notice Internal quote for exact-input swap that mirrors swap() rounding and fee application
    /// @dev Returns amounts consistent with swap() semantics: grossIn includes fees (ceil), amountOut is floored.
    /// @return grossIn amount to transfer in (inclusive of fee), amountOutUint output amount (uint),
    ///         amountInInternalUsed and amountOutInternal (64.64), amountInUintNoFee input amount excluding fee (uint),
    ///         feeUint fee taken from the gross input (uint)
    function _quoteSwapExactIn(
        uint256 inputTokenIndex,
        uint256 outputTokenIndex,
        uint256 maxAmountIn,
        int128 limitPrice
    )
        internal
        view
        returns (
            uint256 grossIn,
            uint256 amountOutUint,
            int128 amountInInternalUsed,
            int128 amountOutInternal,
            uint256 amountInUintNoFee,
            uint256 feeUint
        )
    {
        uint256 n = tokens.length;
        require(inputTokenIndex < n && outputTokenIndex < n, "swap: idx");
        require(maxAmountIn > 0, "swap: input zero");
        require(lmsr.nAssets > 0, "swap: empty pool");

        // Estimate max net input (fee on gross rounded up, then subtract)
        (, uint256 netUintForSwap) = _computeFee(maxAmountIn);

        // Convert to internal (floor)
        int128 deltaInternalI = _uintToInternalFloor(netUintForSwap, bases[inputTokenIndex]);
        require(deltaInternalI > int128(0), "swap: input too small after fee");

        // Compute internal amounts using LMSR (exact-input with price limit)
        // if _stablePair is true, use the optimized path
        console2.log('stablepair optimization?', _stablePair);
        (amountInInternalUsed, amountOutInternal) =
            _stablePair ? LMSRStabilizedBalancedPair.swapAmountsForExactInput(lmsr, inputTokenIndex, outputTokenIndex, deltaInternalI, limitPrice)
                        : lmsr.swapAmountsForExactInput(inputTokenIndex, outputTokenIndex, deltaInternalI, limitPrice);

        // Convert actual used input internal -> uint (ceil)
        amountInUintNoFee = _internalToUintCeil(amountInInternalUsed, bases[inputTokenIndex]);
        require(amountInUintNoFee > 0, "swap: input zero");

        // Compute gross transfer including fee on the used input (ceil)
        feeUint = 0;
        grossIn = amountInUintNoFee;
        if (swapFeePpm > 0) {
            feeUint = _ceilFee(amountInUintNoFee, swapFeePpm);
            grossIn += feeUint;
        }

        // Ensure within user max
        require(grossIn <= maxAmountIn, "swap: transfer exceeds max");

        // Compute output (floor)
        amountOutUint = _internalToUintFloor(amountOutInternal, bases[outputTokenIndex]);
        require(amountOutUint > 0, "swap: output zero");
    }

    /// @notice Internal quote for swap-to-limit that mirrors swapToLimit() rounding and fee application
    /// @dev Computes the input required to reach limitPrice and the resulting output; all rounding matches swapToLimit.
    /// @return grossIn amount to transfer in (inclusive of fee), amountOutUint output amount (uint),
    ///         amountInInternal and amountOutInternal (64.64), amountInUintNoFee input amount excluding fee (uint),
    ///         feeUint fee taken from the gross input (uint)
    function _quoteSwapToLimit(
        uint256 inputTokenIndex,
        uint256 outputTokenIndex,
        int128 limitPrice
    )
        internal
        view
        returns (
            uint256 grossIn,
            uint256 amountOutUint,
            int128 amountInInternal,
            int128 amountOutInternal,
            uint256 amountInUintNoFee,
            uint256 feeUint
        )
    {
        uint256 n = tokens.length;
        require(inputTokenIndex < n && outputTokenIndex < n, "swapToLimit: idx");
        require(limitPrice > int128(0), "swapToLimit: limit <= 0");
        require(lmsr.nAssets > 0, "swapToLimit: pool uninitialized");

        // Compute internal maxima at the price limit
        (amountInInternal, amountOutInternal) = lmsr.swapAmountsForPriceLimit(inputTokenIndex, outputTokenIndex, limitPrice);

        // Convert input to uint (ceil) and output to uint (floor)
        amountInUintNoFee = _internalToUintCeil(amountInInternal, bases[inputTokenIndex]);
        require(amountInUintNoFee > 0, "swapToLimit: input zero");

        feeUint = 0;
        grossIn = amountInUintNoFee;
        if (swapFeePpm > 0) {
            feeUint = _ceilFee(amountInUintNoFee, swapFeePpm);
            grossIn += feeUint;
        }

        amountOutUint = _internalToUintFloor(amountOutInternal, bases[outputTokenIndex]);
        require(amountOutUint > 0, "swapToLimit: output zero");
    }

    /// @notice Compute fee and net amounts for a gross input (fee rounded up to favor the pool).
    /// @return feeUint fee taken (uint) and netUint remaining for protocol use (uint)
    function _computeFee(uint256 gross) internal view returns (uint256 feeUint, uint256 netUint) {
        if (swapFeePpm == 0) {
            return (0, gross);
        }
        feeUint = _ceilFee(gross, swapFeePpm);
        netUint = gross - feeUint;
    }

    /// @notice Convenience: return gross = net + fee(net) using ceiling for fee.
    function _addFee(uint256 netUint) internal view returns (uint256 gross) {
        if (swapFeePpm == 0) return netUint;
        uint256 fee = _ceilFee(netUint, swapFeePpm);
        return netUint + fee;
    }

    // --- New events for single-token mint/burn flows ---
    // Note: events intentionally avoid exposing internal ΔS and avoid duplicating LP mint/burn data
    // which is already present in the standard Mint/Burn events.

    /// @notice Single-token mint: deposit a single token, charge swap-LMSR cost, and mint LP.
    /// @dev swapMint executes as an exact-in planned swap followed by proportional scaling of qInternal.
    ///      The function emits SwapMint (gross, net, fee) and also emits Mint for LP issuance.
    /// @param payer who transfers the input token
    /// @param receiver who receives the minted LP tokens
    /// @param inputTokenIndex index of the input token
    /// @param maxAmountIn maximum uint token input (inclusive of fee)
    /// @param deadline optional deadline
    /// @return lpMinted actual LP minted (uint)
    function swapMint(
        address payer,
        address receiver,
        uint256 inputTokenIndex,
        uint256 maxAmountIn,
        uint256 deadline
    ) external nonReentrant returns (uint256 lpMinted) {
        uint256 n = tokens.length;
        require(inputTokenIndex < n, "swapMint: idx");
        require(maxAmountIn > 0, "swapMint: input zero");
        require(deadline == 0 || block.timestamp <= deadline, "swapMint: deadline");

        // Ensure pool initialized
        require(lmsr.nAssets > 0, "swapMint: uninit pool");

        // compute fee on gross maxAmountIn to get an initial net estimate (we'll recompute based on actual used)
        (, uint256 netUintGuess) = _computeFee(maxAmountIn);

        // Convert the net guess to internal (floor)
        int128 netInternalGuess = _uintToInternalFloor(netUintGuess, bases[inputTokenIndex]);
        require(netInternalGuess > int128(0), "swapMint: input too small after fee");

        // Use LMSR view to determine actual internal consumed and size-increase (ΔS) for mint
        (int128 amountInInternalUsed, int128 sizeIncreaseInternal) = lmsr.swapAmountsForMint(inputTokenIndex, netInternalGuess);

        // amountInInternalUsed may be <= netInternalGuess. Convert to uint (ceil) to determine actual transfer
        uint256 amountInUint = _internalToUintCeil(amountInInternalUsed, bases[inputTokenIndex]);
        require(amountInUint > 0, "swapMint: input zero after internal conversion");

        // Compute fee on the actual used input and total transfer amount (ceiling)
        uint256 feeUintActual = _ceilFee(amountInUint, swapFeePpm);
        uint256 totalTransfer = amountInUint + feeUintActual;
        require(totalTransfer > 0 && totalTransfer <= maxAmountIn, "swapMint: transfer exceeds max");

        // Record pre-balance and transfer tokens from payer, require exact receipt (revert on fee-on-transfer)
        uint256 prevBalI = IERC20(tokens[inputTokenIndex]).balanceOf(address(this));
        tokens[inputTokenIndex].safeTransferFrom(payer, address(this), totalTransfer);
        uint256 balIAfter = IERC20(tokens[inputTokenIndex]).balanceOf(address(this));
        require(balIAfter == prevBalI + totalTransfer, "swapMint: non-standard tokenIn");

        // Update cached uint balances for token inputTokenIndex (only inputTokenIndex changed externally)
        cachedUintBalances[inputTokenIndex] = balIAfter;

        // Compute old and new scaled size metrics to determine LP minted
        int128 oldTotal = _computeSizeMetric(lmsr.qInternal);
        require(oldTotal > int128(0), "swapMint: zero total");
        uint256 oldScaled = ABDKMath64x64.mulu(oldTotal, LP_SCALE);

        int128 newTotal = oldTotal.add(sizeIncreaseInternal);
        uint256 newScaled = ABDKMath64x64.mulu(newTotal, LP_SCALE);

        uint256 actualLpToMint;
        if (totalSupply() == 0) {
            // If somehow supply zero (shouldn't happen as lmsr.nAssets>0), mint newScaled
            actualLpToMint = newScaled;
        } else {
            require(oldScaled > 0, "swapMint: oldScaled zero");
            uint256 delta = (newScaled > oldScaled) ? (newScaled - oldScaled) : 0;
            if (delta > 0) {
                // floor truncation rounds in favor of pool
                actualLpToMint = (totalSupply() * delta) / oldScaled;
            } else {
                actualLpToMint = 0;
            }
        }

        require(actualLpToMint > 0, "swapMint: zero LP minted");

        // Update LMSR internal state: scale qInternal proportionally by newTotal/oldTotal
        int128[] memory newQInternal = new int128[](n);
        for (uint256 idx = 0; idx < n; idx++) {
            // newQInternal[idx] = qInternal[idx] * (newTotal / oldTotal)
            newQInternal[idx] = lmsr.qInternal[idx].mul(newTotal).div(oldTotal);
        }

        // Update cached internal and kappa via updateForProportionalChange
        lmsr.updateForProportionalChange(newQInternal);

        // Note: we updated cachedUintBalances[inputTokenIndex] above via reading balance; other token uint balances did not
        // change externally (they were not transferred in). We keep cachedUintBalances for others unchanged.
        // Mint LP tokens to receiver
        _mint(receiver, actualLpToMint);

        // Emit SwapMint event with gross transfer, net input and fee (planned exact-in)
        emit SwapMint(payer, receiver, inputTokenIndex, totalTransfer, amountInUint, feeUintActual);

        // Emit standard Mint event which records deposit amounts and LP minted
        emit Mint(payer, receiver, new uint256[](n), actualLpToMint);
        // Note: depositAmounts array omitted (empty) since swapMint uses single-token input

        return actualLpToMint;
    }

    /// @notice Burn LP tokens then swap the redeemed proportional basket into a single asset `inputTokenIndex` and send to receiver.
    /// @dev The function burns LP tokens (authorization via allowance if needed), sends the single-asset payout and updates LMSR state.
    /// @param payer who burns LP tokens
    /// @param receiver who receives the single asset
    /// @param lpAmount amount of LP tokens to burn
    /// @param inputTokenIndex index of target asset to receive
    /// @param deadline optional deadline
    /// @return amountOutUint uint amount of asset i sent to receiver
    function burnSwap(
        address payer,
        address receiver,
        uint256 lpAmount,
        uint256 inputTokenIndex,
        uint256 deadline
    ) external nonReentrant returns (uint256 amountOutUint) {
        uint256 n = tokens.length;
        require(inputTokenIndex < n, "burnSwap: idx");
        require(lpAmount > 0, "burnSwap: zero lp");
        require(deadline == 0 || block.timestamp <= deadline, "burnSwap: deadline");

        uint256 supply = totalSupply();
        require(supply > 0, "burnSwap: empty supply");
        require(balanceOf(payer) >= lpAmount, "burnSwap: insufficient LP");

        // alpha = lpAmount / supply as Q64.64
        int128 alpha = ABDKMath64x64.divu(lpAmount, supply);

        // Use LMSR view to compute single-asset payout and burned size-metric
        (int128 payoutInternal, ) = lmsr.swapAmountsForBurn(inputTokenIndex, alpha);

        // Convert payoutInternal -> uint (floor) to favor pool
        amountOutUint = _internalToUintFloor(payoutInternal, bases[inputTokenIndex]);
        require(amountOutUint > 0, "burnSwap: output zero");

        // Transfer the payout to receiver
        tokens[inputTokenIndex].safeTransfer(receiver, amountOutUint);

        // Burn LP tokens from payer (authorization via allowance)
        if (msg.sender != payer) {
            uint256 allowed = allowance(payer, msg.sender);
            require(allowed >= lpAmount, "burnSwap: allowance insufficient");
            _approve(payer, msg.sender, allowed - lpAmount);
        }
        _burn(payer, lpAmount);

        // Update cached balances by reading on-chain balances for all tokens
        int128[] memory newQInternal = new int128[](n);
        for (uint256 idx = 0; idx < n; idx++) {
            uint256 bal = IERC20(tokens[idx]).balanceOf(address(this));
            cachedUintBalances[idx] = bal;
            newQInternal[idx] = _uintToInternalFloor(bal, bases[idx]);
        }

        // Emit BurnSwap with public-facing info only (do not expose ΔS or LP burned)
        emit BurnSwap(payer, receiver, inputTokenIndex, amountOutUint);

        // If entire pool drained, deinit; else update proportionally
        bool allZero = true;
        for (uint256 idx = 0; idx < n; idx++) {
            if (newQInternal[idx] != int128(0)) { allZero = false; break; }
        }
        if (allZero) {
            lmsr.deinit();
        } else {
            lmsr.updateForProportionalChange(newQInternal);
        }

        emit Burn(payer, receiver, new uint256[](n), lpAmount);
        return amountOutUint;
    }


    /// @inheritdoc IPartyPool
    function flashRepaymentAmounts(uint256[] memory loanAmounts) external view
    returns (uint256[] memory repaymentAmounts) {
        repaymentAmounts = new uint256[](tokens.length);
        for (uint256 i = 0; i < tokens.length; i++) {
            uint256 amount = loanAmounts[i];
            if (amount > 0) {
                repaymentAmounts[i] = amount + _ceilFee(amount, flashFeePpm);
            }
        }
    }


    /// @notice Receive token amounts and require them to be repaid plus a fee inside a callback.
    /// @dev The caller must implement IPartyFlashCallback#partyFlashCallback which receives (amounts, repaymentAmounts, data).
    ///      This function verifies that, after the callback returns, the pool's balances have increased by at least the fees
    ///      for each borrowed token. Reverts if repayment (including fee) did not occur.
    /// @param recipient The address which will receive the token amounts
    /// @param amounts The amount of each token to send (array length must equal pool size)
    /// @param data Any data to be passed through to the callback
    function flash(
        address recipient,
        uint256[] memory amounts,
        bytes calldata data
    ) external nonReentrant {
        require(recipient != address(0), "flash: zero recipient");
        require(amounts.length == tokens.length, "flash: amounts length mismatch");

        // Calculate repayment amounts for each token including fee
        uint256[] memory repaymentAmounts = new uint256[](tokens.length);

        // Store initial balances to verify repayment later
        uint256[] memory initialBalances = new uint256[](tokens.length);

        // Track if any token amount is non-zero
        bool hasNonZeroAmount = false;

        // Process each token, skipping those with zero amounts
        for (uint256 i = 0; i < tokens.length; i++) {
            uint256 amount = amounts[i];

            if (amount > 0) {
                hasNonZeroAmount = true;

                // Calculate repayment amount with fee (ceiling)
                repaymentAmounts[i] = amount + _ceilFee(amount, flashFeePpm);

                // Record initial balance
                initialBalances[i] = IERC20(tokens[i]).balanceOf(address(this));

                // Transfer token to recipient
                tokens[i].safeTransfer(recipient, amount);
            }
        }

        // Ensure at least one token is being borrowed
        require(hasNonZeroAmount, "flash: no tokens requested");

        // Call flash callback with expected repayment amounts
        IPartyFlashCallback(msg.sender).partyFlashCallback(amounts, repaymentAmounts, data);

        // Verify repayment amounts for tokens that were borrowed
        for (uint256 i = 0; i < tokens.length; i++) {
            if (amounts[i] > 0) {
                uint256 currentBalance = IERC20(tokens[i]).balanceOf(address(this));

                // Verify repayment: current balance must be at least (initial balance + fee)
                require(
                    currentBalance >= initialBalances[i] + _ceilFee(amounts[i], flashFeePpm),
                    "flash: repayment failed"
                );

                // Update cached balance
                cachedUintBalances[i] = currentBalance;
            }
        }
    }


    /* ----------------------
       Conversion helpers
       ---------------------- */

    // Convert uint token amount -> internal 64.64 (floor). Uses ABDKMath64x64.divu which truncates.
    function _uintToInternalFloor(uint256 amount, uint256 base) internal pure returns (int128) {
        // internal = amount / base  (as Q64.64)
        return ABDKMath64x64.divu(amount, base);
    }

    // Convert internal 64.64 -> uint token amount (floor). Uses ABDKMath64x64.mulu which floors the product.
    function _internalToUintFloor(int128 internalAmount, uint256 base) internal pure returns (uint256) {
        // uint = internal * base (floored)
        return ABDKMath64x64.mulu(internalAmount, base);
    }

    // Convert internal 64.64 -> uint token amount (ceiling). Rounds up to protect the pool.
    function _internalToUintCeil(int128 internalAmount, uint256 base) internal pure returns (uint256) {
        // Get the floor value first
        uint256 floorValue = ABDKMath64x64.mulu(internalAmount, base);

        // Check if there was any fractional part by comparing to a reconstruction of the original
        int128 reconstructed = ABDKMath64x64.divu(floorValue, base);

        // If reconstructed is less than original, there was a fractional part that was truncated
        if (reconstructed < internalAmount) {
            return floorValue + 1;
        }

        return floorValue;
    }

    /// @notice Helper to compute size metric (sum of all asset quantities) from internal balances
    /// @dev Returns the sum of all provided qInternal_ entries as a Q64.64 value.
    function _computeSizeMetric(int128[] memory qInternal_) private pure returns (int128) {
        int128 total = int128(0);
        for (uint i = 0; i < qInternal_.length; ) {
            total = total.add(qInternal_[i]);
            unchecked { i++; }
        }
        return total;
    }

}