move warmup logic outside of staking contract

Signed-off-by: Uncle Fatso <uncle.fatso@ghostchain.io>

src/Staking.sol

@@ -4,6 +4,7 @@ pragma solidity ^0.8.20;
 import {SafeERC20} from "@openzeppelin-contracts/token/ERC20/utils/SafeERC20.sol";
 import {IERC20} from "@openzeppelin-contracts/token/ERC20/IERC20.sol";
 
+import {GhostWarmup} from "./Warmup.sol";
 import {GhostAccessControlled} from "./types/GhostAccessControlled.sol";
 
 import {ISTNK} from "./interfaces/ISTNK.sol";
@@ -12,6 +13,7 @@ import {IDistributor} from "./interfaces/IDistributor.sol";
 import {IStaking} from "./interfaces/IStaking.sol";
 import {IGatekeeper} from "./interfaces/IGatekeeper.sol";
 import {IGhostAuthority} from "./interfaces/IGhostAuthority.sol";
+import {IGhostWarmup} from "./interfaces/IGhostWarmup.sol";
 
 contract GhostStaking is IStaking, GhostAccessControlled {
     using SafeERC20 for IERC20;
@@ -27,9 +29,9 @@ contract GhostStaking is IStaking, GhostAccessControlled {
     Epoch public epoch;
     address public distributor;
     address public gatekeeper;
-    uint256 public sharesInWarmup;
+    address public warmup;
 
-    mapping(address => Claim) public warmupInfo;
+    mapping(address => bool) public locks;
 
     constructor(
         address _ftso,
@@ -63,40 +65,35 @@ contract GhostStaking is IStaking, GhostAccessControlled {
         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);
+            if (locks[to] && to != msg.sender) revert ExternalDepositsLocked();
+            uint48 expiry = epoch.number + warmupPeriod;
+            IGhostWarmup(warmup).addToWarmup(amount, to, expiry);
         }
     }
 
     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);
+        if (locks[to] && to != msg.sender) revert ExternalDepositsLocked();
+        claimedAmount = IGhostWarmup(warmup).claim(to, epoch.number);
+        if (isRebase) {
+            claimedAmount = IGHST(ghst).balanceFrom(claimedAmount);
+            ISTNK(stnk).safeTransfer(to, claimedAmount);
+        } else {
+            IGHST(ghst).mint(to, claimedAmount);
         }
     }
 
