ghost-dao-contracts/src/Staking.sol

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.20;

import "@openzeppelin-contracts/token/ERC20/utils/SafeERC20.sol";
import "@openzeppelin-contracts/token/ERC20/IERC20.sol";

import "./types/GhostAccessControlled.sol";

import "./interfaces/ISTNK.sol";
import "./interfaces/IGHST.sol";
import "./interfaces/IDistributor.sol";
import "./interfaces/IStaking.sol";

contract GhostStaking is IStaking, GhostAccessControlled {
    using SafeERC20 for IERC20;
    using SafeERC20 for ISTNK;
    using SafeERC20 for IGHST;

    address public immutable ftso;
    address public immutable stnk;
    address public immutable ghst;

    uint48 public warmupPeriod;

    Epoch public epoch;
    address public distributor;
    uint256 public sharesInWarmup;

    mapping(address => Claim) public warmupInfo;

    constructor(
        address _ftso,
        address _stnk,
        address _ghst,
        uint48 _epochLength,
        uint48 _firstEpochNumber,
        uint48 _firstEpochTime,
        address _authority
    ) GhostAccessControlled(IGhostAuthority(_authority)) {
        ftso = _ftso;
        stnk = _stnk;
        ghst = _ghst;

        epoch = Epoch({
            length: _epochLength,
            number: _firstEpochNumber,
            end: _firstEpochTime,
            distribute: 0
        });
    }

    function stake(
        uint256 amount,
        address to,
        bool isRebase,
        bool isClaim
    ) external override returns (uint256 returnAmount) {
        returnAmount = amount + rebase();
        IERC20(ftso).safeTransferFrom(msg.sender, address(this), amount);
        if (isClaim && warmupPeriod == 0) {
            returnAmount = _send(returnAmount, to, isRebase);
        } else {
            Claim storage info = warmupInfo[to];
            if (info.lock && to != msg.sender) revert ExternalDepositsLocked();

            info.deposit += returnAmount;
            info.shares += ISTNK(stnk).sharesForBalance(returnAmount);
            info.expiry = epoch.number + warmupPeriod;

            sharesInWarmup += ISTNK(stnk).sharesForBalance(returnAmount);
        }
    }

    function claim(address to, bool isRebase) public override returns (uint256 claimedAmount) {
        Claim memory info = warmupInfo[to];
        if (info.lock && to != msg.sender) revert ExternalDepositsLocked();

        if (epoch.number >= info.expiry && info.expiry > 0) {
            delete warmupInfo[to];
            sharesInWarmup -= info.shares;
            claimedAmount = _send(ISTNK(stnk).balanceForShares(info.shares), to, isRebase);
        }
    }

    function forfeit() external override returns (uint256) {
        Claim memory info = warmupInfo[msg.sender];
        delete warmupInfo[msg.sender];

        sharesInWarmup -= info.shares;

        IERC20(ftso).safeTransfer(msg.sender, info.deposit);
        return info.deposit;
    }

    function toggleLock() external override {
        warmupInfo[msg.sender].lock = !warmupInfo[msg.sender].lock;
    }

    function unstake(
        uint256 amount,
        address to,
        bool isTrigger,
        bool isRebase
    ) external override returns (uint256) {
        amount += isTrigger ? rebase() : 0;
        if (isRebase) {
            ISTNK(stnk).safeTransferFrom(msg.sender, address(this), amount);
        } else {
            IGHST(ghst).burn(msg.sender, amount);
            amount = IGHST(ghst).balanceFrom(amount);
        }

        if (amount > IERC20(ftso).balanceOf(address(this))) revert InsufficientBalance();
        IERC20(ftso).safeTransfer(to, amount);
        return amount;
    }

    function wrap(
        address to,
        uint256 amount
    ) external override returns (uint256 balance) {
        ISTNK(stnk).safeTransferFrom(msg.sender, address(this), amount);
        balance = IGHST(ghst).balanceTo(amount);
        IGHST(ghst).mint(to, balance);
    }

    function unwrap(
        address to,
        uint256 amount
    ) external override returns (uint256 balance) {
        IGHST(ghst).burn(msg.sender, amount);
        balance = IGHST(ghst).balanceFrom(amount);
        ISTNK(stnk).safeTransfer(to, balance);
    }

    function rebase() public override returns (uint256 bounty) {
        if (epoch.end <= block.timestamp) {
            ISTNK(stnk).rebase(epoch.distribute, epoch.number);

            unchecked {
                epoch.end += epoch.length;
                ++epoch.number;
            }

            if (distributor != address(0)) {
                IDistributor(distributor).distribute();
                bounty = IDistributor(distributor).retrieveBounty();
            }

            uint256 balance = IERC20(ftso).balanceOf(address(this));
            uint256 extra = ISTNK(stnk).circulatingSupply() + bounty;

            epoch.distribute = balance > extra ? balance - extra : 0;
        }
    }

    function setDistributor(address _distributor) external onlyGovernor {
        distributor = _distributor;
        emit DistributorSet(_distributor);
    }

    function setWarmupPeriod(uint256 _warmupPeriod) external onlyGovernor {
        warmupPeriod = uint48(_warmupPeriod);
        emit WarmupSet(_warmupPeriod);
    }

    function index() public view override returns (uint256) {
        return ISTNK(stnk).index();
    }

    function supplyInWarmup() public view override returns (uint256) {
        return ISTNK(stnk).balanceForShares(sharesInWarmup);
    }

    function _send(
        uint256 amount,
        address to,
        bool isRebase
    ) internal returns (uint256) {
        if (isRebase) {
            ISTNK(stnk).safeTransfer(to, amount);
            return amount;
        } else {
            uint256 balanceTo = IGHST(ghst).balanceTo(amount);
            IGHST(ghst).mint(to, balanceTo);
            return balanceTo;
        }
    }
}