checkpoint

src/BaseStaking.sol

@@ -207,7 +207,6 @@ abstract contract BaseStaking is OwnableUpgradeable, ERC20VotesUpgradeable {
     //////////////////////////////////////////////////////////////*/
 
     /// @notice Deposit SHU into the contract
-    /// @param user The user address
     /// @param amount The amount of SHU to deposit
     function _deposit(uint256 amount) internal {
         // Calculate the amount of shares to mint

src/Staking.sol

@@ -242,20 +242,8 @@ contract Staking is BaseStaking {
                 );
             }
 
-            uint256 maxWithdraw = keyperStake.amount - minStake;
-            require(amount <= maxWithdraw, WithdrawAmountTooHigh());
-
-            // The unstake can't never result in a keyper SHU staked < minStake
-            require(
-                _maxWithdraw(keyper, keyperStake.amount) >= amount,
-                WithdrawAmountTooHigh()
-            );
-        } else {
-            // doesn't include the min stake and locked staked as the keyper is not a keyper anymore
-            require(
-                convertToAssets(balanceOf(keyper)) >= amount,
-                WithdrawAmountTooHigh()
-            );
+            uint256 maxWithdrawAvailable = keyperStake.amount - minStake;
+            require(amount <= maxWithdrawAvailable, WithdrawAmountTooHigh());
         }
 
         // Decrease the amount from the stake
@@ -307,31 +295,12 @@ contract Staking is BaseStaking {
     /// @return amount The maximum amount of assets that a keyper can withdraw
     function maxWithdraw(
         address keyper
-    ) public view override returns (uint256) {
-        return _maxWithdraw(keyper, 0);
-    }
-
-    /*//////////////////////////////////////////////////////////////
-                          INTERNAL FUNCTIONS
-    //////////////////////////////////////////////////////////////*/
-
-    /// @notice Get the maximum amount of assets that a keyper can withdraw
-    ///         after unlocking a certain amount
-    ///         - if the keyper has no shares, the function will revert
-    ///         - if the keyper sSHU balance is less or equal than the minimum
-    ///         stake or the total locked amount, the function will return 0
-    /// @param keyper The keyper address
-    /// @param unlockedAmount The amount of unlocked assets
-    /// @return amount The maximum amount of assets that a keyper can withdraw after unlocking a certain amount
-    function _maxWithdraw(
-        address keyper,
-        uint256 unlockedAmount
-    ) internal view virtual returns (uint256 amount) {
+    ) public view override returns (uint256 amount) {
         uint256 assets = convertToAssets(balanceOf(keyper));
         require(assets > 0, UserHasNoShares());
 
         unchecked {
-            uint256 locked = totalLocked[keyper] - unlockedAmount;
+            uint256 locked = totalLocked[keyper];
             uint256 compare = locked >= minStake ? locked : minStake;
 
             // need the first branch as convertToAssets rounds down

test/RewardsDistributor.t.sol

@@ -224,7 +224,11 @@ contract OwnableFunctions is RewardsDistributorTest {
 
         uint256 balanceBefore = govToken.balanceOf(_to);
 
-        rewardsDistributor.withdrawFunds(_to, _amount);
+        rewardsDistributor.withdrawFunds(
+            address(rewardsDistributor.rewardToken()),
+            _to,
+            _amount
+        );
 
         assertEq(govToken.balanceOf(_to), balanceBefore + _amount);
     }
@@ -243,7 +247,11 @@ contract OwnableFunctions is RewardsDistributorTest {
             )
         );
         vm.prank(_anyone);
-        rewardsDistributor.withdrawFunds(_to, _amount);
+        rewardsDistributor.withdrawFunds(
+            address(rewardsDistributor.rewardToken()),
+            _to,
+            _amount
+        );
     }
 }
 

test/Staking.t.sol

@@ -712,9 +712,8 @@ contract Stake is StakingTest {
         );
 
         // bob unstake maximum he can unstake
-        uint256 maxBobCanWithdraw = staking.exposed_maxWithdraw(bob, bobStake);
         vm.prank(bob);
-        staking.unstake(bob, bobStakeId, maxBobCanWithdraw);
+        staking.unstake(bob, bobStakeId, bobStake);
 
         uint256 bobBalance = govToken.balanceOf(bob);
 

test/helpers/StakingHarness.sol

@@ -8,13 +8,6 @@ contract StakingHarness is Staking {
         return nextStakeId;
     }
 
-    function exposed_maxWithdraw(
-        address keyper,
-        uint256 unlockedAmount
-    ) external view virtual returns (uint256) {
-        return _maxWithdraw(keyper, unlockedAmount);
-    }
-
     function exposed_previewWithdraw(
         uint256 amount
     ) external view returns (uint256) {