From 6b8a17d8b0f46ef433cfe652e370813e594c48a6 Mon Sep 17 00:00:00 2001 From: zhi Date: Tue, 20 Aug 2024 00:43:40 +0800 Subject: [PATCH] add weight manager --- contracts/WeightManager.sol | 39 +++++++++++++++++++++++++ contracts/factories/GaugeFactory.sol | 2 +- contracts/gauges/DeviceGauge.sol | 18 ++++++++---- contracts/interfaces/IWeightManager.sol | 10 +++++++ contracts/interfaces/IWeightedNFT.sol | 7 ----- contracts/test/TestDeviceNFT.sol | 20 +------------ test/TestIncentive.sol | 12 ++++---- test/TestVoter.t.sol | 1 + 8 files changed, 70 insertions(+), 39 deletions(-) create mode 100644 contracts/WeightManager.sol create mode 100644 contracts/interfaces/IWeightManager.sol delete mode 100644 contracts/interfaces/IWeightedNFT.sol diff --git a/contracts/WeightManager.sol b/contracts/WeightManager.sol new file mode 100644 index 0000000..f520d1f --- /dev/null +++ b/contracts/WeightManager.sol @@ -0,0 +1,39 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.0; + +import "./interfaces/IWeightManager.sol"; +import "@openzeppelin/contracts/access/Ownable.sol"; + +contract WeightManager is Ownable, IWeightManager { + mapping(address => mapping(uint256 => uint256)) public weights; + mapping(address => address) public operators; + + modifier onlyOperator(address _nft) { + if (operators[_nft] != msg.sender) { + revert NotOperator(); + } + _; + } + + function addOperator(address _nft, address _operator) external onlyOwner { + if (operators[_nft] != address(0)) revert DuplicateOperator(); + operators[_nft] = _operator; + emit OperatorSet(_nft, _operator); + } + + function changeOperator(address _nft, address _operator) external onlyOperator(_nft) { + operators[_nft] = _operator; + emit OperatorSet(_nft, _operator); + } + + function setWeight(address _nft, uint256 _tokenId, uint256 _weight) external onlyOperator(_nft) { + weights[_nft][_tokenId] = _weight; + } + + function weight(address _nft, uint256 _tokenId) external view override returns (uint256) { + if (operators[_nft] == address(0)) { + return 1; + } + return weights[_nft][_tokenId]; + } +} diff --git a/contracts/factories/GaugeFactory.sol b/contracts/factories/GaugeFactory.sol index e52de1e..aa503de 100644 --- a/contracts/factories/GaugeFactory.sol +++ b/contracts/factories/GaugeFactory.sol @@ -44,7 +44,7 @@ contract GaugeFactory is IGaugeFactory { gauge = address(new DeviceGauge(_forwarder, _deviceNFT, msg.sender, _incentives)); } - function createWithdrawalGauge(address _guage) internal returns (address gauge) { + function createWithdrawalGauge(address _guage) internal pure returns (address gauge) { gauge = _guage; } } diff --git a/contracts/gauges/DeviceGauge.sol b/contracts/gauges/DeviceGauge.sol index 2f16793..12d4f42 100644 --- a/contracts/gauges/DeviceGauge.sol +++ b/contracts/gauges/DeviceGauge.sol @@ -5,7 +5,7 @@ import {ERC721Holder} from "@openzeppelin/contracts/token/ERC721/utils/ERC721Hol import {IERC721} from "@openzeppelin/contracts/token/ERC721/IERC721.sol"; import {IVoter} from "../interfaces/IVoter.sol"; -import {IWeightedNFT} from "../interfaces/IWeightedNFT.sol"; +import {IWeightManager} from "../interfaces/IWeightManager.sol"; import {RewardGauge} from "./RewardGauge.sol"; import {IIncentive} from "../interfaces/IIncentive.sol"; @@ -16,15 +16,18 @@ contract DeviceGauge is RewardGauge, ERC721Holder { mapping(uint256 => address) public tokenStaker; mapping(uint256 => uint256) public tokenWeight; - address public immutable weightedNFT; + address public weightManager; constructor( address _forwarder, - address _weightedNFT, + address _nft, address _voter, address _incentives - ) RewardGauge(_forwarder, IWeightedNFT(_weightedNFT).nft(), _voter, _incentives) { - weightedNFT = _weightedNFT; + ) RewardGauge(_forwarder, _nft, _voter, _incentives) {} + + function setWeightManager(address _weightManager) external { + if (msg.sender != IVoter(voter).team()) revert NotTeam(); + weightManager = _weightManager; } function _depositFor(uint256 _tokenId, address _recipient) internal override nonReentrant { @@ -35,7 +38,10 @@ contract DeviceGauge is RewardGauge, ERC721Holder { _updateRewards(_recipient); IERC721(stakingToken).safeTransferFrom(sender, address(this), _tokenId); - uint256 _amount = IWeightedNFT(weightedNFT).weight(_tokenId); + uint256 _amount = 1; + if (weightManager != address(0)) { + _amount = IWeightManager(weightManager).weight(stakingToken, _tokenId); + } totalSupply += _amount; balanceOf[_recipient] += _amount; tokenStaker[_tokenId] = _recipient; diff --git a/contracts/interfaces/IWeightManager.sol b/contracts/interfaces/IWeightManager.sol new file mode 100644 index 0000000..797007e --- /dev/null +++ b/contracts/interfaces/IWeightManager.sol @@ -0,0 +1,10 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.0; + +interface IWeightManager { + error DuplicateOperator(); + error NotOperator(); + event OperatorSet(address indexed nft, address indexed operator); + + function weight(address nft, uint256 tokenId) external view returns (uint256); +} diff --git a/contracts/interfaces/IWeightedNFT.sol b/contracts/interfaces/IWeightedNFT.sol deleted file mode 100644 index 3cf63c0..0000000 --- a/contracts/interfaces/IWeightedNFT.sol +++ /dev/null @@ -1,7 +0,0 @@ -// SPDX-License-Identifier: MIT -pragma solidity ^0.8.0; - -interface IWeightedNFT { - function weight(uint256 tokenId) external view returns (uint256); - function nft() external view returns (address); -} diff --git a/contracts/test/TestDeviceNFT.sol b/contracts/test/TestDeviceNFT.sol index 2c52c1b..fec3ee4 100644 --- a/contracts/test/TestDeviceNFT.sol +++ b/contracts/test/TestDeviceNFT.sol @@ -1,34 +1,16 @@ pragma solidity ^0.8.0; import {ERC721} from "@openzeppelin/contracts/token/ERC721/ERC721.sol"; -import {IWeightedNFT} from "../interfaces/IWeightedNFT.sol"; -contract TestDeviceNFT is IWeightedNFT, ERC721 { - mapping(uint256 => uint256) public weightOf; +contract TestDeviceNFT is ERC721 { constructor(string memory name_, string memory symbol_) ERC721(name_, symbol_) { _mint(msg.sender, 1); _mint(msg.sender, 2); _mint(msg.sender, 3); - weightOf[1] = 1 ether; - weightOf[2] = 2 ether; - weightOf[3] = 3 ether; - } - - function weight(uint256 tokenId) external view returns (uint256) { - return weightOf[tokenId]; - } - - function nft() external view override returns (address) { - return address(this); - } - - function setWeight(uint256 _tokenId, uint256 _weight) public { - weightOf[_tokenId] = _weight; } function mint(address to, uint tokenId) external { _mint(to, tokenId); - weightOf[tokenId] = 1 ether; } } diff --git a/test/TestIncentive.sol b/test/TestIncentive.sol index a42dd9c..9014864 100644 --- a/test/TestIncentive.sol +++ b/test/TestIncentive.sol @@ -120,21 +120,21 @@ contract TestIncentive is Test { // 2. deposit lp into gauge and callback to incentives deviceGauge.deposit(1); - assertEq(1 ether, dinti.balanceOf(address(this))); - assertEq(1 ether, dinti.totalSupply()); + assertEq(1, dinti.balanceOf(address(this))); + assertEq(1, dinti.totalSupply()); // 3. again deposit to check increase deviceGauge.deposit(2); - assertEq(3 ether, dinti.balanceOf(address(this))); - assertEq(3 ether, dinti.totalSupply()); + assertEq(2, dinti.balanceOf(address(this))); + assertEq(2, dinti.totalSupply()); // 4. rewardRate should be zero due to not notifyReward. assertEq(0, dinti.rewardRate(rewardTokens[0])); // 5. withdraw deviceGauge.withdraw(1); - assertEq(2 ether, dinti.balanceOf(address(this))); - assertEq(2 ether, dinti.totalSupply()); + assertEq(1, dinti.balanceOf(address(this))); + assertEq(1, dinti.totalSupply()); } function test_notifyReward_Claim() external { diff --git a/test/TestVoter.t.sol b/test/TestVoter.t.sol index fc4f949..26ad051 100644 --- a/test/TestVoter.t.sol +++ b/test/TestVoter.t.sol @@ -6,6 +6,7 @@ import {stdError} from "forge-std/StdError.sol"; import {TestToken} from "../contracts/test/TestToken.sol"; import {Vault} from "../contracts/Vault.sol"; import {Voter} from "../contracts/Voter.sol"; +import {WeightManager} from "../contracts/WeightManager.sol"; import {IVoter} from "../contracts/interfaces/IVoter.sol"; import {IRewardGauge} from "../contracts/interfaces/IRewardGauge.sol"; import {IVault} from "../contracts/interfaces/IVault.sol";