NAY4-4 Incorrect boosted balance (#141)

* chore: reproduce the issue * fix: boosted balance calculation * doc: comment the boosted balance calculation * fix: better boosted balance calculation * fix: boosted balance workaround * chore: remove logging * refactor: revert argument reordering
nayms · Oct 3, 2024 · 93fe561 · 93fe561
1 parent fafdffe
commit 93fe561
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 19 deletions.
diff --git a/src/libs/LibTokenizedVaultStaking.sol b/src/libs/LibTokenizedVaultStaking.sol
@@ -168,18 +168,25 @@ library LibTokenizedVaultStaking {
 
         uint256 boost1 = ((((_getD(_entityId) - ratio) * _amount) / _getD(_entityId)) * _getA(_entityId)) / _getD(_entityId);
         uint256 boost2 = (((ratio * _amount) / _getD(_entityId)) * _getA(_entityId)) / _getD(_entityId);
+
         uint256 balance1 = _amount - (ratio * _amount) / _getD(_entityId);
         uint256 balance2 = (ratio * _amount) / _getD(_entityId);
 
         s.stakeBalance[_vTokenId(_entityId, tokenId, currentInterval + 1)][_stakerId] += balance1 + boost1;
         s.stakeBalance[_vTokenId(_entityId, tokenId, currentInterval + 1)][_entityId] += balance1 + boost1;
 
+        s.stakeBalanceAdded[_vTokenId(_entityId, tokenId, currentInterval + 1)][_stakerId] += balance1;
+        s.stakeBalanceAdded[_vTokenId(_entityId, tokenId, currentInterval + 1)][_entityId] += balance1;
+
         s.stakeBoost[_vTokenId(_entityId, tokenId, currentInterval + 1)][_stakerId] += (boost1 * _getR(_entityId)) / _getD(_entityId) + boost2;
         s.stakeBoost[_vTokenId(_entityId, tokenId, currentInterval + 1)][_entityId] += (boost1 * _getR(_entityId)) / _getD(_entityId) + boost2;
 
         s.stakeBalance[_vTokenId(_entityId, tokenId, currentInterval + 2)][_stakerId] += balance2;
         s.stakeBalance[_vTokenId(_entityId, tokenId, currentInterval + 2)][_entityId] += balance2;
 
+        s.stakeBalanceAdded[_vTokenId(_entityId, tokenId, currentInterval + 2)][_stakerId] += balance2;
+        s.stakeBalanceAdded[_vTokenId(_entityId, tokenId, currentInterval + 2)][_entityId] += balance2;
+
         emit TokenStaked(_stakerId, _entityId, tokenId, _amount);
     }
 
@@ -405,36 +412,25 @@ library LibTokenizedVaultStaking {
         intervalTime_ = _calculateStartTimeOfInterval(_entityId, _currentInterval(_entityId));
     }
 
-    function _stakedAmount(bytes32 _stakerId, bytes32 _entityId) internal view returns (uint256) {
+    function _getStakingAmounts(bytes32 _stakerId, bytes32 _entityId) internal view returns (uint256 stakedBalance_, uint256 boostedBalance_) {
         AppStorage storage s = LibAppStorage.diamondStorage();
 
-        bytes32 tokenId = s.stakingConfigs[_entityId].tokenId;
-        bytes32 vTokenIdMax = _vTokenIdBucket(_entityId, tokenId);
-
-        return s.stakeBalance[vTokenIdMax][_stakerId];
-    }
-
-    function _getStakingAmounts(bytes32 _stakerId, bytes32 _entityId) internal view returns (uint256 stakedBalance_, uint256 boostedBalance_) {
         uint64 currentInterval = _currentInterval(_entityId);
+        bytes32 tokenId = s.stakingConfigs[_entityId].tokenId;
 
-        stakedBalance_ = _stakedAmount(_stakerId, _entityId);
+        stakedBalance_ = s.stakeBalance[_vTokenIdBucket(_entityId, tokenId)][_stakerId];
 
         if (!_isStakingInitialized(_entityId)) {
             // boost is always 1 before init
             return (stakedBalance_, boostedBalance_);
         }
 
-        (StakingState memory state, ) = _getStakingStateWithRewardsBalances(_stakerId, _entityId, currentInterval + 2);
-
-        uint256 boostPrevious = state.boost;
-        uint256 balancePrevious = state.balance;
+        (StakingState memory state, ) = _getStakingStateWithRewardsBalances(_stakerId, _entityId, currentInterval);
 
-        for (uint i = 0; i < 2; i++) {
-            boostPrevious = (boostPrevious * _getD(_entityId)) / _getR(_entityId);
-            balancePrevious = balancePrevious - boostPrevious;
-        }
+        uint256 balance1 = s.stakeBalanceAdded[_vTokenId(_entityId, tokenId, currentInterval + 1)][_stakerId];
+        uint256 balance2 = s.stakeBalanceAdded[_vTokenId(_entityId, tokenId, currentInterval + 2)][_stakerId];
 
-        boostedBalance_ = balancePrevious;
+        boostedBalance_ = state.balance + balance1 + balance2;
 
         if (boostedBalance_ < stakedBalance_) {
             boostedBalance_ = stakedBalance_;

diff --git a/src/shared/AppStorage.sol b/src/shared/AppStorage.sol
@@ -83,12 +83,13 @@ struct AppStorage {
     mapping(address userAddress => EntityApproval) selfOnboarding; // map address => { entityId, roleId }
     /// Staking
     mapping(bytes32 entityId => StakingConfig) stakingConfigs; // StakingConfig for an entity
-    mapping(bytes32 vTokenId => mapping(bytes32 stakerId => uint256 balance)) stakeBalance; // [vTokenId][ownerId] boost at interval
+    mapping(bytes32 vTokenId => mapping(bytes32 stakerId => uint256 balance)) stakeBalance; // [vTokenId][ownerId] balance at interval
     mapping(bytes32 vTokenId => mapping(bytes32 stakerId => uint256 boost)) stakeBoost; // [vTokenId][ownerId] boost at interval
     mapping(bytes32 entityId => mapping(bytes32 stakerId => uint64 interval)) stakeCollected; // last interval reward was collected or pain for a staker in staking entity
     mapping(bytes32 vTokenId => uint256 amount) stakingDistributionAmount; // [vTokenId] Reward at interval
     mapping(bytes32 vTokenId => bytes32 denomination) stakingDistributionDenomination; // [vTokenId] Reward currency
     mapping(bytes32 entityId => mapping(bytes32 stakerId => uint64 interval)) stakingSynced; // last interval when data was synced into storage for staker
+    mapping(bytes32 vTokenId => mapping(bytes32 stakerId => uint256 balance)) stakeBalanceAdded; // raw balance staked at an interval, withouth any boost included, only for reading future intervals (to calculate the total boosted balance)
 }
 
 /// Staking-Related Mappings

diff --git a/test/T06Staking.t.sol b/test/T06Staking.t.sol
@@ -1325,6 +1325,42 @@ contract T06Staking is D03ProtocolDefaults {
         // printAppstorage();
     }
 
+    function test_stake_QS_4_4() public {
+        uint256 startTime = block.timestamp + 1;
+        uint256 stake1Time = startTime + 15 days;
+        uint256 stake2Time = startTime + 31 days;
+
+        initStaking(startTime);
+
+        vm.warp(stake1Time);
+
+        c.log("-- Stake 1 ETH -- ".yellow());
+        startPrank(bob);
+        nayms.stake(nlf.entityId, 1 ether);
+
+        (uint256 stakedBalance, uint256 boostedBalance) = nayms.getStakingAmounts(bob.entityId, nlf.entityId);
+        assertEq(stakedBalance, boostedBalance, "Bob should have no boost".red());
+        printCurrentState(nlf.entityId, bob.entityId, "Bob");
+
+        vm.warp(stake2Time);
+        c.log("-- WARP 31 DAYS -- ".yellow());
+        printCurrentState(nlf.entityId, bob.entityId, "Bob");
+
+        uint256 bobBoost = calculateBoost(startTime, stake2Time, R, I, SCALE_FACTOR);
+        uint256 boostedBalance1 = (0.5 ether * bobBoost) / SCALE_FACTOR / SCALE_FACTOR + 0.5 ether;
+
+        (, uint256 boostedBalance2) = nayms.getStakingAmounts(bob.entityId, nlf.entityId);
+        assertEq(boostedBalance2, boostedBalance1, "Bob should have boost at[1]".red());
+
+        c.log("~~~ Stake 1 ETH -- ".yellow());
+        nayms.stake(nlf.entityId, 1 ether);
+
+        (, uint256 boostedBalance3) = nayms.getStakingAmounts(bob.entityId, nlf.entityId);
+        assertEq(boostedBalance3, boostedBalance1 + 1 ether, "Bob should have boost and more stake".red());
+
+        printCurrentState(nlf.entityId, bob.entityId, "Bob");
+    }
+
     function printAppstorage() public {
         uint64 interval = currentInterval() + 2;