Support staking and receiving rewards in the chain's native token

contracts/mock/tokens/MockERC20.sol

@@ -7,8 +7,9 @@ import { ERC20Mod } from "./ERC20Mod.sol";
 /* solhint-disable */
 contract MockERC20 is ERC20Mod {
     constructor(string memory name, string memory symbol) ERC20Mod(name, symbol) {
-        _mint(msg.sender, 9000000000000 * 10 ** 18);
+        _mint(msg.sender, 999999999999999999 * 10 ** 18);
     }
+
     function mint(address to, uint256 amount) public virtual {
         _mint(to, amount);
     }

contracts/staking/ERC20/IStakingERC20.sol

@@ -34,15 +34,20 @@ interface IStakingERC20 {
      */
     error UnstakeMoreThanStake();
 
-    function stakeWithLock(uint256 amount, uint256 lockDuration) external;
+    /**
+     * @notice Revert when the user is staking an amount inequal to the amount given 
+     */
+    error InsufficientValue();
+
+    function stakeWithLock(uint256 amount, uint256 lockDuration) external payable;
 
-    function stakeWithoutLock(uint256 amount) external;
+    function stakeWithoutLock(uint256 amount) external payable;
 
     function claim() external;
 
-    function unstake(uint256 amount, bool exit) external;
+    function unstake(uint256 amount, bool exit) external payable;
 
-    function unstakeLocked(uint256 amount, bool exit) external;
+    function unstakeLocked(uint256 amount, bool exit) external payable;
 
     function getRemainingLockTime() external view returns (uint256);
 

contracts/staking/ERC20/StakingERC20.sol

@@ -7,7 +7,6 @@ import { IStakingERC20 } from "./IStakingERC20.sol";
 import { StakingBase } from "../StakingBase.sol";
 import { IERC20MintableBurnable } from "../../types/IERC20MintableBurnable.sol";
 
-
 /**
  * @title StakingERC20
  * @notice A staking contract for ERC20 tokens
@@ -38,7 +37,7 @@ contract StakingERC20 is StakingBase, IStakingERC20 {
      * @param amount The amount to stake
      * @param lockDuration The duration of the lock period
      */
-    function stakeWithLock(uint256 amount, uint256 lockDuration) external override {
+    function stakeWithLock(uint256 amount, uint256 lockDuration) external payable override {
         if (lockDuration < config.minimumLockTime) {
             revert LockTimeTooShort();
         }
@@ -52,7 +51,7 @@ contract StakingERC20 is StakingBase, IStakingERC20 {
      *
      * @param amount The amount to stake
      */
-    function stakeWithoutLock(uint256 amount) external override {
+    function stakeWithoutLock(uint256 amount) external payable override {
         _stake(amount, 0);
     }
 
@@ -72,7 +71,7 @@ contract StakingERC20 is StakingBase, IStakingERC20 {
      *
      * @param amount The amount to withdraw
      */
-    function unstake(uint256 amount, bool exit) external override {
+    function unstake(uint256 amount, bool exit) external payable override {
         _unstake(amount, false, exit);
     }
 
@@ -83,7 +82,7 @@ contract StakingERC20 is StakingBase, IStakingERC20 {
      * @param amount The amount to withdraw
      * @param exit Boolean if user wants to forfeit rewards
      */
-    function unstakeLocked(uint256 amount, bool exit) public override {
+    function unstakeLocked(uint256 amount, bool exit) external payable override {
         _unstake(amount, true, exit);
     }
 
@@ -110,14 +109,24 @@ contract StakingERC20 is StakingBase, IStakingERC20 {
             revert ZeroValue();
         }
 
+        // If stakingToken is gas token, `msg.value` must equal `amount`
+        if (config.stakingToken == address(0)) {
+            if (msg.value != amount) {
+                revert InsufficientValue();
+            }
+        }
+
         Staker storage staker = stakers[msg.sender];
 
         _coreStake(staker, amount, lockDuration);
 
         totalStaked += amount;
 
-        // Transfers user's funds to this contract
-        IERC20(config.stakingToken).safeTransferFrom(msg.sender, address(this), amount);
+        // Transfers user's funds to this contract  
+        if (config.stakingToken != address(0)) {
+            IERC20(config.stakingToken).safeTransferFrom(msg.sender, address(this), amount);
+        } // the `else` case is handled by including `recieve` above to accept `msg.value`
+
         // Mint the user's stake as a representative token
         IERC20MintableBurnable(config.stakeRepToken).mint(msg.sender, amount);
 
@@ -150,11 +159,12 @@ contract StakingERC20 is StakingBase, IStakingERC20 {
                 staker.lastTimestampLocked = 0;
                 staker.unlockedTimestamp = 0;
             } else if (_getRemainingLockTime(staker) > 0) {
-                // if still locked and not exiting, revert
+                // if still locked revert
                 revert TimeLockNotPassed();
             } else {
-                // If claims happen after lock period has passed, the lastTimestamp is more accurate
+                // If claims happen after lock period has passed, the lastTimestampLocked is more accurate
                 // but if they don't happen, then lastTimestampLocked may still be the original stake timestamp
+                // and we should use `unlockedTimestamp` instead. 
                 // so we have to calculate which is more recent before calculating rewards
                 uint256 mostRecentTimestamp = _mostRecentTimestamp(staker);
 
@@ -227,14 +237,17 @@ contract StakingERC20 is StakingBase, IStakingERC20 {
         if (rewards > 0) {
             // Transfer the user's rewards
             // Will fail if the contract does not have funding for this
-            config.rewardsToken.safeTransfer(msg.sender, rewards);
+            // If rewards address is `0x0` we use the chain's native token
+            _transferAmount(config.rewardsToken, rewards);
+
             emit Claimed(msg.sender, rewards);
         }
 
         totalStaked -= amount;
 
         // Return the user's initial stake
-        IERC20(config.stakingToken).safeTransfer(msg.sender, amount);
+        _transferAmount(config.stakingToken, amount);
+
         // Burn the user's stake representative token
         IERC20MintableBurnable(config.stakeRepToken).burn(msg.sender, amount);
 
@@ -248,9 +261,15 @@ contract StakingERC20 is StakingBase, IStakingERC20 {
      * return the balance of the rewards token minus the total staked amount
      */
     function _getContractRewardsBalance() internal view override returns (uint256) {
-        uint256 balance = super._getContractRewardsBalance();
+        uint256 balance;
+
+        if (address(config.rewardsToken) == address(0)) {
+            balance = address(this).balance;
+        } else {
+            balance = IERC20(config.rewardsToken).balanceOf(address(this));
+        }
 
-        if (address(config.rewardsToken) == config.stakingToken) {
+        if (config.rewardsToken == config.stakingToken) {
             return balance - totalStaked;
         }
 

contracts/staking/ERC721/IStakingERC721.sol

@@ -15,7 +15,15 @@ interface IStakingERC721 is IERC721Receiver, IStakingBase {
      */
     struct NFTStaker {
         Staker stake;
-        uint256[] tokenIds;
+
+        // A) have array of token ids AND `staked` mapping
+        // This way we can mark tokens as `unstaked` without iterating
+        // `tokenIds` array each time.
+        // B) we can just remove the `unstakeAll` option because the front end could do this
+        // Considering we don't yet have a subgraph for this it might be tricky
+        uint256[] tokenIds; // use sNFT as proof of ownership of stake, and `amountStaked(locked)` as quantity
+        // TODO look at gas costs for this and see if off chain tids is a better solution (with a subgraph)
+        mapping(uint256 tokenId => bool staked) staked;
         mapping(uint256 tokenId => bool locked) locked;
     }
 

contracts/staking/ERC721/StakingERC721.sol

@@ -4,8 +4,6 @@ pragma solidity 0.8.26;
 import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import { IERC721 } from "@openzeppelin/contracts/token/ERC721/IERC721.sol";
 import { SafeERC20 } from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
-import { ERC721 } from "@openzeppelin/contracts/token/ERC721/ERC721.sol";
-import { ERC721URIStorage } from "@openzeppelin/contracts/token/ERC721/extensions/ERC721URIStorage.sol";
 import { IStakingERC721 } from "./IStakingERC721.sol";
 import { StakingBase } from "../StakingBase.sol";
 import { IERC721MintableBurnableURIStorage } from "../../types/IERC721MintableBurnableURIStorage.sol";
@@ -39,7 +37,12 @@ contract StakingERC721 is StakingBase, IStakingERC721 {
         Config memory config
     )
         StakingBase(config)
-    {}
+    {
+        if (config.stakingToken.code.length == 0) {
+            revert InitializedWithZero();
+        }
+
+    }
 
     /**
      * @notice Stake one or more ERC721 tokens with a lock period
@@ -166,8 +169,9 @@ contract StakingERC721 is StakingBase, IStakingERC721 {
                 tokenIds[i]
             );
 
-            // Add to array and to mapping for indexing in unstake
+            // Add to array and to mapping for when unstaking
             nftStaker.tokenIds.push(tokenIds[i]);
+            nftStaker.staked[tokenIds[i]] = true;
             nftStaker.locked[tokenIds[i]] = lockDuration > 0;
 
             // Mint user sNFT
@@ -183,6 +187,10 @@ contract StakingERC721 is StakingBase, IStakingERC721 {
     }
 
     function _unstakeMany(uint256[] memory _tokenIds, bool exit) internal {
+        // its possible that token IDs that are already unstaked are passed here
+        // because removing them from the users tokenIds[] would be gas expensive
+        // and so burning will fail with `non-existent token` error
+        // so we check if the token is owned by the user and if not, skip it
         NFTStaker storage nftStaker = nftStakers[msg.sender];
 
         uint256 rewards;
@@ -198,7 +206,8 @@ contract StakingERC721 is StakingBase, IStakingERC721 {
         uint256 i;
         for (i; i < _tokenIds.length;) {
             if (
-                IERC721(config.stakeRepToken).ownerOf(_tokenIds[i]) == address(0)
+                nftStaker.staked[_tokenIds[i]] == false
+                || IERC721(config.stakeRepToken).ownerOf(_tokenIds[i]) == address(0)
                 || IERC721(config.stakeRepToken).ownerOf(_tokenIds[i]) != msg.sender
             ) {
                 // Either the list of tokenIds contains a non-existent token
@@ -236,9 +245,8 @@ contract StakingERC721 is StakingBase, IStakingERC721 {
                         if (!rewardsGivenLocked) {
                             rewards += nftStaker.stake.owedRewardsLocked;
 
-                            // set to 0 so they don't get it again in future calls
+                            // can only get this once per tx, not each loop, so set to 0
                             nftStaker.stake.owedRewardsLocked = 0;
-
                             rewardsGivenLocked = true;
                         }
                     }
@@ -249,6 +257,7 @@ contract StakingERC721 is StakingBase, IStakingERC721 {
                     // Unstake if they are passed their lock time or exiting
                     _unstake(_tokenIds[i]);
                     --nftStaker.stake.amountStakedLocked;
+                    nftStaker.staked[_tokenIds[i]] = false;
                     isAction = true;
                 } else {
                     // stake is locked and cannot be unstaked
@@ -276,6 +285,7 @@ contract StakingERC721 is StakingBase, IStakingERC721 {
 
                 _unstake(_tokenIds[i]);
                 --nftStaker.stake.amountStaked;
+                nftStaker.staked[_tokenIds[i]] = false;
                 isAction = true;
             }
 
@@ -301,20 +311,19 @@ contract StakingERC721 is StakingBase, IStakingERC721 {
 
         if (!exit) {
             // Transfer the user's rewards
-            // Will fail if the contract does not have funding
-            config.rewardsToken.safeTransfer(msg.sender, rewards);
+            _transferAmount(config.rewardsToken, rewards);
+
             emit Claimed(msg.sender, rewards);
         }
 
         // If a complete withdrawal, delete the staker struct for this user as well
         if (nftStaker.stake.amountStaked == 0 && nftStaker.stake.amountStakedLocked == 0) {
             delete nftStakers[msg.sender];
-        } else if (nftStaker.stake.amountStaked != 0) {
+        } else if (nftStaker.stake.amountStaked != 0 && nftStaker.stake.amountStakedLocked == 0) {
             nftStaker.stake.amountStakedLocked = 0;
             nftStaker.stake.lastTimestampLocked = 0;
             nftStaker.stake.unlockedTimestamp = 0;
-            nftStaker.stake.lockDuration = 0;
-        } else {
+        } else if (nftStaker.stake.amountStaked == 0 && nftStaker.stake.amountStakedLocked != 0) {
             nftStaker.stake.amountStaked = 0;
             nftStaker.stake.lastTimestamp = 0;
         }

contracts/staking/IStakingBase.sol

@@ -1,7 +1,6 @@
 // SPDX-License-Identifier: MIT
 pragma solidity 0.8.26;
 
-import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 
 /**
  * @title IStakingBase
@@ -47,7 +46,7 @@ interface IStakingBase {
     struct Config {
         address stakingToken;
         address contractOwner;
-        IERC20 rewardsToken;
+        address rewardsToken;
         address stakeRepToken;
         uint256 rewardsPerPeriod;
         uint256 periodLength;
@@ -77,6 +76,26 @@ interface IStakingBase {
         uint256 indexed rewards
     );
 
+    /**
+     * @notice Emit when `reqwardsPerPeriod` is set
+     * @param owner The address of the contract owner
+     * @param rewardsPerPeriod The new rewards per period value 
+     */
+    event RewardsPerPeriodSet(
+        address indexed owner,
+        uint256 indexed rewardsPerPeriod
+    );
+
+    /**
+     * @notice Emit when the period length is set
+     * @param owner The address of the contract owner
+     * @param periodLength The new period length value
+     */
+    event PeriodLengthSet(
+        address indexed owner,
+        uint256 indexed periodLength
+    );
+
     /**
      * @notice Emit when the multiplier is set
      * @param owner The address of the contract owner
@@ -97,6 +116,26 @@ interface IStakingBase {
         uint256 indexed minimumLockTime
     );
 
+    /**
+     * @notice Emit when the minimum rewards multiplier is set
+     * @param owner The address of the contract owner
+     * @param minimumRewardsMultiplier The new minimum rewards multiplier
+     */
+    event MinimumRewardsMultiplierSet(
+        address indexed owner,
+        uint256 indexed minimumRewardsMultiplier
+    );
+
+    /**
+     * @notice Emit when the maximum rewards multiplier is set
+     * @param owner The address of the contract owner
+     * @param maximumRewardsMultiplier The new maximum rewards multiplier
+     */
+    event MaximumRewardsMultiplierSet(
+        address indexed owner,
+        uint256 indexed maximumRewardsMultiplier
+    );
+
     /**
      * @notice Emit incoming stake is not valid
      */
@@ -145,6 +184,10 @@ interface IStakingBase {
      */
     error InitializedWithZero();
 
+    receive() external payable;
+
+    fallback() external payable;
+
     function withdrawLeftoverRewards() external;
 
     function setRewardsPerPeriod(uint256 _rewardsPerPeriod) external;
@@ -161,7 +204,7 @@ interface IStakingBase {
 
     function getStakingToken() external view returns(address);
 
-    function getRewardsToken() external view returns(IERC20);
+    function getRewardsToken() external view returns(address);
 
     function getStakeRepToken() external view returns (address);