From 260b919259e887d6c805126d0aca2a1af0940b6b Mon Sep 17 00:00:00 2001
From: Samyak Jain <34437877+KaymasJain@users.noreply.github.com>
Date: Thu, 8 Apr 2021 05:40:38 +0530
Subject: [PATCH] L1 migration connector update

---
 .../senders2/aave-v2-connector/interfaces.sol |  7 ++--
 contracts/senders2/aave-v2-connector/main.sol | 37 +++++--------------
 2 files changed, 13 insertions(+), 31 deletions(-)

diff --git a/contracts/senders2/aave-v2-connector/interfaces.sol b/contracts/senders2/aave-v2-connector/interfaces.sol
index c8eaf5e..2af06ab 100644
--- a/contracts/senders2/aave-v2-connector/interfaces.sol
+++ b/contracts/senders2/aave-v2-connector/interfaces.sol
@@ -72,12 +72,11 @@ interface ATokenInterface {
 }
 
 interface AaveMigratorInterface {
-    function migrate(address, address, address[] calldata, address[] calldata) external;
-    function migrate(address, AaveData memory) external;
+    function migrate(AaveDataRaw memory _data) external;
+    function migrateWithFlash(AaveDataRaw memory _data, uint ethAmt) external;
 }
 
-struct AaveData {
-    bool isFinal;
+struct AaveDataRaw {
     address targetDsa;
     uint[] supplyAmts;
     uint[] variableBorrowAmts;
diff --git a/contracts/senders2/aave-v2-connector/main.sol b/contracts/senders2/aave-v2-connector/main.sol
index 9e91a3f..794af30 100644
--- a/contracts/senders2/aave-v2-connector/main.sol
+++ b/contracts/senders2/aave-v2-connector/main.sol
@@ -2,45 +2,23 @@ pragma solidity ^0.7.0;
 pragma experimental ABIEncoderV2;
 
 import { TokenInterface, AccountInterface } from "../../common/interfaces.sol";
-import { AaveInterface, ATokenInterface, AaveData } from "./interfaces.sol";
+import { AaveInterface, ATokenInterface, AaveDataRaw } from "./interfaces.sol";
 import { Helpers } from "./helpers.sol";
 import { Events } from "./events.sol";
 
 contract AaveMigrateResolver is Helpers, Events {
 
-    // function migrate(
-    //     address targetDsa,
-    //     address[] calldata supplyTokens,
-    //     address[] calldata borrowTokens
-    // ) external payable returns (string memory _eventName, bytes memory _eventParam) {
-    //     require(supplyTokens.length > 0, "0-length-not-allowed");
-    //     require(targetDsa != address(0), "invalid-address");
-
-    //     for (uint i = 0; i < supplyTokens.length; i++) {
-    //         address _token = supplyTokens[i] == ethAddr ? wethAddr : supplyTokens[i];
-    //         (address _aToken, ,) = aaveData.getReserveTokensAddresses(_token);
-    //         ATokenInterface _aTokenContract = ATokenInterface(_aToken);
-    //         _aTokenContract.approve(address(migrator), _aTokenContract.balanceOf(address(this)));
-    //     }
-
-    //     migrator.migrate(msg.sender, targetDsa, supplyTokens, borrowTokens);
-
-    //     _eventName = "LogAaveV2Migrate(address,address,address[],address[])";
-    //     _eventParam = abi.encode(msg.sender, targetDsa, supplyTokens, borrowTokens);
-    // }
-
     function migrate(
-        AaveData calldata _data
+        AaveDataRaw calldata _data,
+        uint ethAmt // if ethAmt is > 0 then use migrateWithflash
     ) external payable returns (string memory _eventName, bytes memory _eventParam) {
         require(_data.supplyTokens.length > 0, "0-length-not-allowed");
         require(_data.supplyTokens.length == _data.supplyAmts.length, "invalid-length");
         require(_data.targetDsa != address(0), "invalid-address");
-        require(!_data.isFinal, "wrong-data");
 
-        AaveData memory data;
+        AaveDataRaw memory data;
 
         data.borrowTokens = _data.borrowTokens;
-        data.isFinal = _data.isFinal;
         data.stableBorrowAmts = _data.stableBorrowAmts;
         data.supplyAmts = _data.supplyAmts;
         data.supplyTokens = _data.supplyTokens;
@@ -60,11 +38,16 @@ contract AaveMigrateResolver is Helpers, Events {
             _aTokenContract.approve(address(migrator), data.supplyAmts[i]);
         }
 
-        migrator.migrate(msg.sender, data);
+        if (ethAmt > 0) {
+            migrator.migrateWithFlash(data, ethAmt);
+        } else {
+            migrator.migrate(data);
+        }
 
         _eventName = "LogAaveV2Migrate(address,address,address[],address[])";
         _eventParam = abi.encode(msg.sender, data.targetDsa, data.supplyTokens, data.borrowTokens);
     }
+
 }
 
 contract AaveV2Migrator is AaveMigrateResolver {