iEpoch.totalSupply underflow, migration fixes, MP formula fix

.gas-snapshot

@@ -1,31 +1,45 @@
 CreateVaultTest:testDeployment() (gas: 9774)
-CreateVaultTest:test_createVault() (gas: 719374)
+CreateVaultTest:test_createVault() (gas: 692913)
-ExecuteAccountTest:testDeployment() (gas: 26401)
+ExecuteAccountTest:testDeployment() (gas: 26495)
-ExecuteAccountTest:test_RevertWhen_InvalidLimitEpoch() (gas: 1098349)
+ExecuteAccountTest:test_ExecuteAccountMintMP() (gas: 4504158)
-LeaveTest:testDeployment() (gas: 26173)
+ExecuteAccountTest:test_MintMPLimit() (gas: 209)
-LeaveTest:test_RevertWhen_NoPendingMigration() (gas: 752180)
+ExecuteAccountTest:test_PayRewards() (gas: 165)
-LeaveTest:test_RevertWhen_SenderIsNotVault() (gas: 10770)
+ExecuteAccountTest:test_RevertWhen_InvalidLimitEpoch() (gas: 1071955)
-LockTest:testDeployment() (gas: 26401)
+ExecuteAccountTest:test_UpdateEpoch() (gas: 186)
-LockTest:test_RevertWhen_InvalidLockupPeriod() (gas: 891958)
+ExecuteEpochTest:testDeployment() (gas: 26495)
-LockTest:test_RevertWhen_SenderIsNotVault() (gas: 10630)
+ExecuteEpochTest:test_ExecuteEpochShouldIncreaseEpoch() (gas: 77749)
-MigrateTest:testDeployment() (gas: 26173)
+ExecuteEpochTest:test_ExecuteEpochShouldIncreasePendingReward() (gas: 239498)
-MigrateTest:test_RevertWhen_NoPendingMigration() (gas: 747548)
+ExecuteEpochTest:test_ExecuteEpochShouldNotIncreaseEpochBeforeEnd() (gas: 17972)
-MigrateTest:test_RevertWhen_SenderIsNotVault() (gas: 10770)
+ExecuteEpochTest:test_ExecuteEpochShouldNotIncreaseEpochInMigration() (gas: 78257)
+LeaveTest:testDeployment() (gas: 26495)
+LeaveTest:test_RevertWhen_NoPendingMigration() (gas: 1072307)
+LeaveTest:test_RevertWhen_SenderIsNotVault() (gas: 10772)
+LockTest:testDeployment() (gas: 26495)
+LockTest:test_RevertWhen_InvalidLockupPeriod() (gas: 865522)
+LockTest:test_RevertWhen_SenderIsNotVault() (gas: 10674)
+MigrateTest:testDeployment() (gas: 26495)
+MigrateTest:test_RevertWhen_NoPendingMigration() (gas: 1069924)
+MigrateTest:test_RevertWhen_SenderIsNotVault() (gas: 10772)
 SetStakeManagerTest:testDeployment() (gas: 9774)
 SetStakeManagerTest:test_RevertWhen_InvalidStakeManagerAddress() (gas: 20481)
 SetStakeManagerTest:test_SetStakeManager() (gas: 19869)
-StakeManagerTest:testDeployment() (gas: 26173)
+StakeManagerTest:testDeployment() (gas: 26267)
-StakeTest:testDeployment() (gas: 26401)
+StakeTest:testDeployment() (gas: 26473)
-StakeTest:test_RevertWhen_InvalidLockupPeriod() (gas: 909883)
+StakeTest:test_RevertWhen_InvalidLockupPeriod() (gas: 883415)
-StakeTest:test_RevertWhen_SenderIsNotVault() (gas: 10672)
+StakeTest:test_RevertWhen_SenderIsNotVault() (gas: 10651)
 StakeTest:test_RevertWhen_StakeTokenTransferFails() (gas: 175040)
-StakeTest:test_StakeWithoutLockUpTimeMintsMultiplierPoints() (gas: 963376)
+StakeTest:test_StakeWithoutLockUpTimeMintsMultiplierPoints() (gas: 963829)
+StakeTest:test_mintBonusMP() (gas: 165)
+StakeTest:test_updateBonusMP() (gas: 210)
+StakeTest:test_updateLockUpTime() (gas: 165)
+StakeTest:test_updateTotalSupplies() (gas: 187)
 StakedTokenTest:testStakeToken() (gas: 7616)
