mv poolBalances to regular mapping and array

pkg/interfaces/contracts/test/IVaultMainMock.sol

@@ -50,7 +50,7 @@ interface IVaultMainMock {
 
     function manualSetStaticSwapFeePercentage(address pool, uint256 value) external;
 
-    function manualSetPoolTokenBalances(address, IERC20[] memory, uint256[] memory, uint256[] memory) external;
+    function manualSetPoolTokensAndBalances(address, IERC20[] memory, uint256[] memory, uint256[] memory) external;
 
     function manualSetPoolConfigBits(address pool, PoolConfigBits config) external;
 

pkg/interfaces/contracts/vault/IVaultErrors.sol

@@ -48,6 +48,12 @@ interface IVaultErrors {
      */
     error TokenAlreadyRegistered(IERC20 token);
 
+    /**
+     * @dev Token were already registered (i.e., it is a duplicate in the pool).
+     * @param tokens The tokens from the storage
+     */
+    error TokensAlreadyRegistered(IERC20[] tokens);
+
     /// @dev The token count is below the minimum allowed.
     error MinTokens();
 

pkg/pool-utils/test/foundry/PoolInfo.t.sol

@@ -53,7 +53,12 @@ contract PoolInfoTest is BaseTest {
     function testGetTokenInfo() public {
         uint256[] memory expectedRawBalances = [uint256(1), uint256(2)].toMemoryArray();
         uint256[] memory expectedLastLiveBalances = [uint256(3), uint256(4)].toMemoryArray();
-        vault.manualSetPoolTokenBalances(address(poolInfo), poolTokens, expectedRawBalances, expectedLastLiveBalances);
+        vault.manualSetPoolTokensAndBalances(
+            address(poolInfo),
+            poolTokens,
+            expectedRawBalances,
+            expectedLastLiveBalances
+        );
 
         TokenInfo[] memory expectedTokenInfo = new TokenInfo[](2);
         expectedTokenInfo[0] = TokenInfo({
@@ -125,7 +130,12 @@ contract PoolInfoTest is BaseTest {
     function testGetCurrentLiveBalances() public {
         uint256[] memory expectedRawBalances = [uint256(12), uint256(34)].toMemoryArray();
         uint256[] memory expectedLastLiveBalances = [uint256(56), uint256(478)].toMemoryArray();
-        vault.manualSetPoolTokenBalances(address(poolInfo), poolTokens, expectedRawBalances, expectedLastLiveBalances);
+        vault.manualSetPoolTokensAndBalances(
+            address(poolInfo),
+            poolTokens,
+            expectedRawBalances,
+            expectedLastLiveBalances
+        );
 
         PoolConfig memory config;
         config.isPoolRegistered = true;

pkg/pool-weighted/test/gas/.hardhat-snapshots/[WeightedPool - Standard] add liquidity proportional

@@ -1 +1 @@
-179.6k
+179.4k

pkg/pool-weighted/test/gas/.hardhat-snapshots/[WeightedPool - Standard] add liquidity single token exact out - warm slots

@@ -1 +1 @@
-197.4k
+196.2k

pkg/pool-weighted/test/gas/.hardhat-snapshots/[WeightedPool - Standard] add liquidity unbalanced - warm slots

@@ -1 +1 @@
-211.7k
+211.6k

pkg/pool-weighted/test/gas/.hardhat-snapshots/[WeightedPool - Standard] remove liquidity proportional

@@ -1 +1 @@
-168.7k
+168.5k

pkg/pool-weighted/test/gas/.hardhat-snapshots/[WeightedPool - Standard] remove liquidity single token exact in - warm slots

@@ -1 +1 @@
-186.9k
+185.7k

pkg/pool-weighted/test/gas/.hardhat-snapshots/[WeightedPool - Standard] remove liquidity single token exact out - warm slots

@@ -1 +1 @@
-208.7k
+207.5k

pkg/pool-weighted/test/gas/.hardhat-snapshots/[WeightedPool - Standard] swap single token exact in with fees - cold slots

@@ -1 +1 @@
-175.1k
+171.2k

pkg/pool-weighted/test/gas/.hardhat-snapshots/[WeightedPool - Standard] swap single token exact in with fees - warm slots

@@ -1 +1 @@
-158.6k
+154.7k

pkg/pool-weighted/test/gas/.hardhat-snapshots/[WeightedPool - With rate] add liquidity proportional

@@ -1 +1 @@
-235.0k
+234.9k

pkg/pool-weighted/test/gas/.hardhat-snapshots/[WeightedPool - With rate] add liquidity single token exact out - warm slots

@@ -1 +1 @@
-199.6k
+198.4k

pkg/pool-weighted/test/gas/.hardhat-snapshots/[WeightedPool - With rate] add liquidity unbalanced - warm slots

@@ -1 +1 @@
-228.4k
+228.3k

pkg/pool-weighted/test/gas/.hardhat-snapshots/[WeightedPool - With rate] remove liquidity proportional

@@ -1 +1 @@
-223.9k
+223.7k

pkg/pool-weighted/test/gas/.hardhat-snapshots/[WeightedPool - With rate] remove liquidity single token exact in - warm slots

@@ -1 +1 @@
-188.8k
+187.6k

pkg/pool-weighted/test/gas/.hardhat-snapshots/[WeightedPool - With rate] remove liquidity single token exact out - warm slots

@@ -1 +1 @@
-210.6k
+209.3k

pkg/pool-weighted/test/gas/.hardhat-snapshots/[WeightedPool - With rate] swap single token exact in with fees - cold slots

@@ -1 +1 @@
-208.4k
+204.5k

pkg/pool-weighted/test/gas/.hardhat-snapshots/[WeightedPool - With rate] swap single token exact in with fees - warm slots

@@ -1 +1 @@
-174.8k
+170.9k

pkg/pool-weighted/test/gas/.hardhat-snapshots/[WeightedPool] donation

@@ -1 +1 @@
-178.3k
+178.1k

pkg/pool-weighted/test/gas/.hardhat-snapshots/[WeightedPool] initialize with ETH

@@ -1 +1 @@
-347.8k
+348.0k

pkg/pool-weighted/test/gas/.hardhat-snapshots/[WeightedPool] initialize without ETH

@@ -1 +1 @@
-335.2k
+335.4k

pkg/vault/.forge-snapshots/forkBoostedPoolEmptyBufferSwapExactIn.snap

@@ -1 +1 @@
-726841
+722920

pkg/vault/.forge-snapshots/forkBoostedPoolEmptyBufferSwapExactOut.snap

@@ -1 +1 @@
-753760
+749817

pkg/vault/.forge-snapshots/forkBoostedPoolSwapExactIn.snap

@@ -1 +1 @@
-289759
+285948

pkg/vault/.forge-snapshots/forkBoostedPoolSwapExactOut.snap

@@ -1 +1 @@
-312573
+308762

pkg/vault/.forge-snapshots/forkBoostedPoolSwapTooLarge-ExactIn.snap

@@ -1 +1 @@
-670336
+666525

pkg/vault/.forge-snapshots/forkBoostedPoolSwapTooLarge-ExactOut.snap

@@ -1 +1 @@
-697239
+693428

pkg/vault/contracts/Vault.sol

@@ -40,7 +40,6 @@ import { PoolDataLib } from "./lib/PoolDataLib.sol";
 import { VaultCommon } from "./VaultCommon.sol";
 
 contract Vault is IVaultMain, VaultCommon, Proxy {
-    using EnumerableMap for EnumerableMap.IERC20ToBytes32Map;
     using PackedTokenBalance for bytes32;
     using InputHelpers for uint256;
     using FixedPoint for *;
@@ -268,29 +267,19 @@ contract Vault is IVaultMain, VaultCommon, Proxy {
     function _loadSwapState(
         SwapParams memory params,
         PoolData memory poolData
-    ) private view returns (SwapState memory state) {
-        // Use the storage map only for translating token addresses to indices. Raw balances can be read from poolData.
-        EnumerableMap.IERC20ToBytes32Map storage poolBalances = _poolTokenBalances[params.pool];
-
-        // EnumerableMap stores indices *plus one* to use the zero index as a sentinel value for non-existence.
-        uint256 indexIn = poolBalances.unchecked_indexOf(params.tokenIn);
-        uint256 indexOut = poolBalances.unchecked_indexOf(params.tokenOut);
+    ) private pure returns (SwapState memory state) {
+        int indexIn = _findTokenIndex(poolData.tokens, params.tokenIn);
+        int indexOut = _findTokenIndex(poolData.tokens, params.tokenOut);
 
         // If either are zero, revert because the token wasn't registered to this pool.
-        if (indexIn == 0 || indexOut == 0) {
+        if (indexIn == -1 || indexOut == -1) {
             // We require the pool to be initialized, which means it's also registered.
             // This can only happen if the tokens are not registered.
             revert TokenNotRegistered();
         }
 
-        // Convert to regular 0-based indices now, since we've established the tokens are valid.
-        unchecked {
-            indexIn -= 1;
-            indexOut -= 1;
-        }
-
-        state.indexIn = indexIn;
-        state.indexOut = indexOut;
+        state.indexIn = uint256(indexIn);
+        state.indexOut = uint(indexOut);
 
         // If the amountGiven is entering the pool math (ExactIn), round down, since a lower apparent amountIn leads
         // to a lower calculated amountOut, favoring the pool.
@@ -463,20 +452,14 @@ contract Vault is IVaultMain, VaultCommon, Proxy {
         );
 
         // 6) Store pool balances, raw and live (only index in and out)
-        EnumerableMap.IERC20ToBytes32Map storage poolBalances = _poolTokenBalances[params.pool];
-        poolBalances.unchecked_setAt(
-            state.indexIn,
-            PackedTokenBalance.toPackedBalance(
-                poolData.balancesRaw[state.indexIn],
-                poolData.balancesLiveScaled18[state.indexIn]
-            )
+        mapping(uint256 => bytes32) storage poolBalances = _poolTokenBalances[params.pool];
+        poolBalances[state.indexIn] = PackedTokenBalance.toPackedBalance(
+            poolData.balancesRaw[state.indexIn],
+            poolData.balancesLiveScaled18[state.indexIn]
         );
-        poolBalances.unchecked_setAt(
-            state.indexOut,
-            PackedTokenBalance.toPackedBalance(
-                poolData.balancesRaw[state.indexOut],
-                poolData.balancesLiveScaled18[state.indexOut]
-            )
+        poolBalances[state.indexOut] = PackedTokenBalance.toPackedBalance(
+            poolData.balancesRaw[state.indexOut],
+            poolData.balancesLiveScaled18[state.indexOut]
         );
 
         // 7) Off-chain events
@@ -1490,17 +1473,14 @@ contract Vault is IVaultMain, VaultCommon, Proxy {
         address pool,
         IERC20 token
     ) external view withRegisteredPool(pool) returns (uint256, uint256) {
-        EnumerableMap.IERC20ToBytes32Map storage poolTokenBalances = _poolTokenBalances[pool];
-        uint256 tokenCount = poolTokenBalances.length();
-        // unchecked indexOf returns index + 1, or 0 if token is not present.
-        uint256 index = poolTokenBalances.unchecked_indexOf(token);
-        if (index == 0) {
+        IERC20[] memory poolTokens = _poolTokens[pool];
+
+        int index = _findTokenIndex(poolTokens, token);
+        if (index == -1) {
             revert TokenNotRegistered();
         }
 
-        unchecked {
-            return (tokenCount, index - 1);
-        }
+        return (poolTokens.length, uint(index));
     }
 
     /*******************************************************************************

pkg/vault/contracts/VaultCommon.sol

@@ -33,7 +33,6 @@ import { PoolDataLib } from "./lib/PoolDataLib.sol";
  * that require storage to work and will be required in both the main Vault and its extension.
  */
 abstract contract VaultCommon is IVaultEvents, IVaultErrors, VaultStorage, ReentrancyGuardTransient, ERC20MultiToken {
-    using EnumerableMap for EnumerableMap.IERC20ToBytes32Map;
     using PackedTokenBalance for bytes32;
     using PoolConfigLib for PoolConfigBits;
     using ScalingHelpers for *;
@@ -256,20 +255,25 @@ abstract contract VaultCommon is IVaultEvents, IVaultErrors, VaultStorage, Reent
      * and poolData.liveBalances in the same storage slot.
      */
     function _writePoolBalancesToStorage(address pool, PoolData memory poolData) internal {
-        EnumerableMap.IERC20ToBytes32Map storage poolBalances = _poolTokenBalances[pool];
+        mapping(uint256 => bytes32) storage poolBalances = _poolTokenBalances[pool];
 
         for (uint256 i = 0; i < poolData.balancesRaw.length; ++i) {
-            // Since we assume all newBalances are properly ordered, we can simply use `unchecked_setAt`
-            // to avoid one less storage read per token.
-            poolBalances.unchecked_setAt(
-                i,
-                PackedTokenBalance.toPackedBalance(poolData.balancesRaw[i], poolData.balancesLiveScaled18[i])
+            // Since we assume all newBalances are properly ordered
+            poolBalances[i] = PackedTokenBalance.toPackedBalance(
+                poolData.balancesRaw[i],
+                poolData.balancesLiveScaled18[i]
             );
         }
     }
 
     function _loadPoolData(address pool, Rounding roundingDirection) internal view returns (PoolData memory poolData) {
-        poolData.load(_poolTokenBalances[pool], _poolConfigBits[pool], _poolTokenInfo[pool], roundingDirection);
+        poolData.load(
+            _poolTokenBalances[pool],
+            _poolConfigBits[pool],
+            _poolTokenInfo[pool],
+            _poolTokens[pool],
+            roundingDirection
+        );
     }
 
     /**
@@ -283,7 +287,13 @@ abstract contract VaultCommon is IVaultEvents, IVaultErrors, VaultStorage, Reent
         Rounding roundingDirection
     ) internal nonReentrant returns (PoolData memory poolData) {
         // Initialize poolData with base information for subsequent calculations.
-        poolData.load(_poolTokenBalances[pool], _poolConfigBits[pool], _poolTokenInfo[pool], roundingDirection);
+        poolData.load(
+            _poolTokenBalances[pool],
+            _poolConfigBits[pool],
+            _poolTokenInfo[pool],
+            _poolTokens[pool],
+            roundingDirection
+        );
 
         PoolDataLib.syncPoolBalancesAndFees(poolData, _poolTokenBalances[pool], _aggregateFeeAmounts[pool]);
     }
@@ -334,6 +344,17 @@ abstract contract VaultCommon is IVaultEvents, IVaultErrors, VaultStorage, Reent
         emit SwapFeePercentageChanged(pool, swapFeePercentage);
     }
 
+    /// @dev Find the index of a token in a token array. Returns -1 if not found.
+    function _findTokenIndex(IERC20[] memory tokens, IERC20 token) internal pure returns (int) {
+        for (uint256 i = 0; i < tokens.length; i++) {
+            if (tokens[i] == token) {
+                return int(i);
+            }
+        }
+
+        return -1;
+    }
+
     /*******************************************************************************
                                     Recovery Mode
     *******************************************************************************/