chore: add rewardTokenInfos

pyk · May 30, 2024 · 1a80de9 · 1a80de9
1 parent 1797d09
commit 1a80de9
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 32 deletions.
diff --git a/src/LlamaLocker.sol b/src/LlamaLocker.sol
@@ -60,7 +60,7 @@ contract LlamaLocker is ERC721Holder, Ownable2Step {
 
     Epoch[] public epochs;
     mapping(uint256 nftId => LockInfo info) public locks;
-    mapping(address token => RewardTokenInfo info) public rewards;
+    mapping(address token => RewardTokenInfo info) public rewardTokenInfos;
     mapping(uint256 nftId => mapping(IERC20 token => NFTYield)) private nftYields;
 
     // TODO(pyk): check variables below
@@ -177,10 +177,10 @@ contract LlamaLocker is ERC721Holder, Ownable2Step {
 
     function _addRewardToken(address token_) private {
         if (address(token_) == address(0)) revert InvalidRewardToken();
-        if (rewards[token_].updatedAt > 0) revert InvalidRewardToken();
+        if (rewardTokenInfos[token_].updatedAt > 0) revert InvalidRewardToken();
         tokens.push(token_);
-        rewards[token_].updatedAt = block.timestamp.toUint48();
-        rewards[token_].epochEndAt = block.timestamp.toUint48();
+        rewardTokenInfos[token_].updatedAt = block.timestamp.toUint48();
+        rewardTokenInfos[token_].epochEndAt = block.timestamp.toUint48();
         emit NewRewardToken(token_);
     }
 
@@ -197,6 +197,29 @@ contract LlamaLocker is ERC721Holder, Ownable2Step {
         count = tokens.length;
     }
 
+    function distributeRewards(IERC20 _token, uint256 _amount) external onlyOwner {
+        // if (rewardTokenInfos[_token].updatedAt == 0) revert RewardTokenNotExists();
+        // if (_amount == 0) revert RewardAmountInvalid();
+        // if (totalLockedNFT == 0) revert NoLockers();
+
+        // _updateRewardStates();
+
+        // RewardState storage rewardState = rewardStates[_token];
+        // if (block.timestamp >= rewardState.epochEndAt) {
+        //     rewardState.rewardPerSecond = (_amount / EPOCH_DURATION).toUint208();
+        // } else {
+        //     uint256 remaining = rewardState.epochEndAt - block.timestamp;
+        //     uint256 leftover = remaining * rewardState.rewardPerSecond;
+        //     rewardState.rewardPerSecond = ((_amount + leftover) / EPOCH_DURATION).toUint208();
+        // }
+
+        // rewardState.updatedAt = block.timestamp.toUint48();
+        // rewardState.epochEndAt = (block.timestamp + EPOCH_DURATION).toUint48();
+
+        // _token.safeTransferFrom(msg.sender, address(this), _amount);
+        // emit RewardDistributed(_token, _amount);
+    }
+
     // function getYieldInfo(IERC20 _token) external view returns (YieldInfo memory info) {
     //     info = yields[_token];
     // }
@@ -240,32 +263,6 @@ contract LlamaLocker is ERC721Holder, Ownable2Step {
     //     }
     // }
 
-    // / @notice Distribute rewards to lockers
-    // / @param _token The reward token address
-    // / @param _amount The amount of reward token
-    // function distribute(IERC20 _token, uint256 _amount) external onlyOwner {
-    //     if (rewardStates[_token].updatedAt == 0) revert RewardTokenNotExists();
-    //     if (_amount == 0) revert RewardAmountInvalid();
-    //     if (totalLockedNFT == 0) revert NoLockers();
-
-    //     _updateRewardStates();
-
-    //     RewardState storage rewardState = rewardStates[_token];
-    //     if (block.timestamp >= rewardState.epochEndAt) {
-    //         rewardState.rewardPerSecond = (_amount / EPOCH_DURATION).toUint208();
-    //     } else {
-    //         uint256 remaining = rewardState.epochEndAt - block.timestamp;
-    //         uint256 leftover = remaining * rewardState.rewardPerSecond;
-    //         rewardState.rewardPerSecond = ((_amount + leftover) / EPOCH_DURATION).toUint208();
-    //     }
-
-    //     rewardState.updatedAt = block.timestamp.toUint48();
-    //     rewardState.epochEndAt = (block.timestamp + EPOCH_DURATION).toUint48();
-
-    //     _token.safeTransferFrom(msg.sender, address(this), _amount);
-    //     emit RewardDistributed(_token, _amount);
-    // }
-
     // function _earned(address _account, IERC20 _token, uint256 _lockedAmount) internal view returns (uint256) {
     //     uint256 rewardPerNFTDiff = _rewardPerNFT(_token) - accountRewardPerNFTPaid[_account][_token];
     //     return (_lockedAmount * rewardPerNFTDiff) + rewards[_account][_token];

diff --git a/test/AddRewardTokens.t.sol b/test/AddRewardTokens.t.sol
@@ -52,13 +52,13 @@ contract AddRewardTokensTest is Test {
 
         // addRewardTokens() should set initial values
         (uint208 amountPerSecond, uint48 epochEndAt, uint208 amountPerNFTStored, uint48 updatedAt) =
-            llama.rewards(address(token0));
+            llama.rewardTokenInfos(address(token0));
         assertEq(amountPerSecond, 0, "invalid token0 amountPerSecond");
         assertEq(epochEndAt, block.timestamp, "invalid token0 epochEndAt");
         assertEq(amountPerNFTStored, 0, "invalid token0 amountPerNFTStored");
         assertEq(updatedAt, block.timestamp, "invalid token0 updatedAt");
 
-        (amountPerSecond, epochEndAt, amountPerNFTStored, updatedAt) = llama.rewards(address(token1));
+        (amountPerSecond, epochEndAt, amountPerNFTStored, updatedAt) = llama.rewardTokenInfos(address(token1));
         assertEq(amountPerSecond, 0, "invalid token1 amountPerSecond");
         assertEq(epochEndAt, block.timestamp, "invalid token1 epochEndAt");
         assertEq(amountPerNFTStored, 0, "invalid token1 amountPerNFTStored");