diff --git a/contracts/gauges/FixedRewardPool.sol b/contracts/gauges/FixedRewardPool.sol index 14951ab..338ad64 100644 --- a/contracts/gauges/FixedRewardPool.sol +++ b/contracts/gauges/FixedRewardPool.sol @@ -32,7 +32,9 @@ contract FixedRewardPool is OwnableUpgradeable, ReentrancyGuardUpgradeable, ERC7 uint256 public accTokenPerShare; // The block number of the last pool update uint256 public lastRewardBlock; - // Reward tokens created per block. + // The block number when mining ends + uint256 public bonusEndBlock; + // Reward tokens created per block uint256 public rewardPerBlock; // Total staked weight uint256 public totalStakedWeight; @@ -44,7 +46,12 @@ contract FixedRewardPool is OwnableUpgradeable, ReentrancyGuardUpgradeable, ERC7 // Mapping from tokenId to weight mapping(uint256 => uint256) public tokenWeight; - function initialize(address _nft, uint256 _startBlock, uint256 _rewardPerBlock) external initializer { + function initialize( + address _nft, + uint256 _startBlock, + uint256 _rewardPerBlock, + uint256 _totalBlocks + ) external initializer { require(_startBlock >= block.number, "invalid start block"); require(_rewardPerBlock > 0, "invalid reward per block"); @@ -53,6 +60,7 @@ contract FixedRewardPool is OwnableUpgradeable, ReentrancyGuardUpgradeable, ERC7 weightNFT = IWeightedNFT(_nft); lastRewardBlock = _startBlock; rewardPerBlock = _rewardPerBlock; + bonusEndBlock = _startBlock + _totalBlocks; } function deposit(uint256 _tokenId, address _recipient) public nonReentrant { @@ -172,7 +180,8 @@ contract FixedRewardPool is OwnableUpgradeable, ReentrancyGuardUpgradeable, ERC7 UserInfo storage user = userInfo[_user]; uint256 _totalStakedWeight = totalStakedWeight; if (block.number > lastRewardBlock && _totalStakedWeight != 0) { - uint256 rewards = (block.number - lastRewardBlock) * rewardPerBlock; + uint256 multiplier = _getMultiplier(lastRewardBlock, block.number); + uint256 rewards = multiplier * rewardPerBlock; uint256 adjustedTokenPerShare = accTokenPerShare + (rewards * PRECISION_FACTOR) / _totalStakedWeight; return (user.amount * adjustedTokenPerShare) / PRECISION_FACTOR - user.rewardDebt; } else { @@ -190,11 +199,22 @@ contract FixedRewardPool is OwnableUpgradeable, ReentrancyGuardUpgradeable, ERC7 return; } - uint256 rewards = (block.number - lastRewardBlock) * rewardPerBlock; + uint256 multiplier = _getMultiplier(lastRewardBlock, block.number); + uint256 rewards = multiplier * rewardPerBlock; accTokenPerShare = accTokenPerShare + (rewards * PRECISION_FACTOR) / totalStakedWeight; lastRewardBlock = block.number; } + function _getMultiplier(uint256 _from, uint256 _to) internal view returns (uint256) { + if (_to <= bonusEndBlock) { + return _to - _from; + } else if (_from >= bonusEndBlock) { + return 0; + } else { + return bonusEndBlock - _from; + } + } + function stakingToken() external view returns (address) { return weightNFT.nft(); } diff --git a/test/FixedRewardPool.t.sol b/test/FixedRewardPool.t.sol index 2164a4f..e6cd0c7 100644 --- a/test/FixedRewardPool.t.sol +++ b/test/FixedRewardPool.t.sol @@ -26,7 +26,7 @@ contract TestFixedRewardPool is Test { testNFT = new TestNFT(); ownedWeightedNFT = new OwnedWeightedNFT(address(testNFT), alice); fixedRewardPool = new FixedRewardPool(); - fixedRewardPool.initialize(address(ownedWeightedNFT), 10, 0.1 ether); + fixedRewardPool.initialize(address(ownedWeightedNFT), 10, 0.1 ether, 3); } function testOwnership() public { @@ -114,5 +114,11 @@ contract TestFixedRewardPool is Test { assertEq(alice.balance, 0.075 ether); assertEq(fixedRewardPool.totalStakedWeight(), 100); assertEq(fixedRewardPool.tokenWeight(1), 0); + + vm.roll(15); + assertEq(fixedRewardPool.pendingReward(bob), 0); + vm.prank(alice); + fixedRewardPool.withdraw(3); + assertEq(alice.balance, 0.075 ether); } }