Added claimAll to AbstractStaking

contracts/finance/staking/AbstractStaking.sol

@@ -89,6 +89,13 @@ abstract contract AbstractStaking is AbstractValueDistributor, Initializable {
         _distributeValue(msg.sender, amount_);
     }
 
+    /**
+     * @notice Claims all the available rewards.
+     */
+    function claimAll() public stakingStarted {
+        _distributeAllValue(msg.sender);
+    }
+
     /**
      * @notice Withdraws all the staked tokens together with rewards.
      *

contracts/finance/staking/AbstractValueDistributor.sol

@@ -107,16 +107,15 @@ abstract contract AbstractValueDistributor {
      * @param amount_ The amount of shares to remove.
      */
     function _removeShares(address user_, uint256 amount_) internal virtual {
+        UserDistribution storage _userDist = _userDistributions[user_];
+
         require(amount_ > 0, "ValueDistributor: amount has to be more than 0");
-        require(
-            amount_ <= _userDistributions[user_].shares,
-            "ValueDistributor: insufficient amount"
-        );
+        require(amount_ <= _userDist.shares, "ValueDistributor: insufficient amount");
 
         _update(user_);
 
         _totalShares -= amount_;
-        _userDistributions[user_].shares -= amount_;
+        _userDist.shares -= amount_;
 
         emit SharesRemoved(user_, amount_);
 
@@ -131,13 +130,32 @@ abstract contract AbstractValueDistributor {
     function _distributeValue(address user_, uint256 amount_) internal virtual {
         _update(user_);
 
+        UserDistribution storage _userDist = _userDistributions[user_];
+
+        require(amount_ > 0, "ValueDistributor: amount has to be more than 0");
+        require(amount_ <= _userDist.owedValue, "ValueDistributor: insufficient amount");
+
+        _userDist.owedValue -= amount_;
+
+        emit ValueDistributed(user_, amount_);
+
+        _afterDistributeValue(user_, amount_);
+    }
+
+    /**
+     * @notice Distributes all the available value to a specific user.
+     * @param user_ The address of the user.
+     */
+    function _distributeAllValue(address user_) internal virtual {
+        _update(user_);
+
+        UserDistribution storage _userDist = _userDistributions[user_];
+
+        uint256 amount_ = _userDist.owedValue;
+
         require(amount_ > 0, "ValueDistributor: amount has to be more than 0");
-        require(
-            amount_ <= _userDistributions[user_].owedValue,
-            "ValueDistributor: insufficient amount"
-        );
 
-        _userDistributions[user_].owedValue -= amount_;
+        _userDist.owedValue -= amount_;
 
         emit ValueDistributed(user_, amount_);
 
@@ -189,12 +207,12 @@ abstract contract AbstractValueDistributor {
         _updatedAt = block.timestamp;
 
         if (user_ != address(0)) {
-            UserDistribution storage userDist = _userDistributions[user_];
+            UserDistribution storage _userDist = _userDistributions[user_];
 
-            userDist.owedValue +=
-                (userDist.shares * (_cumulativeSum - userDist.cumulativeSum)) /
+            _userDist.owedValue +=
+                (_userDist.shares * (_cumulativeSum - _userDist.cumulativeSum)) /
                 PRECISION;
-            userDist.cumulativeSum = _cumulativeSum;
+            _userDist.cumulativeSum = _cumulativeSum;
         }
     }
 

contracts/mock/finance/staking/AbstractValueDistributorMock.sol

@@ -19,6 +19,10 @@ contract AbstractValueDistributorMock is AbstractValueDistributor, Multicall {
         _distributeValue(user_, amount_);
     }
 
+    function distributeAllValue(address user_) external {
+        _distributeAllValue(user_);
+    }
+
     function userShares(address user_) external view returns (uint256) {
         return userDistribution(user_).shares;
     }

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@solarity/solidity-lib",
-  "version": "2.7.8",
+  "version": "2.7.9",
   "license": "MIT",
   "author": "Distributed Lab",
   "readme": "README.md",

test/finance/staking/AbstractStaking.test.ts

@@ -185,6 +185,7 @@ describe("AbstractStaking", () => {
       await expect(abstractStaking.unstake(wei(100, sharesDecimals))).to.be.revertedWith(revertMessage);
       await expect(abstractStaking.withdraw()).to.be.revertedWith(revertMessage);
       await expect(abstractStaking.claim(wei(100, sharesDecimals))).to.be.revertedWith(revertMessage);
+      await expect(abstractStaking.claimAll()).to.be.revertedWith(revertMessage);
     });
 
     it("should work as expected if the staking start time is set to the timestamp in the past", async () => {
@@ -597,6 +598,57 @@ describe("AbstractStaking", () => {
     });
   });
 
+  describe("claimAll()", () => {
+    it("should claim all the rewards correctly", async () => {
+      await performStakingManipulations();
+
+      await abstractStaking.connect(FIRST).claimAll();
+      await abstractStaking.connect(SECOND).claimAll();
+      await abstractStaking.connect(THIRD).claimAll();
+
+      expect(await abstractStaking.getOwedValue(FIRST)).to.equal(0);
+      expect(await abstractStaking.getOwedValue(SECOND)).to.equal(0);
+      expect(await abstractStaking.getOwedValue(THIRD)).to.equal(0);
+
+      expect(await abstractStaking.userOwedValue(FIRST)).to.equal(0);
+      expect(await abstractStaking.userOwedValue(SECOND)).to.equal(0);
+      expect(await abstractStaking.userOwedValue(THIRD)).to.equal(0);
+    });
+
+    it("should transfer tokens correctly on the claim", async () => {
+      await performStakingManipulations();
+
+      const initialRewardsBalance = await rewardsToken.balanceOf(abstractStaking);
+
+      const firstOwed = await abstractStaking.getOwedValue(FIRST);
+      const secondOwed = await abstractStaking.getOwedValue(SECOND);
+      const thirdOwed = await abstractStaking.getOwedValue(THIRD);
+
+      await abstractStaking.connect(FIRST).claimAll();
+      await abstractStaking.connect(SECOND).claimAll();
+
+      await abstractStaking.connect(THIRD).claimAll();
+
+      expect(await rewardsToken.balanceOf(abstractStaking)).to.equal(
+        initialRewardsBalance - (firstOwed + secondOwed + thirdOwed),
+      );
+      expect(await rewardsToken.balanceOf(FIRST)).to.equal(firstOwed);
+      expect(await rewardsToken.balanceOf(SECOND)).to.equal(secondOwed);
+      expect(await rewardsToken.balanceOf(THIRD)).to.equal(thirdOwed);
+    });
+
+    it("should not allow to claim 0 rewards", async () => {
+      await performStakingManipulations();
+
+      await expect(
+        abstractStaking.multicall([
+          abstractStaking.interface.encodeFunctionData("claimAll"),
+          abstractStaking.interface.encodeFunctionData("claimAll"),
+        ]),
+      ).to.be.revertedWith("ValueDistributor: amount has to be more than 0");
+    });
+  });
+
   describe("rate", () => {
     it("should accept 0 as a rate and calculate owed values according to this rate correctly", async () => {
       const AbstractStakingMock = await ethers.getContractFactory("AbstractStakingMock");

test/finance/staking/AbstractValueDistributor.test.ts

@@ -228,6 +228,22 @@ describe("AbstractValueDistributor", () => {
       expect(await abstractValueDistributor.userOwedValue(THIRD)).to.equal(0);
     });
 
+    it("should distribute all the owed values optimally", async () => {
+      await performSharesManipulations();
+
+      await abstractValueDistributor.distributeAllValue(FIRST);
+      await abstractValueDistributor.distributeAllValue(SECOND);
+      await abstractValueDistributor.distributeAllValue(THIRD);
+
+      expect(await abstractValueDistributor.getOwedValue(FIRST)).to.equal(0);
+      expect(await abstractValueDistributor.getOwedValue(SECOND)).to.equal(0);
+      expect(await abstractValueDistributor.getOwedValue(THIRD)).to.equal(0);
+
+      expect(await abstractValueDistributor.userOwedValue(FIRST)).to.equal(0);
+      expect(await abstractValueDistributor.userOwedValue(SECOND)).to.equal(0);
+      expect(await abstractValueDistributor.userOwedValue(THIRD)).to.equal(0);
+    });
+
     it("should correctly distribute owed values partially", async () => {
       await performSharesManipulations();