checkpoint

src/RewardsDistributor.sol

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: MIT
 pragma solidity 0.8.25;
 
+import {console} from "@forge-std/console.sol";
 import {ERC20} from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
 import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
@@ -41,11 +42,11 @@ contract RewardsDistributor is Ownable2StepUpgradeable {
                                  EVENTS
     //////////////////////////////////////////////////////////////*/
 
-    event RewardConfigurationSet(
-        address indexed receiver,
-        address indexed token,
-        uint256 emissionRate
-    );
+    //    event RewardConfigurationSet(
+    //        address indexed receiver,
+    //        address indexed token,
+    //        uint256 emissionRate
+    //    );
 
     event RewardDistributed(
         address indexed receiver,
@@ -89,7 +90,7 @@ contract RewardsDistributor is Ownable2StepUpgradeable {
         );
 
         if (emissionRate == 0) {
-            // remove the token from the list
+            // remove the token
             address[] storage tokens = rewardTokens[receiver];
             for (uint256 i = 0; i < tokens.length; i++) {
                 if (tokens[i] == token) {
@@ -100,7 +101,7 @@ contract RewardsDistributor is Ownable2StepUpgradeable {
             }
         }
 
-        emit RewardConfigurationSet(receiver, token, emissionRate);
+        //  emit RewardConfigurationSet(receiver, token, emissionRate);
     }
 
     /// @notice Distribute rewards to receiver
@@ -147,4 +148,10 @@ contract RewardsDistributor is Ownable2StepUpgradeable {
 
         emit RewardDistributed(receiver, token, reward);
     }
+
+    function getRewardTokens(
+        address receiver
+    ) external view returns (address[] memory) {
+        return rewardTokens[receiver];
+    }
 }

src/Staking.sol

@@ -114,20 +114,43 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
         // Distribute rewards
         rewardsDistributor.distributeRewards();
 
-        if (caller != address(0)) {
-            address[] memory rewardTokens = rewardsDistributor.rewardTokens();
+        address[] memory rewardTokens = rewardsDistributor.getRewardTokens(
+            address(this)
+        );
+
+        for (uint256 i = 0; i < rewardTokens.length; i++) {
+            address token = rewardTokens[i];
 
-            for (uint256 i = 0; i < rewardTokens.length; i++) {
-                address token = rewardTokens[i];
+            uint256 _rewardPerToken = rewardPerToken(token);
+            rewardPerTokenStored[token] = _rewardPerToken;
 
-                uint256 _rewardPerToken = rewardPerToken(token);
+            console.log(
+                "after rewardPerTokenStored[token]",
+                rewardPerTokenStored[token]
+            );
 
-                keyperRewards[caller][token].earned += (balanceOf(caller) *
+            if (caller != address(0)) {
+                console.log("caller balance before earn", balanceOf(caller));
+                uint256 earned = (balanceOf(caller) *
                     (_rewardPerToken -
-                        keyperRewards[caller][token].userRewardPerTokenPaid));
+                        keyperRewards[caller][token].userRewardPerTokenPaid)) /
+                    1e18;
 
+                console.log("earned", earned);
+
+                keyperRewards[caller][token].earned += earned;
                 keyperRewards[caller][token]
                     .userRewardPerTokenPaid = _rewardPerToken;
+
+                // compound
+                if (token == address(stakingToken)) {
+                    uint256 reward = keyperRewards[caller][token].earned;
+
+                    if (reward > 0) {
+                        keyperRewards[caller][token].earned = 0;
+                        _mint(caller, reward);
+                    }
+                }
             }
         }
 
@@ -354,6 +377,8 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
 
             require(maxWithdrawAmount > 0, "No rewards to claim");
 
+            console.log("maxWithdrawAmount", maxWithdrawAmount);
+
             // If the amount is greater than the max withdraw amount, the contract
             // will transfer the maximum amount available not the requested amount
             // If the amount is 0, claim all the rewards
@@ -379,6 +404,8 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
         emit ClaimRewards(keyper, address(rewardToken), amount);
     }
 
+    function compound(address keyper) external updateRewards(keyper) {}
+
     /*//////////////////////////////////////////////////////////////
                          RESTRICTED FUNCTIONS
     //////////////////////////////////////////////////////////////*/
@@ -446,9 +473,13 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
     /// @param keyper The keyper address
     /// @return The maximum amount of assets that a keyper can withdraw
     function maxWithdraw(address keyper) public view virtual returns (uint256) {
+        console.log("maxWithdraw: keyper balance", balanceOf(keyper));
+        console.log("maxWithdraw: total locked", totalLocked[keyper]);
+        console.log("maxWithdraw: min stake", minStake);
+
         return
             balanceOf(keyper) -
-            (totalLocked[keyper] >= minStake ? minStake : totalLocked[keyper]);
+            (totalLocked[keyper] >= minStake ? totalLocked[keyper] : minStake);
     }
 
     function maxWithdraw(
@@ -493,10 +524,15 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable {
             return rewardPerTokenStored[token];
         }
 
-        (, uint256 rewardRate) = rewardsDistributor.rewardConfigurations(
+        (uint256 rewardRate, ) = rewardsDistributor.rewardConfigurations(
             address(this),
             token
         );
+        console.log("reward rate", rewardRate);
+        console.log("rewardPerTokenStored[token]", rewardPerTokenStored[token]);
+        console.log("block.timestamp", block.timestamp);
+        console.log("updatedAt", updatedAt);
+        console.log("supply", supply / 1e18);
 
         return
             rewardPerTokenStored[token] +

src/interfaces/IRewardsDistributor.sol

@@ -12,7 +12,9 @@ interface IRewardsDistributor {
 
     function distributeRewards() external;
 
-    function rewardTokens() external view returns (address[] memory);
+    function getRewardTokens(
+        address receiver
+    ) external view returns (address[] memory);
 
     function rewardConfigurations(
         address receiver,

src/interfaces/IStaking.sol

@@ -51,13 +51,14 @@ interface IStaking {
         address keyper
     ) external view returns (uint256);
 
+    function compound(address keyper) external;
+
     event Staked(
         address indexed user,
         uint256 indexed amount,
-        uint256 indexed shares,
         uint256 lockPeriod
     );
-    event Unstaked(address user, uint256 amount, uint256 shares);
+    event Unstaked(address user, uint256 amount);
     event ClaimRewards(address user, address rewardToken, uint256 rewards);
     event KeyperSet(address keyper, bool isKeyper);
 }

test/unit/StakingUnitTest.t.sol

@@ -17,7 +17,7 @@ contract StakingUnitTest is Test {
     MockShu public shu;
 
     uint256 constant lockPeriod = 60 * 24 * 30 * 6; // 6 months
-    uint256 constant minStake = 50_000 * 1e18; // 50k
+    uint256 constant minStake = 10_000 * 1e18; // 50k
 
     address keyper1 = address(0x1234);
     address keyper2 = address(0x5678);
@@ -69,50 +69,137 @@ contract StakingUnitTest is Test {
                                HAPPY PATHS
     //////////////////////////////////////////////////////////////*/
 
-    function testAddKeyper() public {
+    function testAddKeypers() public {
         vm.expectEmit(address(staking));
         emit IStaking.KeyperSet(keyper1, true);
         staking.setKeyper(keyper1, true);
+
+        emit IStaking.KeyperSet(keyper2, true);
+        staking.setKeyper(keyper2, true);
     }
 
     function testStakeSucceed() public returns (uint256 stakeIndex) {
-        testAddKeyper();
+        testAddKeypers();
+
+        uint256 contractBalanceBefore = shu.balanceOf(address(staking));
+        uint256 keyperBalanceBefore = shu.balanceOf(keyper1);
 
         vm.startPrank(keyper1);
         shu.approve(address(staking), minStake);
 
         vm.expectEmit(true, true, true, true, address(staking));
-        emit IStaking.Staked(
-            keyper1,
-            minStake,
-            minStake, // first stake, shares == amount
-            lockPeriod
-        );
+        emit IStaking.Staked(keyper1, minStake, lockPeriod);
 
         stakeIndex = staking.stake(minStake);
-
         vm.stopPrank();
+
+        uint256 contractBalanceAfter = shu.balanceOf(address(staking));
+        uint256 keyperBalanceAfter = shu.balanceOf(keyper1);
+
+        assertEq(
+            contractBalanceAfter - contractBalanceBefore,
+            minStake,
+            "Wrong contract balance"
+        );
+        assertEq(
+            keyperBalanceBefore - minStake,
+            keyperBalanceAfter,
+            "Wrong keyper balance"
+        );
     }
 
     function testMultipleKeyperStakeSucceed() public {
         testStakeSucceed();
-        vm.warp(block.timestamp + 1000); // 500 seconds later
 
-        staking.setKeyper(keyper2, true);
+        vm.warp(block.timestamp + 1000);
+
+        // TODO move assertions to modifier
+        uint256 contractBalanceBefore = shu.balanceOf(address(staking));
+        uint256 keyperBalanceBefore = shu.balanceOf(keyper2);
 
         vm.startPrank(keyper2);
-        shu.mint(keyper2, 1_000_000 * 1e18);
         shu.approve(address(staking), minStake);
+
         staking.stake(minStake);
         vm.stopPrank();
 
-        vm.warp(block.timestamp + 1000); // 500 seconds later
+        vm.warp(block.timestamp + 1000);
+
+        uint256 contractBalanceAfter = shu.balanceOf(address(staking));
+        uint256 keyperBalanceAfter = shu.balanceOf(keyper2);
+
+        assertEq(
+            contractBalanceAfter - contractBalanceBefore,
+            minStake,
+            "Wrong contract balance"
+        );
+        assertEq(
+            keyperBalanceBefore - minStake,
+            keyperBalanceAfter,
+            "Wrong keyper balance"
+        );
+    }
+
+    function testClaimAllRewardSucceed() public {
+        testStakeSucceed();
+
+        vm.warp(block.timestamp + 100_000);
+
+        uint256 contractBalanceBefore = shu.balanceOf(address(staking));
+        uint256 keyperBalanceBefore = shu.balanceOf(keyper1);
+
+        //vm.expectEmit(true, true, true, true, address(staking));
+        //emit IStaking.ClaimRewards(keyper1, address(shu), claimAmount);
 
         vm.prank(keyper1);
         staking.claimReward(shu, 0);
+
+        uint256 contractBalanceAfter = shu.balanceOf(address(staking));
+        uint256 keyperBalanceAfter = shu.balanceOf(keyper1);
+
+        assertEq(
+            contractBalanceAfter - contractBalanceBefore,
+            0,
+            "All rewards receive must be transferred"
+        );
+        assertEq(
+            contractBalanceAfter,
+            minStake,
+            "Contract balance after must be equal minStake"
+        );
+    }
+
+    function testClaimAllRewardAfterCompoundSucceed() public {
+        testStakeSucceed();
+
+        uint256 contractBalanceBefore = shu.balanceOf(address(staking));
+        uint256 keyperBalanceBefore = shu.balanceOf(keyper1);
+
+        vm.warp(block.timestamp + 50_000);
+
+        staking.compound(keyper1);
+
+        vm.warp(block.timestamp + 50_000);
+
+        vm.prank(keyper1);
+        staking.claimReward(shu, 0);
+
+        uint256 contractBalanceAfter = shu.balanceOf(address(staking));
+        uint256 keyperBalanceAfter = shu.balanceOf(keyper1);
+
+        assertEq(
+            contractBalanceAfter - contractBalanceBefore,
+            0,
+            "All rewards receive must be transferred"
+        );
+        assertEq(
+            contractBalanceAfter,
+            minStake,
+            "Contract balance after must be equal minStake"
+        );
     }
 
-    function testClaimRewardSucceed() public {
+    function testMultipleClaimRewardSucceed() public {
         testStakeSucceed();
 
         vm.warp(block.timestamp + 1000); // 1000 seconds later
@@ -132,7 +219,7 @@ contract StakingUnitTest is Test {
         vm.warp(block.timestamp + 1000); // 1000 seconds later
 
         vm.expectEmit(true, true, true, true, address(staking));
-        emit IStaking.Unstaked(keyper1, minStake, minStake);
+        emit IStaking.Unstaked(keyper1, minStake);
 
         vm.prank(keyper1);
         staking.unstake(keyper1, index, 0);