Add partial migration and token mapping

contracts/receivers/aave-v2-receiver/helpers.sol

@@ -0,0 +1,21 @@
+pragma solidity ^0.7.0;
+
+import { DSMath } from "../../common/math.sol";
+import { TokenMappingInterface, AaveData } from "./interfaces.sol";
+
+abstract contract Helpers is DSMath {
+    // Replace this
+    TokenMappingInterface tokenMapping = TokenMappingInterface(address(2));
+
+    function remapTokens(AaveData memory data) internal returns (AaveData memory) {
+        for (uint i = 0; i < data.supplyTokens.length; i++) {
+            data.supplyTokens[i] = tokenMapping.getMapping(data.supplyTokens[i]);
+        }
+
+        for (uint i = 0; i < data.borrowTokens.length; i++) {
+            data.borrowTokens[i] = tokenMapping.getMapping(data.borrowTokens[i]);
+        }
+
+        return data;
+    }
+}

contracts/receivers/aave-v2-receiver/interfaces.sol

@@ -11,4 +11,18 @@ interface AccountInterface {
         address _origin
     ) external payable returns (bytes32);
     function migrateAave(address) external payable returns (bytes32);
+}
+
+interface TokenMappingInterface {
+    function getMapping(address) external view returns (address);
+}
+
+struct AaveData {
+    bool isFinal;
+    address targetDsa;
+    uint[] supplyAmts;
+    uint[] variableBorrowAmts;
+    uint[] stableBorrowAmts;
+    address[] supplyTokens;
+    address[] borrowTokens;
 }

contracts/receivers/aave-v2-receiver/main.sol

@@ -4,23 +4,13 @@ pragma experimental ABIEncoderV2;
 import { SafeERC20 } from "@openzeppelin/contracts/token/ERC20/SafeERC20.sol";
 import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 
-import { DSMath } from "../../common/math.sol";
+import { AccountInterface, AaveData } from "./interfaces.sol";
-import { AccountInterface } from "./interfaces.sol";
 import { Events } from "./events.sol";
+import { Helpers } from "./helpers.sol";
 
