Add Migrator

contracts/Migrator.sol

@@ -28,13 +28,14 @@ contract Migrator is Governable {
 
     event TokenExchanged(uint256 ogvAmountIn, uint256 ognAmountOut);
     event Decommissioned();
-    event LockupsMigrated(address indexed user, uint256[] ogvLockupIds, uint256 newDuration);
+    event LockupsMigrated(address indexed user, uint256[] ogvLockupIds, uint256 newStakeAmount, uint256 newDuration);
 
     error MigrationAlreadyStarted();
     error MigrationIsInactive();
     error MigrationNotComplete();
-    error BalanceMismatch(uint256 ogvAmountIn, uint256 expectedOgnAmountOut, uint256 actualOgnAmountOut);
+    error ContractInsolvent(uint256 expectedOGN, uint256 availableOGN);
     error LockupIdsRequired();
+    error InvalidStakeAmount();
 
     constructor(address _ogv, address _ogn, address _ogvStaking, address _ognStaking) {
         ogv = ERC20Burnable(_ogv);
@@ -43,12 +44,32 @@ contract Migrator is Governable {
         ognStaking = IStaking(_ognStaking);
     }
 
+    /**
+     * @notice Solvency Checks
+     *
+     * This ensures that the contract always has enough OGN to
+     * continue with the migration.
+     * However, it doesn't revert if the difference is in favour
+     * of the contract (i.e. has more OGN than expected).
+     */
+    modifier isSolvent() {
+        _;
+
+        uint256 availableOGN = ogn.balanceOf(address(this));
+        uint256 totalOGV = ogv.totalSupply() - ogv.balanceOf(address(this));
+        uint256 maxOGNNeeded = (totalOGV * CONVERSION_RATE) / 1 ether;
+
+        if (availableOGN < maxOGNNeeded) {
+            revert ContractInsolvent(maxOGNNeeded, availableOGN);
+        }
+    }
+
     /**
      * @notice Starts the migration and sets it to end after
      *          365 days. Also, approves xOGN to transfer OGN
      *          held in this contract. Can be invoked only once
      */
-    function start() external onlyGovernor {
+    function start() external onlyGovernor isSolvent {
         if (endTime != 0) {
             revert MigrationAlreadyStarted();
         }
@@ -63,7 +84,7 @@ contract Migrator is Governable {
      * @notice Decommissions the contract. Can be called only
      *          after a year since `start()` was invoked. Burns
      *          all OGV held in the contract and transfers OGN
-     *          to address(1).
+     *          to address(0xdead).
      */
     function decommission() external {
         // Only after a year of staking
@@ -87,112 +108,73 @@ contract Migrator is Governable {
             // everything to address(1). The `owner` multisig of
             // OGN token can call `burnFrom(address(1))` later.abi
 
-            ogn.transfer(address(1), ognBalance);
+            ogn.transfer(address(0xdead), ognBalance);
         }
     }
 
     /**
-     * @notice Returns the active status of the migration.
-     * @return True if migration has started and has not ended yet.
+     * @notice Computes the amount of OGN needed for migration
+     *          and if the contract has more OGN than that, it
+     *          transfers it back to the treasury.
+     * @param treasury Address that receives excess OGN
      */
-    function isMigrationActive() public view returns (bool) {
-        return endTime > 0 && block.timestamp < endTime;
+    function transferExcessTokens(address treasury) external onlyGovernor isSolvent {
+        uint256 availableOGN = ogn.balanceOf(address(this));
+        uint256 totalOGV = ogv.totalSupply() - ogv.balanceOf(address(this));
+        uint256 maxOGNNeeded = (totalOGV * CONVERSION_RATE) / 1 ether;
+
+        if (availableOGN > maxOGNNeeded) {
+            ogn.transfer(treasury, availableOGN - maxOGNNeeded);
+        }
     }
 
     /**
-     * @notice Solvency Checks
-     *
-     * This ensures that the contract never transfers more than
-     * desired OGN amount in any case. This takes a balance diff
-     * of OGV and OGN and makes sure that the difference adds up.
-     * However, it doesn't revert if the difference is in favour
-     * of the contract (i.e. less OGN sent than expected).
+     * @notice Returns the active status of the migration.
+     * @return True if migration has started and has not ended yet.
      */
-    modifier netBalanceCheck() {
-        uint256 ogvBalanceBefore = ogv.balanceOf(address(this));
-        uint256 ognBalanceBefore = ogn.balanceOf(address(this));
-
-        _;
-
-        uint256 netOgvAmountIn = ogv.balanceOf(address(this)) - ogvBalanceBefore;
-        uint256 netOgnAmountOut = ognBalanceBefore - ogn.balanceOf(address(this));
-
-        uint256 netExpectedOgnAmountOut = (netOgvAmountIn * CONVERSION_RATE / 1 ether);
-
-        if (netExpectedOgnAmountOut < netOgnAmountOut) {
-            // TODO: Do we need some sort of tolerance (may be 0.01%)?
-            revert BalanceMismatch(netOgvAmountIn, netExpectedOgnAmountOut, netOgnAmountOut);
-        }
+    function isMigrationActive() public view returns (bool) {
+        return endTime > 0 && block.timestamp < endTime;
     }
 
     /**
      * @notice Migrates the specified amount of OGV to OGN
      * @param ogvAmount Amount of OGV to migrate
      * @return ognReceived OGN Received
      */
-    function migrate(uint256 ogvAmount) external netBalanceCheck returns (uint256 ognReceived) {
+    function migrate(uint256 ogvAmount) external isSolvent returns (uint256 ognReceived) {
         return _migrate(ogvAmount, msg.sender);
     }
 
-    /**
-     * @notice Migrates all of user's OGV to OGN
-     * @return ognReceived OGN Received
-     */
-    function migrateAll() external netBalanceCheck returns (uint256 ognReceived) {
-        return _migrate(ogv.balanceOf(msg.sender), msg.sender);
-    }
-
-    /**
-     * @notice Migrates caller's OGV to OGN and sends it to the `receiver`
-     * @return ognReceived OGN Received
-     */
-    function _migrate(uint256 ogvAmount, address receiver) internal returns (uint256 ognReceived) {
-        if (!isMigrationActive()) {
-            revert MigrationIsInactive();
-        }
-
-        ognReceived = ogvAmount * CONVERSION_RATE / 1 ether;
-
-        emit TokenExchanged(ogvAmount, ognReceived);
-
-        ogv.transferFrom(msg.sender, address(this), ogvAmount);
-
-        if (receiver != address(this)) {
-            // When migrating stakes, the contract would directly
-            // stake the balance on behalf of the user. So there's
-            // no need to transfer to self. Transfering to user and then
-            // back to this contract would only increase gas cost (and
-            // an additional tx for the user).
-            ogn.transfer(receiver, ognReceived);
-        }
-    }
-
     /**
      * @notice Migrates OGV stakes to OGN. Can also include unstaked OGN & OGV
      *          balances from the user's wallet (if specified).
      * @param lockupIds OGV Lockup IDs to be migrated
      * @param ogvAmountFromWallet Extra OGV balance from user's wallet to migrate & stake
      * @param ognAmountFromWallet Extra OGN balance from user's wallet to stake
      * @param migrateRewards If true, Migrate & Stake received rewards
+     * @param newStakeAmount Max amount of OGN (from wallet+unstake) to stake
      * @param newStakeDuration Duration of the new stake
      */
     function migrate(
         uint256[] calldata lockupIds,
         uint256 ogvAmountFromWallet,
         uint256 ognAmountFromWallet,
         bool migrateRewards,
+        uint256 newStakeAmount,
         uint256 newStakeDuration
-    ) external netBalanceCheck {
+    ) external isSolvent {
         if (!isMigrationActive()) {
             revert MigrationIsInactive();
         }
 
+        if (newStakeAmount == 0) {
+            revert InvalidStakeAmount();
+        }
+
         if (lockupIds.length == 0) {
             revert LockupIdsRequired();
         }
 
-        // TODO: Migrate delegation
-
         // Unstake
         (uint256 ogvAmountUnlocked, uint256 rewardsCollected) = ogvStaking.unstakeFrom(msg.sender, lockupIds);
 
@@ -211,16 +193,50 @@ contract Migrator is Governable {
         // Migrate OGV to OGN and include that along with existing balance
         ognAmountFromWallet += _migrate(ogvAmountFromWallet, address(this));
 
+        if (ognAmountFromWallet < newStakeAmount) {
+            revert InvalidStakeAmount();
+        }
+
+        uint256 ognToWallet = ognAmountFromWallet - newStakeAmount;
+
+        if (ognToWallet > 0) {
+            ogn.transfer(msg.sender, ognToWallet);
+        }
+
         // Stake it
         ognStaking.stake(
-            ognAmountFromWallet,
+            newStakeAmount,
             newStakeDuration,
             msg.sender,
             false,
             -1 // New stake
         );
 
-        // TODO: Emit new lockupId?
-        emit LockupsMigrated(msg.sender, lockupIds, newStakeDuration);
+        emit LockupsMigrated(msg.sender, lockupIds, newStakeAmount, newStakeDuration);
+    }
+
+    /**
+     * @notice Migrates caller's OGV to OGN and sends it to the `receiver`
+     * @return ognReceived OGN Received
+     */
+    function _migrate(uint256 ogvAmount, address receiver) internal returns (uint256 ognReceived) {
+        if (!isMigrationActive()) {
+            revert MigrationIsInactive();
+        }
+
+        ognReceived = (ogvAmount * CONVERSION_RATE) / 1 ether;
+
+        emit TokenExchanged(ogvAmount, ognReceived);
+
+        ogv.transferFrom(msg.sender, address(this), ogvAmount);
+
+        if (receiver != address(this)) {
+            // When migrating stakes, the contract would directly
+            // stake the balance on behalf of the user. So there's
+            // no need to transfer to self. Transfering to user and then
+            // back to this contract would only increase gas cost (and
+            // an additional tx for the user).
+            ogn.transfer(receiver, ognReceived);
+        }
     }
 }

contracts/OgvStaking.sol

@@ -176,6 +176,8 @@ contract OgvStaking is ERC20Votes {
         // Collect rewards
         rewardCollected = _collectRewards(staker);
 
+        uint256 unstakedPoints = 0;
+
         for (uint256 i = 0; i < lockupIds.length; ++i) {
             uint256 lockupId = lockupIds[i];
             Lockup memory lockup = lockups[staker][lockupId];
@@ -184,6 +186,7 @@ contract OgvStaking is ERC20Votes {
             uint256 points = lockup.points;
 
             unstakedAmount += amount;
+            unstakedPoints += points;
 
             // Make sure it isn't unstaked already
             if (end == 0) {
@@ -192,10 +195,13 @@ contract OgvStaking is ERC20Votes {
 
             delete lockups[staker][lockupId]; // Keeps empty in array, so indexes are stable
 
-            _burn(staker, points);
-            ogv.transfer(staker, amount);
             emit Unstake(staker, lockupId, amount, end, points);
         }
+
+        // Transfer unstaked OGV
+        ogv.transfer(staker, unstakedAmount);
+        // ... and burn veOGV
+        _burn(staker, unstakedPoints);
     }
 
     /// @notice Extend a stake lockup for additional points.
@@ -223,21 +229,15 @@ contract OgvStaking is ERC20Votes {
     /// @return points staking points that would be returned
     /// @return end staking period end date
     function previewPoints(uint256 amount, uint256 duration) public view returns (uint256, uint256) {
-        require(duration >= minStakeDuration, "Staking: Too short");
-        require(duration <= 1461 days, "Staking: Too long");
-        uint256 start = block.timestamp > epoch ? block.timestamp : epoch;
-        uint256 end = start + duration;
-        uint256 endYearpoc = ((end - epoch) * 1e18) / 365 days;
-        uint256 multiplier = PRBMathUD60x18.pow(YEAR_BASE, endYearpoc);
-        return ((amount * multiplier) / 1e18, end);
+        revert StakingDisabled();
     }
 
     // 3. Reward functions
 
     /// @notice Collect all earned OGV rewards.
     /// @return rewardCollected OGV reward amount collected
     function collectRewards() external returns (uint256 rewardCollected) {
-        _collectRewards(msg.sender);
+        return _collectRewards(msg.sender);
     }
 
     /// @notice Shows the amount of OGV a user would receive if they collected

contracts/tests/MockOGVStaking.sol

@@ -14,12 +14,22 @@ contract MockOGVStaking is OgvStaking {
         OgvStaking(ogv_, epoch_, minStakeDuration_, rewardsSource_, migrator_)
     {}
 
+    function _previewPoints(uint256 amount, uint256 duration) internal view returns (uint256, uint256) {
+        require(duration >= minStakeDuration, "Staking: Too short");
+        require(duration <= 1461 days, "Staking: Too long");
+        uint256 start = block.timestamp > epoch ? block.timestamp : epoch;
+        uint256 end = start + duration;
+        uint256 endYearpoc = ((end - epoch) * 1e18) / 365 days;
+        uint256 multiplier = PRBMathUD60x18.pow(YEAR_BASE, endYearpoc);
+        return ((amount * multiplier) / 1e18, end);
+    }
+
     function mockStake(uint256 amountIn, uint256 duration) external {
         Lockup memory lockup;
 
         ogv.transferFrom(msg.sender, address(this), amountIn);
 
-        (uint256 points, uint256 end) = previewPoints(amountIn, duration);
+        (uint256 points, uint256 end) = _previewPoints(amountIn, duration);
         require(points + totalSupply() <= type(uint192).max, "Staking: Max points exceeded");
 
         lockup.end = uint128(end);