improve tests

contracts/gas-snapshots/ccip.gas-snapshot

@@ -684,41 +684,42 @@ TokenPoolFactory_createTokenPool:test_createTokenPool_ExistingRemoteToken_AndPre
 TokenPoolFactory_createTokenPool:test_createTokenPool_WithNoExistingRemoteContracts_predict_Success() (gas: 13556200)
 TokenPoolFactory_createTokenPool:test_createTokenPool_WithNoExistingTokenOnRemoteChain_Success() (gas: 6196728)
 TokenPoolFactory_createTokenPool:test_createTokenPool_WithRemoteTokenAndRemotePool_Success() (gas: 6426549)
-TokenPoolWithAllowList_applyAllowListUpdates:test_AllowListNotEnabled_Revert() (gas: 2730121)
-TokenPoolWithAllowList_applyAllowListUpdates:test_OnlyOwner_Revert() (gas: 12185)
-TokenPoolWithAllowList_applyAllowListUpdates:test_SetAllowListSkipsZero_Success() (gas: 23633)
-TokenPoolWithAllowList_applyAllowListUpdates:test_SetAllowList_Success() (gas: 178556)
-TokenPoolWithAllowList_getAllowList:test_GetAllowList_Success() (gas: 23929)
+TokenPoolWithAllowList_applyAllowListUpdates:test_AllowListNotEnabled_Revert() (gas: 2750102)
+TokenPoolWithAllowList_applyAllowListUpdates:test_OnlyOwner_Revert() (gas: 12119)
+TokenPoolWithAllowList_applyAllowListUpdates:test_SetAllowListSkipsZero_Success() (gas: 23477)
+TokenPoolWithAllowList_applyAllowListUpdates:test_SetAllowList_Success() (gas: 178290)
+TokenPoolWithAllowList_getAllowList:test_GetAllowList_Success() (gas: 23884)
 TokenPoolWithAllowList_getAllowListEnabled:test_GetAllowListEnabled_Success() (gas: 8408)
-TokenPoolWithAllowList_setRouter:test_SetRouter_Success() (gas: 24983)
-TokenPoolWithAllowList_setRouter:test_ZeroAddressNotAllowed_Revert() (gas: 10663)
-TokenPool_addRemotePool:test_NonExistentChain_Revert() (gas: 14186)
-TokenPool_addRemotePool:test_PoolAlreadyAdded_Revert() (gas: 124357)
-TokenPool_addRemotePool:test_ZeroAddressNotAllowed_Revert() (gas: 14201)
-TokenPool_addRemotePool:test_ZeroLengthAddressNotAllowed_Revert() (gas: 14097)
-TokenPool_addRemotePool:test_addRemotePool_MultipleActive() (gas: 150422)
-TokenPool_addRemotePool:test_addRemotePool_Success() (gas: 160757)
-TokenPool_applyChainUpdates:test_applyChainUpdates_InvalidRateLimitRate_Revert() (gas: 704340)
+TokenPoolWithAllowList_setRouter:test_SetRouter_Success() (gas: 25049)
+TokenPoolWithAllowList_setRouter:test_ZeroAddressNotAllowed_Revert() (gas: 10685)
+TokenPool_addRemotePool:test_NonExistentChain_Revert() (gas: 14208)
+TokenPool_addRemotePool:test_PoolAlreadyAdded_Revert() (gas: 124401)
+TokenPool_addRemotePool:test_ZeroAddressNotAllowed_Revert() (gas: 14223)
+TokenPool_addRemotePool:test_ZeroLengthAddressNotAllowed_Revert() (gas: 14119)
+TokenPool_addRemotePool:test_addRemotePool_MultipleActive() (gas: 422223)
+TokenPool_addRemotePool:test_addRemotePool_Success() (gas: 160801)
+TokenPool_applyChainUpdates:test_applyChainUpdates_InvalidRateLimitRate_Revert() (gas: 704516)
 TokenPool_applyChainUpdates:test_applyChainUpdates_NonExistentChain_Revert() (gas: 14999)
 TokenPool_applyChainUpdates:test_applyChainUpdates_OnlyCallableByOwner_Revert() (gas: 11841)
