✨ Add user convenience features to ValidatorStaking

src/external/interfaces/ITMITO.sol

@@ -0,0 +1,162 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.28;
+
+import { IERC20 } from '@oz/interfaces/IERC20.sol';
+
+interface ITMITO is IERC20 {
+  // =========================== EVENTS =========================== //
+
+  /**
+   * @notice Emitted when extra rewards are added to the contract.
+   * @param amount The amount of MITO added as extra rewards
+   */
+  event ExtraRewardsAdded(uint256 amount);
+
+  /**
+   * @notice Emitted when the MITO:TMITO ratio is finalized.
+   * @param totalMITOAmount The total MITO amount at finalization
+   * @param totalTMITOAmount The total TMITO amount at finalization
+   */
+  event RatioFinalized(uint256 totalMITOAmount, uint256 totalTMITOAmount);
+
+  /**
+   * @notice Emitted when MITO is converted to TMITO after ratio is fixed.
+   * @param account The address that converted the MITO to TMITO
+   * @param to The address that received the TMITO
+   * @param mitoAmount The amount of MITO converted
+   * @param tmitoAmount The amount of TMITO minted
+   */
+  event ConvertedMITOToTMITO(
+    address indexed account, address indexed to, uint256 mitoAmount, uint256 tmitoAmount
+  );
+
+  /**
+   * @notice Emitted when tMITO tokens are minted.
+   * @param to The address that received the tokens
+   * @param amount The amount of tokens minted
+   */
+  event Minted(address indexed to, uint256 amount);
+
+  /**
+   * @notice Emitted when tMITO tokens are redeemed for MITO and extra rewards.
+   * @param account The address that redeemed the tokens
+   * @param to The address that received the MITO
+   * @param tmitoAmount The amount of tMITO tokens burned
+   * @param baseAmount The amount of base MITO received (1:1)
+   * @param extraRewardAmount The amount of extra MITO received
+   * @param totalAmount The total amount of MITO received (base + extra)
+   */
+  event Redeemed(
+    address indexed account,
+    address indexed to,
+    uint256 tmitoAmount,
+    uint256 baseAmount,
+    uint256 extraRewardAmount,
+    uint256 totalAmount
+  );
+
+  /**
+   * @notice Emitted when the lockup end time is set.
+   * @param lockupEndTime The new lockup end timestamp
+   */
+  event LockupEndTimeSet(uint48 lockupEndTime);
+
+  // =========================== ERRORS =========================== //
+
+  error TMITO__LockupNotEnded();
+  error TMITO__ZeroAmount();
+  error TMITO__ZeroAddress();
+  error TMITO__RatioAlreadyFinalized();
+  error TMITO__RatioNotFinalized();
+
+  // =========================== VIEW FUNCTIONS =========================== //
+
+  /**
+   * @notice Returns the total amount of extra MITO rewards.
+   * @return The total amount of extra MITO rewards
+   */
+  function totalExtraRewards() external view returns (uint256);
+
+  /**
+   * @notice Returns the lockup end timestamp.
+   * @return The timestamp when lockup ends
+   */
+  function lockupEndTime() external view returns (uint48);
+
+  /**
+   * @notice Previews the total MITO amount a user would receive when redeeming tMITO.
+   * @param tmitoAmount The amount of tMITO tokens to redeem
+   * @return baseAmount The base MITO amount (1:1 ratio)
+   * @return extraRewardAmount The extra MITO reward amount
+   * @return totalAmount The total MITO amount (base + extra rewards)
+   */
+  function previewRedeem(uint256 tmitoAmount)
+    external
+    view
+    returns (uint256 baseAmount, uint256 extraRewardAmount, uint256 totalAmount);
+
+  /**
+   * @notice Calculates the extra MITO rewards a user would receive for a given tMITO amount.
+   * @param tmitoAmount The amount of tMITO tokens
+   * @return The amount of extra MITO rewards
+   */
+  function previewExtraRewards(uint256 tmitoAmount) external view returns (uint256);
+
+  /**
+   * @notice Previews the amount of TMITO tokens that would be minted for a given MITO amount.
+   * @param mitoAmount The amount of MITO to convert
+   * @return tmitoAmount The amount of TMITO that would be minted
+   */
+  function previewConvert(uint256 mitoAmount) external view returns (uint256 tmitoAmount);
+
+  /**
+   * @notice Checks if the lockup period has ended.
+   * @return Whether the lockup period has ended
+   */
+  function isLockupEnded() external view returns (bool);
+
+  /**
+   * @notice Returns whether the MITO:TMITO ratio has been finalized.
+   * @return True if the ratio has been finalized, false otherwise
+   */
+  function ratioFinalized() external view returns (bool);
+
+  // =========================== MUTATIVE FUNCTIONS =========================== //
+
+  /**
+   * @notice Mints tMITO tokens to a user.
+   * @dev Only addresses with MINTER_ROLE can call this function.
+   * @param to The address to mint tokens to
+   */
+  function mint(address to) external payable;
+
+  /**
+   * @notice Adds extra MITO rewards to the contract.
+   * @dev Only addresses with REWARD_MANAGER_ROLE can call this function.
+   */
+  function addExtraRewards() external payable;
+
+  /**
+   * @notice Converts MITO to TMITO based on finalized ratio.
+   * @dev Can only be called after ratio is finalized.
+   * @param to The address to receive the TMITO
+   */
+  function convertMITOToTMITO(address to) external payable;
+
+  /**
+   * @notice Redeems tMITO tokens for MITO and extra rewards.
+   * @dev Can only be called after lockup period ends.
+   * @param to The address to receive the MITO
+   * @param tmitoAmount The amount of tMITO tokens to redeem
+   */
+  function redeem(address to, uint256 tmitoAmount) external;
+
+  // =========================== ADMIN FUNCTIONS =========================== //
+
+  /**
+   * @notice Sets the lockup end time.
+   * @dev Only addresses with DEFAULT_ADMIN_ROLE can call this function.
+   * @param lockupEndTime_ The new lockup end timestamp
+   */
+  function setLockupEndTime(uint48 lockupEndTime_) external;
+}

