feature(RewardsStreamerMP): add rewardAccrued to account to allow acc…

…ounts to stake again multiple times
vacp2p · Jan 14, 2025 · e095752 · e095752
1 parent 26f76d4
commit e095752
Show file tree

Hide file tree

Showing 2 changed files with 79 additions and 1 deletion.
diff --git a/src/RewardsStreamerMP.sol b/src/RewardsStreamerMP.sol
@@ -60,6 +60,7 @@ contract RewardsStreamerMP is
         uint256 stakedBalance;
         uint256 accountRewardIndex;
         uint256 mpAccrued;
+        uint256 rewardAccrued;
         uint256 maxMP;
         uint256 lastMPUpdateTime;
         uint256 lockUntil;
@@ -201,6 +202,8 @@ contract RewardsStreamerMP is
         _updateAccountMP(msg.sender, true);
 
         Account storage account = accounts[msg.sender];
+        _updateAccountReward(msg.sender);
+
         if (account.lockUntil != 0 && account.lockUntil > block.timestamp) {
             revert StakingManager__CannotRestakeWithLockedFunds();
         }
@@ -231,6 +234,45 @@ contract RewardsStreamerMP is
         account.accountRewardIndex = rewardIndex;
     }
 
+    function currentRewardIndex() external view returns (uint256) {
+        (, uint256 newRewardIndex) = _pendingRewardIndex();
+        return newRewardIndex;
+    }
+
+    function accountAccruedRewards(address accountAddress) external view returns (uint256) {
+        return accounts[accountAddress].rewardAccrued;
+    }
+
+    function accountRewardIndex(address accountAddress) external view returns (uint256) {
+        return accounts[accountAddress].accountRewardIndex;
+    }
+
+    function _updateAccountReward(address accountAddress) internal {
+        Account storage account = accounts[accountAddress];
+        (, uint256 currentRewardIndex) = _pendingRewardIndex();
+
+        uint256 accountWeight = account.stakedBalance + _mpBalanceOf(accountAddress);
+        if (accountWeight == 0) {
+            account.accountRewardIndex = currentRewardIndex;
+            return;
+        }
+
+        // If there's no difference in indexes, do nothing
+        if (currentRewardIndex == account.accountRewardIndex) {
+            return;
+        }
+
+        // Calculate user’s pending reward, e.g.:
+        uint256 deltaIndex = currentRewardIndex - account.accountRewardIndex;
+        uint256 pending = (accountWeight * deltaIndex) / SCALE_FACTOR;
+
+        if (pending > 0) {
+            account.rewardAccrued += pending;
+        }
+
+        account.accountRewardIndex = currentRewardIndex;
+    }
+
     function lock(uint256 lockPeriod)
         external
         onlyTrustedCodehash
@@ -299,6 +341,7 @@ contract RewardsStreamerMP is
 
         account.stakedBalance -= amount;
         account.mpAccrued -= mpToReduce;
+        // FIXME: update account.rewardAccrued
         account.maxMP -= maxMPToReduce;
         account.accountRewardIndex = rewardIndex;
         totalMPAccrued -= mpToReduce;
@@ -528,7 +571,7 @@ contract RewardsStreamerMP is
         uint256 accountWeight = account.stakedBalance + _mpBalanceOf(accountAddress);
         uint256 deltaRewardIndex = newRewardIndex - account.accountRewardIndex;
 
-        return (accountWeight * deltaRewardIndex) / SCALE_FACTOR;
+        return account.rewardAccrued + (accountWeight * deltaRewardIndex) / SCALE_FACTOR;
     }
 
     function rewardsBalanceOfUser(address user) external view returns (uint256) {

diff --git a/test/RewardsStreamerMP.t.sol b/test/RewardsStreamerMP.t.sol
@@ -2133,6 +2133,41 @@ contract RewardsStreamerMP_RewardsTest is RewardsStreamerMPTest {
         assertEq(streamer.rewardsBalanceOf(vaults[alice]), secondLiveBalanceBeforeGlobalUpdate);
         assertApproxEqAbs(streamer.rewardsBalanceOf(vaults[alice]), 1000e18, tolerance);
     }
+
+    function testRewardsBalanceOfStakingAgain() public {
+        assertEq(streamer.totalRewardsSupply(), 0);
+
+        uint256 initialTime = vm.getBlockTimestamp();
+
+        _stake(alice, 100e18, 0);
+        _stake(bob, 100e18, 0);
+        assertEq(streamer.rewardsBalanceOf(vaults[alice]), 0);
+
+        vm.prank(admin);
+        streamer.setReward(2000e18, 10 days);
+        assertEq(streamer.rewardsBalanceOf(vaults[alice]), 0);
+        assertEq(streamer.rewardsBalanceOf(vaults[bob]), 0);
+
+        vm.warp(initialTime + 5 days);
+
+        uint256 tolerance = 300; // 300 wei
+        assertApproxEqAbs(streamer.rewardsBalanceOf(vaults[alice]), 500e18, tolerance);
+        assertApproxEqAbs(streamer.accountAccruedRewards(vaults[alice]), 0, tolerance);
+        assertApproxEqAbs(streamer.rewardsBalanceOf(vaults[bob]), 500e18, tolerance);
+        assertApproxEqAbs(streamer.accountAccruedRewards(vaults[bob]), 0, tolerance);
+
+        _stake(alice, 100e18, 0);
+
+        assertApproxEqAbs(streamer.rewardsBalanceOf(vaults[alice]), 500e18, tolerance);
+        assertApproxEqAbs(streamer.accountAccruedRewards(vaults[alice]), 500e18, tolerance);
+        assertApproxEqAbs(streamer.rewardsBalanceOf(vaults[bob]), 500e18, tolerance);
+        assertApproxEqAbs(streamer.accountAccruedRewards(vaults[bob]), 0, tolerance);
+
+        vm.warp(initialTime + 10 days);
+
+        assertApproxEqAbs(streamer.rewardsBalanceOf(vaults[alice]), 1_166_666_666_666_666_666_666, tolerance);
+        assertApproxEqAbs(streamer.rewardsBalanceOf(vaults[alice]), 833_333_333_333_333_333_333, tolerance);
+    }
 }
 
 contract MultipleVaultsStakeTest is RewardsStreamerMPTest {