lmsr-amm/src/PartyPool.sol

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

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 "./IPartyPool.sol";
import "./IPartyFlashCallback.sol";

/// @title PartyPool - LMSR-backed multi-asset pool with LP ERC20 token
/// @notice Uses LMSRStabilized library; stores per-token uint bases to convert to/from 64.64 fixed point.
/// - Caches qInternal[] (int128 64.64) and cachedUintBalances[] to minimize balanceOf() calls.
/// - swap and swapToLimit mimic core lib; mint/burn call updateForProportionalChange() and manage LP tokens.
contract PartyPool is IPartyPool, ERC20, ReentrancyGuard {
    using ABDKMath64x64 for int128;
    using LMSRStabilized for LMSRStabilized.State;
    using SafeERC20 for IERC20;


    //
    // Immutable pool configuration
    //

    address[] public tokens; // effectively immutable since there is no interface to change the tokens
    function numTokens() external view returns (uint256) { return tokens.length; }
    function allTokens() external view returns (address[] 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.
    int128 public immutable tradeFrac; // slippage target trade size as a fraction of one asset's inventory
    int128 public immutable targetSlippage; // target slippage applied to that trade size

    // fee in parts-per-million (ppm), taken from inputs before swaps
    uint256 public immutable swapFeePpm;

    // flash loan fee in parts-per-million (ppm)
    uint256 public immutable flashFeePpm;

    //
    // Internal state
    //
    
    LMSRStabilized.State internal lmsr;

    // Cached on-chain balances (uint) and internal 64.64 representation
    // balance / base = internal
    uint256[] internal cachedUintBalances;
    uint256[] internal bases; // per-token uint base used to scale token amounts <-> internal
    
    mapping(address=>uint) public tokenAddressToIndexPlusOne; // Uses index+1 so a result of 0 indicates a failed lookup
    
    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
    constructor(
        string memory name_,
        string memory symbol_,
        address[] memory _tokens,
        uint256[] memory _bases,
        int128 _tradeFrac,
        int128 _targetSlippage,
        uint256 _swapFeePpm,
        uint256 _flashFeePpm
    ) 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;

        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)
       ---------------------- */

    /// @notice Calculate the proportional deposit amounts required for a given LP token amount
    /// @param lpTokenAmount The amount of LP tokens desired
    /// @return depositAmounts Array of token amounts to deposit (rounded up)
    function computeMintAmounts(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 Calculate the proportional withdrawal amounts for a given LP token amount
    /// @param lpTokenAmount The amount of LP tokens to burn
    /// @return withdrawAmounts Array of token amounts to withdraw (rounded down)
    function computeBurnAmounts(uint256 lpTokenAmount) external view returns (uint256[] memory withdrawAmounts) {
        return _computeBurnAmounts(lpTokenAmount);
    }

    function _computeBurnAmounts(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 Proportional mint (or initial supply if first call).
    /// For initial supply: assumes tokens have already been transferred to the pool
    /// For subsequent mints: payer must approve tokens beforehand, receiver gets the LP tokens
    /// @param payer address that provides the input tokens (ignored for initial deposit)
    /// @param receiver address that receives the LP tokens
    /// @param lpTokenAmount desired amount of LP tokens to mint (ignored for initial deposit)
    /// @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 {
        require(deadline == 0 || block.timestamp <= deadline, "mint: deadline exceeded");
        uint256 n = tokens.length;
        // Check if this is initial deposit
        bool isInitialDeposit = totalSupply() == 0 || lmsr.nAssets == 0;

        require(lpTokenAmount > 0 || isInitialDeposit, "mint: zero LP amount");

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

        // For non-initial deposits, transfer tokens from payer
        uint256[] memory depositAmounts = new uint256[](n);

        if (!isInitialDeposit) {
            // Calculate required deposit amounts for the desired LP tokens
            depositAmounts = computeMintAmounts(lpTokenAmount);

            // Transfer in all token amounts
            for (uint i = 0; i < n; ) {
                if (depositAmounts[i] > 0) {
                    _safeTransferFrom(tokens[i], 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]);

            // For initial deposit, record the actual deposited amounts
            if (isInitialDeposit) {
                depositAmounts[i] = bal;
            }

            unchecked { i++; }
        }

        // If first time, call init, otherwise update proportional change.
        if (isInitialDeposit) {
            // Initialize the stabilized LMSR state
            lmsr.init(newQInternal, tradeFrac, targetSlippage);
        } else {
            // 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;

        if (isInitialDeposit) {
            // Initial provisioning: mint newScaled (as LP units)
            actualLpToMint = newScaled;
        } else {
            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;
            }
        }

        // For subsequent mints, ensure the calculated LP amount is not too different from requested
        if (!isInitialDeposit) {
            // Allow for some rounding error but ensure we're not far off from requested amount
            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");
        }

        console2.log('actualLpToMint', actualLpToMint);
        require( actualLpToMint > 0, "mint: zero LP amount");
        _mint(receiver, actualLpToMint);
        emit Mint(payer, receiver, depositAmounts, actualLpToMint);
    }

    /// @notice Burn LP tokens and withdraw the proportional basket to receiver.
    /// Payer must own the LP tokens; withdraw amounts are computed from current proportions.
    /// @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 = _computeBurnAmounts(lpAmount);

        // Transfer underlying tokens out to receiver according to computed proportions
        for (uint i = 0; i < n; ) {
            if (withdrawAmounts[i] > 0) {
                _safeTransfer(tokens[i], 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
       ---------------------- */

    /// @notice Swap input token i -> token j. Payer must approve token i.
    /// @param payer address of the account that pays for the swap
    /// @param receiver address that will receive the output tokens
    /// @param i index of input asset
    /// @param j 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)
    function swap(
        address payer,
        address receiver,
        uint256 i,
        uint256 j,
        uint256 maxAmountIn,
        int128 limitPrice,
        uint256 deadline
    ) external nonReentrant returns (uint256 amountIn, uint256 amountOut) {
        uint256 n = tokens.length;
        require(i < n && j < 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[i]).balanceOf(address(this));
        uint256 prevBalJ = IERC20(tokens[j]).balanceOf(address(this));

        // Calculate fee (ceiling) and net amount
        (, uint256 netUintForSwap) = _computeFee(maxAmountIn);

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

        // Make sure LMSR state exists
        require(lmsr.nAssets > 0, "swap: empty pool");

        // Compute swap amounts in internal space using exact-input logic (with limitPrice)
        (int128 amountInInternalUsed, int128 amountOutInternal) = lmsr.swapAmountsForExactInput(
            i,
            j,
            deltaInternalI,
            limitPrice
        );

        // Convert actual used input internal -> uint (ceiling to protect the pool)
        uint256 amountInUint = _internalToUintCeil(amountInInternalUsed, bases[i]);

        // Total transfer amount includes fee calculated on the actual used input (ceiling)
        uint256 totalTransferAmount = amountInUint;
        if (swapFeePpm > 0) {
            uint256 feeOnUsed = _ceilFee(amountInUint, swapFeePpm);
            totalTransferAmount += feeOnUsed;
        }

        // Ensure we do not attempt to transfer more than the caller specified as maximum
        require(totalTransferAmount > 0, 'swap: input zero');
        require(totalTransferAmount <= maxAmountIn, "swap: transfer exceeds max");

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

        // Compute output uint amount (floor)
        uint256 amountOutUint = _internalToUintFloor(amountOutInternal, bases[j]);
        require(amountOutUint > 0, "swap: output zero");

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

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

        // Apply swap to LMSR state with the internal amounts actually used 
        // (fee is already accounted for in the reduced input amount)
        lmsr.applySwap(i, j, amountInInternalUsed, amountOutInternal);

        emit Swap(payer, receiver, tokens[i], tokens[j], totalTransferAmount, amountOutUint);

        return (totalTransferAmount, amountOutUint);
    }

    /// @notice Swap up to the price limit; computes max input to reach limit then performs swap.
    /// If the pool can't fill entirely because of balances, it caps appropriately and returns actuals.
    /// Payer must approve token i for the exact computed input amount.
    /// @param deadline timestamp after which the transaction will revert. Pass 0 to ignore.
    function swapToLimit(
        address payer,
        address receiver,
        uint256 i,
        uint256 j,
        int128 limitPrice,
        uint256 deadline
    ) external returns (uint256 amountInUsed, uint256 amountOut) {
        uint256 n = tokens.length;
        require(i < n && j < n, "swapToLimit: idx");
        require(limitPrice > int128(0), "swapToLimit: limit <= 0");
        require(deadline == 0 || block.timestamp <= deadline, "swapToLimit: deadline exceeded");

        // Ensure LMSR state exists
        require(lmsr.nAssets > 0, "swapToLimit: pool uninitialized");

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

        // Compute maxima in internal space using library
        (int128 amountInInternalMax, int128 amountOutInternal) = lmsr.swapAmountsForPriceLimit(i, j, limitPrice);

        // Calculate how much input will be needed with fee included (ceiling to protect the pool)
        uint256 amountInUsedUint = _internalToUintCeil(amountInInternalMax, bases[i]);
        require(amountInUsedUint > 0, "swapToLimit: input zero");

        // Total transfer amount is the input amount including what will be taken as fee (ceiling)
        uint256 totalTransferAmount = amountInUsedUint;

        if (swapFeePpm > 0) {
            uint256 feeOnUsed = _ceilFee(amountInUsedUint, swapFeePpm);
            totalTransferAmount += feeOnUsed;
        }

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

        // Compute output amount (floor)
        uint256 amountOutUint = _internalToUintFloor(amountOutInternal, bases[j]);
        require(amountOutUint > 0, "swapToLimit: output zero");

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

        // Update caches to actual balances
        cachedUintBalances[i] = balIAfter;
        cachedUintBalances[j] = balJAfter;

        // Apply swap to LMSR state with the internal amounts
        // (fee is already part of the reduced effective input)
        lmsr.applySwap(i, j, amountInInternalMax, amountOutInternal);

        emit Swap(payer, receiver, tokens[i], tokens[j], amountInUsedUint, amountOutUint);

        return (amountInUsedUint, amountOutUint);
    }

    /// @notice Ceiling fee helper: computes ceil(x * feePpm / 1_000_000)
    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 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.
    /// @param payer who transfers the input token
    /// @param receiver who receives the minted LP tokens
    /// @param i 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 i,
        uint256 maxAmountIn,
        uint256 deadline
    ) external nonReentrant returns (uint256 lpMinted) {
        uint256 n = tokens.length;
        require(i < 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[i]);
        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(i, netInternalGuess);

        // amountInInternalUsed may be <= netInternalGuess. Convert to uint (ceil) to determine actual transfer
        uint256 amountInUint = _internalToUintCeil(amountInInternalUsed, bases[i]);
        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[i]).balanceOf(address(this));
        _safeTransferFrom(tokens[i], payer, address(this), totalTransfer);
        uint256 balIAfter = IERC20(tokens[i]).balanceOf(address(this));
        require(balIAfter == prevBalI + totalTransfer, "swapMint: non-standard tokenIn");

        // Update cached uint balances for token i (only i changed externally)
        cachedUintBalances[i] = 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[i] 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, i, 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 `i` and send to receiver.
    /// @param payer who burns LP tokens
    /// @param receiver who receives the single asset
    /// @param lpAmount amount of LP tokens to burn
    /// @param i 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 i,
        uint256 deadline
    ) external nonReentrant returns (uint256 amountOutUint) {
        uint256 n = tokens.length;
        require(i < 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(i, alpha);

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

        // Transfer the payout to receiver
        _safeTransfer(tokens[i], 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, i, 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;
    }


    function computeFlashRepaymentAmounts(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 token0 and/or token1 and pay it back, plus a fee, in the callback
    /// @dev The caller of this method receives a callback in the form of IPartyFlashCallback#partyFlashCallback
    /// @param recipient The address which will receive the token amounts
    /// @param amounts The amount of each token to send
    /// @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
                _safeTransfer(tokens[i], 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;
    }

    /* ----------------------
       ERC20 helpers (minimal)
       ---------------------- */

    function _safeTransferFrom(address token, address from, address to, uint256 amt) internal {
        IERC20(token).safeTransferFrom(from, to, amt);
    }

    function _safeTransfer(address token, address to, uint256 amt) internal {
        IERC20(token).safeTransfer(to, amt);
    }

    /// @notice Helper to compute size metric (sum of all asset quantities) from internal balances
    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;
    }

}