update StakingContractMainnet to support vault tokens (#11)

* add incentive flag for rounding down rewards

* move reward rounding logic to claim handler

* add unit test

* preserve reward delta

* add comments

* update deploy script

* feat: add batch subscription (#12)

* feat: add NftVaultManager (#13)

* feat: add NftVaultManager

* add NftVaultManager deploy script

* feat: add depositBatch to NftVaultManager

---------

Co-authored-by: Alec Ananian <[email protected]>

* rerun fmt after rebase

* fix unit tests failing with too many rejects

* run ci

* add `isRewardRounded` to `IncentiveCreated` event

---------

Co-authored-by: crisis <[email protected]>
Co-authored-by: Mark Vujevits <[email protected]>

.github/workflows/test.yml

@@ -0,0 +1,43 @@
+name: test
+
+on:
+    workflow_dispatch:
+    pull_request:
+    push:
+        branches:
+        - "main"
+
+env:
+  FOUNDRY_PROFILE: ci
+
+jobs:
+  foundry:
+    strategy:
+      fail-fast: true
+
+    name: Foundry project
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+        with:
+          submodules: recursive
+
+      - name: Install Foundry
+        uses: foundry-rs/foundry-toolchain@v1
+        with:
+          version: nightly
+
+      - name: Run Forge build
+        run: |
+          forge --version
+          forge build --sizes
+        id: build
+
+      - name: Run Forge tests
+        run: |
+          forge test -vvv
+        id: forge-test
+
+      - name: Forge style
+        run: |
+          forge fmt --check

script/NftVaultManager.s.sol

@@ -0,0 +1,15 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity >=0.8.13;
+
+import "forge-std/Script.sol";
+import "../src/Vault/NftVaultManager.sol";
+
+contract NftVaultManagerScript is Script {
+    function run() public {
+        vm.startBroadcast();
+
+        new NftVaultManager();
+
+        vm.stopBroadcast();
+    }
+}

script/StakingContract.s.sol

@@ -6,8 +6,7 @@ import "../src/Rewards/StakingContractMainnet.sol";
 
 contract StakingContractScript is Script {
     function run() public {
-        uint256 deployerPrivateKey = vm.envUint("PRIVATE_KEY");
-        vm.startBroadcast(deployerPrivateKey);
+        vm.startBroadcast();
 
         new StakingContractMainnet();
 

sh/deployArbitrum.sh

@@ -3,3 +3,4 @@ source .env
 
 # To deploy and verify our contract
 forge script script/MagicswapV2.s.sol:MagicswapV2Script --aws --rpc-url $ARBITRUM_RPC --broadcast --verify -vvvv
+forge script script/StakingContract.s.sol:StakingContractScript --aws --rpc-url $ARBITRUM_RPC --broadcast --verify -vvvv

sh/deployArbitrumSepolia.sh

@@ -3,3 +3,5 @@ source .env
 
 # To deploy and verify our contract
 forge script script/MagicswapV2.s.sol:MagicswapV2Script --aws --rpc-url $ARBITRUM_SEPOLIA_RPC --broadcast --verify -vvvv
+forge script script/StakingContract.s.sol:StakingContractScript --aws --rpc-url $ARBITRUM_SEPOLIA_RPC --broadcast --verify -vvvv
+forge script script/NftVaultManager.s.sol:NftVaultManagerScript --aws --rpc-url $ARBITRUM_SEPOLIA_RPC --broadcast --verify -vvvv

src/Rewards/StakingContractMainnet.sol

@@ -20,6 +20,7 @@ contract StakingContractMainnet is ReentrancyGuard {
         address token; // 2nd slot
         address rewardToken; // 3rd slot
         uint32 endTime; // 3rd slot
+        bool isRewardRounded; // 3rd slot
         uint256 rewardPerLiquidity; // 4th slot
         uint32 lastRewardTime; // 5th slot
         uint112 rewardRemaining; // 5th slot
@@ -65,7 +66,8 @@ contract StakingContractMainnet is ReentrancyGuard {
         uint256 id,
         uint256 amount,
         uint256 startTime,
-        uint256 endTime
+        uint256 endTime,
+        bool isRewardRounded
     );
     event IncentiveUpdated(uint256 indexed id, int256 changeAmount, uint256 newStartTime, uint256 newEndTime);
     event Stake(address indexed token, address indexed user, uint256 amount);
@@ -74,11 +76,14 @@ contract StakingContractMainnet is ReentrancyGuard {
     event Unsubscribe(uint256 indexed id, address indexed user);
     event Claim(uint256 indexed id, address indexed user, uint256 amount);
 
-    function createIncentive(address token, address rewardToken, uint112 rewardAmount, uint32 startTime, uint32 endTime)
-        external
-        nonReentrant
-        returns (uint256 incentiveId)
-    {
+    function createIncentive(
+        address token,
+        address rewardToken,
+        uint112 rewardAmount,
+        uint32 startTime,
+        uint32 endTime,
+        bool isRewardRounded
+    ) external nonReentrant returns (uint256 incentiveId) {
         if (rewardAmount <= 0) revert InvalidInput();
 
         if (startTime < block.timestamp) startTime = uint32(block.timestamp);
@@ -99,13 +104,16 @@ contract StakingContractMainnet is ReentrancyGuard {
             rewardToken: rewardToken,
             lastRewardTime: startTime,
             endTime: endTime,
+            isRewardRounded: isRewardRounded,
             rewardRemaining: rewardAmount,
             liquidityStaked: 0,
             // Initial value of rewardPerLiquidity can be arbitrarily set to a non-zero value.
             rewardPerLiquidity: type(uint256).max / 2
         });
 
-        emit IncentiveCreated(token, rewardToken, msg.sender, incentiveId, rewardAmount, startTime, endTime);
+        emit IncentiveCreated(
+            token, rewardToken, msg.sender, incentiveId, rewardAmount, startTime, endTime, isRewardRounded
+        );
     }
 
     function updateIncentive(uint256 incentiveId, int112 changeAmount, uint32 newStartTime, uint32 newEndTime)
@@ -119,18 +127,24 @@ contract StakingContractMainnet is ReentrancyGuard {
         _accrueRewards(incentive);
 
         if (newStartTime != 0) {
-            if (newStartTime < block.timestamp) newStartTime = uint32(block.timestamp);
+            if (newStartTime < block.timestamp) {
+                newStartTime = uint32(block.timestamp);
+            }
 
             incentive.lastRewardTime = newStartTime;
         }
 
         if (newEndTime != 0) {
-            if (newEndTime < block.timestamp) newEndTime = uint32(block.timestamp);
+            if (newEndTime < block.timestamp) {
+                newEndTime = uint32(block.timestamp);
+            }
 
             incentive.endTime = newEndTime;
         }
 
-        if (incentive.lastRewardTime >= incentive.endTime) revert InvalidTimeFrame();
+        if (incentive.lastRewardTime >= incentive.endTime) {
+            revert InvalidTimeFrame();
+        }
 
         if (changeAmount > 0) {
             incentive.rewardRemaining += uint112(changeAmount);
@@ -139,7 +153,9 @@ contract StakingContractMainnet is ReentrancyGuard {
         } else if (changeAmount < 0) {
             uint112 transferOut = uint112(-changeAmount);
 
-            if (transferOut > incentive.rewardRemaining) transferOut = incentive.rewardRemaining;
+            if (transferOut > incentive.rewardRemaining) {
+                transferOut = incentive.rewardRemaining;
+            }
 
             unchecked {
                 incentive.rewardRemaining -= transferOut;
@@ -229,14 +245,28 @@ contract StakingContractMainnet is ReentrancyGuard {
         emit Unstake(token, msg.sender, amount);
     }
 
+    function subscribeToIncentives(uint256[] memory incentiveIds) external {
+        uint256 n = incentiveIds.length;
+
+        for (uint256 i = 0; i < n; i = _increment(i)) {
+            subscribeToIncentive(incentiveIds[i]);
+        }
+    }
+
     function subscribeToIncentive(uint256 incentiveId) public nonReentrant {
-        if (incentiveId > incentiveCount || incentiveId <= 0) revert InvalidInput();
+        if (incentiveId > incentiveCount || incentiveId <= 0) {
+            revert InvalidInput();
+        }
 
-        if (rewardPerLiquidityLast[msg.sender][incentiveId] != 0) revert AlreadySubscribed();
+        if (rewardPerLiquidityLast[msg.sender][incentiveId] != 0) {
+            revert AlreadySubscribed();
+        }
 
         Incentive storage incentive = incentives[incentiveId];
 
-        if (userStakes[msg.sender][incentive.token].liquidity <= 0) revert NotStaked();
+        if (userStakes[msg.sender][incentive.token].liquidity <= 0) {
+            revert NotStaked();
+        }
 
         _accrueRewards(incentive);
 
@@ -262,14 +292,18 @@ contract StakingContractMainnet is ReentrancyGuard {
 
         uint256 incentiveId = userStake.subscribedIncentiveIds.getUint24ValueAt(incentiveIndex);
 
-        if (rewardPerLiquidityLast[msg.sender][incentiveId] == 0) revert AlreadyUnsubscribed();
+        if (rewardPerLiquidityLast[msg.sender][incentiveId] == 0) {
+            revert AlreadyUnsubscribed();
+        }
 
         Incentive storage incentive = incentives[incentiveId];
 
         _accrueRewards(incentive);
 
         /// In case there is a token specific issue we can ignore rewards.
-        if (!ignoreRewards) _claimReward(incentive, incentiveId, userStake.liquidity);
+        if (!ignoreRewards) {
+            _claimReward(incentive, incentiveId, userStake.liquidity);
+        }
 
         rewardPerLiquidityLast[msg.sender][incentiveId] = 0;
 
@@ -281,7 +315,9 @@ contract StakingContractMainnet is ReentrancyGuard {
     }
 
     function accrueRewards(uint256 incentiveId) external nonReentrant {
-        if (incentiveId > incentiveCount || incentiveId <= 0) revert InvalidInput();
+        if (incentiveId > incentiveCount || incentiveId <= 0) {
+            revert InvalidInput();
+        }
 
         _accrueRewards(incentives[incentiveId]);
     }
@@ -292,7 +328,9 @@ contract StakingContractMainnet is ReentrancyGuard {
         rewards = new uint256[](n);
 
         for (uint256 i = 0; i < n; i = _increment(i)) {
-            if (incentiveIds[i] > incentiveCount || incentiveIds[i] <= 0) revert InvalidInput();
+            if (incentiveIds[i] > incentiveCount || incentiveIds[i] <= 0) {
+                revert InvalidInput();
+            }
 
             Incentive storage incentive = incentives[incentiveIds[i]];
 
@@ -302,6 +340,26 @@ contract StakingContractMainnet is ReentrancyGuard {
         }
     }
 
+    /// @dev Claims rewards for all incentives in the list, skipping reward rounding.
+    function claimAllRewards(uint256[] calldata incentiveIds)
+        external
+        nonReentrant
+        returns (uint256[] memory rewards)
+    {
+        uint256 n = incentiveIds.length;
+        rewards = new uint256[](n);
+        for (uint256 i = 0; i < n; i = _increment(i)) {
+            if (incentiveIds[i] > incentiveCount || incentiveIds[i] <= 0) {
+                revert InvalidInput();
+            }
+
+            Incentive storage incentive = incentives[incentiveIds[i]];
+            _accrueRewards(incentive);
+            rewards[i] =
+                _claimReward(incentive, incentiveIds[i], userStakes[msg.sender][incentive.token].liquidity, true);
+        }
+    }
+
     function _accrueRewards(Incentive storage incentive) internal {
         uint256 lastRewardTime = incentive.lastRewardTime;
 
@@ -315,10 +373,10 @@ contract StakingContractMainnet is ReentrancyGuard {
 
                 uint256 passedTime = maxTime - lastRewardTime;
 
-                uint256 reward = uint256(incentive.rewardRemaining) * passedTime / totalTime;
+                uint256 reward = (uint256(incentive.rewardRemaining) * passedTime) / totalTime;
 
                 // Increments of less than type(uint224).max - overflow is unrealistic.
-                incentive.rewardPerLiquidity += reward * type(uint112).max / incentive.liquidityStaked;
+                incentive.rewardPerLiquidity += (reward * type(uint112).max) / incentive.liquidityStaked;
 
                 incentive.rewardRemaining -= uint112(reward);
 
@@ -332,13 +390,31 @@ contract StakingContractMainnet is ReentrancyGuard {
     function _claimReward(Incentive storage incentive, uint256 incentiveId, uint112 usersLiquidity)
         internal
         returns (uint256 reward)
+    {
+        return _claimReward(incentive, incentiveId, usersLiquidity, false);
+    }
+
+    function _claimReward(Incentive storage incentive, uint256 incentiveId, uint112 usersLiquidity, bool skipRounding)
+        internal
+        returns (uint256 reward)
     {
         reward = _calculateReward(incentive, incentiveId, usersLiquidity);
 
-        rewardPerLiquidityLast[msg.sender][incentiveId] = incentive.rewardPerLiquidity;
+        uint256 rewardDelta;
+        // Check if the reward should be rounded
+        if (!skipRounding && incentive.isRewardRounded) {
+            uint8 decimals = ERC20(incentive.rewardToken).decimals();
+            uint256 roundedReward = (reward / 10 ** decimals) * 10 ** decimals;
+            // Delta of rewards to be left claimable for the user in the future
+            rewardDelta = reward - roundedReward;
+            reward = roundedReward;
+        }
 
-        ERC20(incentive.rewardToken).safeTransfer(msg.sender, reward);
+        // Calculate the reward per liquidity delta based on actual rewards given
+        uint256 rewardPerLiquidityDelta = (rewardDelta * type(uint112).max) / usersLiquidity;
+        rewardPerLiquidityLast[msg.sender][incentiveId] = incentive.rewardPerLiquidity - rewardPerLiquidityDelta;
 
+        ERC20(incentive.rewardToken).safeTransfer(msg.sender, reward);
         emit Claim(incentiveId, msg.sender, reward);
     }
 
@@ -349,7 +425,7 @@ contract StakingContractMainnet is ReentrancyGuard {
     {
         reward = _calculateReward(incentive, incentiveId, usersLiquidity);
 
-        uint256 rewardPerLiquidityDelta = reward * type(uint112).max / newLiquidity;
+        uint256 rewardPerLiquidityDelta = (reward * type(uint112).max) / newLiquidity;
 
         rewardPerLiquidityLast[msg.sender][incentiveId] = incentive.rewardPerLiquidity - rewardPerLiquidityDelta;
     }

src/Rewards/test/StakingContractMainnet.t.sol

@@ -5,11 +5,11 @@ import "./TestSetup.sol";
 
 contract CreateIncentiveTest is TestSetup {
     function testCreateIncentive(uint112 amount, uint32 startTime, uint32 endTime) public {
-        _createIncentive(address(tokenA), address(tokenB), amount, startTime, endTime);
+        _createIncentive(address(tokenA), address(tokenB), amount, startTime, endTime, false);
     }
 
     function testFailCreateIncentiveInvalidRewardToken(uint32 startTime, uint32 endTime) public {
-        _createIncentive(address(tokenA), zeroAddress, 1, startTime, endTime);
+        _createIncentive(address(tokenA), zeroAddress, 1, startTime, endTime, false);
     }
 
     function testUpdateIncentive(
@@ -147,6 +147,45 @@ contract CreateIncentiveTest is TestSetup {
         assertEqInexact(reward0 + reward1 + soloReward, totalReward, 10);
     }
 
+    function testClaimRoundedRewards() public {
+        uint112 amount = testIncentiveAmount;
+        uint256 duration = testIncentiveDuration;
+        uint256 incentiveId = _createIncentive(
+            address(tokenA), address(tokenB), amount, uint32(block.timestamp), uint32(block.timestamp + duration), true
+        );
+        uint256[] memory incentiveIds = new uint256[](1);
+        incentiveIds[0] = incentiveId;
+        StakingContractMainnet.Incentive memory incentive = _getIncentive(incentiveId);
+
+        // 2 users stake and subscribe
+        _stake(address(tokenA), 1, johnDoe, true);
+        _stake(address(tokenA), 1, janeDoe, true);
+        _subscribeToIncentive(incentiveId, johnDoe);
+        _subscribeToIncentive(incentiveId, janeDoe);
+
+        // 1/30 the time has passed
+        vm.warp(incentive.lastRewardTime + 86400);
+
+        // Each user got 1/60 of the total reward amount
+        (,, uint256 johnDoeReward) = _calculateReward(incentiveId, johnDoe);
+        (,, uint256 janeDoeReward) = _calculateReward(incentiveId, janeDoe);
+        assertEq(johnDoeReward, 16666666666666666666);
+        assertEq(janeDoeReward, 16666666666666666666);
+
+        // 1 user claims
+        vm.prank(johnDoe);
+        uint256[] memory johnDoeClaimed = stakingContract.claimRewards(incentiveIds);
+        assertEq(johnDoeClaimed[0], 16000000000000000000);
+
+        // User still has some rewards pending
+        (,, johnDoeReward) = _calculateReward(incentiveId, johnDoe);
+        assertEq(johnDoeReward, 666666666666666666);
+
+        // Other user still has the same reward
+        (,, janeDoeReward) = _calculateReward(incentiveId, janeDoe);
+        assertEq(janeDoeReward, 16666666666666666666);
+    }
+
     function testUnstakeSaveRewards() public {
         _stake(address(tokenA), 1, johnDoe, true);
         _subscribeToIncentive(ongoingIncentive, johnDoe);