-contract MigrateResolver is DSMath, Events {
+contract MigrateResolver is Helpers, Events {
     using SafeERC20 for IERC20;
 
-    struct AaveData {
-        bool isFinal;
-        address targetDsa;
-        uint[] supplyAmts;
-        uint[] variableBorrowAmts;
-        uint[] stableBorrowAmts;
-        address[] supplyTokens;
-        address[] borrowTokens;
-    }
-
     mapping (address => AaveData) public positions;
     mapping(address => mapping(address => uint)) deposits;
 
@@ -108,7 +98,7 @@ contract AaveV2Migrator is MigrateResolver {
         lastStateId = stateId;
 
         (address owner, AaveData memory data) = abi.decode(receivedData, (address, AaveData));
-        positions[owner] = data;
+        positions[owner] = remapTokens(data);
 
         if (canMigrate(owner)) {
             _migratePosition(owner);

contracts/senders/aave-v2-connector/interfaces.sol

@@ -73,9 +73,15 @@ interface ATokenInterface {
 
 interface AaveMigratorInterface {
     function migrate(address, address, address[] calldata, address[] calldata) external;
+    function migrate(address, AaveData memory) external;
 }
 
-// interface StateSenderInterface {
+struct AaveData {
-//     function syncState(address receiver, bytes calldata data) external;
+    bool isFinal;
-//     function register(address sender, address receiver) external;
+    address targetDsa;
-// }
+    uint[] supplyAmts;
+    uint[] variableBorrowAmts;
+    uint[] stableBorrowAmts;
+    address[] supplyTokens;
+    address[] borrowTokens;
+}

contracts/senders/aave-v2-connector/main.sol

@@ -2,31 +2,68 @@ pragma solidity ^0.7.0;
 pragma experimental ABIEncoderV2;
 
 import { TokenInterface, AccountInterface } from "../../common/interfaces.sol";
-import { AaveInterface, ATokenInterface } from "./interfaces.sol";
+import { AaveInterface, ATokenInterface, AaveData } from "./interfaces.sol";
 import { Helpers } from "./helpers.sol";
 import { Events } from "./events.sol";
 
 contract AaveMigrateResolver is Helpers, Events {
 
-    function migrate(
+    // function migrate(
-        address targetDsa,
+    //     address targetDsa,
-        address[] calldata supplyTokens,
+    //     address[] calldata supplyTokens,
-        address[] calldata borrowTokens
+    //     address[] calldata borrowTokens
-    ) external payable returns (string memory _eventName, bytes memory _eventParam) {
+    // ) external payable returns (string memory _eventName, bytes memory _eventParam) {
-        require(supplyTokens.length > 0, "0-length-not-allowed");
+    //     require(supplyTokens.length > 0, "0-length-not-allowed");
-        require(targetDsa != address(0), "invalid-address");
+    //     require(targetDsa != address(0), "invalid-address");
 
-        for (uint i = 0; i < supplyTokens.length; i++) {
+    //     for (uint i = 0; i < supplyTokens.length; i++) {
-            address _token = supplyTokens[i] == ethAddr ? wethAddr : supplyTokens[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
+    ) 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;
+
+        data.borrowTokens = _data.borrowTokens;
+        data.isFinal = _data.isFinal;
+        data.stableBorrowAmts = _data.stableBorrowAmts;
+        data.supplyAmts = _data.supplyAmts;
+        data.supplyTokens = _data.supplyTokens;
+        data.targetDsa = _data.targetDsa;
+        data.variableBorrowAmts = _data.variableBorrowAmts;
+
+        for (uint i = 0; i < data.supplyTokens.length; i++) {
+            address _token = data.supplyTokens[i] == ethAddr ? wethAddr : data.supplyTokens[i];
+            data.supplyTokens[i] = _token;
             (address _aToken, ,) = aaveData.getReserveTokensAddresses(_token);
             ATokenInterface _aTokenContract = ATokenInterface(_aToken);
-            _aTokenContract.approve(address(migrator), _aTokenContract.balanceOf(address(this)));
+
+            if (data.supplyAmts[i] == uint(-1)) {
+                data.supplyAmts[i] = _aTokenContract.balanceOf(address(this));
+            }
+
+            _aTokenContract.approve(address(migrator), data.supplyAmts[i]);
         }
 
-        migrator.migrate(msg.sender, targetDsa, supplyTokens, borrowTokens);
+        migrator.migrate(msg.sender, data);
 
         _eventName = "LogAaveV2Migrate(address,address,address[],address[])";
-        _eventParam = abi.encode(msg.sender, targetDsa, supplyTokens, borrowTokens);
+        _eventParam = abi.encode(msg.sender, data.targetDsa, data.supplyTokens, data.borrowTokens);
     }
 }
 

contracts/senders/aave-v2-migrator/main.sol

@@ -88,67 +88,64 @@ contract MigrateResolver is LiquidityResolver {
 
     mapping (address => AaveData) public positions;
 
-    function migrate(
+    function migrate(address owner, AaveData calldata _data) external {
-        address owner,
+        require(_data.supplyTokens.length > 0, "0-length-not-allowed");
-        address targetDsa,
+        require(_data.targetDsa != address(0), "invalid-address");
-        address[] calldata supplyTokens,
+        require(_data.supplyTokens.length == _data.supplyAmts.length, "invalid-length");
-        address[] calldata borrowTokens
+        require(
-    ) external {
+            _data.borrowTokens.length == _data.variableBorrowAmts.length &&
-        require(supplyTokens.length > 0, "0-length-not-allowed");
+            _data.borrowTokens.length == _data.stableBorrowAmts.length,
-        require(targetDsa != address(0), "invalid-address");
+            "invalid-length"
+        );
+
+        AaveData memory data;
+
+        data.borrowTokens = _data.borrowTokens;
+        data.isFinal = _data.isFinal;
+        data.stableBorrowAmts = _data.stableBorrowAmts;
+        data.supplyAmts = _data.supplyAmts;
+        data.supplyTokens = _data.supplyTokens;
+        data.targetDsa = _data.targetDsa;
+        data.variableBorrowAmts = _data.variableBorrowAmts;
 
         address sourceDsa = msg.sender;
 
         AaveInterface aave = AaveInterface(aaveProvider.getLendingPool());
 
-        AaveData memory data;
-
         (,,,,,uint healthFactor) = aave.getUserAccountData(sourceDsa);
         require(healthFactor > 1e18, "position-not-safe");
 
-        data.supplyAmts = new uint[](supplyTokens.length);
+        for (uint i = 0; i < data.supplyTokens.length; i++) {
-        data.supplyTokens = new address[](supplyTokens.length);
+            (address _aToken, ,) = aaveData.getReserveTokensAddresses(data.supplyTokens[i]);
-        data.targetDsa = targetDsa;
-
-        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);
 
-            data.supplyTokens[i] = _token;
-            data.supplyAmts[i] = aTokenContract.balanceOf(sourceDsa);
-
             aTokenContract.transferFrom(msg.sender, address(this), data.supplyAmts[i]);
         }
 
-        if (borrowTokens.length > 0) {
+        for (uint i = 0; i < data.borrowTokens.length; i++) {
-            data.variableBorrowAmts = new uint[](borrowTokens.length);
+            address _token = data.borrowTokens[i] == ethAddr ? wethAddr : data.borrowTokens[i];
-            data.stableBorrowAmts = new uint[](borrowTokens.length);
+            data.borrowTokens[i] = _token;
 
-            for (uint i = 0; i < borrowTokens.length; i++) {
+            (
-                address _token = borrowTokens[i] == ethAddr ? wethAddr : borrowTokens[i];
+                ,
-                data.borrowTokens[i] = _token;
+                uint stableDebt,
+                uint variableDebt,
+                ,,,,,
+            ) = aaveData.getUserReserveData(_token, sourceDsa);
 
-                (
+            data.stableBorrowAmts[i] = data.stableBorrowAmts[i] == uint(-1) ? stableDebt : data.stableBorrowAmts[i];
-                    ,
+            data.variableBorrowAmts[i] = data.variableBorrowAmts[i] == uint(-1) ? variableDebt : data.variableBorrowAmts[i];
-                    data.stableBorrowAmts[i],
-                    data.variableBorrowAmts[i],
-                    ,,,,,
-                ) = aaveData.getUserReserveData(_token, sourceDsa);
 
-                uint totalBorrowAmt = add(data.stableBorrowAmts[i], data.variableBorrowAmts[i]);
 
-                if (totalBorrowAmt > 0) {
+            uint totalBorrowAmt = add(data.stableBorrowAmts[i], data.variableBorrowAmts[i]);
-                    IERC20(_token).safeApprove(address(aave), totalBorrowAmt);
+            if (totalBorrowAmt > 0) {
-                }
+                IERC20(_token).safeApprove(address(aave), totalBorrowAmt);
             }
-
-            _PaybackStable(borrowTokens.length, aave, data.borrowTokens, data.stableBorrowAmts, sourceDsa);
-            _PaybackVariable(borrowTokens.length, aave, data.borrowTokens, data.variableBorrowAmts, sourceDsa);
         }
 
-        _Withdraw(supplyTokens.length, aave, data.supplyTokens, data.supplyAmts);
+        _PaybackStable(data.borrowTokens.length, aave, data.borrowTokens, data.stableBorrowAmts, sourceDsa);
+        _PaybackVariable(data.borrowTokens.length, aave, data.borrowTokens, data.variableBorrowAmts, sourceDsa);
+        _Withdraw(data.supplyTokens.length, aave, data.supplyTokens, data.supplyAmts);
 
         positions[owner] = data;
         bytes memory positionData = abi.encode(owner, data);
@@ -156,12 +153,89 @@ contract MigrateResolver is LiquidityResolver {
 
         emit LogAaveV2Migrate(
             msg.sender,
-            targetDsa,
+            data.targetDsa,
-            supplyTokens,
+            data.supplyTokens,
-            borrowTokens,
+            data.borrowTokens,
             data.supplyAmts,
             data.variableBorrowAmts,
             data.stableBorrowAmts
         );
     }
+
+    // function migrate(
+    //     address owner,
+    //     address targetDsa,
+    //     address[] calldata supplyTokens,
+    //     address[] calldata borrowTokens
+    // ) external {
+    //     require(supplyTokens.length > 0, "0-length-not-allowed");
+    //     require(targetDsa != address(0), "invalid-address");
+
+    //     address sourceDsa = msg.sender;
+
+    //     AaveInterface aave = AaveInterface(aaveProvider.getLendingPool());
+
+    //     AaveData memory data;
+
+    //     (,,,,,uint healthFactor) = aave.getUserAccountData(sourceDsa);
+    //     require(healthFactor > 1e18, "position-not-safe");
+
+    //     data.supplyAmts = new uint[](supplyTokens.length);
+    //     data.supplyTokens = new address[](supplyTokens.length);
+    //     data.targetDsa = targetDsa;
+
+    //     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);
+
+    //         data.supplyTokens[i] = _token;
+    //         data.supplyAmts[i] = aTokenContract.balanceOf(sourceDsa);
+
+    //         aTokenContract.transferFrom(msg.sender, address(this), data.supplyAmts[i]);
+    //     }
+
+    //     if (borrowTokens.length > 0) {
+    //         data.variableBorrowAmts = new uint[](borrowTokens.length);
+    //         data.stableBorrowAmts = new uint[](borrowTokens.length);
+
+    //         for (uint i = 0; i < borrowTokens.length; i++) {
+    //             address _token = borrowTokens[i] == ethAddr ? wethAddr : borrowTokens[i];
+    //             data.borrowTokens[i] = _token;
+
+    //             (
+    //                 ,
+    //                 data.stableBorrowAmts[i],
+    //                 data.variableBorrowAmts[i],
+    //                 ,,,,,
+    //             ) = aaveData.getUserReserveData(_token, sourceDsa);
+
+    //             uint totalBorrowAmt = add(data.stableBorrowAmts[i], data.variableBorrowAmts[i]);
+
+    //             if (totalBorrowAmt > 0) {
+    //                 IERC20(_token).safeApprove(address(aave), totalBorrowAmt);
+    //             }
+    //         }
+
+    //         _PaybackStable(borrowTokens.length, aave, data.borrowTokens, data.stableBorrowAmts, sourceDsa);
+    //         _PaybackVariable(borrowTokens.length, aave, data.borrowTokens, data.variableBorrowAmts, sourceDsa);
+    //     }
+
+    //     _Withdraw(supplyTokens.length, aave, data.supplyTokens, data.supplyAmts);
+
+    //     positions[owner] = data;
+    //     bytes memory positionData = abi.encode(owner, data);
+    //     stateSender.syncState(polygonReceiver, positionData);
+
+    //     emit LogAaveV2Migrate(
+    //         msg.sender,
+    //         targetDsa,
+    //         supplyTokens,
+    //         borrowTokens,
+    //         data.supplyAmts,
+    //         data.variableBorrowAmts,
+    //         data.stableBorrowAmts
+    //     );
+    // }
 }