-UnstakeTest:testDeployment() (gas: 26423)
+UnstakeTest:testDeployment() (gas: 26517)
-UnstakeTest:test_RevertWhen_FundsLocked() (gas: 1098058)
+UnstakeTest:test_RevertWhen_FundsLocked() (gas: 1072013)
-UnstakeTest:test_RevertWhen_SenderIsNotVault() (gas: 10587)
+UnstakeTest:test_RevertWhen_SenderIsNotVault() (gas: 10653)
-UnstakeTest:test_UnstakeShouldReturnFunds() (gas: 961933)
+UnstakeTest:test_UnstakeShouldReturnFund_NoLockUp() (gas: 961835)
-UserFlowsTest:testDeployment() (gas: 26401)
+UnstakeTest:test_UnstakeShouldReturnFund_WithLockUp() (gas: 1062238)
-UserFlowsTest:test_StakeWithLockUpTimeLocksStake() (gas: 1025811)
+UserFlowsTest:testDeployment() (gas: 26495)
-UserFlowsTest:test_StakedSupplyShouldIncreaseAndDecreaseAgain() (gas: 1843385)
+UserFlowsTest:test_StakeWithLockUpTimeLocksStake() (gas: 1062311)
+UserFlowsTest:test_StakedSupplyShouldIncreaseAndDecreaseAgain() (gas: 1857520)
 VaultFactoryTest:testDeployment() (gas: 9774)

contracts/StakeManager.sol

