Merge pull request #435 from push-protocol/fix-reviews-on-staking

Fix reviews on staking

contracts/PushCore/PushCoreStorageV2.sol

@@ -15,6 +15,9 @@ contract PushCoreStorageV2 {
     /**
      * Staking V2 state variables *
      */
+
+    //UNUSED STATE VARIABLES START HERE
+
     mapping(address => uint256) public usersRewardsClaimed;
 
     uint256 public genesisEpoch; // Block number at which Stakig starts
@@ -25,11 +28,13 @@ contract PushCoreStorageV2 {
     uint256 public constant epochDuration = 21 * 7156; // 21 * number of blocks per day(7156) ~ 20 day approx
 
     /// @notice Stores all the individual epoch rewards
-    mapping(uint256 => uint256) public epochRewards;
+    mapping(uint256 => uint256) public epochRewards; 
     /// @notice Stores User's Fees Details
-    mapping(address => StakingTypes.UserFeesInfo) public userFeesInfo;
+    mapping(address => StakingTypes.UserFeesInfo) public userFeesInfo; 
     /// @notice Stores the total staked weight at a specific epoch.
-    mapping(uint256 => uint256) public epochToTotalStakedWeight;
+    mapping(uint256 => uint256) public epochToTotalStakedWeight; 
+
+    //UNUSED STATE VARIABLES END HERE
 
     /**
      * Handling bridged information *

contracts/PushStaking/PushStaking.sol

@@ -9,15 +9,17 @@ import { Errors } from "../libraries/Errors.sol";
 import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import { SafeERC20 } from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import { Initializable } from "@openzeppelin/contracts/proxy/utils/Initializable.sol";
+import { GenericTypes } from "../libraries/DataTypes.sol";
+import { BaseHelper } from "../libraries/BaseHelper.sol";
 
 contract PushStaking is Initializable, PushStakingStorage {
     using SafeERC20 for IERC20;
 
     event Staked(address indexed user, uint256 indexed amountStaked);
     event Unstaked(address indexed user, uint256 indexed amountUnstaked);
     event RewardsHarvested(address indexed user, uint256 indexed rewardAmount, uint256 fromEpoch, uint256 tillEpoch);
-    event NewSharesIssued(address indexed Wallet, uint256 indexed Shares);
-    event SharesRemoved(address indexed Wallet, uint256 indexed Shares);
+    event NewSharesIssued(address indexed wallet, uint256 indexed shares);
+    event SharesRemoved(address indexed wallet, uint256 indexed shares);
     event SharesDecreased(address indexed Wallet, uint256 indexed oldShares, uint256 newShares);
 
     function initialize(address _pushChannelAdmin, address _core, address _pushToken) public initializer {
@@ -61,6 +63,10 @@ contract PushStaking is Initializable, PushStakingStorage {
         removeWalletShare(oldFoundation);
     }
 
+    function getEpochToWalletShare(address wallet, uint epoch) public view returns(uint){
+        return walletShareInfo[wallet].epochToWalletShares[epoch];
+    }
+
     function initializeStake(uint256 _walletTotalShares) external {
         require(genesisEpoch == 0, "PushCoreV2::initializeStake: Already Initialized");
         genesisEpoch = block.number;
@@ -78,11 +84,9 @@ contract PushStaking is Initializable, PushStakingStorage {
         emit NewSharesIssued(FOUNDATION, sharesToBeAllocated);
     }
 
-    /*
-     * Fetching shares for a wallet 
-     * 1. Internal helper function
-     * 2. helps in getting the shares to be assigned for a wallet based on params passed in this function
-    */
+    /**
+     * @notice Calcultes the share amount based on requested shares and total shares 
+     */
     function getSharesAmount(
         uint256 _totalShares,
         GenericTypes.Percentage memory _percentage
@@ -97,18 +101,13 @@ contract PushStaking is Initializable, PushStakingStorage {
         sharesToBeAllocated = (_percentage.percentageNumber * _totalShares)
             / ((100 * (10 ** _percentage.decimalPlaces)) - _percentage.percentageNumber);
     }
+    /**
+     * @notice allows Governance to add/increase wallet shares. 
+     * @notice If a wallet has already has a share, then it acts as a "increase share" function, given that the percenatge passed
+     * @notice should be greater than the already assigned percentge.
+     * Emits NewSharesIssued
+     */
 
-    /*
-     * Adding Wallet Share to a Wallet
-     * 1. addWalletShare(address wallet, uint256 percentageOfShares)
-     * 2. Can be called by governance.
-     * 3. Uses the formulae to derive the percent of shares to be assigned to a specific wallet
-     * 4. Updates WALLET_TOTAL_SHARES
-     * 5. Updates walletShareInfo mapping
-     * 6. Emits out an event.
-    * 7. If a wallet has already has a share, then it acts as a "increase share" function. And the percenatge passed
-     *    should be greater than the already assigned percentge.
-    */
     function addWalletShare(address _walletAddress, GenericTypes.Percentage memory _percentage) public onlyGovernance {
         if(_walletAddress == address(0)){
             revert Errors.InvalidArgument_WrongAddress(_walletAddress);
@@ -129,18 +128,20 @@ contract PushStaking is Initializable, PushStakingStorage {
         WALLET_TOTAL_SHARES = TotalShare + sharesToBeAllocated;
         emit NewSharesIssued(_walletAddress, sharesToBeAllocated);
     }
-    /*
-     * Removing Wallet Share from a Wallet
-     * 1. removes the shares from a wallet completely
-     * 2. Can be called by governance.
-     * 3. Updates WALLET_TOTAL_SHARES
-     * 4. Emits out an event.
-    */
+
+    /**
+     * @notice allows Governance to remove wallet shares. 
+     * @notice shares to be removed are given back to FOUNDATION
+     * Emits SharesRemoved
+     */
 
     function removeWalletShare(address _walletAddress) public onlyGovernance {
         if(_walletAddress == address(0) || _walletAddress == FOUNDATION) {
             revert Errors.InvalidArgument_WrongAddress(_walletAddress);
         }
+        if (block.number <= walletShareInfo[_walletAddress].lastStakedBlock + epochDuration) {
+            revert Errors.PushStaking_InvalidEpoch_LessThanExpected();
+        }
         uint256 sharesToBeRemoved = walletShareInfo[_walletAddress].walletShare;
         _adjustWalletAndTotalStake(_walletAddress, 0, sharesToBeRemoved);
         _adjustWalletAndTotalStake(FOUNDATION, sharesToBeRemoved, 0);
@@ -149,31 +150,40 @@ contract PushStaking is Initializable, PushStakingStorage {
     }
 
     function decreaseWalletShare(
-        address _walletAddress,
-        GenericTypes.Percentage memory _percentage
-    )
-        external
-        onlyGovernance
-    {
-        uint256 sharesToBeRemoved = walletShareInfo[_walletAddress].walletShare;
-        removeWalletShare(_walletAddress);
-        addWalletShare(_walletAddress, _percentage);
-        emit SharesDecreased(_walletAddress, sharesToBeRemoved, walletShareInfo[_walletAddress].walletShare);
-    }
+            address _walletAddress,
+            GenericTypes.Percentage memory _percentage
+        )
+            external
+            onlyGovernance
+        {
+            if(_walletAddress == address(0) || _walletAddress == FOUNDATION) {
+               revert Errors.InvalidArgument_WrongAddress(_walletAddress);
+            }
+
+            uint currentEpoch = lastEpochRelative(genesisEpoch,block.number);
+
+            uint256 currentShares = walletShareInfo[_walletAddress].walletShare;
+            uint256 sharesToBeAllocated = BaseHelper.calcPercentage(WALLET_TOTAL_SHARES, _percentage);
 
-    /*
-     *Reward Calculation for a Given Wallet
-     * 1. calculateWalletRewards(address wallet)
-     * 2. public helper function
-     * 3. Helps in calculating rewards for a specific wallet based on
-       * a. Their Wallet Share 
-        * b. Total Wallet Shares
-        * c. Total Rewards available in WALLET_TOTAL_SHARES
-     * 4. Once calculated rewards for a sepcific wallet can be updated in WalletToRewards mapping
-     * 5. Reward can be calculated for wallets similar they are calculated for Token holders in a specific epoch.
-        * Reward for a Given Wallet X = ( Wallet Share of X / WALLET_TOTAL_SHARES) * WALLET_FEE_POOL
-    */
-    //TODO logic yet to be finalized
+            if(sharesToBeAllocated >= currentShares){
+                revert Errors.InvalidArg_MoreThanExpected(currentShares, sharesToBeAllocated);
+            }
+            uint256 sharesToBeRemoved = currentShares - sharesToBeAllocated;
+            walletShareInfo[_walletAddress].walletShare = sharesToBeAllocated;
+            walletShareInfo[_walletAddress].epochToWalletShares[currentEpoch] = sharesToBeAllocated;
+
+            walletShareInfo[FOUNDATION].walletShare += sharesToBeRemoved;
+            walletShareInfo[FOUNDATION].epochToWalletShares[currentEpoch] += sharesToBeRemoved;
+
+            emit SharesDecreased(_walletAddress, currentShares, sharesToBeAllocated);
+       }
+
+    /**
+     * @notice calculates rewards for share holders, for any given epoch. 
+     * @notice The rewards are calcluated based on -Their Wallet Share in that epoch, Total Wallet Shares in that epoch,
+     * @notice Total Rewards available in that epoch
+     * @dev Reward for a Given Wallet X in epoch i = ( Wallet Share of X in epoch i  * Rewards in epoch i) / WALLET_TOTAL_SHARES in epoch i
+     */
     function calculateWalletRewards(address _wallet, uint256 _epochId) public view returns (uint256) {
         return (walletShareInfo[_wallet].epochToWalletShares[_epochId] * epochRewardsForWallets[_epochId])
             / epochToTotalShares[_epochId];
@@ -192,7 +202,7 @@ contract PushStaking is Initializable, PushStakingStorage {
             rewards = rewards + claimableReward;
         }
 
-        usersRewardsClaimed[msg.sender] = usersRewardsClaimed[msg.sender] + rewards;
+        walletRewardsClaimed[msg.sender] = walletRewardsClaimed[msg.sender] + rewards;
         // set the lastClaimedBlock to blocknumer at the end of `_tillEpoch`
         uint256 _epoch_to_block_number = genesisEpoch + _tillEpoch * epochDuration;
         walletShareInfo[msg.sender].lastClaimedBlock = _epoch_to_block_number;

contracts/PushStaking/PushStakingStorage.sol

@@ -1,40 +1,45 @@
 pragma solidity ^0.8.20;
 
-import { StakingTypes, GenericTypes } from "../libraries/DataTypes.sol";
+import { StakingTypes } from "../libraries/DataTypes.sol";
 
 contract PushStakingStorage {
-    /**
-     * Staking V2 state variables *
-     */
-    mapping(address => uint256) public usersRewardsClaimed;
-
+    //State variables for Stakers 
     uint256 public genesisEpoch; // Block number at which Stakig starts
     uint256 lastEpochInitialized; // The last EPOCH ID initialized with the respective epoch rewards
     uint256 lastTotalStakeEpochInitialized; // The last EPOCH ID initialized with the respective total staked weight
+    uint256 public previouslySetEpochRewards; // Amount of rewards set in last initialized epoch
+    uint256 public totalStakedAmount; // Total token weight staked in Protocol at any given time
+
+    // State variables for Share holders 
     uint256 walletLastEpochInitialized; // todo new variable
     uint256 walletLastTotalStakeEpochInitialized;//todo new variable
-    uint256 public totalStakedAmount; // Total token weight staked in Protocol at any given time
-    uint256 public previouslySetEpochRewards; // Amount of rewards set in last initialized epoch
     uint256 public walletPreviouslySetEpochRewards; //todo new variable
-    uint256 public constant epochDuration = 21 * 7156; // 21 * number of blocks per day(7156) ~ 20 day approx
     uint256 public WALLET_TOTAL_SHARES; //Total Shares
 
+    uint256 public constant epochDuration = 21 * 7156; // 21 * number of blocks per day(7156) ~ 20 day approx
     address public pushChannelAdmin;
     address public PUSH_TOKEN_ADDRESS;
     address public governance;
     address public core;
     address public FOUNDATION;
 
+    //Mappings for Stakers
+    ///@notice stores total rewards claimed by a user
+    mapping(address => uint256) public usersRewardsClaimed;
     /// @notice Stores all the individual epoch rewards for stakers
     mapping(uint256 => uint256) public epochRewardsForStakers;
-    /// @notice Stores all the individual epoch rewards for Wallet share holders
-    mapping(uint256 => uint256) public epochRewardsForWallets;
     /// @notice Stores User's Fees Details
     mapping(address => StakingTypes.UserFeesInfo) public userFeesInfo;
-    ///@notice stores Wallet share details for a given address
-    mapping(address => StakingTypes.WalletShareInfo) public walletShareInfo;
     /// @notice Stores the total staked weight at a specific epoch.
     mapping(uint256 => uint256) public epochToTotalStakedWeight;
+
+    //Mappings for Share Holders
+    ///@notice stores total reward claimed by a wallet
+    mapping(address => uint256) public walletRewardsClaimed;
+    /// @notice Stores all the individual epoch rewards for Wallet share holders
+    mapping(uint256 => uint256) public epochRewardsForWallets;
+    ///@notice stores Wallet share details for a given address
+    mapping(address => StakingTypes.WalletShareInfo) public walletShareInfo;
     ///@notice stores the total shares in a specific epoch
     mapping(uint256 => uint256) public epochToTotalShares;
 }

test/PushStaking/WalletSharesStaking/BaseWalletSharesStaking.t.sol

@@ -48,7 +48,7 @@ contract BaseWalletSharesStaking is BasePushStaking {
             aliceWalletShares +
             charlieWalletShares +
             tonyWalletShares;
-        assertEq(walletTotalShares, totalSharesSum);
+        assertEq(walletTotalShares, totalSharesSum,"wallet Share Sum");
     }
 
     /**
@@ -57,7 +57,7 @@ contract BaseWalletSharesStaking is BasePushStaking {
     function _validateEpochShares() internal {
         uint256 walletTotalShares = pushStaking.WALLET_TOTAL_SHARES();
         for (uint256 i=genesisEpoch; i<=getCurrentEpoch(); ) {
-            assertLe(pushStaking.epochToTotalShares(i), walletTotalShares);
+            assertLe(pushStaking.epochToTotalShares(i), walletTotalShares, "Epoch Shares");
             unchecked {
                 i++;
             }

test/PushStaking/WalletSharesStaking/unit_tests/ClaimRewards/ClaimRewards.t.sol

@@ -47,7 +47,7 @@ contract ClaimRewardsTest is BaseWalletSharesStaking {
         assertEq(bobStakedBlockBefore, bobStakedBlockAfter, "StakedBlock");
         assertEq(bobClaimedBlockAfter, genesisEpoch + (getCurrentEpoch() - 1) * epochDuration, "ClaimedBlock");
 
-        uint256 claimedRewards = pushStaking.usersRewardsClaimed(actor.bob_channel_owner);
+        uint256 claimedRewards = pushStaking.walletRewardsClaimed(actor.bob_channel_owner);
         assertEq(balanceBobBefore + expectedRewards, pushToken.balanceOf(actor.bob_channel_owner), "Balance");
         assertEq(expectedRewards, claimedRewards);
 
@@ -94,7 +94,7 @@ contract ClaimRewardsTest is BaseWalletSharesStaking {
         assertEq(bobStakedBlockBefore, bobStakedBlockAfter, "StakedBlock");
         assertEq(bobClaimedBlockAfter, genesisEpoch + (getCurrentEpoch() - 1) * epochDuration, "ClaimedBlock");
 
-        uint256 claimedRewards = pushStaking.usersRewardsClaimed(actor.bob_channel_owner);
+        uint256 claimedRewards = pushStaking.walletRewardsClaimed(actor.bob_channel_owner);
         assertEq(balanceBobBefore + expectedRewards, pushToken.balanceOf(actor.bob_channel_owner), "Balance");
         assertEq(expectedRewards, claimedRewards);
 
@@ -106,7 +106,7 @@ contract ClaimRewardsTest is BaseWalletSharesStaking {
         test_WhenUserClaims_in_DifferentEpoch();
 
         uint256 balanceBobBefore = pushToken.balanceOf(actor.bob_channel_owner);
-        uint256 claimedRewardsBefore = pushStaking.usersRewardsClaimed(actor.bob_channel_owner);
+        uint256 claimedRewardsBefore = pushStaking.walletRewardsClaimed(actor.bob_channel_owner);
         (uint256 bobWalletSharesBefore, uint256 bobStakedBlockBefore, uint256 bobLastClaimedBlock) =
             pushStaking.walletShareInfo(actor.bob_channel_owner);
 
@@ -129,7 +129,7 @@ contract ClaimRewardsTest is BaseWalletSharesStaking {
         assertEq(bobStakedBlockBefore, bobStakedBlockAfter2, "StakedBlock");
         assertEq(bobClaimedBlockAfter2, genesisEpoch + (getCurrentEpoch() - 1) * epochDuration, "ClaimedBlock");
 
-        uint256 claimedRewards = pushStaking.usersRewardsClaimed(actor.bob_channel_owner);
+        uint256 claimedRewards = pushStaking.walletRewardsClaimed(actor.bob_channel_owner);
         assertEq(balanceBobBefore, pushToken.balanceOf(actor.bob_channel_owner), "Balance");
         assertEq(claimedRewardsBefore, claimedRewards);
     }
@@ -184,7 +184,7 @@ contract ClaimRewardsTest is BaseWalletSharesStaking {
         assertEq(bobStakedBlockBefore, bobStakedBlockAfter, "StakedBlock");
         assertEq(bobClaimedBlockAfter2, genesisEpoch + (getCurrentEpoch() - 1) * epochDuration, "ClaimedBlock");
 
-        uint256 claimedRewards = pushStaking.usersRewardsClaimed(actor.bob_channel_owner);
+        uint256 claimedRewards = pushStaking.walletRewardsClaimed(actor.bob_channel_owner);
 
         assertEq(balanceBobBefore + expectedRewards, pushToken.balanceOf(actor.bob_channel_owner), "Balance");
         assertEq(expectedRewards, claimedRewards);
@@ -228,7 +228,7 @@ contract ClaimRewardsTest is BaseWalletSharesStaking {
         assertEq(adminStakedBlockBefore, adminStakedBlockAfter, "StakedBlock");
         assertEq(adminClaimedBlockAfter, genesisEpoch + (getCurrentEpoch() - 1) * epochDuration, "ClaimedBlock");
 
-        uint256 claimedRewards = pushStaking.usersRewardsClaimed(actor.admin);
+        uint256 claimedRewards = pushStaking.walletRewardsClaimed(actor.admin);
         assertEq(balanceAdminBefore + expectedRewards, pushToken.balanceOf(actor.admin), "Balance");
         assertEq(expectedRewards, claimedRewards);
 
@@ -282,7 +282,7 @@ contract ClaimRewardsTest is BaseWalletSharesStaking {
         assertEq(bobStakedBlockBefore, bobStakedBlockAfter, "StakedBlock");
         assertEq(bobClaimedBlockAfter, genesisEpoch + (getCurrentEpoch() - 1) * epochDuration, "ClaimedBlock");
 
-        uint256 claimedRewardsBob = pushStaking.usersRewardsClaimed(actor.bob_channel_owner);
+        uint256 claimedRewardsBob = pushStaking.walletRewardsClaimed(actor.bob_channel_owner);
 
         uint256 expectedRewardsBob = (coreProxy.WALLET_FEE_POOL() * 10) / 100;
         assertEq(balanceBobBefore + expectedRewardsBob, pushToken.balanceOf(actor.bob_channel_owner), "balanceBob");
@@ -292,7 +292,7 @@ contract ClaimRewardsTest is BaseWalletSharesStaking {
         assertEq(aliceStakedBlockBefore, aliceStakedBlockAfter, "StakedBlock");
         assertEq(aliceClaimedBlockAfter, genesisEpoch + (getCurrentEpoch() - 1) * epochDuration, "ClaimedBlock");
 
-        uint256 claimedRewardsAlice = pushStaking.usersRewardsClaimed(actor.alice_channel_owner);
+        uint256 claimedRewardsAlice = pushStaking.walletRewardsClaimed(actor.alice_channel_owner);
 
         uint256 expectedRewardsAlice = coreProxy.WALLET_FEE_POOL() * 50 / 100;
         assertEq(