add weight manager

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];
+    }
+}

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;
   }
 }

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;

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);
+}

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 {

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 {

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);

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();

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,