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

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

Hide file tree

Showing 2 changed files with 101 additions and 2 deletions.
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,43 @@ 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 (currentRewardIndex == account.accountRewardIndex) {
+            return;
+        }
+
+        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 +339,7 @@ contract RewardsStreamerMP is
 
         account.stakedBalance -= amount;
         account.mpAccrued -= mpToReduce;
+        // FIXME: update account.rewardAccrued
         account.maxMP -= maxMPToReduce;
         account.accountRewardIndex = rewardIndex;
         totalMPAccrued -= mpToReduce;
@@ -528,7 +569,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
@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.26;
 
-import { Test } from "forge-std/Test.sol";
+import { Test, console } from "forge-std/Test.sol";
 import { Math } from "@openzeppelin/contracts/utils/math/Math.sol";
 import { DeployRewardsStreamerMPScript } from "../script/DeployRewardsStreamerMP.s.sol";
 import { UpgradeRewardsStreamerMPScript } from "../script/UpgradeRewardsStreamerMP.s.sol";
@@ -2133,6 +2133,64 @@ 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);
+        // console.log("global index", streamer.currentRewardIndex());
+        // console.log("alice index", streamer.accountRewardIndex(vaults[alice]));
+
+        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);
+
+        console.log("alice rewards balance", streamer.rewardsBalanceOf(vaults[alice]));
+
+        vm.warp(initialTime + 100 days);
+        console.log("alice rewards balance", streamer.rewardsBalanceOf(vaults[alice]));
+
+        vm.warp(initialTime + 200 days);
+        console.log("alice rewards balance", streamer.rewardsBalanceOf(vaults[alice]));
+
+        vm.warp(initialTime + 10_000 days);
+        console.log("alice rewards balance", streamer.rewardsBalanceOf(vaults[alice]));
+
+        // assertApproxEqAbs(
+        //     streamer.rewardsBalanceOf(vaults[alice]) + streamer.rewardsBalanceOf(vaults[bob]), 2000e18, tolerance
+        // );
+
+        // console.log("rewards total supply", streamer.totalRewardsSupply());
+        // console.log("alice rewards balance", streamer.rewardsBalanceOf(vaults[alice]));
+        // console.log("alice rewards accrued", streamer.accountAccruedRewards(vaults[alice]));
+        // console.log("bob rewards balance", streamer.rewardsBalanceOf(vaults[bob]));
+        // console.log("bob rewards accrued", streamer.accountAccruedRewards(vaults[bob]));
+
+        // 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 {