staking/contracts/StakeManager.sol

// SPDX-License-Identifier: MIT

pragma solidity ^0.8.18;

import { ERC20 } from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
import { Ownable } from "@openzeppelin/contracts/access/Ownable.sol";
import { Math } from "@openzeppelin/contracts/utils/math/Math.sol";

import { StakeVault } from "./StakeVault.sol";

contract StakeManager is Ownable {
    error StakeManager__SenderIsNotVault();
    error StakeManager__FundsLocked();
    error StakeManager__InvalidLockTime();
    error StakeManager__NoPendingMigration();
    error StakeManager__PendingMigration();
    error StakeManager__SenderIsNotPreviousStakeManager();
    error StakeManager__InvalidLimitEpoch();
    error StakeManager__AccountNotInitialized();
    error StakeManager__InvalidMigration();
    error StakeManager__AlreadyProcessedEpochs();
    error StakeManager__InsufficientFunds();

    struct Account {
        address rewardAddress;
        uint256 balance;
        uint256 initialMP;
        uint256 currentMP;
        uint256 lastMint;
        uint256 lockUntil;
        uint256 epoch;
    }

    struct Epoch {
        uint256 startTime;
        uint256 epochReward;
        uint256 totalSupply;
    }

    uint256 public constant EPOCH_SIZE = 1 weeks;
    uint256 public constant YEAR = 365 days;
    uint256 public constant MIN_LOCKUP_PERIOD = 2 weeks;
    uint256 public constant MAX_LOCKUP_PERIOD = 4 * YEAR; // 4 years
    uint256 public constant MP_APY = 1;
    uint256 public constant MAX_BOOST = 4;

    mapping(address index => Account value) public accounts;
    mapping(uint256 index => Epoch value) public epochs;
    mapping(bytes32 codehash => bool approved) public isVault;

    uint256 public currentEpoch;
    uint256 public pendingReward;
    uint256 public totalSupplyMP;
    uint256 public totalSupplyBalance;
    StakeManager public migration;
    StakeManager public immutable previousManager;
    ERC20 public immutable stakedToken;

    /**
     * @notice Only callable by vaults
     */
    modifier onlyVault() {
        if (!isVault[msg.sender.codehash]) {
            revert StakeManager__SenderIsNotVault();
        }
        _;
    }

    modifier onlyAccountInitialized(address account) {
        if (accounts[account].lockUntil == 0) {
            revert StakeManager__AccountNotInitialized();
        }
        _;
    }

    /**
     * @notice Only callable when migration is not initialized.
     */
    modifier noPendingMigration() {
        if (address(migration) != address(0)) {
            revert StakeManager__PendingMigration();
        }
        _;
    }

    /**
     * @notice Only callable when migration is initialized.
     */
    modifier onlyPendingMigration() {
        if (address(migration) == address(0)) {
            revert StakeManager__NoPendingMigration();
        }
        _;
    }

    /**
     * @notice Only callable from old manager.
     */
    modifier onlyPreviousManager() {
        if (msg.sender != address(previousManager)) {
            revert StakeManager__SenderIsNotPreviousStakeManager();
        }
        _;
    }

    /**
     * @notice Process epoch if it has ended
     */
    modifier finalizeEpoch() {
        if (block.timestamp >= epochEnd() && address(migration) == address(0)) {
            //finalize current epoch
            epochs[currentEpoch].epochReward = epochReward();
            epochs[currentEpoch].totalSupply = totalSupply();
            pendingReward += epochs[currentEpoch].epochReward;
            //create new epoch
            currentEpoch++;
            epochs[currentEpoch].startTime = block.timestamp;
        }
        _;
    }

    constructor(address _stakedToken, address _previousManager) {
        epochs[0].startTime = block.timestamp;
        previousManager = StakeManager(_previousManager);
        stakedToken = ERC20(_stakedToken);
    }

    /**
     * Increases balance of msg.sender;
     * @param _amount Amount of balance to be decreased.
     * @param _timeToIncrease Seconds to increase in locked time. If stake is unlocked, increases from block.timestamp.
     *
     * @dev Reverts when resulting locked time is not in range of [MIN_LOCKUP_PERIOD, MAX_LOCKUP_PERIOD]
     */
    function stake(uint256 _amount, uint256 _timeToIncrease) external onlyVault noPendingMigration finalizeEpoch {
        Account storage account = accounts[msg.sender];

        if (account.lockUntil == 0) {
            // account not initialized
            account.lockUntil = block.timestamp;
            account.epoch = currentEpoch; //starts in current epoch
            account.rewardAddress = StakeVault(msg.sender).owner();
        } else {
            _processAccount(account, currentEpoch);
        }

        uint256 deltaTime = 0;

        if (_timeToIncrease > 0) {
            uint256 lockUntil = account.lockUntil + _timeToIncrease;
            if (lockUntil < block.timestamp) {
                revert StakeManager__InvalidLockTime();
            }

            deltaTime = lockUntil - block.timestamp;
            if (deltaTime < MIN_LOCKUP_PERIOD || deltaTime > MAX_LOCKUP_PERIOD) {
                revert StakeManager__InvalidLockTime();
            }
        }
        _mintInitialMP(account, deltaTime, _amount);

        //update storage
        totalSupplyBalance += _amount;
        account.balance += _amount;
        account.lockUntil += _timeToIncrease;
    }

    /**
     * leaves the staking pool and withdraws all funds;
     */
    function unstake(uint256 _amount)
        external
        onlyVault
        onlyAccountInitialized(msg.sender)
        noPendingMigration
        finalizeEpoch
    {
        Account storage account = accounts[msg.sender];
        if (_amount > account.balance) {
            revert StakeManager__InsufficientFunds();
        }
        if (account.lockUntil > block.timestamp) {
            revert StakeManager__FundsLocked();
        }
        _processAccount(account, currentEpoch);

        uint256 reducedMP = Math.mulDiv(_amount, account.currentMP, account.balance);
        uint256 reducedInitialMP = Math.mulDiv(_amount, account.initialMP, account.balance);

        //update storage
        account.balance -= _amount;
        account.initialMP -= reducedInitialMP;
        account.currentMP -= reducedMP;
        totalSupplyBalance -= _amount;
        totalSupplyMP -= reducedMP;
    }

    /**
     * @notice Locks entire balance for more amount of time.
     * @param _timeToIncrease Seconds to increase in locked time. If stake is unlocked, increases from block.timestamp.
     *
     * @dev Reverts when resulting locked time is not in range of [MIN_LOCKUP_PERIOD, MAX_LOCKUP_PERIOD]
     */
    function lock(uint256 _timeToIncrease)
        external
        onlyVault
        onlyAccountInitialized(msg.sender)
        noPendingMigration
        finalizeEpoch
    {
        Account storage account = accounts[msg.sender];
        _processAccount(account, currentEpoch);
        uint256 lockUntil = account.lockUntil;
        uint256 deltaTime;
        if (lockUntil < block.timestamp) {
            lockUntil = block.timestamp + _timeToIncrease;
            deltaTime = _timeToIncrease;
        } else {
            lockUntil += _timeToIncrease;
            deltaTime = lockUntil - block.timestamp;
        }
        if (deltaTime < MIN_LOCKUP_PERIOD || deltaTime > MAX_LOCKUP_PERIOD) {
            revert StakeManager__InvalidLockTime();
        }
        _mintInitialMP(account, _timeToIncrease, 0);
        //update account storage
        account.lockUntil = lockUntil;
    }

    /**
     * @notice Release rewards for current epoch and increase epoch.
     * @dev only executes the prerequisite modifier finalizeEpoch
     */
    function executeEpoch() external noPendingMigration finalizeEpoch {
        return; //see modifier finalizeEpoch
    }

    /**
     * @notice Execute rewards for account until limit has reached
     * @param _vault Referred account
     * @param _limitEpoch Until what epoch it should be executed
     */
    function executeAccount(
        address _vault,
        uint256 _limitEpoch
    )
        external
        onlyAccountInitialized(_vault)
        finalizeEpoch
    {
        _processAccount(accounts[_vault], _limitEpoch);
    }

    /**
     * @notice Enables a contract class to interact with staking functions
     * @param _codehash bytecode hash of contract
     */
    function setVault(bytes32 _codehash) external onlyOwner {
        isVault[_codehash] = true;
    }

    /**
     * @notice starts migration to new StakeManager
     * @param _migration new StakeManager
     */
    function startMigration(StakeManager _migration) external onlyOwner noPendingMigration finalizeEpoch {
        if (_migration == this || address(_migration) == address(0)) {
            revert StakeManager__InvalidMigration();
        }
        migration = _migration;
        stakedToken.transfer(address(migration), epochReward());
        migration.migrationInitialize(currentEpoch, totalSupplyMP, totalSupplyBalance, epochs[currentEpoch].startTime);
    }

    /**
     * @dev Callable automatically from old StakeManager.startMigration(address)
     * @notice Initilizes migration process
     * @param _currentEpoch epoch of old manager
     * @param _totalSupplyMP MP supply on old manager
     * @param _totalSupplyBalance stake supply on old manager
     * @param _epochStartTime epoch start time of old manager
     */
    function migrationInitialize(
        uint256 _currentEpoch,
        uint256 _totalSupplyMP,
        uint256 _totalSupplyBalance,
        uint256 _epochStartTime
    )
        external
        onlyPreviousManager
    {
        if (address(migration) != address(0)) {
            revert StakeManager__PendingMigration();
        }
        if (currentEpoch > 0) {
            revert StakeManager__AlreadyProcessedEpochs();
        }
        currentEpoch = _currentEpoch;
        totalSupplyMP = _totalSupplyMP;
        totalSupplyBalance = _totalSupplyBalance;
        epochs[currentEpoch].startTime = _epochStartTime;
    }

    /**
     * @notice Transfer current epoch funds for migrated manager
     */
    function transferNonPending() external onlyPendingMigration {
        stakedToken.transfer(address(migration), epochReward());
    }

    /**
     * @notice Migrate account to new manager.
     * @param _acceptMigration true if wants to migrate, false if wants to leave
     */
    function migrateTo(bool _acceptMigration)
        external
        onlyVault
        onlyAccountInitialized(msg.sender)
        onlyPendingMigration
        finalizeEpoch
        returns (StakeManager newManager)
    {
        _processAccount(accounts[msg.sender], currentEpoch);
        Account memory account = accounts[msg.sender];
        totalSupplyMP -= account.currentMP;
        totalSupplyBalance -= account.balance;
        delete accounts[msg.sender];
        migration.migrateFrom(msg.sender, _acceptMigration, account);
        return migration;
    }

    /**
     * @dev Only callable from old manager.
     * @notice Migrate account from old manager
     * @param _vault Account address
     * @param _account Account data
     * @param _acceptMigration If account should be stored or its MP/balance supply reduced
     */
    function migrateFrom(address _vault, bool _acceptMigration, Account memory _account) external onlyPreviousManager {
        if (_acceptMigration) {
            accounts[_vault] = _account;
        } else {
            totalSupplyMP -= _account.currentMP;
            totalSupplyBalance -= _account.balance;
        }
    }

    /**
     * @dev Only callable from old manager.
     * @notice Increase total MP from old manager
     * @param _amount amount MP increased on account after migration initialized
     */
    function increaseTotalMP(uint256 _amount) external onlyPreviousManager {
        totalSupplyMP += _amount;
    }

    /**
     * @notice Process account until limit has reached
     * @param account Account to process
     * @param _limitEpoch Until what epoch it should be executed
     */
    function _processAccount(Account storage account, uint256 _limitEpoch) private {
        if (_limitEpoch > currentEpoch) {
            revert StakeManager__InvalidLimitEpoch();
        }
        uint256 userReward;
        uint256 userEpoch = account.epoch;
        uint256 mpDifference = account.currentMP;
        for (Epoch storage iEpoch = epochs[userEpoch]; userEpoch < _limitEpoch; userEpoch++) {
            //mint multiplier points to that epoch
            _mintMP(account, iEpoch.startTime + EPOCH_SIZE, iEpoch);
            uint256 userSupply = account.balance + account.currentMP;
            uint256 userEpochReward = Math.mulDiv(userSupply, iEpoch.epochReward, iEpoch.totalSupply);

            userReward += userEpochReward;
            iEpoch.epochReward -= userEpochReward;
            iEpoch.totalSupply -= userSupply;
        }
        account.epoch = userEpoch;
        if (userReward > 0) {
            pendingReward -= userReward;
            stakedToken.transfer(account.rewardAddress, userReward);
        }
        mpDifference = account.currentMP - mpDifference;
        if (address(migration) != address(0)) {
            migration.increaseTotalMP(mpDifference);
        } else if (userEpoch == currentEpoch) {
            _mintMP(account, block.timestamp, epochs[currentEpoch]);
        }
    }

    /**
     * @notice Mint initial multiplier points for given staking amount and time
     * @dev if amount is greater 0, it increases difference of amount for current remaining lock time
     * @dev if increased lock time, increases difference of total new balance for increased lock time
     * @param account Account to mint multiplier points
     * @param increasedLockTime increased lock time
     * @param amount amount to stake
     */
    function _mintInitialMP(Account storage account, uint256 increasedLockTime, uint256 amount) private {
        uint256 mpToMint;
        if (amount > 0) {
            mpToMint += amount; //initial multiplier points
            if (block.timestamp < account.lockUntil) {
                //increasing balance on locked account?
                //bonus for remaining previously locked time of new balance.
                mpToMint += _getMPToMint(amount, account.lockUntil - block.timestamp);
            }
        }
        if (increasedLockTime > 0) {
            //bonus for increased lock time
            mpToMint += _getMPToMint(account.balance + amount, increasedLockTime);
        }
        //update storage
        totalSupplyMP += mpToMint;
        account.initialMP += mpToMint;
        account.currentMP += mpToMint;
        account.lastMint = block.timestamp;
    }

    /**
     * @notice Mint multiplier points for given account and epoch
     * @param account Account earning multiplier points
     * @param processTime amount of time of multiplier points
     * @param epoch Epoch to increment total supply
     */
    function _mintMP(Account storage account, uint256 processTime, Epoch storage epoch) private {
        uint256 increasedMP = _getMaxMPToMint( //check for MAX_BOOST
            _getMPToMint(account.balance, processTime - account.lastMint),
            account.balance,
            account.initialMP,
            account.currentMP
        );

        //update storage
        account.lastMint = processTime;
        account.currentMP += increasedMP;
        totalSupplyMP += increasedMP;
        epoch.totalSupply += increasedMP;
    }

    /**
     * @notice Calculates maximum multiplier point increase for given balance
     * @param _mpToMint tested value
     * @param _balance balance of account
     * @param _currentMP current multiplier point of the account
     * @param _initialMP initial multiplier point of the account
     * @return _maxToIncrease maximum multiplier point increase
     */
    function _getMaxMPToMint(
        uint256 _mpToMint,
        uint256 _balance,
        uint256 _initialMP,
        uint256 _currentMP
    )
        private
        pure
        returns (uint256 _maxToIncrease)
    {
        // Maximum multiplier point for given balance
        _maxToIncrease = _getMPToMint(_balance, MAX_BOOST * YEAR) + _initialMP;
        if (_mpToMint + _currentMP > _maxToIncrease) {
            //reached cap when increasing MP
            return _maxToIncrease - _currentMP; //how much left to reach cap
        } else {
            //not reached capw hen increasing MP
            return _mpToMint; //just return tested value
        }
    }

    /**
     * @notice Calculates multiplier points to mint for given balance and time
     * @param _balance balance of account
     * @param _deltaTime time difference
     * @return multiplier points to mint
     */
    function _getMPToMint(uint256 _balance, uint256 _deltaTime) private pure returns (uint256) {
        return Math.mulDiv(_balance, _deltaTime, YEAR) * MP_APY;
    }

    /**
     * @notice Returns total of multiplier points and balance
     * @return _totalSupply current total supply
     */
    function totalSupply() public view returns (uint256 _totalSupply) {
        return totalSupplyMP + totalSupplyBalance;
    }

    /**
     * @notice Returns funds available for current epoch
     * @return _epochReward current epoch reward
     */
    function epochReward() public view returns (uint256 _epochReward) {
        return stakedToken.balanceOf(address(this)) - pendingReward;
    }

    /**
     * @notice Returns end time of current epoch
     * @return _epochEnd end time of current epoch
     */
    function epochEnd() public view returns (uint256 _epochEnd) {
        return epochs[currentEpoch].startTime + EPOCH_SIZE;
    }
}