refactor(StakeManager): use custom errors everywhere (#30)

.gas-snapshot

@@ -1,7 +1,11 @@
+ExecuteAccountTest:testDeployment() (gas: 26400)
+ExecuteAccountTest:test_RevertWhen_InvalidLimitEpoch() (gas: 974502)
 LeaveTest:testDeployment() (gas: 26172)
+LeaveTest:test_RevertWhen_NoPendingMigration() (gas: 662520)
 LeaveTest:test_RevertWhen_SenderIsNotVault() (gas: 10562)
-LockTest:testDeployment() (gas: 26172)
+LockTest:testDeployment() (gas: 26400)
-LockTest:test_RevertWhen_SenderIsNotVault() (gas: 10573)
+LockTest:test_RevertWhen_DecreasingLockTime() (gas: 976972)
+LockTest:test_RevertWhen_SenderIsNotVault() (gas: 10607)
 MigrateTest:testDeployment() (gas: 26172)
 MigrateTest:test_RevertWhen_SenderIsNotVault() (gas: 10629)
 StakeManagerTest:testDeployment() (gas: 26172)

contracts/StakeManager.sol

@@ -9,6 +9,11 @@ import { StakeVault } from "./StakeVault.sol";
 contract StakeManager is Ownable {
     error StakeManager__SenderIsNotVault();
     error StakeManager__FundsLocked();
+    error StakeManager__DecreasingLockTime();
+    error StakeManager__NoPendingMigration();
+    error StakeManager__PendingMigration();
+    error StakeManager__SenderIsNotPreviousStakeManager();
+    error StakeManager__InvalidLimitEpoch();
 
     struct Account {
         uint256 lockUntil;
@@ -93,7 +98,9 @@ contract StakeManager is Ownable {
     function lock(uint256 _time) external onlyVault {
         Account storage account = accounts[msg.sender];
         processAccount(account, currentEpoch);
-        require(block.timestamp + _time > account.lockUntil, "Cannot decrease lock time");
+        if (block.timestamp + _time < account.lockUntil) {
+            revert StakeManager__DecreasingLockTime();
+        }
         mintIntialMultiplier(account, _time, account.balance, 0);
     }
 
@@ -101,7 +108,9 @@ contract StakeManager is Ownable {
      * @notice leave without processing account
      */
     function leave() external onlyVault {
-        require(address(migration) != address(0), "Leave only during migration");
+        if (address(migration) == address(0)) {
+            revert StakeManager__NoPendingMigration();
+        }
         Account memory account = accounts[msg.sender];
         delete accounts[msg.sender];
         multiplierSupply -= account.multiplier;
@@ -136,7 +145,9 @@ contract StakeManager is Ownable {
      */
 
     function migrate() external onlyVault returns (StakeManager newManager) {
-        require(address(migration) != address(0), "Migration not available");
+        if (address(migration) == address(0)) {
+            revert StakeManager__NoPendingMigration();
+        }
         Account storage account = accounts[msg.sender];
         stakedToken.approve(address(migration), account.balance);
         migration.migrate(msg.sender, account);
@@ -151,7 +162,9 @@ contract StakeManager is Ownable {
      * @param _account Account data
      */
     function migrate(address _vault, Account memory _account) external {
-        require(msg.sender == address(oldManager), "Unauthorized");
+        if (msg.sender != address(oldManager)) {
+            revert StakeManager__SenderIsNotPreviousStakeManager();
+        }
         stakedToken.transferFrom(address(oldManager), address(this), _account.balance);
         accounts[_vault] = _account;
     }
@@ -190,8 +203,12 @@ contract StakeManager is Ownable {
 
     function processAccount(Account storage account, uint256 _limitEpoch) private {
         processEpoch();
-        require(address(migration) == address(0), "Contract ended, please migrate");
+        if (address(migration) != address(0)) {
-        require(_limitEpoch <= currentEpoch, "Non-sense call");
+            revert StakeManager__PendingMigration();
+        }
+        if (_limitEpoch > currentEpoch) {
+            revert StakeManager__InvalidLimitEpoch();
+        }
         uint256 userReward;
         uint256 userEpoch = account.epoch;
         for (Epoch memory iEpoch = epochs[userEpoch]; userEpoch < _limitEpoch; userEpoch++) {

test/StakeManager.t.sol

@@ -89,6 +89,22 @@ contract LockTest is StakeManagerTest {
         vm.expectRevert(StakeManager.StakeManager__SenderIsNotVault.selector);
         stakeManager.lock(100);
     }
+
+    function test_RevertWhen_DecreasingLockTime() public {
+        // ensure user has funds
+        deal(stakeToken, testUser, 1000);
+        StakeVault userVault = _createTestVault(testUser);
+
+        vm.startPrank(testUser);
+        // ensure user vault can spend user tokens
+        ERC20(stakeToken).approve(address(userVault), 100);
+
+        uint256 lockTime = 1 days;
+        userVault.stake(100, lockTime);
+
+        vm.expectRevert(StakeManager.StakeManager__DecreasingLockTime.selector);
+        userVault.lock(1);
+    }
 }
 
 contract LeaveTest is StakeManagerTest {
@@ -100,6 +116,13 @@ contract LeaveTest is StakeManagerTest {
         vm.expectRevert(StakeManager.StakeManager__SenderIsNotVault.selector);
         stakeManager.leave();
     }
+
+    function test_RevertWhen_NoPendingMigration() public {
+        StakeVault userVault = _createTestVault(testUser);
+        vm.prank(testUser);
+        vm.expectRevert(StakeManager.StakeManager__NoPendingMigration.selector);
+        userVault.leave();
+    }
 }
 
 contract MigrateTest is StakeManagerTest {
@@ -111,4 +134,33 @@ contract MigrateTest is StakeManagerTest {
         vm.expectRevert(StakeManager.StakeManager__SenderIsNotVault.selector);
         stakeManager.migrate();
     }
+
+    function test_RevertWhen_NoPendingMigration() public {
+        StakeVault userVault = _createTestVault(testUser);
+        vm.prank(testUser);
+        vm.expectRevert(StakeManager.StakeManager__NoPendingMigration.selector);
+        userVault.updateManager();
+    }
+}
+
+contract ExecuteAccountTest is StakeManagerTest {
+    function setUp() public override {
+        StakeManagerTest.setUp();
+    }
+
+    function test_RevertWhen_InvalidLimitEpoch() public {
+        deal(stakeToken, testUser, 1000);
+        StakeVault userVault = _createTestVault(testUser);
+
+        vm.startPrank(testUser);
+        ERC20(stakeToken).approve(address(userVault), 100);
+
+        uint256 lockTime = 1 days;
+        userVault.stake(100, lockTime);
+
+        uint256 currentEpoch = stakeManager.currentEpoch();
+
+        vm.expectRevert(StakeManager.StakeManager__InvalidLimitEpoch.selector);
+        stakeManager.executeAccount(address(userVault), currentEpoch + 1);
+    }
 }