mini-swap/contracts/MiniSwapAMM.sol

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.19;

import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
import "@openzeppelin/contracts/token/ERC20/ERC20.sol";
import "@openzeppelin/contracts/utils/ReentrancyGuard.sol";
import "@openzeppelin/contracts/utils/math/Math.sol";

/**
 * @title MiniSwapAMM
 * @dev 简单的自动做市商 (AMM) 实现，基于 x * y = k 公式
 */
contract MiniSwapAMM is ERC20, ReentrancyGuard {
    IERC20 public immutable tokenA;
    IERC20 public immutable tokenB;
    
    uint256 public reserveA;
    uint256 public reserveB;
    
    uint256 private constant MINIMUM_LIQUIDITY = 10**3;
    
    event AddLiquidity(
        address indexed provider,
        uint256 amountA,
        uint256 amountB,
        uint256 liquidity
    );
    
    event RemoveLiquidity(
        address indexed provider,
        uint256 amountA,
        uint256 amountB,
        uint256 liquidity
    );
    
    event Swap(
        address indexed trader,
        address tokenIn,
        address tokenOut,
        uint256 amountIn,
        uint256 amountOut
    );

    constructor(address _tokenA, address _tokenB) ERC20("MiniSwap LP Token", "MSLP") {
        tokenA = IERC20(_tokenA);
        tokenB = IERC20(_tokenB);
    }

    /**
     * @dev 添加流动性
     * @param amountADesired 期望添加的 tokenA 数量
     * @param amountBDesired 期望添加的 tokenB 数量
     * @param amountAMin 最小添加的 tokenA 数量
     * @param amountBMin 最小添加的 tokenB 数量
     * @return amountA 实际添加的 tokenA 数量
     * @return amountB 实际添加的 tokenB 数量
     * @return liquidity 获得的 LP 代币数量
     */
    function addLiquidity(
        uint256 amountADesired,
        uint256 amountBDesired,
        uint256 amountAMin,
        uint256 amountBMin
    ) external nonReentrant returns (uint256 amountA, uint256 amountB, uint256 liquidity) {
        (amountA, amountB) = _calculateOptimalAmounts(
            amountADesired,
            amountBDesired,
            amountAMin,
            amountBMin
        );

        // 转移代币到合约
        tokenA.transferFrom(msg.sender, address(this), amountA);
        tokenB.transferFrom(msg.sender, address(this), amountB);

        liquidity = _mintLiquidity(amountA, amountB);

        emit AddLiquidity(msg.sender, amountA, amountB, liquidity);
    }

    /**
     * @dev 移除流动性
     * @param liquidity 要移除的 LP 代币数量
     * @param amountAMin 最小获得的 tokenA 数量
     * @param amountBMin 最小获得的 tokenB 数量
     * @return amountA 获得的 tokenA 数量
     * @return amountB 获得的 tokenB 数量
     */
    function removeLiquidity(
        uint256 liquidity,
        uint256 amountAMin,
        uint256 amountBMin
    ) external nonReentrant returns (uint256 amountA, uint256 amountB) {
        require(liquidity > 0, "MiniSwapAMM: INSUFFICIENT_LIQUIDITY");

        uint256 totalSupply = totalSupply();
        amountA = (liquidity * reserveA) / totalSupply;
        amountB = (liquidity * reserveB) / totalSupply;

        require(amountA >= amountAMin, "MiniSwapAMM: INSUFFICIENT_A_AMOUNT");
        require(amountB >= amountBMin, "MiniSwapAMM: INSUFFICIENT_B_AMOUNT");

        _burn(msg.sender, liquidity);

        tokenA.transfer(msg.sender, amountA);
        tokenB.transfer(msg.sender, amountB);

        _updateReserves();

        emit RemoveLiquidity(msg.sender, amountA, amountB, liquidity);
    }

    /**
     * @dev 交换代币 A 到代币 B
     * @param amountAIn 输入的 tokenA 数量
     * @param amountBOutMin 最小输出的 tokenB 数量
     * @return amountBOut 实际输出的 tokenB 数量
     */
    function swapAForB(uint256 amountAIn, uint256 amountBOutMin) external nonReentrant returns (uint256 amountBOut) {
        require(amountAIn > 0, "MiniSwapAMM: INSUFFICIENT_INPUT_AMOUNT");
        
        amountBOut = getAmountOut(amountAIn, reserveA, reserveB);
        require(amountBOut >= amountBOutMin, "MiniSwapAMM: INSUFFICIENT_OUTPUT_AMOUNT");

        tokenA.transferFrom(msg.sender, address(this), amountAIn);
        tokenB.transfer(msg.sender, amountBOut);

        _updateReserves();

        emit Swap(msg.sender, address(tokenA), address(tokenB), amountAIn, amountBOut);
    }

    /**
     * @dev 交换代币 B 到代币 A
     * @param amountBIn 输入的 tokenB 数量
     * @param amountAOutMin 最小输出的 tokenA 数量
     * @return amountAOut 实际输出的 tokenA 数量
     */
    function swapBForA(uint256 amountBIn, uint256 amountAOutMin) 
        external nonReentrant returns (uint256 amountAOut) {
        require(amountBIn > 0, "MiniSwapAMM: INSUFFICIENT_INPUT_AMOUNT");
        
        amountAOut = getAmountOut(amountBIn, reserveB, reserveA);
        require(amountAOut >= amountAOutMin, "MiniSwapAMM: INSUFFICIENT_OUTPUT_AMOUNT");

        tokenB.transferFrom(msg.sender, address(this), amountBIn);
        tokenA.transfer(msg.sender, amountAOut);

        _updateReserves();

        emit Swap(msg.sender, address(tokenB), address(tokenA), amountBIn, amountAOut);
    }

    /**
     * @dev 根据输入数量计算输出数量 (基于 x * y = k 公式)
     * @param amountIn 输入数量
     * @param reserveIn 输入代币储备量
     * @param reserveOut 输出代币储备量
     * @return amountOut 输出数量
     */
    function getAmountOut(uint256 amountIn, uint256 reserveIn, uint256 reserveOut) 
        public pure returns (uint256 amountOut) {
        require(amountIn > 0, "MiniSwapAMM: INSUFFICIENT_INPUT_AMOUNT");
        require(reserveIn > 0 && reserveOut > 0, "MiniSwapAMM: INSUFFICIENT_LIQUIDITY");

        // 考虑 0.3% 的交易费用
        uint256 amountInWithFee = amountIn * 997;
        uint256 numerator = amountInWithFee * reserveOut;
        uint256 denominator = (reserveIn * 1000) + amountInWithFee;
        amountOut = numerator / denominator;
    }

    /**
     * @dev 获取当前储备量
     */
    function getReserves() external view returns (uint256 _reserveA, uint256 _reserveB) {
        return (reserveA, reserveB);
    }

    /**
     * @dev 计算最优的代币数量比例
     */
    function _calculateOptimalAmounts(
        uint256 amountADesired,
        uint256 amountBDesired,
        uint256 amountAMin,
        uint256 amountBMin
    ) internal view returns (uint256 amountA, uint256 amountB) {
        if (reserveA == 0 && reserveB == 0) {
            // 首次添加流动性
            return (amountADesired, amountBDesired);
        } else {
            uint256 amountBOptimal = (amountADesired * reserveB) / reserveA;
            if (amountBOptimal <= amountBDesired) {
                require(amountBOptimal >= amountBMin, "MiniSwapAMM: INSUFFICIENT_B_AMOUNT");
                return (amountADesired, amountBOptimal);
            } else {
                uint256 amountAOptimal = (amountBDesired * reserveA) / reserveB;
                require(amountAOptimal <= amountADesired && amountAOptimal >= amountAMin, 
                    "MiniSwapAMM: INSUFFICIENT_A_AMOUNT");
                return (amountAOptimal, amountBDesired);
            }
        }
    }

    /**
     * @dev 铸造流动性代币
     */
    function _mintLiquidity(uint256 amountA, uint256 amountB) internal returns (uint256 liquidity) {
        uint256 totalSupply = totalSupply();
        
        if (totalSupply == 0) {
            // 首次添加流动性
            liquidity = Math.sqrt(amountA * amountB);
            require(liquidity > MINIMUM_LIQUIDITY, "MiniSwapAMM: INSUFFICIENT_LIQUIDITY_MINTED");
            // 从流动性中减去最小流动性，永久锁定
            liquidity = liquidity - MINIMUM_LIQUIDITY;
        } else {
            liquidity = Math.min(
                (amountA * totalSupply) / reserveA,
                (amountB * totalSupply) / reserveB
            );
        }

        require(liquidity > 0, "MiniSwapAMM: INSUFFICIENT_LIQUIDITY_MINTED");
        _mint(msg.sender, liquidity);

        _updateReserves();
    }

    /**
     * @dev 更新储备量
     */
    function _updateReserves() internal {
        reserveA = tokenA.balanceOf(address(this));
        reserveB = tokenB.balanceOf(address(this));
    }
}