diff --git a/src/GMXAutomationBase.sol b/src/GMXAutomationBase.sol index deb10f5..bad31c5 100644 --- a/src/GMXAutomationBase.sol +++ b/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 @@ -32,8 +32,8 @@ contract GMXAutomationBase is Ownable2Step { address public s_forwarderAddress; // 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; + // mapping (tokenAddress => uint256(feedId)) + EnumerableMap.AddressToUintMap internal s_marketTokenToFeedId; /// @param dataStore the DataStore contract address - immutable /// @param reader the Reader contract address - immutable @@ -83,42 +83,42 @@ contract GMXAutomationBase is Ownable2Step { 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); + 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 (!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 (!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 + /// @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_feedIdToMarketTokenMap.length(); + uint256 length = s_marketTokenToFeedId.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; + 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++; } } diff --git a/test/GMXAutomationBase.t.sol b/test/GMXAutomationBase.t.sol index 55c253d..76d7598 100644 --- a/test/GMXAutomationBase.t.sol +++ b/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); } } diff --git a/test/helpers/GMXAutomationBaseHelper.sol b/test/helpers/GMXAutomationBaseHelper.sol index 44ea0c2..91b02dc 100644 --- a/test/helpers/GMXAutomationBaseHelper.sol +++ b/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(); } }