From 0192f7d0a490311c82414069a93e5caf5ba0ad9c Mon Sep 17 00:00:00 2001
From: zhi <thecoderzhi@gmail.com>
Date: Fri, 9 Aug 2024 21:19:12 +0800
Subject: [PATCH] update device gauge

---
 contracts/gauges/DeviceGauge.sol                | 17 +++++++++++------
 .../{IDeviceNFT.sol => IWeightedNFT.sol}        |  5 ++---
 contracts/test/TestDeviceNFT.sol                |  8 ++++++--
 test/TestVoter.t.sol                            |  3 ++-
 4 files changed, 21 insertions(+), 12 deletions(-)
 rename contracts/interfaces/{IDeviceNFT.sol => IWeightedNFT.sol} (53%)

diff --git a/contracts/gauges/DeviceGauge.sol b/contracts/gauges/DeviceGauge.sol
index f28e37e..3dc8af8 100644
--- a/contracts/gauges/DeviceGauge.sol
+++ b/contracts/gauges/DeviceGauge.sol
@@ -2,9 +2,10 @@
 pragma solidity ^0.8.0;
 
 import {ERC721Holder} from "@openzeppelin/contracts/token/ERC721/utils/ERC721Holder.sol";
+import {IERC721} from "@openzeppelin/contracts/token/ERC721/IERC721.sol";
 
 import {IVoter} from "../interfaces/IVoter.sol";
-import {IDeviceNFT} from "../interfaces/IDeviceNFT.sol";
+import {IWeightedNFT} from "../interfaces/IWeightedNFT.sol";
 import {RewardGauge} from "./RewardGauge.sol";
 import {IIncentive} from "../interfaces/IIncentive.sol";
 
@@ -15,12 +16,16 @@ contract DeviceGauge is RewardGauge, ERC721Holder {
   mapping(uint256 => address) public tokenStaker;
   mapping(uint256 => uint256) public tokenWeight;
 
+  address public immutable weightedNFT;
+
   constructor(
     address _forwarder,
-    address _stakingToken,
+    address _weightedNFT,
     address _voter,
     address _incentives
-  ) RewardGauge(_forwarder, _stakingToken, _voter, _incentives) {}
+  ) RewardGauge(_forwarder, IWeightedNFT(_weightedNFT).nft(), _voter, _incentives) {
+    weightedNFT = _weightedNFT;
+  }
 
   function _depositFor(uint256 _tokenId, address _recipient) internal override nonReentrant {
     if (_tokenId == 0) revert ZeroAmount();
@@ -29,8 +34,8 @@ contract DeviceGauge is RewardGauge, ERC721Holder {
     address sender = _msgSender();
     _updateRewards(_recipient);
 
-    IDeviceNFT(stakingToken).safeTransferFrom(sender, address(this), _tokenId);
-    uint256 _amount = IDeviceNFT(stakingToken).weight(_tokenId);
+    IERC721(stakingToken).safeTransferFrom(sender, address(this), _tokenId);
+    uint256 _amount = IWeightedNFT(weightedNFT).weight(_tokenId);
     totalSupply += _amount;
     balanceOf[_recipient] += _amount;
     tokenStaker[_tokenId] = _recipient;
@@ -50,7 +55,7 @@ contract DeviceGauge is RewardGauge, ERC721Holder {
     uint256 _amount = tokenWeight[_tokenId];
     totalSupply -= _amount;
     balanceOf[sender] -= _amount;
-    IDeviceNFT(stakingToken).safeTransferFrom(address(this), sender, _tokenId);
+    IERC721(stakingToken).safeTransferFrom(address(this), sender, _tokenId);
     delete tokenStaker[_tokenId];
     delete tokenWeight[_tokenId];
     updateWeightBalance(sender);
diff --git a/contracts/interfaces/IDeviceNFT.sol b/contracts/interfaces/IWeightedNFT.sol
similarity index 53%
rename from contracts/interfaces/IDeviceNFT.sol
rename to contracts/interfaces/IWeightedNFT.sol
index 9f4ea9f..3cf63c0 100644
--- a/contracts/interfaces/IDeviceNFT.sol
+++ b/contracts/interfaces/IWeightedNFT.sol
@@ -1,8 +1,7 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.0;
 
-import {IERC721} from "@openzeppelin/contracts/token/ERC721/IERC721.sol";
-
-interface IDeviceNFT is IERC721 {
+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 380d14e..af171f5 100644
--- a/contracts/test/TestDeviceNFT.sol
+++ b/contracts/test/TestDeviceNFT.sol
@@ -1,9 +1,9 @@
 pragma solidity ^0.8.0;
 
 import {ERC721} from "@openzeppelin/contracts/token/ERC721/ERC721.sol";
-import {IDeviceNFT} from "../interfaces/IDeviceNFT.sol";
+import {IWeightedNFT} from "../interfaces/IWeightedNFT.sol";
 
-contract TestDeviceNFT is IDeviceNFT, ERC721 {
+contract TestDeviceNFT is IWeightedNFT, ERC721 {
   mapping(uint256 => uint256) public weightOf;
 
   constructor(string memory name_, string memory symbol_) ERC721(name_, symbol_) {
@@ -19,6 +19,10 @@ contract TestDeviceNFT is IDeviceNFT, ERC721 {
     return weightOf[tokenId];
   }
 
+  function nft() external view override returns (address) {
+    return address(this);
+  }
+
   function setWeight(uint256 _tokenId, uint256 _weight) public {
     weightOf[_tokenId] = _weight;
   }
diff --git a/test/TestVoter.t.sol b/test/TestVoter.t.sol
index 9bfd033..8a9e3b7 100644
--- a/test/TestVoter.t.sol
+++ b/test/TestVoter.t.sol
@@ -10,6 +10,7 @@ import {IVoter} from "../contracts/interfaces/IVoter.sol";
 import {IRewardGauge} from "../contracts/interfaces/IRewardGauge.sol";
 import {IVault} from "../contracts/interfaces/IVault.sol";
 import {DAOForwarder} from "../contracts/DAOForwarder.sol";
+import {TestDeviceNFT} from "../contracts/test/TestDeviceNFT.sol";
 import {TestStrategyManager} from "../contracts/test/TestStrategyManager.sol";
 import {FactoryRegistry} from "../contracts/factories/FactoryRegistry.sol";
 import {GaugeFactory} from "../contracts/factories/GaugeFactory.sol";
@@ -162,7 +163,7 @@ contract TestVoter is Test {
     voter.createGauge(poolFactory, address(pool), 0);
 
     // 3. create NFT gauge
-    address deviceNFT = address(10);
+    address deviceNFT = address(new TestDeviceNFT("name", "symbol"));
     voter.createGauge(poolFactory, deviceNFT, 1);
     gauge = voter.gauges(address(deviceNFT));
     assertTrue(gauge != address(0));