@@ -16,6 +16,7 @@ contract StakeManager is Ownable {
     error StakeManager__InvalidLimitEpoch();
     error StakeManager__InvalidLockupPeriod();
     error StakeManager__AccountNotInitialized();
+    error StakeManager__InvalidMigration();
 
     struct Account {
         address rewardAddress;
@@ -39,6 +40,7 @@ contract StakeManager is Ownable {
     uint256 public constant MAX_LOCKUP_PERIOD = 4 * YEAR; // 4 years
     uint256 public constant MP_APY = 1;
     uint256 public constant MAX_BOOST = 4;
+    uint256 public constant PRECISION = 1000000;
 
     mapping(address index => Account value) public accounts;
     mapping(uint256 index => Epoch value) public epochs;
@@ -223,6 +225,9 @@ contract StakeManager is Ownable {
      * @param _migration new StakeManager
      */
     function startMigration(StakeManager _migration) external onlyOwner noMigration processEpoch {
+        if (_migration == this || address(_migration) == address(0)) {
+            revert StakeManager__InvalidMigration();
+        }
         migration = _migration;
         stakedToken.transfer(address(migration), epochReward());
         migration.migrationInitialize(currentEpoch, totalSupplyMP, totalSupplyBalance, epochs[currentEpoch].startTime);
@@ -320,19 +325,19 @@ contract StakeManager is Ownable {
             //mint multiplier points to that epoch
             _mintMP(account, iEpoch.startTime + EPOCH_SIZE, iEpoch);
             uint256 userSupply = account.balance + account.currentMP;
-            uint256 userShare = (userSupply / iEpoch.totalSupply) * 100;
+            uint256 userShare = (userSupply / iEpoch.totalSupply) * PRECISION;
-            uint256 userEpochReward = (userShare * iEpoch.epochReward) / 100;
+            uint256 userEpochReward = (userShare * iEpoch.epochReward) / PRECISION;
             userReward += userEpochReward;
-            iEpoch.epochReward -= userEpochReward;
+            //iEpoch.epochReward -= userEpochReward;
-            iEpoch.totalSupply -= userSupply;
+            //iEpoch.totalSupply -= userSupply;
         }
         account.epoch = userEpoch;
         if (userReward > 0) {
             pendingReward -= userReward;
             stakedToken.transfer(account.rewardAddress, userReward);
         }
-        mpDifference = account.currentMP - mpDifference;
         if (address(migration) != address(0)) {
+            mpDifference = account.currentMP - mpDifference;
             migration.increaseMPFromMigration(mpDifference);
         } else if (userEpoch == currentEpoch) {
             _mintMP(account, block.timestamp, epochs[currentEpoch]);
@@ -384,7 +389,6 @@ contract StakeManager is Ownable {
             account.initialMP,
             account.currentMP
         );
-
         //update storage
         account.lastMint = processTime;
         account.currentMP += increasedMP;
@@ -428,7 +432,7 @@ contract StakeManager is Ownable {
      */
 
     function _getIncreasedMP(uint256 _balance, uint256 _deltaTime) private pure returns (uint256 _increasedMP) {
-        return _balance * ((MP_APY / YEAR) * _deltaTime);
+        return (_balance * (((_deltaTime * PRECISION) / YEAR) * MP_APY)) / PRECISION;
     }
 
     /**
@@ -454,4 +458,8 @@ contract StakeManager is Ownable {
     function epochEnd() public view returns (uint256 _epochEnd) {
         return epochs[currentEpoch].startTime + EPOCH_SIZE;
     }
+
+    function getAccount(address index) public view returns (Account memory value) {
+        return accounts[index];
+    }
 }

script/Deploy.s.sol

@@ -7,15 +7,16 @@ import { StakeManager } from "../contracts/StakeManager.sol";
 import { VaultFactory } from "../contracts/VaultFactory.sol";
 
 contract Deploy is BaseScript {
-    function run() public returns (VaultFactory, StakeManager, DeploymentConfig) {
+    function run() public returns (VaultFactory, StakeManager, StakeManager, DeploymentConfig) {
         DeploymentConfig deploymentConfig = new DeploymentConfig(broadcaster);
         (, address token) = deploymentConfig.activeNetworkConfig();
 
         vm.startBroadcast(broadcaster);
         StakeManager stakeManager = new StakeManager(token, address(0));
         VaultFactory vaultFactory = new VaultFactory(address(stakeManager));
+        StakeManager migrationStakeManager = new StakeManager(token, address(stakeManager));
         vm.stopBroadcast();
 
-        return (vaultFactory, stakeManager, deploymentConfig);
+        return (vaultFactory, stakeManager, migrationStakeManager, deploymentConfig);
     }
 }

test/StakeManager.t.sol

@@ -3,7 +3,7 @@ pragma solidity ^0.8.19;
 
 import { ERC20 } from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
 
-import { Test } from "forge-std/Test.sol";
+import { Test, console2 } from "forge-std/Test.sol";
 import { Deploy } from "../script/Deploy.s.sol";
 import { DeploymentConfig } from "../script/DeploymentConfig.s.sol";
 import { StakeManager } from "../contracts/StakeManager.sol";
@@ -14,6 +14,7 @@ contract StakeManagerTest is Test {
     DeploymentConfig internal deploymentConfig;
     StakeManager internal stakeManager;
     VaultFactory internal vaultFactory;
+    StakeManager internal migrationStakeManager;
 
     address internal stakeToken;
     address internal deployer;
@@ -22,7 +23,7 @@ contract StakeManagerTest is Test {
 
     function setUp() public virtual {
         Deploy deployment = new Deploy();
-        (vaultFactory, stakeManager, deploymentConfig) = deployment.run();
+        (vaultFactory, stakeManager, migrationStakeManager, deploymentConfig) = deployment.run();
         (deployer, stakeToken) = deploymentConfig.activeNetworkConfig();
     }
 
@@ -100,6 +101,14 @@ contract StakeTest is StakeManagerTest {
         assertEq(stakeManager.totalSupplyMP(), 0);
         assertEq(currentMP, 0);
     }
+
+    function test_updateLockUpTime() public { }
+
+    function test_mintBonusMP() public { }
+
+    function test_updateBonusMP() public { }
+
+    function test_updateTotalSupplies() public { }
 }
 
 contract UnstakeTest is StakeManagerTest {
@@ -127,7 +136,7 @@ contract UnstakeTest is StakeManagerTest {
         userVault.unstake(1);
     }
 
-    function test_UnstakeShouldReturnFunds() public {
+    function test_UnstakeShouldReturnFund_NoLockUp() public {
         // ensure user has funds
         deal(stakeToken, testUser, 1000);
         StakeVault userVault = _createTestVault(testUser);
@@ -144,6 +153,27 @@ contract UnstakeTest is StakeManagerTest {
         assertEq(ERC20(stakeToken).balanceOf(address(userVault)), 0);
         assertEq(ERC20(stakeToken).balanceOf(testUser), 1000);
     }
+
+    function test_UnstakeShouldReturnFund_WithLockUp() public {
+        // ensure user has funds
+        deal(stakeToken, testUser, 1000);
+        StakeVault userVault = _createTestVault(testUser);
+
+        vm.startPrank(testUser);
+        ERC20(stakeToken).approve(address(userVault), 100);
+
+        uint256 lockTime = stakeManager.MIN_LOCKUP_PERIOD();
+        userVault.stake(100, lockTime);
+        assertEq(ERC20(stakeToken).balanceOf(testUser), 900);
+
+        vm.warp(block.timestamp + lockTime + 1);
+
+        userVault.unstake(100);
+
+        assertEq(stakeManager.totalSupplyBalance(), 0);
+        assertEq(ERC20(stakeToken).balanceOf(address(userVault)), 0);
+        assertEq(ERC20(stakeToken).balanceOf(testUser), 1000);
+    }
 }
 
 contract LockTest is StakeManagerTest {
@@ -223,6 +253,8 @@ contract MigrateTest is StakeManagerTest {
         userVault.acceptMigration();
         vm.stopPrank();
     }
+
+    function increaseEpoch(uint256 epochNumber) internal { }
 }
 
 contract ExecuteAccountTest is StakeManagerTest {
@@ -245,6 +277,31 @@ contract ExecuteAccountTest is StakeManagerTest {
         vm.expectRevert(StakeManager.StakeManager__InvalidLimitEpoch.selector);
         stakeManager.executeAccount(address(userVault), currentEpoch + 1);
     }
+
+    function test_ExecuteAccountMintMP() public {
+        uint256 stakeAmount = 10_000_000;
+        deal(stakeToken, testUser, stakeAmount);
+        StakeVault userVault = _createTestVault(testUser);
+        vm.startPrank(testUser);
+        ERC20(stakeToken).approve(address(userVault), stakeAmount);
+        userVault.stake(stakeAmount, 0);
+        
+        vm.warp(stakeManager.epochEnd()-1);
+        stakeManager.executeAccount(address(userVault), stakeManager.currentEpoch());
+
+        
+        for (uint256 i = 0; i < 54; i++) {
+            vm.warp(stakeManager.epochEnd());
+            stakeManager.executeAccount(address(userVault), stakeManager.currentEpoch());
+            console2.log("current epoch", stakeManager.currentEpoch());
+            console2.log("account MP", stakeManager.getAccount(address(userVault)).currentMP);
+        }
+    }
+
+    function test_UpdateEpoch() public { }
+    function test_PayRewards() public { }
+
+    function test_MintMPLimit() public { }
 }
 
 contract UserFlowsTest is StakeManagerTest {
@@ -311,3 +368,54 @@ contract UserFlowsTest is StakeManagerTest {
         assertEq(stakeManager.totalSupplyBalance(), 0);
     }
 }
+
+contract ExecuteEpochTest is StakeManagerTest {
+    function setUp() public override {
+        StakeManagerTest.setUp();
+    }
+
+    //currentEpoch can only increase if time stakeManager.epochEnd().
+    function test_ExecuteEpochShouldNotIncreaseEpochBeforeEnd() public {
+        assertEq(stakeManager.currentEpoch(), 0);
+
+        vm.warp(stakeManager.epochEnd() - 1);
+        stakeManager.executeEpoch();
+        assertEq(stakeManager.currentEpoch(), 0);
+    }
+
+    function test_ExecuteEpochShouldNotIncreaseEpochInMigration() public {
+        assertEq(stakeManager.currentEpoch(), 0);
+
+        assertEq(address(stakeManager.migration()), address(0));
+        vm.prank(deployer);
+        stakeManager.startMigration(migrationStakeManager);
+        assertEq(address(stakeManager.migration()), address(migrationStakeManager));
+
+        vm.warp(stakeManager.epochEnd());
+        vm.expectRevert(StakeManager.StakeManager__PendingMigration.selector);
+        stakeManager.executeEpoch();
+        assertEq(stakeManager.currentEpoch(), 0);
+    }
+
+    //currentEpoch can only increase.
+    function test_ExecuteEpochShouldIncreaseEpoch() public {
+        assertEq(stakeManager.currentEpoch(), 0);
+
+        vm.warp(stakeManager.epochEnd());
+        stakeManager.executeEpoch();
+        assertEq(stakeManager.currentEpoch(), 1);
+    }
+
+    //invariant: stakeManager balanceOf stakeToken > pendingReward
+    function test_ExecuteEpochShouldIncreasePendingReward() public {
+        assertEq(stakeManager.pendingReward(), 0);
+        assertEq(stakeManager.epochReward(), 0);
+        deal(stakeToken, address(stakeManager), 1);
+        assertEq(stakeManager.pendingReward(), 0);
+        assertEq(stakeManager.epochReward(), 1);
+        vm.warp(stakeManager.epochEnd());
+        stakeManager.executeEpoch();
+        assertEq(stakeManager.pendingReward(), 1);
+        assertEq(stakeManager.epochReward(), 0);
+    }
+}

test/StakeVault.t.sol

@@ -28,7 +28,7 @@ contract StakeVaultTest is Test {
 
     function setUp() public virtual {
         Deploy deployment = new Deploy();
-        (vaultFactory, stakeManager, deploymentConfig) = deployment.run();
+        (vaultFactory, stakeManager,, deploymentConfig) = deployment.run();
         (deployer, stakeToken) = deploymentConfig.activeNetworkConfig();
 
         vm.prank(testUser);

test/VaultFactory.t.sol

@@ -16,6 +16,8 @@ contract VaultFactoryTest is Test {
 
     VaultFactory internal vaultFactory;
 
+    StakeManager internal migrationStakeManager;
+
     address internal deployer;
 
     address internal stakedToken;
@@ -24,7 +26,7 @@ contract VaultFactoryTest is Test {
 
     function setUp() public virtual {
         Deploy deployment = new Deploy();
-        (vaultFactory, stakeManager, deploymentConfig) = deployment.run();
+        (vaultFactory, stakeManager, migrationStakeManager, deploymentConfig) = deployment.run();
         (deployer, stakedToken) = deploymentConfig.activeNetworkConfig();
     }