src/hub/validator/ValidatorStaking.sol

@@ -36,6 +36,8 @@ contract ValidatorStakingStorageV1 {
     mapping(address valAddr => Checkpoints.Trace208) validatorTotal;
     mapping(address valAddr => mapping(address staker => Checkpoints.Trace208)) staked;
     mapping(address staker => mapping(address valAddr => uint256)) lastRedelegationTime;
+    mapping(address valAddr => mapping(address owner => mapping(address spender => uint256))) stakingAllowances;
+    address migrationAgent;
   }
 
   string private constant _NAMESPACE = 'mitosis.storage.ValidatorStaking.v1';
@@ -123,7 +125,7 @@ contract ValidatorStaking is
   }
 
   /// @inheritdoc IValidatorStaking
-  function staked(address valAddr, address staker, uint48 timestamp) external view virtual returns (uint256) {
+  function staked(address valAddr, address staker, uint48 timestamp) public view virtual returns (uint256) {
     return _getStorageV1().staked[valAddr][staker].upperLookupRecent(timestamp);
   }
 
@@ -188,6 +190,16 @@ contract ValidatorStaking is
     return _getStorageV1().lastRedelegationTime[staker][valAddr];
   }
 
+  /// @inheritdoc IValidatorStaking
+  function stakingAllowance(address valAddr, address owner, address spender) external view returns (uint256) {
+    return _getStorageV1().stakingAllowances[valAddr][owner][spender];
+  }
+
+  /// @notice Returns the migration agent address.
+  function migrationAgent() external view returns (address) {
+    return _getStorageV1().migrationAgent;
+  }
+
   // ===================================== MUTATIVE FUNCTIONS ===================================== //
 
   /// @inheritdoc IValidatorStaking
@@ -210,6 +222,52 @@ contract ValidatorStaking is
     return _redelegate(_getStorageV1(), _msgSender(), fromValAddr, toValAddr, amount);
   }
 