-    function forfeit() external override returns (uint256) {
-        Claim memory info = warmupInfo[msg.sender];
-        delete warmupInfo[msg.sender];
+    function breakout(bytes32 receiver, uint256 amount) public override {
+        IGhostWarmup(warmup).breakout(msg.sender, amount);
+        IGatekeeper(gatekeeper).ghost(receiver, amount);
+    }
 
-        sharesInWarmup -= info.shares;
-
-        IERC20(ftso).safeTransfer(msg.sender, info.deposit);
-        return info.deposit;
+    function forfeit() external override returns (uint256 deposit) {
+        deposit = IGhostWarmup(warmup).forfeit(msg.sender);
+        IERC20(ftso).safeTransfer(msg.sender, deposit);
     }
 
     function toggleLock() external override {
-        warmupInfo[msg.sender].lock = !warmupInfo[msg.sender].lock;
+        locks[msg.sender] = !locks[msg.sender];
     }
 
     function unstake(
@@ -183,6 +180,10 @@ contract GhostStaking is IStaking, GhostAccessControlled {
     function setWarmupPeriod(uint256 _warmupPeriod) external onlyGovernor {
         // forge-lint: disable-next-line(unsafe-typecast)
         warmupPeriod = uint48(_warmupPeriod);
+        if (warmup == address(0)) {
+            GhostWarmup newWarmup = new GhostWarmup(ghst);
+            warmup = address(newWarmup);
+        }
         emit WarmupSet(_warmupPeriod);
     }
 
@@ -196,7 +197,10 @@ contract GhostStaking is IStaking, GhostAccessControlled {
     }
 
     function supplyInWarmup() public view override returns (uint256) {
-        return ISTNK(stnk).balanceForShares(sharesInWarmup);
+        if (warmup == address(0)) {
+            return 0;
+        }
+        return IGHST(ghst).balanceFrom(IGhostWarmup(warmup).ghstInWarmup());
     }
 
     function ghostedSupply() public view override returns (uint256 amount) {
@@ -205,6 +209,12 @@ contract GhostStaking is IStaking, GhostAccessControlled {
         }
     }
 
+    function warmupInfo(address who) external view returns (uint256, uint256, uint48, bool) {
+        bool lock = locks[who];
+        (uint256 deposit, uint256 payout, uint48 expiry) = IGhostWarmup(warmup).warmupInfo(who);
+        return (deposit, payout, expiry, lock);
+    }
+
     function _send(
         uint256 amount,
         address to,

src/Warmup.sol

@@ -0,0 +1,80 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.20;
+
+import {IGhostWarmup} from "./interfaces/IGhostWarmup.sol";
+import {IGHST} from "./interfaces/IGHST.sol";
+
+import {FullMath} from "./libraries/FullMath.sol";
+
+contract GhostWarmup is IGhostWarmup {
+    address public immutable STAKING;
+    address public immutable GHST;
+
+    uint256 private _ghstInWarmup;
+    mapping(address => Claim) private _warmupInfo;
+
+    constructor(address ghst) {
+        STAKING = msg.sender;
+        GHST = ghst;
+    }
+
+    function addToWarmup(
+        uint256 payout,
+        address who,
+        uint48 expiry
+    ) external override {
+        if (msg.sender != STAKING) revert NotStakingContract();
+
+        Claim storage info = _warmupInfo[who];
+        uint256 ghstPayout = IGHST(GHST).balanceTo(payout);
+
+        info.deposit += payout;
+        info.payout += ghstPayout;
+        info.expiry = expiry;
+
+        _ghstInWarmup += ghstPayout;
+    }
+
+    function claim(
+        address to,
+        uint256 epochNumber
+    ) external override returns (uint256 claimedAmount) {
+        if (msg.sender != STAKING) revert NotStakingContract();
+        Claim memory info = _warmupInfo[to];
+
+        if (epochNumber >= info.expiry && info.expiry > 0) {
+            delete _warmupInfo[to];
+            _ghstInWarmup -= info.payout;
+            claimedAmount = info.payout;
+        }
+    }
+
+    function breakout(address who, uint256 payout) external override {
+        if (msg.sender != STAKING) revert NotStakingContract();
+
+        Claim storage info = _warmupInfo[who];
+        uint256 depositReduction = FullMath.mulDiv(info.deposit, payout, info.payout);
+
+        info.deposit -= depositReduction;
+        info.payout -= payout;
+        _ghstInWarmup -= payout;
+    }
+
+    function forfeit(address who) external override returns (uint256) {
+        if (msg.sender != STAKING) revert NotStakingContract();
+
+        Claim memory info = _warmupInfo[who];
+        delete _warmupInfo[who];
+        _ghstInWarmup -= info.payout;
+        return info.deposit;
+    }
+
+    function ghstInWarmup() external view override returns (uint256) {
+        return _ghstInWarmup;
+    }
+
+    function warmupInfo(address who) external view override returns (uint256, uint256, uint48) {
+        Claim memory info = _warmupInfo[who];
+        return (info.deposit, info.payout, info.expiry);
+    }
+}

src/interfaces/IGhostWarmup.sol

@@ -0,0 +1,20 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.20;
+
+interface IGhostWarmup {
+    struct Claim {
+        uint256 deposit;
+        uint256 payout;
+        uint48 expiry;
+    }
+
+    error NotStakingContract();
+    error ExternalDepositsLocked();
+
+    function addToWarmup(uint256 payout, address who, uint48 expiry) external;
+    function claim(address who, uint256 epochNumber) external returns (uint256);
+    function breakout(address who, uint256 amount) external;
+    function forfeit(address who) external returns (uint256);
+    function ghstInWarmup() external view returns (uint256);
+    function warmupInfo(address who) external view returns (uint256, uint256, uint48);
+}

src/interfaces/INoteKeeper.sol

@@ -21,7 +21,13 @@ interface INoteKeeper {
         uint256[] memory _indexes
     ) external returns (uint256);
 
+    function forceRedeem(
+        bytes32 _receiver,
+        uint256[] memory _indexes
+    ) external returns (uint256);
+
     function redeemAll(address _user, bool _sendGhst) external returns (uint256);
+    function forceRedeemAll(bytes32 _receiver) external returns (uint256);
     function pushNote(address to, uint256 index) external;
     function pullNote(address from, uint256 index) external returns (uint256 newIndex_);
     function indexesFor(address _user) external view returns (uint256[] memory);

src/interfaces/IStaking.sol

@@ -32,6 +32,7 @@ interface IStaking {
     ) external returns (uint256);
 
     function claim(address _recipient, bool _rebasing) external returns (uint256);
+    function breakout(bytes32 _receiver, uint256 _amount) external;
     function forfeit() external returns (uint256);
     function toggleLock() external;
 

src/types/NoteKeeper.sol

@@ -93,6 +93,27 @@ abstract contract NoteKeeper is INoteKeeper, FrontEndRewarder {
         else _STAKING.unwrap(user, payout);
     }
 
+    function forceRedeem(
+        bytes32 receiver,
+        uint256[] memory indexes
+    ) public override returns (uint256 payout) {
+        address user = msg.sender;
+        uint48 time = uint48(block.timestamp);
+        uint256 i;
+
+        for (; i < indexes.length; ) {
+            (uint256 pay, bool matured) = pendingFor(user, indexes[i]);
+            if (matured) {
+                _pendingIndexes[user].remove(indexes[i]);
+                notes[user][indexes[i]].redeemed = time;
+                payout += pay;
+            }
+            unchecked { ++i; }
+        }
+
+        _STAKING.breakout(receiver, payout);
+    }
+
     function redeemAll(
         address user,
         bool sendGhst
@@ -100,6 +121,10 @@ abstract contract NoteKeeper is INoteKeeper, FrontEndRewarder {
         return redeem(user, sendGhst, indexesFor(user));
     }
 
+    function forceRedeemAll(bytes32 receiver) external override returns (uint256) {
+        return forceRedeem(receiver, indexesFor(msg.sender));
+    }
+
     function pushNote(address to, uint256 index) external override {
         if (notes[msg.sender][index].created == 0) revert NoteNotFound(msg.sender, index);
         _noteTransfers[msg.sender][index] = to;