add end blocks

iotexproject · Oct 24, 2024 · 0624c30 · 0624c30
1 parent ebdf607
commit 0624c30
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 5 deletions.
diff --git 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
@@ -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);
   }
 }