diff --git a/src/RewardsDistributor.sol b/src/RewardsDistributor.sol index 3c938f0..23cbf1b 100644 --- a/src/RewardsDistributor.sol +++ b/src/RewardsDistributor.sol @@ -24,39 +24,29 @@ contract RewardsDistributor is Ownable2StepUpgradeable { /// @notice the reward configuration struct RewardConfiguration { - address token; // the reward token uint256 emissionRate; // emission per second uint256 lastUpdate; // last update timestamp } /*////////////////////////////////////////////////////////////// - MAPPINGS + MAPPINGS/ARRAYS //////////////////////////////////////////////////////////////*/ - /// @notice reward configurations - mapping(address receiver => RewardConfiguration[]) + mapping(address receiver => mapping(address token => RewardConfiguration configuration)) public rewardConfigurations; - mapping(address receiver => mapping(address token => uint256 id)) - public rewardConfigurationsIds; + mapping(address receiver => address[] rewardsTokens) public rewardTokens; /*////////////////////////////////////////////////////////////// EVENTS //////////////////////////////////////////////////////////////*/ - event RewardConfigurationAdded( + event RewardConfigurationSet( address indexed receiver, address indexed token, uint256 emissionRate ); - event RewardConfigurationUpdated( - address indexed receiver, - address indexed token, - uint256 oldEmissionRate, - uint256 newEmissionRate - ); - event RewardDistributed( address indexed receiver, address indexed token, @@ -81,7 +71,7 @@ contract RewardsDistributor is Ownable2StepUpgradeable { /// @param receiver The receiver of the rewards /// @param token The reward token /// @param emissionRate The emission rate - function addRewardConfiguration( + function setRewardConfiguration( address receiver, address token, uint256 emissionRate @@ -89,81 +79,42 @@ contract RewardsDistributor is Ownable2StepUpgradeable { require(token != address(0), "No native rewards allowed"); require(emissionRate > 0, "Emission rate must be greater than 0"); - rewardConfigurations[receiver].push( - RewardConfiguration(token, emissionRate, block.timestamp) - ); - - rewardConfigurationsIds[receiver][token] = rewardConfigurations[ - receiver - ].length; - - emit RewardConfigurationAdded(receiver, token, emissionRate); - } - - /// @notice Update the emission rate of a reward configuration - /// @param receiver The receiver of the rewards - /// @param token The reward token - /// @param emissionRate The new emission rate - /// @dev set the emission rate to 0 to stop the rewards - function updateEmissonRate( - address receiver, - address token, - uint256 emissionRate - ) external onlyOwner { - uint256 id = rewardConfigurationsIds[receiver][token]; - require( - rewardConfigurations[receiver].length > 0 && id > 0, - "No reward configuration found" + rewardConfigurations[receiver][token] = RewardConfiguration( + emissionRate, + block.timestamp ); - // index is always 1 less than the id - RewardConfiguration storage conf = rewardConfigurations[receiver][ - id - 1 - ]; - - uint256 oldEmissionRate = conf.emissionRate; - - conf.emissionRate = emissionRate; + // TODO check if token already exists + rewardTokens[receiver].push(token); - emit RewardConfigurationUpdated( - receiver, - token, - oldEmissionRate, - emissionRate - ); + emit RewardConfigurationSet(receiver, token, emissionRate); } /// @notice Distribute rewards to receiver /// @param token The reward token function distributeReward(address token) external { - address receiver = msg.sender; - - uint256 id = rewardConfigurationsIds[receiver][token]; - - require(id > 0, "No reward configuration found"); - - distributeRewardInternal(receiver, id - 1); + distributeRewardInternal(msg.sender, token); } /// @notice Distribute rewards to all tokens function distributeRewards() external { address receiver = msg.sender; - for (uint256 i = 0; i < rewardConfigurations[receiver].length; i++) { - distributeRewardInternal(receiver, i); + for (uint256 i = 0; i < rewardTokens[receiver].length; i++) { + distributeRewardInternal(receiver, rewardTokens[receiver][i]); } } - /// @notice Distribute rewards to token at index + /// @notice Distribute rewards to token /// @param receiver The receiver of the rewards - /// @param index The index of the reward configuration + /// @param token The reward token function distributeRewardInternal( address receiver, - uint256 index + address token ) internal { RewardConfiguration storage rewardConfiguration = rewardConfigurations[ receiver - ][index]; + ][token]; // difference in time since last update uint256 timeDelta = block.timestamp - rewardConfiguration.lastUpdate; diff --git a/src/Staking.sol b/src/Staking.sol index ac40771..6ade9a0 100644 --- a/src/Staking.sol +++ b/src/Staking.sol @@ -45,6 +45,9 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable { /// @dev only owner can change uint256 public minStake; + /// @notice the last time the contract was updated + uint256 updatedAt; + /*////////////////////////////////////////////////////////////// STRUCTS //////////////////////////////////////////////////////////////*/ @@ -58,6 +61,11 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable { uint256 lockPeriod; } + struct Reward { + uint256 earned; + uint256 userRewardPerTokenPaid; + } + /*////////////////////////////////////////////////////////////// MAPPINGS/ARRAYS //////////////////////////////////////////////////////////////*/ @@ -66,12 +74,19 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable { mapping(address keyper => Stake[]) public keyperStakes; /// TODO when remove keyper also unstake the first stake + /// @notice the keypers mapping mapping(address keyper => bool isKeyper) public keypers; /// @notice how many SHU a keyper has locked mapping(address keyper => uint256 totalLocked) public totalLocked; - address[] public rewardTokenList; + mapping(address token => uint256 rewardPerTokenStored) + public rewardPerTokenStored; + + mapping(address keyper => mapping(address token => Reward keyperRewards)) + public keyperRewards; + + Reward[] public rewardTokenList; /*////////////////////////////////////////////////////////////// EVENTS @@ -98,7 +113,7 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable { } /// @notice Update rewards for a keyper - /// @param keyper The keyper address + /// @param caller The keyper address modifier updateRewards(address caller) { // Calculate current assets before distributing rewards uint256 assetsBefore = convertToAssets(balanceOf(caller)); @@ -106,6 +121,8 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable { // Distribute rewards rewardsDistributor.distributeRewards(); + if (caller != address(0)) {} + // If the caller has no assets or is the zero address, skip compound if (assetsBefore == 0 || caller == address(0)) { _; @@ -170,6 +187,8 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable { function stake( uint256 amount ) external onlyKeyper updateRewards(msg.sender) { + /////////////////////////// CHECKS /////////////////////////////// + address keyper = msg.sender; // Get the keyper stakes @@ -183,6 +202,8 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable { ); } + /////////////////////////// EFFECTS /////////////////////////////// + uint256 sharesToMint = convertToShares(amount); // Update the keyper's SHU balance @@ -191,6 +212,8 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable { // Mint the shares _mint(keyper, sharesToMint); + /////////////////////////// INTERACTIONS /////////////////////////// + // Lock the SHU in the contract stakingToken.safeTransferFrom(keyper, address(this), amount); @@ -381,7 +404,7 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable { /// @param _rewardsDistributor The address of the rewards distributor contract function setRewardsDistributor( IRewardsDistributor _rewardsDistributor - ) external onlyOwner updateRewards(0) { + ) external onlyOwner updateRewards(address(0)) { rewardsDistributor = _rewardsDistributor; } @@ -389,7 +412,7 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable { /// @param _lockPeriod The lock period in seconds function setLockPeriod( uint256 _lockPeriod - ) external onlyOwner updateRewards(0) { + ) external onlyOwner updateRewards(address(0)) { lockPeriod = _lockPeriod; } @@ -397,7 +420,7 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable { /// @param _minStake The minimum stake amount function setMinStake( uint256 _minStake - ) external onlyOwner updateRewards(0) { + ) external onlyOwner updateRewards(address(0)) { minStake = _minStake; } @@ -407,7 +430,7 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable { function setKeyper( address keyper, bool isKeyper - ) external onlyOwner updateRewards(0) { + ) external onlyOwner updateRewards(address(0)) { keypers[keyper] = isKeyper; emit KeyperSet(keyper, isKeyper); @@ -419,7 +442,7 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable { function setKeypers( address[] memory _keypers, bool isKeyper - ) external onlyOwner updateRewards(0) { + ) external onlyOwner updateRewards(address(0)) { for (uint256 i = 0; i < _keypers.length; i++) { keypers[_keypers[i]] = isKeyper; @@ -427,18 +450,6 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable { } } - /// @notice Add a reward token - /// @dev Only the rewards distributor can add reward tokens - /// @param rewardToken The address of the reward token - function addRewardToken(address rewardToken) external updateRewards(0) { - require( - msg.sender == address(rewardsDistributor), - "Only rewards distributor can add reward tokens" - ); - - rewardTokenList.push(rewardToken); - } - /*////////////////////////////////////////////////////////////// VIEW FUNCTIONS //////////////////////////////////////////////////////////////*/ @@ -528,4 +539,22 @@ contract Staking is ERC20VotesUpgradeable, Ownable2StepUpgradeable { ) public pure override returns (bool) { revert("Transfer is disabled"); } + + function rewardPerToken(address token) private view returns (uint256) { + uint256 supply = totalSupply(); // Saves an extra SLOAD if totalSupply is non-zero. + + if (supply == 0) { + return rewardPerTokenStored[token]; + } + + (, uint256 rewardRate) = rewardsDistributor.rewardConfigurations( + address(this), + token + ); + + return + rewardPerTokenStored[token] + + (rewardRate * (block.timestamp - updatedAt) * 1e18) / + supply; + } } diff --git a/src/interfaces/IRewardsDistributor.sol b/src/interfaces/IRewardsDistributor.sol index 9a77b04..0d608c4 100644 --- a/src/interfaces/IRewardsDistributor.sol +++ b/src/interfaces/IRewardsDistributor.sol @@ -2,19 +2,7 @@ pragma solidity 0.8.25; interface IRewardsDistributor { - struct RewardConfiguration { - address token; // the reward token - uint256 emissionRate; // emission per second - uint256 lastUpdate; // last update timestamp - } - - function addRewardConfiguration( - address receiver, - address token, - uint256 emissionRate - ) external; - - function updateEmissonRate( + function setRewardConfiguration( address receiver, address token, uint256 emissionRate @@ -25,15 +13,7 @@ interface IRewardsDistributor { function distributeRewards() external; function rewardConfigurations( - address receiver, - uint256 index - ) - external - view - returns (address token, uint256 emissionRate, uint256 lastUpdate); - - function rewardConfigurationsIds( address receiver, address token - ) external view returns (uint256); + ) external view returns (uint256 emissionRate, uint256 lastUpdate); } diff --git a/test/unit/StakingUnitTest.t.sol b/test/unit/StakingUnitTest.t.sol index 163482c..193c67f 100644 --- a/test/unit/StakingUnitTest.t.sol +++ b/test/unit/StakingUnitTest.t.sol @@ -19,13 +19,15 @@ contract StakingUnitTest is Test { uint256 constant lockPeriod = 60 * 24 * 30 * 6; // 6 months uint256 constant minStake = 50_000 * 1e18; // 50k - address keyper = address(0x1234); + address keyper1 = address(0x1234); + address keyper2 = address(0x5678); function setUp() public { // deploy mock shu shu = new MockShu(); - shu.mint(keyper, 1_000_000 * 1e18); + shu.mint(keyper1, 1_000_000 * 1e18); + shu.mint(keyper2, 1_000_000 * 1e18); // deploy rewards distributor address rewardsDistributionProxy = address( @@ -53,7 +55,7 @@ contract StakingUnitTest is Test { staking = IStaking(stakingProxy); - IRewardsDistributor(rewardsDistributionProxy).addRewardConfiguration( + IRewardsDistributor(rewardsDistributionProxy).setRewardConfiguration( stakingProxy, address(shu), 1e18 @@ -69,19 +71,19 @@ contract StakingUnitTest is Test { function testAddKeyper() public { vm.expectEmit(address(staking)); - emit IStaking.KeyperSet(keyper, true); - staking.setKeyper(keyper, true); + emit IStaking.KeyperSet(keyper1, true); + staking.setKeyper(keyper1, true); } function testStakeSucceed() public { testAddKeyper(); - vm.startPrank(keyper); + vm.startPrank(keyper1); shu.approve(address(staking), minStake); vm.expectEmit(true, true, true, true, address(staking)); emit IStaking.Staked( - keyper, + keyper1, minStake, minStake, // first stake, shares == amount lockPeriod @@ -100,9 +102,53 @@ contract StakingUnitTest is Test { uint256 claimAmount = 1_000e18; // 1 SHU per second is distributed vm.expectEmit(true, true, true, true, address(staking)); - emit IStaking.ClaimRewards(keyper, address(shu), claimAmount); + emit IStaking.ClaimRewards(keyper1, address(shu), claimAmount); - vm.prank(keyper); + vm.prank(keyper1); staking.claimReward(shu, claimAmount); } + + function testMultipleKeyperStakeSucceed() public { + testStakeSucceed(); + vm.warp(block.timestamp + 1000); // 500 seconds later + + staking.setKeyper(keyper2, true); + + 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.prank(keyper1); + staking.claimReward(shu, 0); + } + + function testHarvestAndClaimRewardSucceed() public { + testStakeSucceed(); + + vm.warp(block.timestamp + 500); // 500 seconds later + + staking.setKeyper(keyper2, true); + + vm.startPrank(keyper2); + shu.mint(keyper2, 1_000_000 * 1e18); + shu.approve(address(staking), minStake); + staking.stake(minStake); + vm.stopPrank(); + + staking.harvest(keyper1); + + vm.warp(block.timestamp + 500); // 500 seconds later + + uint256 claimAmount = 1_000e18; // 1 SHU per second is distributed + + vm.expectEmit(true, true, true, true, address(staking)); + emit IStaking.ClaimRewards(keyper1, address(shu), claimAmount); + + vm.prank(keyper1); + staking.claimReward(shu, 0); + } }