PegPool.sol

// SPDX-License-Identifier: MIT
pragma solidity 0.8.36;

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

/// @notice Full-range constant-product laboratory AMM. All swap fees stay with LPs.
/// @dev Not for rebasing or fee-on-transfer assets. No oracle, lending or protocol fee.
contract PegPool is ERC20, ReentrancyGuard {
    using SafeERC20 for IERC20;
    uint256 public constant BPS = 10_000;
    uint256 public constant SWAP_FEE_BPS = 30;
    uint256 public constant MINIMUM_LIQUIDITY = 1_000;
    address public constant LOCKED_LIQUIDITY = address(1);
    IERC20 public immutable token0;
    IERC20 public immutable token1;
    error InvalidConfiguration();
    error InvalidAmount();
    error InvalidRecipient();
    error InvalidToken();
    error DeadlineExpired();
    error SlippageExceeded();
    error InsufficientLiquidity();
    error UnsupportedAsset();
    event LiquidityAdded(address indexed sender, address indexed recipient, uint256 amount0, uint256 amount1, uint256 shares);
    event LiquidityRemoved(address indexed sender, address indexed recipient, uint256 amount0, uint256 amount1, uint256 shares);
    event Swap(address indexed sender, address indexed tokenIn, address indexed recipient, uint256 amountIn, uint256 amountOut);

    constructor(address token0_, address token1_, string memory symbol_) ERC20("Peg LP", symbol_) {
        if (token0_.code.length == 0 || token1_.code.length == 0 || token0_ == token1_) revert InvalidConfiguration();
        token0 = IERC20(token0_);
        token1 = IERC20(token1_);
    }

    function getReserves() public view returns (uint256 reserve0, uint256 reserve1) {
        return (token0.balanceOf(address(this)), token1.balanceOf(address(this)));
    }

    function quoteLiquidity(uint256 max0, uint256 max1) public view returns (uint256 amount0, uint256 amount1, uint256 shares) {
        uint256 supply = totalSupply();
        if (max0 == 0 || max1 == 0) revert InvalidAmount();
        if (supply == 0) {
            // Bounds keep the initial sqrt multiplication safe and practical in this lab.
            if (max0 > type(uint128).max || max1 > type(uint128).max) revert InvalidAmount();
            uint256 root = Math.sqrt(max0 * max1);
            if (root <= MINIMUM_LIQUIDITY) revert InsufficientLiquidity();
            return (max0, max1, root - MINIMUM_LIQUIDITY);
        }
        (uint256 reserve0, uint256 reserve1) = getReserves();
        if (reserve0 == 0 || reserve1 == 0) revert InsufficientLiquidity();
        uint256 optimal1 = Math.mulDiv(max0, reserve1, reserve0);
        if (optimal1 <= max1) { amount0 = max0; amount1 = optimal1; }
        else { amount0 = Math.mulDiv(max1, reserve0, reserve1); amount1 = max1; }
        shares = Math.min(Math.mulDiv(amount0, supply, reserve0), Math.mulDiv(amount1, supply, reserve1));
        if (amount0 == 0 || amount1 == 0 || shares == 0) revert InvalidAmount();
    }

    function addLiquidity(uint256 max0, uint256 max1, uint256 minShares, address recipient, uint256 deadline)
        external nonReentrant returns (uint256 shares)
    {
        _validate(recipient, deadline);
        (uint256 amount0, uint256 amount1, uint256 quotedShares) = quoteLiquidity(max0, max1);
        shares = quotedShares;
        if (shares < minShares) revert SlippageExceeded();
        if (totalSupply() == 0) _mint(LOCKED_LIQUIDITY, MINIMUM_LIQUIDITY);
        _receiveExact(token0, amount0);
        _receiveExact(token1, amount1);
        _mint(recipient, shares);
        emit LiquidityAdded(msg.sender, recipient, amount0, amount1, shares);
    }

    function removeLiquidity(uint256 shares, uint256 min0, uint256 min1, address recipient, uint256 deadline)
        external nonReentrant returns (uint256 amount0, uint256 amount1)
    {
        _validate(recipient, deadline);
        if (shares == 0) revert InvalidAmount();
        (uint256 reserve0, uint256 reserve1) = getReserves();
        uint256 supply = totalSupply();
        if (supply == 0) revert InsufficientLiquidity();
        amount0 = Math.mulDiv(shares, reserve0, supply);
        amount1 = Math.mulDiv(shares, reserve1, supply);
        if (amount0 == 0 || amount1 == 0) revert InvalidAmount();
        if (amount0 < min0 || amount1 < min1) revert SlippageExceeded();
        _burn(msg.sender, shares);
        token0.safeTransfer(recipient, amount0);
        token1.safeTransfer(recipient, amount1);
        emit LiquidityRemoved(msg.sender, recipient, amount0, amount1, shares);
    }

    function quoteSwap(address tokenIn, uint256 amountIn) public view returns (uint256 amountOut) {
        if (amountIn == 0) revert InvalidAmount();
        if (tokenIn != address(token0) && tokenIn != address(token1)) revert InvalidToken();
        (uint256 reserve0, uint256 reserve1) = getReserves();
        (uint256 reserveIn, uint256 reserveOut) = tokenIn == address(token0) ? (reserve0, reserve1) : (reserve1, reserve0);
        if (reserveIn == 0 || reserveOut == 0) revert InsufficientLiquidity();
        uint256 netInput = Math.mulDiv(amountIn, BPS - SWAP_FEE_BPS, BPS);
        amountOut = Math.mulDiv(netInput, reserveOut, reserveIn + netInput);
        if (amountOut == 0) revert InvalidAmount();
    }

    function swapExactInput(address tokenIn, uint256 amountIn, uint256 minOut, address recipient, uint256 deadline)
        external nonReentrant returns (uint256 amountOut)
    {
        _validate(recipient, deadline);
        amountOut = quoteSwap(tokenIn, amountIn);
        if (amountOut < minOut) revert SlippageExceeded();
        _receiveExact(IERC20(tokenIn), amountIn);
        IERC20 tokenOut = tokenIn == address(token0) ? token1 : token0;
        tokenOut.safeTransfer(recipient, amountOut);
        emit Swap(msg.sender, tokenIn, recipient, amountIn, amountOut);
    }

    function _receiveExact(IERC20 token, uint256 amount) private {
        uint256 beforeBalance = token.balanceOf(address(this));
        token.safeTransferFrom(msg.sender, address(this), amount);
        if (token.balanceOf(address(this)) - beforeBalance != amount) revert UnsupportedAsset();
    }

    function _validate(address recipient, uint256 deadline) private view {
        if (block.timestamp > deadline) revert DeadlineExpired();
        if (recipient == address(0) || recipient == address(this)) revert InvalidRecipient();
    }
}