-TokenPool_applyChainUpdates:test_applyChainUpdates_Success() (gas: 608018)
-TokenPool_applyChainUpdates:test_applyChainUpdates_ZeroAddressNotAllowed_Revert() (gas: 158371)
-TokenPool_constructor:test_ZeroAddressNotAllowed_Revert() (gas: 71389)
-TokenPool_constructor:test_immutableFields_Success() (gas: 20718)
-TokenPool_getRemotePool:test_getRemotePools() (gas: 332279)
-TokenPool_onlyOffRamp:test_CallerIsNotARampOnRouter_Revert() (gas: 21504)
-TokenPool_onlyOffRamp:test_ChainNotAllowed_Revert() (gas: 335819)
-TokenPool_onlyOffRamp:test_onlyOffRamp_Success() (gas: 94291)
-TokenPool_onlyOnRamp:test_CallerIsNotARampOnRouter_Revert() (gas: 21134)
-TokenPool_onlyOnRamp:test_ChainNotAllowed_Revert() (gas: 299725)
-TokenPool_onlyOnRamp:test_onlyOnRamp_Success() (gas: 49216)
-TokenPool_removeRemotePool:test_InvalidRemotePoolForChain_Revert() (gas: 17860)
-TokenPool_removeRemotePool:test_NonExistentChain_Revert() (gas: 14277)
-TokenPool_removeRemotePool:test_removeRemotePool_Success() (gas: 136446)
+TokenPool_applyChainUpdates:test_applyChainUpdates_Success() (gas: 608388)
+TokenPool_applyChainUpdates:test_applyChainUpdates_UpdatesRemotePoolHashes() (gas: 1332044)
+TokenPool_applyChainUpdates:test_applyChainUpdates_ZeroAddressNotAllowed_Revert() (gas: 158393)
+TokenPool_constructor:test_ZeroAddressNotAllowed_Revert() (gas: 71413)
+TokenPool_constructor:test_immutableFields_Success() (gas: 20762)
+TokenPool_getRemotePool:test_getRemotePools() (gas: 332345)
+TokenPool_onlyOffRamp:test_CallerIsNotARampOnRouter_Revert() (gas: 21439)
+TokenPool_onlyOffRamp:test_ChainNotAllowed_Revert() (gas: 335698)
+TokenPool_onlyOffRamp:test_onlyOffRamp_Success() (gas: 94226)
+TokenPool_onlyOnRamp:test_CallerIsNotARampOnRouter_Revert() (gas: 21156)
+TokenPool_onlyOnRamp:test_ChainNotAllowed_Revert() (gas: 299744)
+TokenPool_onlyOnRamp:test_onlyOnRamp_Success() (gas: 49238)
+TokenPool_removeRemotePool:test_InvalidRemotePoolForChain_Revert() (gas: 17882)
+TokenPool_removeRemotePool:test_NonExistentChain_Revert() (gas: 14299)
+TokenPool_removeRemotePool:test_removeRemotePool_Success() (gas: 136499)
 TokenPool_setChainRateLimiterConfig:test_NonExistentChain_Revert() (gas: 17169)
 TokenPool_setChainRateLimiterConfig:test_OnlyOwnerOrRateLimitAdmin_Revert() (gas: 15262)
-TokenPool_setRateLimitAdmin:test_SetRateLimitAdmin_Revert() (gas: 11002)
-TokenPool_setRateLimitAdmin:test_SetRateLimitAdmin_Success() (gas: 37562)
+TokenPool_setRateLimitAdmin:test_SetRateLimitAdmin_Revert() (gas: 11024)
+TokenPool_setRateLimitAdmin:test_SetRateLimitAdmin_Success() (gas: 37628)
 USDCBridgeMigrator_BurnLockedUSDC:test_PrimaryMechanism_Success() (gas: 135933)
 USDCBridgeMigrator_BurnLockedUSDC:test_WhileMigrationPause_Revert() (gas: 109772)
 USDCBridgeMigrator_BurnLockedUSDC:test_invalidPermissions_Revert() (gas: 39427)

contracts/src/v0.8/ccip/pools/TokenPool.sol

