Feat/adjustments (#107)

* return rewards while claim all * made is whitelisted methods internal * bumped version
dl-solarity · Jul 3, 2024 · a372599 · a372599
1 parent 7dc446e
commit a372599
Show file tree

Hide file tree

Showing 9 changed files with 54 additions and 24 deletions.
diff --git a/contracts/access/MerkleWhitelisted.sol b/contracts/access/MerkleWhitelisted.sol
@@ -24,27 +24,35 @@ abstract contract MerkleWhitelisted {
 
     modifier onlyWhitelisted(bytes memory data_, bytes32[] memory merkleProof_) {
         require(
-            isWhitelisted(keccak256(data_), merkleProof_),
+            _isWhitelisted(keccak256(data_), merkleProof_),
             "MerkleWhitelisted: not whitelisted"
         );
         _;
     }
 
     modifier onlyWhitelistedUser(address user_, bytes32[] memory merkleProof_) {
-        require(isWhitelistedUser(user_, merkleProof_), "MerkleWhitelisted: not whitelisted");
+        require(_isWhitelistedUser(user_, merkleProof_), "MerkleWhitelisted: not whitelisted");
         _;
     }
 
+    /**
+     * @notice The function to get the current Merkle root
+     * @return the current Merkle root or zero bytes if it has not been set
+     */
+    function getMerkleRoot() public view returns (bytes32) {
+        return _merkleRoot;
+    }
+
     /**
      * @notice The function to check if the leaf belongs to the Merkle tree
      * @param leaf_ the leaf to be checked
      * @param merkleProof_ the path from the leaf to the Merkle tree root
      * @return true if the leaf belongs to the Merkle tree, false otherwise
      */
-    function isWhitelisted(
+    function _isWhitelisted(
         bytes32 leaf_,
         bytes32[] memory merkleProof_
-    ) public view returns (bool) {
+    ) internal view returns (bool) {
         return merkleProof_.verify(_merkleRoot, leaf_);
     }
 
@@ -54,19 +62,11 @@ abstract contract MerkleWhitelisted {
      * @param merkleProof_ the path from the user to the Merkle tree root
      * @return true if the user belongs to the Merkle tree, false otherwise
      */
-    function isWhitelistedUser(
+    function _isWhitelistedUser(
         address user_,
         bytes32[] memory merkleProof_
-    ) public view returns (bool) {
-        return isWhitelisted(keccak256(abi.encodePacked(user_)), merkleProof_);
-    }
-
-    /**
-     * @notice The function to get the current Merkle root
-     * @return the current Merkle root or zero bytes if it has not been set
-     */
-    function getMerkleRoot() public view returns (bytes32) {
-        return _merkleRoot;
+    ) internal view returns (bool) {
+        return _isWhitelisted(keccak256(abi.encodePacked(user_)), merkleProof_);
     }
 
     /**

diff --git a/contracts/finance/staking/AbstractStaking.sol b/contracts/finance/staking/AbstractStaking.sol
@@ -91,9 +91,10 @@ abstract contract AbstractStaking is AbstractValueDistributor, Initializable {
 
     /**
      * @notice Claims all the available rewards.
+     * @return The total value of the rewards claimed.
      */
-    function claimAll() public stakingStarted {
-        _distributeAllValue(msg.sender);
+    function claimAll() public stakingStarted returns (uint256) {
+        return _distributeAllValue(msg.sender);
     }
 
     /**

diff --git a/contracts/finance/staking/AbstractValueDistributor.sol b/contracts/finance/staking/AbstractValueDistributor.sol
@@ -145,8 +145,9 @@ abstract contract AbstractValueDistributor {
     /**
      * @notice Distributes all the available value to a specific user.
      * @param user_ The address of the user.
+     * @return The amount of value distributed.
      */
-    function _distributeAllValue(address user_) internal virtual {
+    function _distributeAllValue(address user_) internal virtual returns (uint256) {
         _update(user_);
 
         UserDistribution storage _userDist = _userDistributions[user_];
@@ -155,11 +156,13 @@ abstract contract AbstractValueDistributor {
 
         require(amount_ > 0, "ValueDistributor: amount has to be more than 0");
 
-        _userDist.owedValue -= amount_;
+        delete _userDist.owedValue;
 
         emit ValueDistributed(user_, amount_);
 
         _afterDistributeValue(user_, amount_);
+
+        return amount_;
     }
 
     /**

diff --git a/contracts/mock/access/MerkleWhitelistedMock.sol b/contracts/mock/access/MerkleWhitelistedMock.sol
@@ -20,6 +20,20 @@ contract MerkleWhitelistedMock is MerkleWhitelisted {
         emit WhitelistedUser();
     }
 
+    function isWhitelisted(
+        bytes32 leaf_,
+        bytes32[] memory merkleProof_
+    ) internal view returns (bool) {
+        return _isWhitelisted(leaf_, merkleProof_);
+    }
+
+    function isWhitelistedUser(
+        address user_,
+        bytes32[] memory merkleProof_
+    ) internal view returns (bool) {
+        return _isWhitelistedUser(user_, merkleProof_);
+    }
+
     function setMerkleRoot(bytes32 merkleRoot_) external {
         _setMerkleRoot(merkleRoot_);
     }

diff --git a/contracts/mock/finance/staking/AbstractValueDistributorMock.sol b/contracts/mock/finance/staking/AbstractValueDistributorMock.sol
@@ -19,8 +19,8 @@ contract AbstractValueDistributorMock is AbstractValueDistributor, Multicall {
         _distributeValue(user_, amount_);
     }
 
-    function distributeAllValue(address user_) external {
-        _distributeAllValue(user_);
+    function distributeAllValue(address user_) external returns (uint256) {
+        return _distributeAllValue(user_);
     }
 
     function userShares(address user_) external view returns (uint256) {

diff --git a/package-lock.json b/package-lock.json
diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@solarity/solidity-lib",
-  "version": "2.7.9",
+  "version": "2.7.10",
   "license": "MIT",
   "author": "Distributed Lab",
   "readme": "README.md",

diff --git a/test/finance/staking/AbstractStaking.test.ts b/test/finance/staking/AbstractStaking.test.ts
@@ -624,6 +624,10 @@ describe("AbstractStaking", () => {
       const secondOwed = await abstractStaking.getOwedValue(SECOND);
       const thirdOwed = await abstractStaking.getOwedValue(THIRD);
 
+      expect(await abstractStaking.connect(FIRST).claimAll.staticCall()).to.eq(firstOwed);
+      expect(await abstractStaking.connect(SECOND).claimAll.staticCall()).to.eq(secondOwed);
+      expect(await abstractStaking.connect(THIRD).claimAll.staticCall()).to.eq(thirdOwed);
+
       await abstractStaking.connect(FIRST).claimAll();
       await abstractStaking.connect(SECOND).claimAll();
 

diff --git a/test/finance/staking/AbstractValueDistributor.test.ts b/test/finance/staking/AbstractValueDistributor.test.ts
@@ -231,6 +231,14 @@ describe("AbstractValueDistributor", () => {
     it("should distribute all the owed values optimally", async () => {
       await performSharesManipulations();
 
+      const firstOwed = await abstractValueDistributor.getOwedValue(FIRST);
+      const secondOwed = await abstractValueDistributor.getOwedValue(SECOND);
+      const thirdOwed = await abstractValueDistributor.getOwedValue(THIRD);
+
+      expect(await abstractValueDistributor.distributeAllValue.staticCall(FIRST)).to.eq(firstOwed);
+      expect(await abstractValueDistributor.distributeAllValue.staticCall(SECOND)).to.eq(secondOwed);
+      expect(await abstractValueDistributor.distributeAllValue.staticCall(THIRD)).to.eq(thirdOwed);
+
       await abstractValueDistributor.distributeAllValue(FIRST);
       await abstractValueDistributor.distributeAllValue(SECOND);
       await abstractValueDistributor.distributeAllValue(THIRD);