Reverse market address -> feedId mapping update

src/GMXAutomationBase.sol

@@ -14,7 +14,7 @@ import {Reader} from "gmx-synthetics/reader/Reader.sol";
 /// @title Base Automation Contract for GMX Automation Contracts
 /// @author Alex Roan - Cyfrin (@alexroan)
 contract GMXAutomationBase is Ownable2Step {
-    using EnumerableMap for EnumerableMap.UintToAddressMap;
+    using EnumerableMap for EnumerableMap.AddressToUintMap;
     using SafeERC20 for IERC20;
 
     // ERRORS
@@ -33,7 +33,7 @@ contract GMXAutomationBase is Ownable2Step {
 
     // This should be empty after every transaction. It is filled and cleared each time checkLog is called.
     // mapping (uint256(feedId) => tokenAddress)
-    EnumerableMap.UintToAddressMap internal s_feedIdToMarketTokenMap;
+    EnumerableMap.AddressToUintMap internal s_marketTokenToFeedId;
 
     /// @param dataStore the DataStore contract address - immutable
     /// @param reader the Reader contract address - immutable
@@ -75,53 +75,66 @@ contract GMXAutomationBase is Ownable2Step {
     // INTERNAL FUNCTIONS
     ///////////////////////////
 
-    /// @notice Pushes the feedIds for marketProps: indexToken, longToken and shortToken to the feedIdToMarketTokenMap
+    /// @notice Pushes the feedIds for marketProps: indexToken, longToken and shortToken to the marketTokenToFeedId
     /// @dev Does not allow for duplicate feedIds or zero address feedIds
     /// @dev Does not push the Props.marketToken feedId to the feedIdToMarketTokenMap
     /// @param marketProps the Market Props struct to retrieve the feedIds from
-    function _addPropsToMapping(Market.Props memory marketProps) internal {
-        if (marketProps.indexToken != address(0)) {
-            uint256 indexTokenFeedId = uint256(i_dataStore.getBytes32(Keys.realtimeFeedIdKey(marketProps.indexToken)));
-            if (indexTokenFeedId == 0) revert GMXAutomationBase_ZeroIndexTokenFeedId();
-            if (!s_feedIdToMarketTokenMap.contains(indexTokenFeedId)) {
-                s_feedIdToMarketTokenMap.set(indexTokenFeedId, marketProps.indexToken);
-            }
+function _addPropsToMapping(Market.Props memory marketProps) internal {
+    if (marketProps.indexToken != address(0)) {
+        uint256 indexTokenFeedId = uint256(i_dataStore.getBytes32(Keys.realtimeFeedIdKey(marketProps.indexToken)));
+        if (indexTokenFeedId == 0) revert GMXAutomationBase_ZeroIndexTokenFeedId();
+        if (!s_marketTokenToFeedId.contains(marketProps.indexToken)) {
+            s_marketTokenToFeedId.set(marketProps.indexToken, indexTokenFeedId);
         }
+    }
 
-        if (marketProps.longToken != address(0)) {
-            uint256 longTokenFeedId = uint256(i_dataStore.getBytes32(Keys.realtimeFeedIdKey(marketProps.longToken)));
-            if (longTokenFeedId == 0) revert GMXAutomationBase_ZeroLongTokenFeedId();
-            if (!s_feedIdToMarketTokenMap.contains(longTokenFeedId)) {
-                s_feedIdToMarketTokenMap.set(longTokenFeedId, marketProps.longToken);
-            }
+    if (marketProps.longToken != address(0)) {
+        uint256 longTokenFeedId = uint256(i_dataStore.getBytes32(Keys.realtimeFeedIdKey(marketProps.longToken)));
+        if (longTokenFeedId == 0) revert GMXAutomationBase_ZeroLongTokenFeedId();
+        if (!s_marketTokenToFeedId.contains(marketProps.longToken)) {
+            s_marketTokenToFeedId.set(marketProps.longToken, longTokenFeedId);
         }
+    }
 
-        if (marketProps.shortToken != address(0)) {
-            uint256 shortTokenFeedId = uint256(i_dataStore.getBytes32(Keys.realtimeFeedIdKey(marketProps.shortToken)));
-            if (shortTokenFeedId == 0) revert GMXAutomationBase_ZeroShortTokenFeedId();
-            if (!s_feedIdToMarketTokenMap.contains(shortTokenFeedId)) {
-                s_feedIdToMarketTokenMap.set(shortTokenFeedId, marketProps.shortToken);
-            }
+    if (marketProps.shortToken != address(0)) {
+        uint256 shortTokenFeedId = uint256(i_dataStore.getBytes32(Keys.realtimeFeedIdKey(marketProps.shortToken)));
+        if (shortTokenFeedId == 0) revert GMXAutomationBase_ZeroShortTokenFeedId();
+        if (!s_marketTokenToFeedId.contains(marketProps.shortToken)) {
+            s_marketTokenToFeedId.set(marketProps.shortToken, shortTokenFeedId);
         }
     }
+}
 
-    /// @notice Returns all values from and clears the s_feedIdToMarketTokenMap
-    /// @dev Iterates over the feedIdToMarketTokenMap, and removes each feedId and returns them as an array
-    /// @return feedIds the feedIds that were in the feedIdToMarketTokenMap
-    /// @return addresses the addresses that were in the feedIdToMarketTokenMap
-    function _flushMapping() internal returns (string[] memory feedIds, address[] memory addresses) {
-        uint256 length = s_feedIdToMarketTokenMap.length();
-        feedIds = new string[](length);
-        addresses = new address[](length);
-        uint256 count = 0;
-        while (s_feedIdToMarketTokenMap.length() > 0) {
-            (uint256 uintKey, address value) = s_feedIdToMarketTokenMap.at(s_feedIdToMarketTokenMap.length() - 1);
-            s_feedIdToMarketTokenMap.remove(uintKey);
-            feedIds[count] = _toHexString(bytes32(uintKey));
-            addresses[count] = value;
-            count++;
+
+function _isAddressInMap(address addr) private view returns (bool) {
+    for (uint i = 0; i < s_marketTokenToFeedId.length(); i++) {
+        (address currentAddress, ) = s_marketTokenToFeedId.at(i);
+        if (currentAddress == addr) {
+            return true;
         }
     }
+    return false;
+}
+
+
+/// @notice Returns all values from and clears the s_marketTokenToFeedId
+/// @dev Iterates over the addressToMarketTokenMap, and removes each address and returns them as an array along with the corresponding feedIds
+/// @return feedIds the feedIds that were in the addressToMarketTokenMap mapped with respective token addresses
+/// @return addresses the addresses that were in the addressToMarketTokenMap mapped with respective feedIds
+    function _flushMapping() internal returns (string[] memory feedIds, address[] memory addresses) {
+    uint256 length = s_marketTokenToFeedId.length();
+    feedIds = new string[](length);
+    addresses = new address[](length);
+    uint256 count = 0;
+    while (s_marketTokenToFeedId.length() > 0) {
+        (address addressKey, uint256 uintValue) = s_marketTokenToFeedId.at(s_marketTokenToFeedId.length() - 1);
+        s_marketTokenToFeedId.remove(addressKey);
+        feedIds[count] = _toHexString(bytes32(uintValue));
+        addresses[count] = addressKey;
+        count++;
+    }
+}
+
 
     /// @notice Converts a bytes buffer to a hexadecimal string
     /// @param value the bytes32 value to convert

test/GMXAutomationBase.t.sol

@@ -121,10 +121,10 @@ contract GMXAutomationBaseTest__addPropsToMapping is Test {
         );
         s_gmxAutomation.addPropsToMapping(marketProps);
 
-        assertEq(s_gmxAutomation.feedIdToMarketTokenMapLength(), 3);
-        assertEq(s_gmxAutomation.feedIdToMarketTokenMapGet(uint256(indexTokenFeedId)), marketProps.indexToken);
-        assertEq(s_gmxAutomation.feedIdToMarketTokenMapGet(uint256(longTokenFeedId)), marketProps.longToken);
-        assertEq(s_gmxAutomation.feedIdToMarketTokenMapGet(uint256(shortTokenFeedId)), marketProps.shortToken);
+        assertEq(s_gmxAutomation.marketTokenToFeedIdMapLength(), 3);
+        assertEq(s_gmxAutomation.marketTokenToFeedIdMapGet(marketProps.indexToken), uint256(indexTokenFeedId));
+        assertEq(s_gmxAutomation.marketTokenToFeedIdMapGet(marketProps.longToken), uint256(longTokenFeedId));
+        assertEq(s_gmxAutomation.marketTokenToFeedIdMapGet(marketProps.shortToken), uint256(shortTokenFeedId));
     }
 
     function test__addPropsToMapping_ZeroIndexTokenFeedId_reverts() public {
@@ -231,7 +231,7 @@ contract GMXAutomationBaseTest__flushMapping is Test {
         addresses[1] = address(2);
         addresses[2] = address(3);
         for (uint256 i = 0; i < feedIds.length; i++) {
-            s_gmxAutomation.feedIdToMarketTokenMapSet(uint256(feedIds[i]), addresses[i]);
+            s_gmxAutomation.marketTokenToFeedIdMapSet(addresses[i],uint256(feedIds[i]));
         }
         s_feedIdToAddress["0x14e044f932bb959cc2aa8dc1ba110c09224e639aae00264c1ffc2a0830904a3c"] = addresses[0];
         s_feedIdToAddress["0x4ce52cf28e49f4673198074968aeea280f13b5f897c687eb713bcfc1eeab89ba"] = addresses[1];
@@ -246,10 +246,10 @@ contract GMXAutomationBaseTest__flushMapping is Test {
     function test__flushMapping_shouldBeEmptyAfter(address[] memory addresses) public {
         for (uint256 i = 0; i < addresses.length; i++) {
             if (addresses[i] == address(0)) continue;
-            s_gmxAutomation.feedIdToMarketTokenMapSet(uint256(keccak256(abi.encode(addresses[i]))), addresses[i]);
+            s_gmxAutomation.marketTokenToFeedIdMapSet(addresses[i], uint256(keccak256(abi.encode(addresses[i]))));
         }
         s_gmxAutomation.flushMapping();
-        assertEq(s_gmxAutomation.feedIdToMarketTokenMapLength(), 0);
+        assertEq(s_gmxAutomation.marketTokenToFeedIdMapLength(), 0);
     }
 }
 

test/helpers/GMXAutomationBaseHelper.sol

@@ -11,7 +11,7 @@ import {Market} from "gmx-synthetics/market/Market.sol";
 import {EnumerableMap} from "openzeppelin/utils/structs/EnumerableMap.sol";
 
 contract GMXAutomationBaseHelper is GMXAutomationBase {
-    using EnumerableMap for EnumerableMap.UintToAddressMap;
+    using EnumerableMap for EnumerableMap.AddressToUintMap;
 
     constructor(DataStore dataStore, Reader reader) GMXAutomationBase(dataStore, reader) {}
 
@@ -29,27 +29,27 @@ contract GMXAutomationBaseHelper is GMXAutomationBase {
         return _toHexString(value);
     }
 
-    function feedIdToMarketTokenMapSet(uint256 feedId, address addr) public {
-        s_feedIdToMarketTokenMap.set(feedId, addr);
+    function marketTokenToFeedIdMapSet(address addr, uint256 feedId) public {
+        s_marketTokenToFeedId.set(addr, feedId);
     }
 
-    function feedIdToMarketTokenMapLength() public view returns (uint256) {
-        return s_feedIdToMarketTokenMap.length();
+    function marketTokenToFeedIdMapLength() public view returns (uint256) {
+        return s_marketTokenToFeedId.length();
     }
 
-    function feedIdToMarketTokenMapContains(uint256 feedId) public view returns (bool) {
-        return s_feedIdToMarketTokenMap.contains(feedId);
+    function marketTokenToFeedIdMapContains(address addr) public view returns (bool) {
+        return s_marketTokenToFeedId.contains(addr);
     }
 
-    function feedIdToMarketTokenMapGet(uint256 feedId) public view returns (address) {
-        return s_feedIdToMarketTokenMap.get(feedId);
+    function marketTokenToFeedIdMapGet(address addr) public view returns (uint) {
+        return s_marketTokenToFeedId.get(addr);
     }
 
-    function feedIdToMarketTokenMapAt(uint256 index) public view returns (uint256, address) {
-        return s_feedIdToMarketTokenMap.at(index);
+    function marketTokenToFeedIdMapAt(uint256 index) public view returns (address, uint256) {
+        return s_marketTokenToFeedId.at(index);
     }
 
-    function feedIdToMarketTokenMapKeys() public view returns (uint256[] memory) {
-        return s_feedIdToMarketTokenMap.keys();
+    function marketTokenToFeedIdMapKeys() public view returns (address[] memory) {
+        return s_marketTokenToFeedId.keys();
     }
 }