Merge pull request #1 from zarkk01/fix-s_tokenToPoolbug

Fix bug that made s_tokenToPool to store only the last created pool
mgnfy-view · Mar 24, 2024 · 90f27ae · 90f27ae
2 parents a80884a + 2a80f57
commit 90f27ae
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 14 deletions.
diff --git a/src/ThunderSwapPoolFactory.sol b/src/ThunderSwapPoolFactory.sol
@@ -6,9 +6,10 @@ import { Ownable } from "@openzeppelin/contracts/access/Ownable.sol";
 import { ERC20 } from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
 
 contract ThunderSwapPoolFactory is Ownable {
-    mapping(address => bool) private s_supportedTokens;
-    mapping(address => address) private s_tokenToPool;
-    mapping(address => address[]) private s_poolToTokens;
+    mapping(address token => bool supported) private s_supportedTokens;
+    mapping(address token1 => mapping(address token2 => address pool)) private s_pairings;
+    mapping(address token => address[] pools) private s_tokenToPools;
+    mapping(address pool => address[] tokens) private s_poolToTokens;
 
     event PoolCreated(address newPool, address token1, address token2);
     event SupportedToken(address supportedToken);
@@ -42,13 +43,11 @@ contract ThunderSwapPoolFactory is Ownable {
         external
         returns (ThunderSwapPool)
     {
-        address poolFromToken1 = s_tokenToPool[_token1];
         if (!s_supportedTokens[_token1]) revert TokenNotSupported(_token1);
         if (!s_supportedTokens[_token2]) revert TokenNotSupported(_token2);
         if (_token1 == _token2) revert PoolCannotHaveTwoTokensOfTheSameType();
-        if (poolFromToken1 != address(0) && poolFromToken1 == s_tokenToPool[_token2]) {
-            revert PoolAlreadyExists(poolFromToken1);
-        }
+        address _pool = s_pairings[_token1][_token2];
+        if (_pool != address(0)) revert PoolAlreadyExists(_pool);
 
         string memory poolName =
             string.concat("ThunderSwap", ERC20(_token1).name(), ERC20(_token2).name());
@@ -57,8 +56,12 @@ contract ThunderSwapPoolFactory is Ownable {
         ThunderSwapPool newPool = new ThunderSwapPool(_token1, _token2, poolName, poolSymbol);
 
         s_poolToTokens[address(newPool)] = [_token1, _token2];
-        s_tokenToPool[_token1] = address(newPool);
-        s_tokenToPool[_token2] = address(newPool);
+
+        s_tokenToPools[_token1].push(address(newPool));
+        s_tokenToPools[_token2].push(address(newPool));
+
+        s_pairings[_token1][_token2] = address(newPool);
+        s_pairings[_token2][_token1] = address(newPool);
 
         emit PoolCreated(address(newPool), _token1, _token2);
 
@@ -79,8 +82,8 @@ contract ThunderSwapPoolFactory is Ownable {
      * @param _token The token to find a pool for
      * @return The thunder swap pool address for the given token
      */
-    function getPoolFromToken(address _token) external view returns (address) {
-        return s_tokenToPool[_token];
+    function getPoolsFromToken(address _token) external view returns (address[] memory) {
+        return s_tokenToPools[_token];
     }
 
     /**
@@ -91,4 +94,14 @@ contract ThunderSwapPoolFactory is Ownable {
     function getPoolTokens(address _pool) external view returns (address[] memory) {
         return s_poolToTokens[_pool];
     }
+
+    /**
+     * @notice Returns the address of the pool that two tokens have.
+     * @param _token1 The address of the first token.
+     * @param _token2 The address of the second token.
+     * @return The address of the pool for _token1 and _token2.
+     */
+    function getPairing(address _token1, address _token2) external view returns (address) {
+        return s_pairings[_token1][_token2];
+    }
 }
diff --git a/test/unit/miscellaneousTests/miscellaneous.t.sol b/test/unit/miscellaneousTests/miscellaneous.t.sol
@@ -115,9 +115,13 @@ contract MiscellaneousTest is UniversalHelper {
         assert(thunderSwapPoolFactory.isTokenSupported(supportToken));
     }
 
-    function testGetPoolFromToken() public view {
-        assertEq(thunderSwapPoolFactory.getPoolFromToken(address(tokenA)), address(thunderSwapPool));
-        assertEq(thunderSwapPoolFactory.getPoolFromToken(address(tokenB)), address(thunderSwapPool));
+    function testGetPoolsFromToken() public view {
+        assertEq(
+            thunderSwapPoolFactory.getPoolsFromToken(address(tokenA))[0], address(thunderSwapPool)
+        );
+        assertEq(
+            thunderSwapPoolFactory.getPoolsFromToken(address(tokenB))[0], address(thunderSwapPool)
+        );
     }
 
     function testGetPoolTokensFromThunderSwapPool() public view {
@@ -126,4 +130,9 @@ contract MiscellaneousTest is UniversalHelper {
         assertEq(poolTokens[0], address(tokenA));
         assertEq(poolTokens[1], address(tokenB));
     }
+
+    function testGetPoolFromPairings() public view {
+        assertEq(thunderSwapPoolFactory.getPairing(address(tokenA),address(tokenB)), address(thunderSwapPool));
+         assertEq(thunderSwapPoolFactory.getPairing(address(tokenA),address(tokenB)), address(thunderSwapPool));
+    }
 }