+  /// @inheritdoc IValidatorStaking
+  function approveStakingOwnership(address valAddr, address spender, uint256 amount) external {
+    require(_manager.isValidator(valAddr), IValidatorStaking__NotValidator(valAddr));
+
+    _getStorageV1().stakingAllowances[valAddr][_msgSender()][spender] = amount;
+
+    emit StakingOwnershipApproved(valAddr, _msgSender(), spender, amount);
+  }
+
+  /// @inheritdoc IValidatorStaking
+  function transferStakingOwnershipFrom(address valAddr, address from, address to, uint256 amount)
+    external
+    virtual
+    nonReentrant
+    returns (uint256)
+  {
+    require(_msgSender() == to, StdError.InvalidParameter('to'));
+    require(_manager.isValidator(valAddr), IValidatorStaking__NotValidator(valAddr));
+
+    StorageV1 storage $ = _getStorageV1();
+    require(amount >= $.minStakingAmount, IValidatorStaking__InsufficientMinimumAmount($.minStakingAmount));
+    require(amount > 0, StdError.ZeroAmount());
+
+    uint48 _now = Time.timestamp();
+
+    uint256 _staked = staked(valAddr, from, _now);
+    uint256 allowance = $.stakingAllowances[valAddr][from][to];
+
+    require(_staked >= amount, IValidatorStaking__InsufficientStakedAmount(amount, _staked));
+    require(allowance >= amount, IValidatorStaking__InsufficientAllowance(amount, allowance));
+
+    $.stakingAllowances[valAddr][from][to] -= amount;
+
+    {
+      uint208 amount208 = amount.toUint208();
+      _storeUnstake($, _now, valAddr, from, amount208);
+      _storeStake($, _now, valAddr, to, amount208);
+    }
+
+    _hub.notifyUnstake(valAddr, from, amount);
+    _hub.notifyStake(valAddr, to, amount);
+
+    emit StakingOwnershipTransferred(valAddr, from, to, amount);
+    return amount;
+  }
+
   /// @inheritdoc IValidatorStaking
   function setMinStakingAmount(uint256 minAmount) external onlyOwner {
     _setMinStakingAmount(_getStorageV1(), minAmount);
@@ -230,6 +288,19 @@ contract ValidatorStaking is
     _setRedelegationCooldown(_getStorageV1(), redelegationCooldown_);
   }
 
+  /// @notice Sets the migration agent address. Only the migration agent can call claimUnstakeForMigration.
+  function setMigrationAgent(address agent) external onlyOwner {
+    address previousAgent = _getStorageV1().migrationAgent;
+    _getStorageV1().migrationAgent = agent;
+    emit MigrationAgentSet(previousAgent, agent);
+  }
+
+  /// @notice Claims unstaked tokens immediately, bypassing cooldown. Only callable by migration agent.
+  function claimUnstakeForMigration() external nonReentrant returns (uint256) {
+    require(_msgSender() == _getStorageV1().migrationAgent, StdError.Unauthorized());
+    return _claimUnstakeForMigration(_getStorageV1(), _msgSender());
+  }
+
   // ===================================== INTERNAL FUNCTIONS ===================================== //
 
   function _assertUnstakeAmountCondition(StorageV1 storage $, address valAddr, address staker, uint256 amount)
@@ -307,11 +378,22 @@ contract ValidatorStaking is
     return reqId;
   }
 
