Feature/migrate user token

contracts/StakingPoolV2.sol

@@ -2,6 +2,7 @@
 pragma solidity 0.8.4;
 import './libraries/StakingPoolLogicV2.sol';
 import './interface/IStakingPoolV2.sol';
+import './interface/INextStakingPool.sol';
 import './token/StakedElyfiToken.sol';
 import '@openzeppelin/contracts/token/ERC20/IERC20.sol';
 import '@openzeppelin/contracts/access/Ownable.sol';
@@ -35,6 +36,8 @@ contract StakingPoolV2 is IStakingPoolV2, StakedElyfiToken, Ownable {
     bool isFinished;
   }
 
+  address internal _admin;
+  address internal _nextContractAddr;
   IERC20 public stakingAsset;
   IERC20 public rewardAsset;
   PoolData internal _poolData;
@@ -129,8 +132,37 @@ contract StakingPoolV2 is IStakingPoolV2, StakedElyfiToken, Ownable {
     _claim(msg.sender);
   }
 
-  // TODO Implement `migrate` function to send an asset to the next staking contract
-
+  /// @notice Migrate the amount of principal to the current round and transfer the rest principal to the caller
+  function migrate() external override {
+    if (_poolData.isOpened == true) revert Opened();
+    if (_nextContractAddr == address(0)) revert NotSetContractAddr();
+    if (_poolData.userPrincipal[msg.sender] == 0) revert ZeroPrincipal();
+    uint256 amount = _poolData.userPrincipal[msg.sender];
+
+
+    // Claim reward
+    if (_poolData.getUserReward(msg.sender) != 0) {
+      _claim(msg.sender);
+    }
+
+    _poolData.updateStakingPool(msg.sender);
+
+    // Migrate user, total principal 
+    _poolData.userPrincipal[msg.sender] -= amount;
+    _poolData.totalPrincipal -= amount;
+
+    // Migrate
+    _migrate(_nextContractAddr, amount);
+
+    // Call next contract
+    INextStakingPool nextContract = INextStakingPool(_nextContractAddr);
+
+    // Migrate next contract of user, total principal
+    nextContract.setPreviousPoolData(msg.sender, amount);
+
+    emit Migrate(msg.sender, amount);
+  }
+
   /***************** Internal Functions ******************/
 
   function _withdraw(uint256 amount) internal {
@@ -199,6 +231,7 @@ contract StakingPoolV2 is IStakingPoolV2, StakedElyfiToken, Ownable {
 
   function closePool() external onlyOwner {
     if (_poolData.isOpened == false) revert Closed();
+   // _poolData.rewardIndex = _poolData.getRewardIndex();
     _poolData.endTimestamp = block.timestamp;
     _poolData.isOpened = false;
     _poolData.isFinished = true;

contracts/interface/INextStakingPool.sol

@@ -0,0 +1,13 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.4;
+
+interface INextStakingPool {
+  error OnlyAdmin();
+  function setPreviousPoolData(address user, uint256 amount) external;
+
+  function initNewPool(
+    uint256 rewardPerSecond,
+    uint256 startTimestamp,
+    uint32 duration
+  ) external;
+}

contracts/interface/IStakingPoolV2.sol

@@ -8,8 +8,10 @@ interface IStakingPoolV2 {
   error OnlyAdmin();
   error NotEnoughPrincipal(uint256 principal);
   error ZeroPrincipal();
+  error NotSetContractAddr();
   error Finished();
   error Closed();
+  error Opened();
 
   event Stake(
     address indexed user,
@@ -33,12 +35,16 @@ interface IStakingPoolV2 {
     uint256 endTimestamp
   );
 
+  event Migrate(address user, uint256 amount);
+
   function stake(uint256 amount) external;
 
   function claim() external;
 
   function withdraw(uint256 amount) external;
 
+  function migrate() external;
+
   function getRewardIndex() external view returns (uint256);
 
   function getUserReward(address user) external view returns (uint256);

contracts/test/StakingPoolLogicV3.sol

@@ -0,0 +1,82 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.4;
+import './StakingPoolV3.sol';
+import '@openzeppelin/contracts/token/ERC20/IERC20.sol';
+import '@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol';
+
+library StakingPoolLogicV3 {
+  using StakingPoolLogicV3 for StakingPoolV3.PoolData;
+
+  event UpdateStakingPool(
+    address indexed user,
+    uint256 newRewardIndex,
+    uint256 totalPrincipal
+  );
+
+  function getRewardIndex(StakingPoolV3.PoolData storage poolData) internal view returns (uint256) {
+    uint256 currentTimestamp = block.timestamp < poolData.endTimestamp
+      ? block.timestamp
+      : poolData.endTimestamp;
+    uint256 timeDiff = currentTimestamp - poolData.lastUpdateTimestamp;
+    uint256 totalPrincipal = poolData.totalPrincipal;
+
+    if (timeDiff == 0) {
+      return poolData.rewardIndex;
+    }
+
+    if (totalPrincipal == 0) {
+      return poolData.rewardIndex;
+    }
+
+    uint256 rewardIndexDiff = (timeDiff * poolData.rewardPerSecond * 1e9) / totalPrincipal;
+
+    return poolData.rewardIndex + rewardIndexDiff;
+  }
+
+  function getUserReward(StakingPoolV3.PoolData storage poolData, address user)
+    internal
+    view
+    returns (uint256)
+  {
+    if (poolData.userIndex[user] == 0) {
+      return 0;
+    }
+
+    uint256 indexDiff = getRewardIndex(poolData) - poolData.userIndex[user];
+    uint256 balance = poolData.userPrincipal[user];
+    uint256 result = poolData.userReward[user] + (balance * indexDiff) / 1e9;
+
+    return result;
+  }
+
+
+  function updateStakingPool(
+    StakingPoolV3.PoolData storage poolData,
+    address user
+  ) internal {
+    poolData.userReward[user] = getUserReward(poolData, user);
+    poolData.rewardIndex = poolData.userIndex[user] = getRewardIndex(poolData);
+    poolData.lastUpdateTimestamp = block.timestamp < poolData.endTimestamp
+      ? block.timestamp
+      : poolData.endTimestamp;
+    emit UpdateStakingPool(msg.sender, poolData.rewardIndex, poolData.totalPrincipal);
+  }
+
+  function initRound(
+    StakingPoolV3.PoolData storage poolData,
+    uint256 rewardPerSecond,
+    uint256 roundStartTimestamp,
+    uint32 duration
+  ) internal returns (uint256, uint256) {
+    poolData.rewardPerSecond = rewardPerSecond;
+    poolData.startTimestamp = roundStartTimestamp;
+    poolData.endTimestamp = roundStartTimestamp + duration;
+    poolData.lastUpdateTimestamp = roundStartTimestamp;
+    poolData.rewardIndex = 1e18;
+    poolData.duration = duration;
+    return (poolData.startTimestamp, poolData.endTimestamp);
+  }
+
+
+
+}

contracts/test/StakingPoolV3.sol

@@ -0,0 +1,163 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.4;
+import './StakingPoolLogicV3.sol';
+import '../interface/INextStakingPool.sol';
+import '../token/StakedElyfiToken.sol';
+import '@openzeppelin/contracts/token/ERC20/IERC20.sol';
+
+
+contract StakingPoolV3 is INextStakingPool, StakedElyfiToken {
+  using StakingPoolLogicV3 for PoolData;
+  constructor(IERC20 stakingAsset_, IERC20 rewardAsset_) StakedElyfiToken(stakingAsset_) {
+    stakingAsset = stakingAsset_;
+    rewardAsset = rewardAsset_;
+    _admin = msg.sender;
+  }
+
+  struct PoolData {
+    uint32 duration;
+    uint256 rewardPerSecond;
+    uint256 rewardIndex;
+    uint256 startTimestamp;
+    uint256 endTimestamp;
+    uint256 lastUpdateTimestamp;
+    uint256 totalPrincipal;
+    uint256 nextRewardAmount;
+    mapping(address => uint256) userIndex;
+    mapping(address => uint256) userReward;
+    mapping(address => uint256) userPrincipal;
+    bool isOpened;
+    bool isFinished;
+  }
+  address internal _admin;
+  IERC20 public stakingAsset;
+  IERC20 public rewardAsset;
+  PoolData internal _poolData;
+
+
+
+
+  //=================== view ==================
+  /// @notice Returns reward index of the round
+  function getRewardIndex() external view returns (uint256) {
+    return _poolData.getRewardIndex();
+  }
+
+  /// @notice Returns user accrued reward index of the round
+  /// @param user The user address
+  function getUserReward(address user) external view returns (uint256) {
+    return _poolData.getUserReward(user);
+  }
+
+  function getPoolData()
+    external
+    view
+    returns (
+      uint256 rewardPerSecond,
+      uint256 rewardIndex,
+      uint256 startTimestamp,
+      uint256 endTimestamp,
+      uint256 totalPrincipal,
+      uint256 lastUpdateTimestamp
+    )
+  {
+    return (
+      _poolData.rewardPerSecond,
+      _poolData.rewardIndex,
+      _poolData.startTimestamp,
+      _poolData.endTimestamp,
+      _poolData.totalPrincipal,
+      _poolData.lastUpdateTimestamp
+    );
+  }
+
+  function getUserData(address user)
+    external
+    view
+    returns (
+      uint256 userIndex,
+      uint256 userReward,
+      uint256 userPrincipal
+    )
+  {
+    return (_poolData.userIndex[user], _poolData.userReward[user], _poolData.userPrincipal[user]);
+  }
+
+  function setPreviousPoolData(address user, uint256 amount) external override {
+    _poolData.updateStakingPool(user);
+    _poolData.userPrincipal[user] += amount;
+    _poolData.totalPrincipal += amount;
+  }
+
+  function getContractBalance() external view returns(uint256) {
+    return stakingAsset.balanceOf(address(this));
+  }
+
+
+  // ================   main  ==============
+  function stake(uint256 amount) external {
+    _poolData.updateStakingPool(msg.sender);
+    _depositFor(msg.sender, amount);
+
+    _poolData.userPrincipal[msg.sender] += amount;
+    _poolData.totalPrincipal += amount;
+  }
+
+  function withdraw(uint256 amount) external  {
+    _withdraw(amount);
+  }
+
+  /// @notice Transfer accrued reward to msg.sender. User accrued reward will be reset and user reward index will be set to the current reward index.
+  function claim() external {
+    _claim(msg.sender);
+  }
+
+  function _withdraw(uint256 amount) internal {
+    uint256 amountToWithdraw = amount;
+
+    if (amount == type(uint256).max) {
+      amountToWithdraw = _poolData.userPrincipal[msg.sender];
+    }
+
+    _poolData.updateStakingPool(msg.sender);
+
+    _poolData.userPrincipal[msg.sender] -= amountToWithdraw;
+    _poolData.totalPrincipal -= amountToWithdraw;
+
+    _withdrawTo(msg.sender, amountToWithdraw);
+  }
+
+  function _claim(address user) internal {
+    uint256 reward = _poolData.getUserReward(user);
+
+    _poolData.userReward[user] = 0;
+    _poolData.userIndex[user] = _poolData.getRewardIndex();
+
+    SafeERC20.safeTransfer(rewardAsset, user, reward);
+
+    uint256 rewardLeft = rewardAsset.balanceOf(address(this));
+
+  }
+
+
+  function initNewPool(
+    uint256 rewardPerSecond,
+    uint256 startTimestamp,
+    uint32 duration
+  ) external override {
+    (uint256 newRoundStartTimestamp, uint256 newRoundEndTimestamp) = _poolData.initRound(
+      rewardPerSecond,
+      startTimestamp,
+      duration
+    );
+    _poolData.isOpened = true;
+
+    SafeERC20.safeTransferFrom(rewardAsset, msg.sender, address(this), duration * rewardPerSecond);
+  }
+
+  modifier onlyAdmin() {
+    if (msg.sender != _admin) revert OnlyAdmin();
+    _;
+  }
+
+}

contracts/token/StakedElyfiToken.sol

@@ -121,6 +121,13 @@ contract StakedElyfiToken is ERC20, ERC20Permit, ERC20Votes {
     SafeERC20.safeTransfer(underlying, account, amount);
     return true;
   }
+
+  /// @dev Allow a user to migrate underlying tokens to next new contract
+  /// @notice This function is based on the openzeppelin ERC20Wrapper
+  function _migrate(address toContractAddr, uint256 amount) internal virtual returns (bool) {
+    SafeERC20.safeTransfer(underlying, toContractAddr, amount);
+    return true;
+  }
 
   /// @notice The following functions are overrides required by Solidity.
   function _afterTokenTransfer(

deploy/elyfiPool.ts

@@ -20,6 +20,6 @@ const elyfiPool: DeployFunction = async function (hre: HardhatRuntimeEnvironment
     network: hre.network.name,
   });
 };
-elyfiPool.tags = ['elyfiPool'];
+elyfiPool.tags = ['elyfiPool', 'test'];
 
 export default elyfiPool;

test/claim.test.ts

@@ -6,7 +6,7 @@ import { setTestEnv } from './utils/testEnv';
 import { advanceTime, advanceTimeTo, getTimestamp, toTimestamp } from './utils/time';
 import { expectDataAfterClaim } from './utils/expect';
 import { getPoolData, getUserData } from './utils/helpers';
-import TestEnv from './types/TestEnv';
+import {TestEnv} from './types/TestEnv';
 
 const { loadFixture } = waffle;