@@ -276,6 +276,7 @@ abstract contract TokenPool is IPoolV1, Ownable2StepMsgSender {
     bytes32 remotePoolHash = keccak256(remotePoolAddress);
     for (uint256 i = 0; i < remotePoolAddresses.length; ++i) {
       if (keccak256(remotePoolAddresses[i]) == remotePoolHash) {
+        // Swap the last element with the element to remove and then pop the last element.
         s_remoteChainConfigs[remoteChainSelector].remotePoolAddresses[i] =
           remotePoolAddresses[remotePoolAddresses.length - 1];
         s_remoteChainConfigs[remoteChainSelector].remotePoolAddresses.pop();

contracts/src/v0.8/ccip/test/helpers/TokenPoolHelper.sol

@@ -5,15 +5,22 @@ import {Pool} from "../../libraries/Pool.sol";
 import {TokenPool} from "../../pools/TokenPool.sol";
 
 import {IERC20} from "../../../vendor/openzeppelin-solidity/v4.8.3/contracts/token/ERC20/IERC20.sol";
+import {EnumerableSet} from "../../../vendor/openzeppelin-solidity/v5.0.2/contracts/utils/structs/EnumerableSet.sol";
 
 contract TokenPoolHelper is TokenPool {
+  using EnumerableSet for EnumerableSet.Bytes32Set;
+
   constructor(
     IERC20 token,
     address[] memory allowlist,
     address rmnProxy,
     address router
   ) TokenPool(token, allowlist, rmnProxy, router) {}
 
+  function getRemotePoolHashes() external view returns (bytes32[] memory) {
+    return s_remotePoolHashes.values();
+  }
+
   function lockOrBurn(
     Pool.LockOrBurnInV1 calldata lockOrBurnIn
   ) external override returns (Pool.LockOrBurnOutV1 memory) {

contracts/src/v0.8/ccip/test/pools/TokenPool/TokenPool.addRemotePool.t.sol

@@ -44,6 +44,9 @@ contract TokenPool_addRemotePool is TokenPoolSetup {
     vm.expectRevert(abi.encodeWithSelector(TokenPool.InvalidSourcePoolAddress.selector, remotePools[0]));
     s_tokenPool.releaseOrMint(_getReleaseOrMintInV1(remotePools[0]));
 
+    // There's already one pool setup through the test setup
+    assertEq(s_tokenPool.getRemotePoolHashes().length, 1);
+
     vm.startPrank(OWNER);
     s_tokenPool.addRemotePool(DEST_CHAIN_SELECTOR, remotePools[0]);
 
@@ -55,22 +58,33 @@ contract TokenPool_addRemotePool is TokenPoolSetup {
     s_tokenPool.addRemotePool(DEST_CHAIN_SELECTOR, remotePools[1]);
 
     // Both should now work
+    assertEq(s_tokenPool.getRemotePoolHashes().length, 3);
     vm.startPrank(fakeOffRamp);
     s_tokenPool.releaseOrMint(_getReleaseOrMintInV1(remotePools[0]));
     s_tokenPool.releaseOrMint(_getReleaseOrMintInV1(remotePools[1]));
 
     // Adding a third pool, and removing the first one
     vm.startPrank(OWNER);
     s_tokenPool.addRemotePool(DEST_CHAIN_SELECTOR, remotePools[2]);
+    assertEq(s_tokenPool.getRemotePoolHashes().length, 4);
     s_tokenPool.removeRemotePool(DEST_CHAIN_SELECTOR, remotePools[0]);
+    assertEq(s_tokenPool.getRemotePoolHashes().length, 3);
 
     // Only the last two should work
     vm.startPrank(fakeOffRamp);
     vm.expectRevert(abi.encodeWithSelector(TokenPool.InvalidSourcePoolAddress.selector, remotePools[0]));
     s_tokenPool.releaseOrMint(_getReleaseOrMintInV1(remotePools[0]));
-
     s_tokenPool.releaseOrMint(_getReleaseOrMintInV1(remotePools[1]));
     s_tokenPool.releaseOrMint(_getReleaseOrMintInV1(remotePools[2]));
+
+    // Removing the chain removes all associated pool hashes
+    vm.startPrank(OWNER);
+
+    uint64[] memory chainSelectorsToRemove = new uint64[](1);
+    chainSelectorsToRemove[0] = DEST_CHAIN_SELECTOR;
+    s_tokenPool.applyChainUpdates(chainSelectorsToRemove, new TokenPool.ChainUpdate[](0));
+
+    assertEq(s_tokenPool.getRemotePoolHashes().length, 0);
   }
 
   function _getReleaseOrMintInV1(

contracts/src/v0.8/ccip/test/pools/TokenPool/TokenPool.applyChainUpdates.t.sol

@@ -128,6 +128,75 @@ contract TokenPool_applyChainUpdates is RouterSetup {
     s_tokenPool.applyChainUpdates(new uint64[](0), singleChainConfigured);
   }
 
+  function test_applyChainUpdates_UpdatesRemotePoolHashes() public {
+    assertEq(s_tokenPool.getRemotePoolHashes().length, 0);
+
+    uint64 selector1 = 789;
+    uint64 selector2 = 123;
+    uint64 selector3 = 456;
+
+    bytes memory pool1 = abi.encode(makeAddr("pool1"));
+    bytes memory pool2 = abi.encode(makeAddr("pool2"));
+    bytes memory pool3 = abi.encode(makeAddr("pool3"));
+
+    TokenPool.ChainUpdate[] memory chainUpdates = new TokenPool.ChainUpdate[](3);
+    chainUpdates[0] = TokenPool.ChainUpdate({
+      remoteChainSelector: selector1,
+      remotePoolAddress: pool1,
+      remoteTokenAddress: pool1,
+      outboundRateLimiterConfig: _getOutboundRateLimiterConfig(),
+      inboundRateLimiterConfig: _getInboundRateLimiterConfig()
+    });
+    chainUpdates[1] = TokenPool.ChainUpdate({
+      remoteChainSelector: selector2,
+      remotePoolAddress: pool2,
+      remoteTokenAddress: pool2,
+      outboundRateLimiterConfig: _getOutboundRateLimiterConfig(),
+      inboundRateLimiterConfig: _getInboundRateLimiterConfig()
+    });
+    chainUpdates[2] = TokenPool.ChainUpdate({
+      remoteChainSelector: selector3,
+      remotePoolAddress: pool3,
+      remoteTokenAddress: pool3,
+      outboundRateLimiterConfig: _getOutboundRateLimiterConfig(),
+      inboundRateLimiterConfig: _getInboundRateLimiterConfig()
+    });
+
+    s_tokenPool.applyChainUpdates(new uint64[](0), chainUpdates);
+
+    assertEq(s_tokenPool.getRemotePoolHashes().length, 3);
+
+    // This adds 3 for the first chain, 2 for the second, and 1 for the third for a total of 6.
+    // Since each chain already had one, the totals are 4 + 3 + 2
+    for (uint256 i = 0; i < chainUpdates.length; ++i) {
+      for (uint256 j = i; j < chainUpdates.length; ++j) {
+        s_tokenPool.addRemotePool(chainUpdates[i].remoteChainSelector, abi.encode(i, j));
+      }
+    }
+
+    assertEq(s_tokenPool.getRemotePoolHashes().length, 4 + 3 + 2);
+
+    // Removing a chain should remove all associated pool hashes
+    uint64[] memory chainRemoves = new uint64[](1);
+    chainRemoves[0] = selector1;
+
+    s_tokenPool.applyChainUpdates(chainRemoves, new TokenPool.ChainUpdate[](0));
+
+    assertEq(s_tokenPool.getRemotePoolHashes().length, 3 + 2);
+
+    chainRemoves[0] = selector2;
+
+    s_tokenPool.applyChainUpdates(chainRemoves, new TokenPool.ChainUpdate[](0));
+
+    assertEq(s_tokenPool.getRemotePoolHashes().length, 2);
+
+    chainRemoves[0] = selector3;
+
+    s_tokenPool.applyChainUpdates(chainRemoves, new TokenPool.ChainUpdate[](0));
+
+    assertEq(s_tokenPool.getRemotePoolHashes().length, 0);
+  }
+
   // Reverts
 
   function test_applyChainUpdates_OnlyCallableByOwner_Revert() public {