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

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

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;

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

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 {

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