LP token update

Author: Vivkzz

src/Factory.sol

@@ -28,7 +28,7 @@ contract Factory {
         }
 
         getPair[token0][token1] = pair;
-        getPair[token1][token0] = pair; // optional: not needed
+        getPair[token1][token0] = pair; // Easy to find in router
         allPair.push(pair);
         emit PairCreated(token0, token1, pair, allPair.length);
     }

src/LPToken.sol

@@ -0,0 +1,25 @@
+// SPDX-License-Identifier: SEE LICENSE IN LICENSE
+pragma solidity 0.8;
+
+import {ERC20} from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
+
+contract LPToken is ERC20 {
+    address public pair;
+
+    constructor(string memory name, string memory symbol, address _pair) ERC20(name, symbol) {
+        pair = _pair;
+    }
+
+    modifier onlyPair() {
+        require(msg.sender == pair, "Only pair contract can access it");
+        _;
+    }
+
+    function mint(address to, uint256 amount) external onlyPair {
+        _mint(to, amount);
+    }
+
+    function burn(address from, uint256 amount) external onlyPair {
+        _burn(from, amount);
+    }
+}

src/Pair.sol

@@ -2,8 +2,10 @@
 pragma solidity 0.8;
 
 import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import {LPToken} from "./LPToken.sol";
 
 contract Pair {
+    LPToken public lptoken;
     address public token0;
     address public token1;
 
@@ -13,18 +15,46 @@ contract Pair {
     constructor(address _token0, address _token1) {
         token0 = _token0;
         token1 = _token1;
+
+        lptoken = new LPToken("LPTOKEN", "LPTK", address(this));
     }
 
     function getReserve() public view returns (uint256, uint256) {
         return (reserve0, reserve1);
     }
 
-    function addLiquidity(uint256 amount0, uint256 amount1) external {
+    function addLiquidity(uint256 amount0, uint256 amount1, address to) external {
         IERC20(address(token0)).transferFrom(msg.sender, address(this), amount0);
         IERC20(address(token1)).transferFrom(msg.sender, address(this), amount1);
 
         reserve0 += amount0;
         reserve1 += amount1;
+
+        uint256 liquidity;
+        if (lptoken.totalSupply() == 0) {
+            liquidity = sqrt(amount0 * amount1);
+        } else {
+            liquidity = min((amount0 * lptoken.totalSupply()) / reserve0, (amount1 * lptoken.totalSupply()) / reserve1);
+        }
+        //potential error if we use msg.sender it will be router so we have taken to
+        lptoken.mint(to, liquidity);
+    }
+
+    function min(uint256 x, uint256 y) internal pure returns (uint256) {
+        return x < y ? x : y;
+    }
+
+    function sqrt(uint256 y) internal pure returns (uint256 z) {
+        if (y > 3) {
+            z = y;
+            uint256 x = y / 2 + 1;
+            while (x < z) {
+                z = x;
+                x = (y / x + x) / 2;
+            }
+        } else if (y != 0) {
+            z = 1;
+        }
     }
 
     function swap(address tokenIn, uint256 amountIn, address to) external {

src/Router.sol

@@ -3,12 +3,13 @@ pragma solidity 0.8;
 
 import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import {Pair} from "./Pair.sol";
+import {Factory} from "./Factory.sol";
 
 contract Router {
-    address public immutable i_pair;
+    Factory public factory;
 
-    constructor(address _pair) {
-        i_pair = _pair;
+    constructor(address _factory) {
+        factory = Factory(factory);
     }
 
     function sortToken(address tokenA, address tokenB) public pure returns (address, address) {
@@ -20,18 +21,27 @@ contract Router {
     }
 
     function addLiquidity(address tokenA, address tokenB, uint256 amountA, uint256 amountB) external {
+        address pair = factory.getPair(tokenA, tokenB);
+
+        if (pair == address(0)) {
+            factory.createPair(tokenA, tokenB);
+            pair = factory.getPair(tokenA, tokenB);
+        }
         IERC20(tokenA).transferFrom(msg.sender, address(this), amountA);
         IERC20(tokenB).transferFrom(msg.sender, address(this), amountB);
 
-        IERC20(tokenA).approve(i_pair, amountA);
-        IERC20(tokenB).approve(i_pair, amountB);
+        IERC20(tokenA).approve(pair, amountA);
+        IERC20(tokenB).approve(pair, amountB);
 
-        Pair(i_pair).addLiquidity(amountA, amountB);
+        Pair(pair).addLiquidity(amountA, amountB, msg.sender);
     }
 
-    function swapExactToken(address tokenIn, uint256 amountIn) external {
+    function swapExactToken(address tokenIn, address tokenOut, uint256 amountIn) external {
+        address pair = factory.getPair(tokenIn, tokenOut);
+        require(pair != address(0), "Pair doesnt exist");
+
         IERC20(tokenIn).transferFrom(msg.sender, address(this), amountIn);
-        IERC20(tokenIn).approve(i_pair, amountIn);
-        Pair(i_pair).swap(tokenIn, amountIn, msg.sender);
+        IERC20(tokenIn).approve(pair, amountIn);
+        Pair(pair).swap(tokenIn, amountIn, msg.sender);
     }
 }