+  function _claimUnstakeForMigration(StorageV1 storage $, address receiver) internal virtual returns (uint256) {
+    return _claimUnstakeForRequestsUntil($, receiver, Time.timestamp());
+  }
+
   function _claimUnstake(StorageV1 storage $, address receiver) internal virtual returns (uint256) {
+    return _claimUnstakeForRequestsUntil($, receiver, Time.timestamp() - $.unstakeCooldown);
+  }
+
+  function _claimUnstakeForRequestsUntil(StorageV1 storage $, address receiver, uint48 eligibleUntil)
+    private
+    returns (uint256)
+  {
     LibQueue.Trace208OffsetQueue storage queue = $.unstakeQueue[receiver];
 
     uint48 now_ = Time.timestamp();
-    (uint32 reqIdFrom, uint32 reqIdTo) = queue.solveByKey(now_ - $.unstakeCooldown);
+    (uint32 reqIdFrom, uint32 reqIdTo) = queue.solveByKey(eligibleUntil);
     uint256 claimed;
     {
       uint256 fromValue = reqIdFrom == 0 ? 0 : queue.valueAt(reqIdFrom - 1);
@@ -323,7 +405,6 @@ contract ValidatorStaking is
     if (baseAsset_ == NATIVE_TOKEN) receiver.safeTransferETH(claimed);
     else baseAsset_.safeTransfer(receiver, claimed);
 
-    // apply to state
     _push($.totalUnstaking, now_, claimed.toUint208(), _opSub);
 
     emit UnstakeClaimed(receiver, claimed, reqIdFrom, reqIdTo);

src/hub/validator/ValidatorStakingGovMITO.sol

@@ -63,6 +63,11 @@ contract ValidatorStakingGovMITO is ValidatorStaking, SudoVotes {
     return stakerTotal(account, now_) + totalUnstakingAmount;
   }
 
+  /// @dev GovMITO staking is non-transferable; ownership transfer would desync voting power.
+  function transferStakingOwnershipFrom(address, address, address, uint256) external pure override returns (uint256) {
+    revert ValidatorStakingGovMITO__NonTransferable();
+  }
+
   /// @dev Mints voting units to the recipient. No need to care about the validator
   function _stake(StorageV1 storage $, address valAddr, address payer, address recipient, uint256 amount)
     internal

src/hub/validator/ValidatorStakingMigration.sol

@@ -0,0 +1,83 @@
+// SPDX-License-Identifier: Apache-2.0
+pragma solidity ^0.8.28;
+
+import { AccessControlEnumerableUpgradeable } from '@ozu/access/extensions/AccessControlEnumerableUpgradeable.sol';
+import { UUPSUpgradeable } from '@ozu/proxy/utils/UUPSUpgradeable.sol';
+import { ReentrancyGuardUpgradeable } from '@ozu/utils/ReentrancyGuardUpgradeable.sol';
+
+import { ITMITO } from '../../external/interfaces/ITMITO.sol';
+import { IValidatorStaking } from '../../interfaces/hub/validator/IValidatorStaking.sol';
+import { IValidatorStakingMigration } from '../../interfaces/hub/validator/IValidatorStakingMigration.sol';
+import { StdError } from '../../lib/StdError.sol';
+
+/// @title ValidatorStakingMigration
+/// @notice Enables migration of staking between ValidatorStaking_MITO and ValidatorStaking_TMITO.
+/// @dev Uses approveStakingOwnership and transferStakingOwnershipFrom.
+///      Requires ValidatorStaking_TMITO owner to set this contract as migrationAgent for instant migration.
+contract ValidatorStakingMigration is
+  IValidatorStakingMigration,
+  ReentrancyGuardUpgradeable,
+  AccessControlEnumerableUpgradeable,
+  UUPSUpgradeable
+{
+  IValidatorStaking public immutable tmitoValidatorStaking;
+  IValidatorStaking public immutable mitoValidatorStaking;
+  ITMITO public immutable tMITO;
+
+  constructor(IValidatorStaking tmitoValidatorStaking_, IValidatorStaking mitoValidatorStaking_, ITMITO tMITO_) {
+    _disableInitializers();
+    tmitoValidatorStaking = tmitoValidatorStaking_;
+    mitoValidatorStaking = mitoValidatorStaking_;
+    tMITO = tMITO_;
+  }
+
+  function initialize(address admin) external initializer {
+    __ReentrancyGuard_init();
+    __AccessControlEnumerable_init();
+    __UUPSUpgradeable_init();
+    _grantRole(DEFAULT_ADMIN_ROLE, admin);
+  }
+
+  /// @notice Instant migration: TMITO -> MITO in a single tx (bypasses unstake cooldown).
+  /// @dev Requires: 1) User approved this contract on tmitoValidatorStaking.approveStakingOwnership
+  ///                2) ValidatorStaking_TMITO owner set this contract as migrationAgent
+  /// @param valAddr Validator address
+  /// @param amount Amount of TMITO staking to migrate
+  /// @param recipient Address to receive MITO staking
+  function migrateFromTMITOToMITO(address valAddr, uint256 amount, address recipient)
+    external
+    nonReentrant
+    returns (uint256)
+  {
+    require(amount > 0, StdError.ZeroAmount());
+    require(recipient != address(0), StdError.ZeroAddress('recipient'));
+    require(tMITO.isLockupEnded(), ValidatorStakingMigration__LockupNotEnded());
+    require(tmitoValidatorStaking.migrationAgent() == address(this), ValidatorStakingMigration__NotMigrationAgent());
+
+    // 1. Transfer staking ownership from user to this contract
+    tmitoValidatorStaking.transferStakingOwnershipFrom(valAddr, _msgSender(), address(this), amount);
+
+    // 2. Request unstake (adds to queue)
+    tmitoValidatorStaking.requestUnstake(valAddr, address(this), amount);
+
+    // 3. Claim immediately, bypassing cooldown (requires migrationAgent role)
+    uint256 tmitoAmount = tmitoValidatorStaking.claimUnstakeForMigration();
+
+    // 4. Redeem TMITO to MITO
+    uint256 balanceBefore = address(this).balance;
+    tMITO.redeem(address(this), tmitoAmount);
+    uint256 mitoAmount = address(this).balance - balanceBefore;
+
+    // 5. Stake MITO in ValidatorStaking_MITO for recipient
+    mitoValidatorStaking.stake{ value: mitoAmount }(valAddr, recipient, mitoAmount);
+
+    emit InstantMigrationCompleted(_msgSender(), valAddr, recipient, mitoAmount);
+
+    return mitoAmount;
+  }
+
+  function _authorizeUpgrade(address) internal override onlyRole(DEFAULT_ADMIN_ROLE) { }
+
+  /// @dev Required to receive native MITO from tMITO.redeem()
+  receive() external payable { }
+}