Change the Pool assets mapping in SwapOperator from by address to by …

…asset ID

contracts/modules/capital/Pool.sol

@@ -31,7 +31,7 @@ contract Pool is IPool, MasterAwareV2, ReentrancyGuard {
 
   // SwapOperator assets
   uint32 internal assetsInSwapOperatorBitmap;
-  mapping(address => uint) public assetsInSwapOperator;
+  mapping(uint => uint) public assetsInSwapOperator;
 
   /* constants */
 
@@ -142,18 +142,17 @@ contract Pool is IPool, MasterAwareV2, ReentrancyGuard {
         continue;
       }
 
-      address assetAddress = asset.assetAddress;
-
       uint assetAmountInSwapOperator = isAssetInSwapOperator(i, _assetsInSwapOperatorBitmap)
-        ? assetsInSwapOperator[assetAddress]
+        ? assetsInSwapOperator[i]
         : 0;
 
-      if (assetAddress == ETH) {
+      // check if the asset is ETH and skip the oracle call
+      if (i == 0) {
         total += assetAmountInSwapOperator;
         continue;
       }
 
-      total += getAssetValueInEth(assetAddress, assetAmountInSwapOperator);
+      total += getAssetValueInEth(asset.assetAddress, assetAmountInSwapOperator);
     }
 
     return total;
@@ -319,14 +318,14 @@ contract Pool is IPool, MasterAwareV2, ReentrancyGuard {
 
   function setSwapAssetAmount(address assetAddress, uint value) external onlySwapOperator whenNotPaused {
 
-    assetsInSwapOperator[assetAddress] = value;
+    uint assetId = getAssetId(assetAddress);
+    assetsInSwapOperator[assetId] = value;
 
     if (value > 0) {
       if (assetsInSwapOperatorBitmap != 0) {
         revert OrderInProgress();
       }
 
-      uint assetId = getAssetId(assetAddress);
       assetsInSwapOperatorBitmap = uint32(1 << assetId);
     } else {
       assetsInSwapOperatorBitmap = 0;

lib/constants.js

@@ -103,6 +103,7 @@ const PoolAsset = {
   ETH: 0,
   DAI: 1,
   stETH: 2,
+  USDC: 6,
   unknown: '115792089237316195423570985008687907853269984665640564039457584007913129639935',
 };
 

test/unit/Pool/getPoolValueInEth.js

@@ -8,7 +8,8 @@ const utils = require('../utils');
 const {
   helpers: { toBytes8 },
   constants: {
-    Assets: { ETH },
+    PoolAsset: { ETH },
+    Assets: { ETH: ETH_ADDRESS },
   },
 } = utils;
 
@@ -83,7 +84,7 @@ describe('getPoolValueInEth', function () {
     const oldPoolValue = await pool.getPoolValueInEth();
 
     await pool.connect(governance).updateAddressParameters(toBytes8('SWP_OP'), defaultSender.address);
-    await pool.setSwapAssetAmount(ETH, parseEther('1'));
+    await pool.setSwapAssetAmount(ETH_ADDRESS, parseEther('1'));
 
     const swapValue = await pool.assetsInSwapOperator(ETH);
     expect(swapValue.toString()).to.eq(parseEther('1').toString());

test/unit/Pool/setSwapAssetAmount.js

@@ -7,7 +7,8 @@ const { BigNumber } = ethers;
 const utils = require('../utils');
 
 const { toBytes8 } = utils.helpers;
-const { ETH } = utils.constants.Assets;
+const { ETH } = utils.constants.PoolAsset;
+const { ETH: ETH_ADDRESS } = utils.constants.Assets;
 
 describe('setSwapAssetAmount', function () {
   it('is only callabe by swap operator', async function () {
@@ -19,13 +20,15 @@ describe('setSwapAssetAmount', function () {
     } = fixture.accounts;
 
     // Not calling from swap operator reverts
-    await expect(pool.setSwapAssetAmount(ETH, BigNumber.from('123'))).to.be.revertedWith('Pool: Not swapOperator');
+    await expect(pool.setSwapAssetAmount(ETH_ADDRESS, BigNumber.from('123'))).to.be.revertedWith(
+      'Pool: Not swapOperator',
+    );
 
     // Set swap operator
     await pool.connect(governance).updateAddressParameters(toBytes8('SWP_OP'), swapOperator.address);
 
     // Call should succeed
-    await pool.connect(swapOperator).setSwapAssetAmount(ETH, BigNumber.from('123'));
+    await pool.connect(swapOperator).setSwapAssetAmount(ETH_ADDRESS, BigNumber.from('123'));
   });
 
   it('sets the swapValue value', async function () {
@@ -39,7 +42,7 @@ describe('setSwapAssetAmount', function () {
     expect(await pool.assetsInSwapOperator(ETH)).to.eq(0);
     // Set swap operator and set swap value
     await pool.connect(governance).updateAddressParameters(toBytes8('SWP_OP'), swapOperator.address);
-    await pool.connect(swapOperator).setSwapAssetAmount(ETH, 123);
+    await pool.connect(swapOperator).setSwapAssetAmount(ETH_ADDRESS, 123);
 
     expect(await pool.assetsInSwapOperator(ETH)).to.eq(123);
   });

test/unit/SwapOperator/closeOrder.js

@@ -16,7 +16,7 @@ const setup = require('./setup');
 const utils = require('../utils');
 
 const { setEtherBalance, setNextBlockTime, revertToSnapshot, takeSnapshot, increaseTime, mineNextBlock } = utils.evm;
-const { ETH } = utils.constants.Assets;
+const { ETH } = utils.constants.PoolAsset;
 
 const {
   utils: { parseEther, hexZeroPad },

test/unit/SwapOperator/placeOrder.js

@@ -16,7 +16,8 @@ const setup = require('./setup');
 const utils = require('../utils');
 
 const { setEtherBalance, setNextBlockTime } = utils.evm;
-const { ETH } = utils.constants.Assets;
+const { ETH, stETH } = utils.constants.PoolAsset;
+const { ETH: ETH_ADDRESS } = utils.constants.Assets;
 
 const { parseEther, hexZeroPad, hexlify, randomBytes } = ethers.utils;
 
@@ -262,7 +263,7 @@ describe('placeOrder', function () {
     const { contracts, order, domain } = await loadFixture(placeSellWethOrderSetup);
     const { swapOperator } = contracts;
 
-    const badOrder = createContractOrder(domain, order, { sellToken: ETH });
+    const badOrder = createContractOrder(domain, order, { sellToken: ETH_ADDRESS });
     const placeOrder = swapOperator.placeOrder(badOrder.contractOrder, badOrder.orderUID);
     await expect(placeOrder).to.be.revertedWithCustomError(swapOperator, 'InvalidTokenAddress').withArgs('sellToken');
   });
@@ -271,7 +272,7 @@ describe('placeOrder', function () {
     const { contracts, order, domain } = await loadFixture(placeSellDaiOrderSetup);
     const { swapOperator } = contracts;
 
-    const badOrder = createContractOrder(domain, order, { buyToken: ETH });
+    const badOrder = createContractOrder(domain, order, { buyToken: ETH_ADDRESS });
     const placeOrder = swapOperator.placeOrder(badOrder.contractOrder, badOrder.orderUID);
     await expect(placeOrder).to.be.revertedWithCustomError(swapOperator, 'InvalidTokenAddress').withArgs('buyToken');
   });
@@ -961,16 +962,17 @@ describe('placeOrder', function () {
   it('should set totalOutAmount in ETH as pool.swapValue when selling non-ETH assets', async function () {
     const orderSetupsToTest = [placeSellDaiOrderSetup, placeNonEthOrderSellStethSetup, placeNonEthOrderSellDaiSetup];
     for (const orderSetup of orderSetupsToTest) {
-      const { contracts, order, contractOrder, orderUID } = await loadFixture(orderSetup);
+      const { contracts, order, contractOrder, orderUID, poolAssetAddressIdMapping } = await loadFixture(orderSetup);
       const { swapOperator, pool } = contracts;
 
-      expect(await pool.assetsInSwapOperator(order.sellToken)).to.eq(0);
+      const assetId = poolAssetAddressIdMapping[order.sellToken];
+      expect(await pool.assetsInSwapOperator(assetId)).to.eq(0);
 
       await swapOperator.placeOrder(contractOrder, orderUID);
 
       // convert non-ETH sellAmount + fee to ETH
       const { sellAmount, feeAmount } = contractOrder;
-      expect(await pool.assetsInSwapOperator(order.sellToken)).to.be.equal(sellAmount.add(feeAmount));
+      expect(await pool.assetsInSwapOperator(assetId)).to.be.equal(sellAmount.add(feeAmount));
     }
   });
 
@@ -979,13 +981,14 @@ describe('placeOrder', function () {
     const { swapOperator, pool, priceFeedOracle, stEth } = contracts;
     const { sellAmount, feeAmount } = order;
 
+    console.log(order.sellToken);
     expect(await pool.assetsInSwapOperator(order.sellToken)).to.eq(0);
 
     await swapOperator.placeOrder(contractOrder, orderUID);
 
     // convert stETH sellAmount + fee to ETH
     const totalOutAmountInEth = await priceFeedOracle.getEthForAsset(stEth.address, sellAmount.add(feeAmount));
-    expect(await pool.assetsInSwapOperator(order.sellToken)).to.eq(totalOutAmountInEth);
+    expect(await pool.assetsInSwapOperator(stETH)).to.eq(totalOutAmountInEth);
   });
 
   it('approves CoW vault relayer to spend exactly sellAmount + fee', async function () {

test/unit/SwapOperator/setup.js

@@ -167,6 +167,11 @@ async function setup() {
       ETH_ADDRESS: '0xEeeeeEeeeEeEeeEeEeEeeEEEeeeeEeeeeeeeEEeE',
       DAI_ADDRESS: dai.address,
     },
+    poolAssetAddressIdMapping: {
+      [dai.address]: 1,
+      [stEth.address]: 2,
+      [usdc.address]: 6,
+    },
   };
 }