lmsr-amm/src/PartyPoolMintImpl.sol

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

import "@abdk/ABDKMath64x64.sol";
import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
import "./PartyPoolBase.sol";
import "./LMSRStabilized.sol";

/// @title PartyPoolMintImpl - Implementation contract for mint and burn functions
/// @notice This contract contains the mint and burn implementation that will be called via delegatecall
/// @dev This contract inherits from PartyPoolBase to access storage and internal functions
contract PartyPoolMintImpl is PartyPoolBase {
    using ABDKMath64x64 for int128;
    using LMSRStabilized for LMSRStabilized.State;
    using SafeERC20 for IERC20;

    // Events that mirror the main contract events
    event Mint(address indexed payer, address indexed receiver, uint256[] depositAmounts, uint256 lpMinted);
    event Burn(address indexed payer, address indexed receiver, uint256[] withdrawAmounts, uint256 lpBurned);

    constructor() PartyPoolBase('','') {}

    /// @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 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 = mintAmounts(lpTokenAmount, lmsr.nAssets, totalSupply());

        // 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;
    }

    /// @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
    returns (uint256[] memory withdrawAmounts) {
        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)
        withdrawAmounts = burnAmounts(lpAmount, lmsr.nAssets, totalSupply(), cachedUintBalances);

        // 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);
    }

    function mintAmounts(uint256 lpTokenAmount, uint256 numAssets, uint256 totalSupply) public view
    returns (uint256[] memory depositAmounts) {
        depositAmounts = new uint256[](numAssets);

        // 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 || numAssets == 0) {
            return depositAmounts; // Return zeros, initial deposit handled differently
        }

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

        return depositAmounts;
    }

    function burnAmounts(uint256 lpTokenAmount,
        uint256 numAssets, uint256 totalSupply, uint256[] memory cachedUintBalances) public view
    returns (uint256[] memory withdrawAmounts) {
        withdrawAmounts = new uint256[](numAssets);

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

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

        return withdrawAmounts;
    }
}