FIX: make breakout and partial claim automatically

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

src/Staking.sol

@@ -66,7 +66,7 @@ contract GhostStaking is IStaking, GhostAccessControlled {
         returnAmount = amount + rebase();
         IERC20(ftso).safeTransferFrom(msg.sender, address(this), amount);
         if (isClaim && warmupPeriod == 0) {
-            returnAmount = _send(returnAmount, to, isRebase);
+            returnAmount = _sendStnkBased(returnAmount, to, isRebase);
         } else {
             if (locks[to] && to != msg.sender) revert ExternalDepositsLocked();
             uint48 expiry = epoch.number + warmupPeriod;
@@ -75,20 +75,28 @@ contract GhostStaking is IStaking, GhostAccessControlled {
         emit Staked(msg.sender, to, amount, isRebase, isClaim);
     }
 
-    function claim(address to, bool isRebase) public override returns (uint256 claimedAmount) {
+    function claim(address to, bool isRebase) public override returns (uint256) {
         if (locks[to] && to != msg.sender) revert ExternalDepositsLocked();
-        claimedAmount = IGhostWarmup(warmup).claim(to, epoch.number);
+        uint256 claimedAmount = IGhostWarmup(warmup).claim(to, epoch.number);
-        if (isRebase) {
+        return _sendGhstBased(claimedAmount, to, isRebase);
-            claimedAmount = IGHST(ghst).balanceFrom(claimedAmount);
-            ISTNK(stnk).safeTransfer(to, claimedAmount);
-        } else {
-            IGHST(ghst).mint(to, claimedAmount);
-        }
     }
 
-    function breakout(bytes32 receiver, uint256 amount) public override {
+    function claimByAmount(
-        IGhostWarmup(warmup).breakout(msg.sender, amount);
+        address to,
-        IGatekeeper(gatekeeper).ghost(receiver, amount);
+        uint256 amount,
+        bool isRebase
+    ) public override returns (uint256) {
+        if (locks[to] && to != msg.sender) revert ExternalDepositsLocked();
+        uint256 claimedAmount = IGhostWarmup(warmup).claimByAmount(to, amount, epoch.number);
+        return _sendGhstBased(claimedAmount, to, isRebase);
+    }
+
+    function breakout(
+        bytes32 receiver,
+        uint256 amount
+    ) public override returns (uint256 claimedAmount) {
+        claimedAmount = IGhostWarmup(warmup).breakout(msg.sender, amount);
+        IGatekeeper(gatekeeper).ghost(receiver, claimedAmount);
     }
 
     function forfeit() external override returns (uint256 deposit) {
@@ -213,7 +221,7 @@ contract GhostStaking is IStaking, GhostAccessControlled {
         return (deposit, payout, expiry, lock);
     }
 
-    function _send(
+    function _sendStnkBased(
         uint256 amount,
         address to,
         bool isRebase
@@ -227,4 +235,19 @@ contract GhostStaking is IStaking, GhostAccessControlled {
             return balanceTo;
         }
     }
+
+    function _sendGhstBased(
+        uint256 amount,
+        address to,
+        bool isRebase
+    ) internal returns (uint256) {
+        if (isRebase) {
+            uint256 convertedAmount = IGHST(ghst).balanceFrom(amount);
+            ISTNK(stnk).safeTransfer(to, convertedAmount);
+            return convertedAmount;
+        } else {
+            IGHST(ghst).mint(to, amount);
+            return amount;
+        }
+    }
 }

src/Warmup.sol

@@ -49,17 +49,21 @@ contract GhostWarmup is IGhostWarmup {
         }
     }
 
-    function breakout(address who, uint256 payout) external override {
+    function claimByAmount(
+        address who,
+        uint256 payout,
+        uint256 epochNumber
+    ) external override returns (uint256) {
         if (msg.sender != STAKING) revert NotStakingContract();
+        return _reduceWarmupInfo(who, payout, epochNumber);
+    }
 
-        Claim storage info = _warmupInfo[who];
+    function breakout(
-        uint256 mm = mulmod(info.deposit, payout, info.payout);
+        address who,
-        uint256 depositReduction = FullMath.mulDiv(info.deposit, payout, info.payout);
+        uint256 payout
-        if (mm > 0) depositReduction += 1;
+    ) external override returns (uint256) {
-
+        if (msg.sender != STAKING) revert NotStakingContract();
-        info.deposit -= depositReduction;
+        return _reduceWarmupInfo(who, payout, type(uint256).max);
-        info.payout -= payout;
-        _ghstInWarmup -= payout;
     }
 
     function forfeit(address who) external override returns (uint256) {
@@ -79,4 +83,28 @@ contract GhostWarmup is IGhostWarmup {
         Claim memory info = _warmupInfo[who];
         return (info.deposit, info.payout, info.expiry);
     }
+
+    function _reduceWarmupInfo(
+        address who,
+        uint256 amount,
+        uint256 epochNumber
+    ) private returns (uint256) {
+        Claim storage info = _warmupInfo[who];
+
+        if (epochNumber >= info.expiry && info.expiry > 0) {
+            uint256 mm = mulmod(info.deposit, amount, info.payout);
+            uint256 depositReduction = FullMath.mulDiv(info.deposit, amount, info.payout);
+            if (mm > 0) depositReduction += 1;
+
+            info.deposit -= depositReduction;
+            info.payout -= amount;
+            _ghstInWarmup -= amount;
+
+            if (info.payout == 0) delete _warmupInfo[who];
+
+            return amount;
+        }
+
+        revert LockedInWarmupPeriod();
+    }
 }

src/interfaces/IGhostWarmup.sol

@@ -10,10 +10,16 @@ interface IGhostWarmup {
 
     error NotStakingContract();
     error ExternalDepositsLocked();
+    error LockedInWarmupPeriod();
 
     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 claimByAmount(
+        address who,
+        uint256 amount,
+        uint256 epochNumber
+    ) external returns (uint256);
+    function breakout(address who, uint256 amount) external returns (uint256);
     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

@@ -6,6 +6,7 @@ interface INoteKeeper {
     error NoteNotFound(address from, uint256 index);
     error TransferNotFound(address from, uint256 index);
     error AlreadyRedeemed(address from, uint256 index);
+    error IncompleteRedeemPayout();
 
     struct Note {
         uint256 payout;

src/interfaces/IStaking.sol

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

src/types/NoteKeeper.sol

@@ -74,20 +74,15 @@ abstract contract NoteKeeper is INoteKeeper, FrontEndRewarder {
         bool sendGhst,
         uint256[] memory indexes
     ) public override returns (uint256 payout) {
-        _STAKING.claim(address(this), false);
+        uint256 expected = _maturedPayout(msg.sender, indexes);
+        uint256 balance = _GHST.balanceOf(address(this));
 
-        uint48 time = uint48(block.timestamp);
+        if (balance < expected) {
-        uint256 i;
+            uint256 deficit = expected - balance;
+            payout = balance + _STAKING.claimByAmount(address(this), deficit, false);
+        } else payout = expected;
 
-        for (; i < indexes.length; ) {
+        if (payout < expected) revert IncompleteRedeemPayout();
-            (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; }
-        }
 
         if (sendGhst) _GHST.safeTransfer(user, payout);
         else _STAKING.unwrap(user, payout);
@@ -96,22 +91,22 @@ abstract contract NoteKeeper is INoteKeeper, FrontEndRewarder {
     function forceRedeem(
         bytes32 receiver,
         uint256[] memory indexes
-    ) public override returns (uint256 payout) {
+    ) public override returns (uint256 expected) {
-        address user = msg.sender;
+        uint256 payout = _maturedPayout(msg.sender, indexes);
-        uint48 time = uint48(block.timestamp);
+        uint256 balance = _GHST.balanceOf(address(this));
-        uint256 i;
+        expected = payout;
 
-        for (; i < indexes.length; ) {
+        if (balance > 0) {
-            (uint256 pay, bool matured) = pendingFor(user, indexes[i]);
+            uint256 toGhost = payout > balance ? balance : payout;
-            if (matured) {
+            _STAKING.ghost(receiver, toGhost);
-                _pendingIndexes[user].remove(indexes[i]);
+            payout -= toGhost;
-                notes[user][indexes[i]].redeemed = time;
-                payout += pay;
-            }
-            unchecked { ++i; }
         }
 
-        _STAKING.breakout(receiver, payout);
+        if (payout > 0) {
+            payout -= _STAKING.breakout(receiver, payout);
+        }
+
+        if (payout > 0) revert IncompleteRedeemPayout();
     }
 
     function redeemAll(
@@ -167,4 +162,18 @@ abstract contract NoteKeeper is INoteKeeper, FrontEndRewarder {
         payout = note.payout;
         matured = _pendingIndexes[user].contains(index) && note.matured <= block.timestamp && payout > 0;
     }
+
+    function _maturedPayout(address user, uint256[] memory indexes) private returns (uint256 payout) {
+        uint48 time = uint48(block.timestamp);
+        uint256 i = 0;
+        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; }
+        }
+    }
 }