Refactor to avoid leftovers on _swapAndRepay with flash loan

contracts/adapters/UniswapRepayAdapter.sol

@@ -123,8 +123,8 @@ contract UniswapRepayAdapter is BaseUniswapAdapter, IFlashLoanReceiver {
   /**
    * @dev Perform the repay of the debt, pulls the initiator collateral and swaps to repay the flash loan
    *
-   * @param assetFrom Address of token to be swapped
-   * @param assetTo Address of debt token to be received from the swap
+   * @param collateralAsset Address of token to be swapped
+   * @param debtAsset Address of debt token to be received from the swap
    * @param amount Amount of the debt to be repaid
    * @param collateralAmount Amount of the reserve to be swapped
    * @param rateMode Rate mode of the debt to be repaid
@@ -133,8 +133,8 @@ contract UniswapRepayAdapter is BaseUniswapAdapter, IFlashLoanReceiver {
    * @param permitSignature struct containing the permit signature
    */
   function _swapAndRepay(
-    address assetFrom,
-    address assetTo,
+    address collateralAsset,
+    address debtAsset,
     uint256 amount,
     uint256 collateralAmount,
     uint256 rateMode,
@@ -142,25 +142,48 @@ contract UniswapRepayAdapter is BaseUniswapAdapter, IFlashLoanReceiver {
     uint256 premium,
     PermitSignature memory permitSignature
   ) internal {
+    ReserveLogic.ReserveData memory collateralReserveData = _getReserveData(collateralAsset);
+
     // Repay debt
-    IERC20(assetTo).approve(address(POOL), amount);
-    POOL.repay(assetTo, amount, rateMode, initiator);
-    uint256 debtRepayLeftovers = IERC20(assetTo).balanceOf(address(this));
+    IERC20(debtAsset).approve(address(POOL), amount);
+    uint256 repaidAmount = IERC20(debtAsset).balanceOf(address(this));
+    POOL.repay(debtAsset, amount, rateMode, initiator);
+    repaidAmount = repaidAmount.sub(IERC20(debtAsset).balanceOf(address(this)));
 
-    uint256 flashLoanDebt = amount.add(premium);
-    uint256 neededForFlashLoanDebt = flashLoanDebt.sub(debtRepayLeftovers);
+    if (collateralAsset != debtAsset) {
+      uint256 maxCollateralToSwap = collateralAmount;
+      if (repaidAmount < amount) {
+        maxCollateralToSwap = maxCollateralToSwap.mul(repaidAmount).div(amount);
+      }
 
-    // Pull aTokens from user
-    ReserveLogic.ReserveData memory reserveData = _getReserveData(assetFrom);
-    _pullAToken(assetFrom, reserveData.aTokenAddress, initiator, collateralAmount, permitSignature);
+      uint256 neededForFlashLoanDebt = repaidAmount.add(premium);
+      uint256[] memory amounts = _getAmountsIn(collateralAsset, debtAsset, neededForFlashLoanDebt);
+      require(amounts[0] <= maxCollateralToSwap, 'slippage too high');
 
-    uint256 amountSwapped = _swapTokensForExactTokens(assetFrom, assetTo, collateralAmount, neededForFlashLoanDebt);
+      // Pull aTokens from user
+      _pullAToken(
+        collateralAsset,
+        collateralReserveData.aTokenAddress,
+        initiator,
+        amounts[0],
+        permitSignature
+      );
 
-    // Send collateral leftovers from swap to the user
-    _sendLeftovers(assetFrom, initiator);
+      // Swap collateral asset to the debt asset
+      _swapTokensForExactTokens(collateralAsset, debtAsset, amounts[0], neededForFlashLoanDebt);
+    } else {
+      // Pull aTokens from user
+      _pullAToken(
+        collateralAsset,
+        collateralReserveData.aTokenAddress,
+        initiator,
+        repaidAmount.add(premium),
+        permitSignature
+      );
+    }
 
     // Repay flashloan
-    IERC20(assetTo).approve(address(POOL), flashLoanDebt);
+    IERC20(debtAsset).approve(address(POOL), amount.add(premium));
   }
 
   /**
@@ -201,17 +224,4 @@ contract UniswapRepayAdapter is BaseUniswapAdapter, IFlashLoanReceiver {
       )
     );
   }
-
-  /**
-  * @dev Transfers the balance of the adapter to the user, as there shouldn't be any leftover in the adapter
-  * @param asset address of the asset
-  * @param user address
-  */
-  function _sendLeftovers(address asset, address user) internal {
-    uint256 assetLeftover = IERC20(asset).balanceOf(address(this));
-
-    if (assetLeftover > 0) {
-      IERC20(asset).transfer(user, assetLeftover);
-    }
-  }
 }

contracts/mocks/swap/MockUniswapV2Router02.sol

@@ -10,6 +10,7 @@ contract MockUniswapV2Router02 is IUniswapV2Router02 {
   mapping(address => uint256) internal _amountToSwap;
   mapping(address => mapping(address => mapping(uint256 => uint256))) internal _amountsIn;
   mapping(address => mapping(address => mapping(uint256 => uint256))) internal _amountsOut;
+  uint256 internal defaultMockValue;
 
   function setAmountToReturn(address reserve, uint256 amount) public {
     _amountToReturn[reserve] = amount;
@@ -61,16 +62,20 @@ contract MockUniswapV2Router02 is IUniswapV2Router02 {
     _amountsIn[reserveIn][reserveOut][amountOut] = amountIn;
   }
 
+  function setDefaultMockValue(uint value) public {
+    defaultMockValue = value;
+  }
+
   function getAmountsOut(uint amountIn, address[] calldata path) external view override returns (uint[] memory) {
     uint256[] memory amounts = new uint256[](2);
     amounts[0] = amountIn;
-    amounts[1] = _amountsOut[path[0]][path[1]][amountIn];
+    amounts[1] = _amountsOut[path[0]][path[1]][amountIn] > 0 ? _amountsOut[path[0]][path[1]][amountIn] : defaultMockValue;
     return amounts;
   }
 
   function getAmountsIn(uint amountOut, address[] calldata path) external view override returns (uint[] memory) {
     uint256[] memory amounts = new uint256[](2);
-    amounts[0] = _amountsIn[path[0]][path[1]][amountOut];
+    amounts[0] = _amountsIn[path[0]][path[1]][amountOut] > 0 ? _amountsIn[path[0]][path[1]][amountOut] : defaultMockValue;
     amounts[1] = amountOut;
     return amounts;
   }

test/uniswapAdapters.spec.ts

@@ -2091,6 +2091,13 @@ makeSuite('Uniswap adapters', (testEnv: TestEnv) => {
           .multipliedBy(1.0009)
           .toFixed(0);
 
+        await mockUniswapRouter.setAmountIn(
+          flashLoanDebt,
+          weth.address,
+          dai.address,
+          liquidityToSwap
+        );
+
         const params = buildRepayAdapterParams(
           weth.address,
           liquidityToSwap,
@@ -2198,6 +2205,13 @@ makeSuite('Uniswap adapters', (testEnv: TestEnv) => {
           .multipliedBy(1.0009)
           .toFixed(0);
 
+        await mockUniswapRouter.setAmountIn(
+          flashLoanDebt,
+          weth.address,
+          dai.address,
+          liquidityToSwap
+        );
+
         const params = buildRepayAdapterParams(
           weth.address,
           liquidityToSwap,
@@ -2401,6 +2415,17 @@ makeSuite('Uniswap adapters', (testEnv: TestEnv) => {
 
         await mockUniswapRouter.connect(user).setAmountToSwap(weth.address, bigMaxAmountToSwap);
 
+        const flashLoanDebt = new BigNumber(expectedDaiAmount.toString())
+          .multipliedBy(1.0009)
+          .toFixed(0);
+
+        await mockUniswapRouter.setAmountIn(
+          flashLoanDebt,
+          weth.address,
+          dai.address,
+          bigMaxAmountToSwap
+        );
+
         const params = buildRepayAdapterParams(
           weth.address,
           bigMaxAmountToSwap,
@@ -2472,14 +2497,19 @@ makeSuite('Uniswap adapters', (testEnv: TestEnv) => {
           .multipliedBy(0.995)
           .toFixed(0);
 
-        const leftOverWeth = new BigNumber(liquidityToSwap.toString()).minus(actualWEthSwapped);
-
         await mockUniswapRouter.connect(user).setAmountToSwap(weth.address, actualWEthSwapped);
 
         const flashLoanDebt = new BigNumber(expectedDaiAmount.toString())
           .multipliedBy(1.0009)
           .toFixed(0);
 
+        await mockUniswapRouter.setAmountIn(
+          flashLoanDebt,
+          weth.address,
+          dai.address,
+          actualWEthSwapped
+        );
+
         const params = buildRepayAdapterParams(
           weth.address,
           liquidityToSwap,
@@ -2520,8 +2550,8 @@ makeSuite('Uniswap adapters', (testEnv: TestEnv) => {
         expect(userDaiStableDebtAmountBefore).to.be.gte(expectedDaiAmount);
         expect(userDaiStableDebtAmount).to.be.lt(expectedDaiAmount);
         expect(userAEthBalance).to.be.lt(userAEthBalanceBefore);
-        expect(userAEthBalance).to.be.eq(userAEthBalanceBefore.sub(liquidityToSwap));
-        expect(userWethBalance).to.be.gte(userWethBalanceBefore.add(leftOverWeth.toString()));
+        expect(userAEthBalance).to.be.eq(userAEthBalanceBefore.sub(actualWEthSwapped));
+        expect(userWethBalance).to.be.eq(userWethBalanceBefore);
       });
 
       it('should correctly swap tokens and repay the whole stable debt with no leftovers', async () => {
@@ -2560,7 +2590,11 @@ makeSuite('Uniswap adapters', (testEnv: TestEnv) => {
 
         const userDaiStableDebtAmountBefore = await daiStableDebtContract.balanceOf(userAddress);
 
-        const liquidityToSwap = amountWETHtoSwap;
+        // Add a % to repay on top of the debt
+        const liquidityToSwap = new BigNumber(amountWETHtoSwap.toString())
+          .multipliedBy(1.1)
+          .toFixed(0);
+
         await aWETH.connect(user).approve(uniswapRepayAdapter.address, liquidityToSwap);
         const userAEthBalanceBefore = await aWETH.balanceOf(userAddress);
 
@@ -2569,7 +2603,8 @@ makeSuite('Uniswap adapters', (testEnv: TestEnv) => {
           .multipliedBy(1.1)
           .toFixed(0);
 
-        await mockUniswapRouter.connect(user).setAmountToSwap(weth.address, liquidityToSwap);
+        await mockUniswapRouter.connect(user).setAmountToSwap(weth.address, amountWETHtoSwap);
+        await mockUniswapRouter.setDefaultMockValue(amountWETHtoSwap);
 
         const params = buildRepayAdapterParams(
           weth.address,
@@ -2647,7 +2682,11 @@ makeSuite('Uniswap adapters', (testEnv: TestEnv) => {
           userAddress
         );
 
-        const liquidityToSwap = amountWETHtoSwap;
+        // Add a % to repay on top of the debt
+        const liquidityToSwap = new BigNumber(amountWETHtoSwap.toString())
+          .multipliedBy(1.1)
+          .toFixed(0);
+
         await aWETH.connect(user).approve(uniswapRepayAdapter.address, liquidityToSwap);
         const userAEthBalanceBefore = await aWETH.balanceOf(userAddress);
 
@@ -2656,7 +2695,8 @@ makeSuite('Uniswap adapters', (testEnv: TestEnv) => {
           .multipliedBy(1.1)
           .toFixed(0);
 
-        await mockUniswapRouter.connect(user).setAmountToSwap(weth.address, liquidityToSwap);
+        await mockUniswapRouter.connect(user).setAmountToSwap(weth.address, amountWETHtoSwap);
+        await mockUniswapRouter.setDefaultMockValue(amountWETHtoSwap);
 
         const params = buildRepayAdapterParams(
           weth.address,