From f6f936dc914d5a85dd989df5cad6398429c38637 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 | 9 ++++-- contracts/gauges/DeviceGauge.sol | 13 +++++---- contracts/interfaces/IWeightManager.sol | 10 +++++++ contracts/interfaces/IWeightedNFT.sol | 7 ----- contracts/test/TestDeviceNFT.sol | 19 +----------- test/TestIncentive.sol | 17 ++++++----- test/TestVault.t.sol | 5 +++- test/TestVoter.t.sol | 5 +++- test/flow.test.ts | 3 +- 10 files changed, 84 insertions(+), 43 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..0975828 100644 --- a/contracts/factories/GaugeFactory.sol +++ b/contracts/factories/GaugeFactory.sol @@ -9,6 +9,11 @@ contract GaugeFactory is IGaugeFactory { uint8 constant Erc20Gauge = 0; uint8 constant DeviceNFTGauge = 1; uint8 constant WithdrawGauge = 2; + address private weightManager; + + constructor(address _weightManager) { + weightManager = _weightManager; + } function createGauge( address _forwarder, @@ -41,10 +46,10 @@ contract GaugeFactory is IGaugeFactory { address _deviceNFT, address _incentives ) internal returns (address gauge) { - gauge = address(new DeviceGauge(_forwarder, _deviceNFT, msg.sender, _incentives)); + gauge = address(new DeviceGauge(_forwarder, _deviceNFT, msg.sender, weightManager, _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..97c8d62 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,16 @@ contract DeviceGauge is RewardGauge, ERC721Holder { mapping(uint256 => address) public tokenStaker; mapping(uint256 => uint256) public tokenWeight; - address public immutable weightedNFT; + IWeightManager public immutable weightManager; constructor( address _forwarder, - address _weightedNFT, + address _nft, address _voter, + address _weightManager, address _incentives - ) RewardGauge(_forwarder, IWeightedNFT(_weightedNFT).nft(), _voter, _incentives) { - weightedNFT = _weightedNFT; + ) RewardGauge(_forwarder, _nft, _voter, _incentives) { + weightManager = IWeightManager(_weightManager); } function _depositFor(uint256 _tokenId, address _recipient) internal override nonReentrant { @@ -35,7 +36,7 @@ contract DeviceGauge is RewardGauge, ERC721Holder { _updateRewards(_recipient); IERC721(stakingToken).safeTransferFrom(sender, address(this), _tokenId); - uint256 _amount = IWeightedNFT(weightedNFT).weight(_tokenId); + uint256 _amount = 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..a561bc5 100644 --- a/contracts/test/TestDeviceNFT.sol +++ b/contracts/test/TestDeviceNFT.sol @@ -1,30 +1,13 @@ 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 { diff --git a/test/TestIncentive.sol b/test/TestIncentive.sol index a42dd9c..d400db6 100644 --- a/test/TestIncentive.sol +++ b/test/TestIncentive.sol @@ -4,6 +4,7 @@ pragma solidity ^0.8.13; import {Test, console} from "forge-std/Test.sol"; import {ERC20Gauge} from "../contracts/gauges/ERC20Gauge.sol"; import {Voter} from "../contracts/Voter.sol"; +import {WeightManager} from "../contracts/WeightManager.sol"; import {IRewardGauge} from "../contracts/interfaces/IRewardGauge.sol"; import {DAOForwarder} from "../contracts/DAOForwarder.sol"; import {TestToken} from "../contracts/test/TestToken.sol"; @@ -18,6 +19,7 @@ contract TestIncentive is Test { address[] public rewardTokens; DAOForwarder public forwarder; Voter public voter; + WeightManager public wm; TestStrategyManager public strategyManager; address public factoryRegistry; Incentives public inti; @@ -49,6 +51,7 @@ contract TestIncentive is Test { lp = new TestToken("lp", "test-lp"); voter = new Voter(address(forwarder), address(strategyManager), factoryRegistry); + wm = new WeightManager(); inti = new Incentives(address(forwarder), address(voter), rewardTokens); gauge = new ERC20Gauge(address(forwarder), address(lp), address(voter), address(inti)); vm.prank(address(voter)); @@ -57,7 +60,7 @@ contract TestIncentive is Test { dlp = new TestDeviceNFT("device nft", "device"); dinti = new Incentives(address(forwarder), address(voter), rewardTokens); - deviceGauge = new DeviceGauge(address(forwarder), address(dlp), address(voter), address(dinti)); + deviceGauge = new DeviceGauge(address(forwarder), address(dlp), address(voter), address(wm), address(dinti)); vm.prank(address(voter)); dinti.setGauge(address(deviceGauge)); voter.reviveGauge(address(deviceGauge)); @@ -120,21 +123,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(3, dinti.balanceOf(address(this))); + assertEq(3, 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(2, dinti.balanceOf(address(this))); + assertEq(2, dinti.totalSupply()); } function test_notifyReward_Claim() external { diff --git a/test/TestVault.t.sol b/test/TestVault.t.sol index 89eb685..61bb869 100644 --- a/test/TestVault.t.sol +++ b/test/TestVault.t.sol @@ -4,6 +4,7 @@ pragma solidity ^0.8.13; import {Test, console} from "forge-std/Test.sol"; import {Vault} from "../contracts/Vault.sol"; import {Voter} from "../contracts/Voter.sol"; +import {WeightManager} from "../contracts/WeightManager.sol"; import {IVault} from "../contracts/interfaces/IVault.sol"; import {DAOForwarder} from "../contracts/DAOForwarder.sol"; import {TestToken} from "../contracts/test/TestToken.sol"; @@ -16,6 +17,7 @@ import {IncentivesFactory} from "../contracts/factories/IncentivesFactory.sol"; contract TestVault is Test { Vault public vault; Voter public voter; + WeightManager public weightManager; DAOForwarder public forwarder; address public poolFactory; GaugeFactory public gaugeFactory; @@ -25,7 +27,8 @@ contract TestVault is Test { function setUp() public { forwarder = new DAOForwarder(); - gaugeFactory = new GaugeFactory(); + weightManager = new WeightManager(); + gaugeFactory = new GaugeFactory(address(weightManager)); incentiveFactory = new IncentivesFactory(); strategyManager = new TestStrategyManager(); strategyManager.setShare(address (this), 100); diff --git a/test/TestVoter.t.sol b/test/TestVoter.t.sol index fc4f949..ffb0084 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"; @@ -21,6 +22,7 @@ import "../contracts/libraries/ProtocolTimeLibrary.sol"; contract TestVoter is Test { Voter public voter; Vault public vault; + WeightManager public weightManager; TestToken public pool; address public poolFactory; DAOForwarder public forwarder; @@ -31,8 +33,9 @@ contract TestVoter is Test { function setUp() public { forwarder = new DAOForwarder(); + weightManager = new WeightManager(); pool = new TestToken("test-pool", "pool"); - gaugeFactory = new GaugeFactory(); + gaugeFactory = new GaugeFactory(address(weightManager)); incentiveFactory = new IncentivesFactory(); poolFactory = address(1); strategyManager = new TestStrategyManager(); diff --git a/test/flow.test.ts b/test/flow.test.ts index 586cb80..0fc88fb 100644 --- a/test/flow.test.ts +++ b/test/flow.test.ts @@ -8,7 +8,8 @@ describe('Flow', function () { const forwarder = await ethers.deployContract('Forwarder'); const poolFactory = await ethers.deployContract('EmptyPoolFactory', []); const incentiveFactory = await ethers.deployContract('IncentivesFactory', []); - const gaugeFactory = await ethers.deployContract('GaugeFactory', []); + const weightManager = await ethers.deployContract('WeightManager', []); + const gaugeFactory = await ethers.deployContract('GaugeFactory', [weightManager.target]); const factoryRegistry = await ethers.deployContract('FactoryRegistry', [ poolFactory.target, incentiveFactory.target,