rm shares

src/collator/CollatorStaking.sol

@@ -1,90 +1,155 @@
-// SPDX-License-Identifier: MIT
-pragma solidity 0.8.17;
+pragma solidity ^0.5.16;
 
-import "@openzeppelin/contracts@4.9.6/utils/math/Math.sol";
+import "openzeppelin-solidity-2.3.0/contracts/math/Math.sol";
+import "openzeppelin-solidity-2.3.0/contracts/math/SafeMath.sol";
+import "openzeppelin-solidity-2.3.0/contracts/token/ERC20/ERC20Detailed.sol";
+import "openzeppelin-solidity-2.3.0/contracts/token/ERC20/SafeERC20.sol";
+import "openzeppelin-solidity-2.3.0/contracts/utils/ReentrancyGuard.sol";
 
-contract CollatorStaking {
-    using Math for uint256;
+// Inheritance
+import "./interfaces/IStakingRewards.sol";
+import "./RewardsDistributionRecipient.sol";
+
+contract StakingRewards is IStakingRewards, RewardsDistributionRecipient, ReentrancyGuard {
+    using SafeMath for uint256;
+    using SafeERC20 for IERC20;
+
+    /* ========== STATE VARIABLES ========== */
 
-    address public hub;
     address public operator;
+    uint256 public periodFinish = 0;
+    uint256 public rewardRate = 0;
+    uint256 public rewardsDuration = 7 days;
+    uint256 public lastUpdateTime;
+    uint256 public rewardPerTokenStored;
 
-    uint256 public totalAssets;
-    uint256 public totalShares;
+    mapping(address => uint256) public userRewardPerTokenPaid;
+    mapping(address => uint256) public rewards;
 
-    mapping(address => uint256) sharesOf;
+    uint256 private _totalSupply;
+    mapping(address => uint256) private _balances;
 
-    modifier onlyHub() {
-        require(msg.sender == hub);
-        _;
-    }
+    /* ========== CONSTRUCTOR ========== */
 
-    constructor(address operator_) {
-        hub = msg.sender;
+    constructor(address operator_, address wring) public {
+        rewardsDistribution = msg.sender;
         operator = operator_;
     }
 
-    function assetsOf(address account) public view returns (uint256) {
-        return _convertToAssets(sharesOf[account]);
-    }
+    /* ========== VIEWS ========== */
 
-    function previewDeposit(uint256 assets) public view returns (uint256) {
-        return _convertToShares(assets, Math.Rounding.Down);
+    function totalSupply() external view returns (uint256) {
+        return _totalSupply;
     }
 
-    function previewRedeem(uint256 shares) public view returns (uint256) {
-        return _convertToAssets(shares, Math.Rounding.Down);
+    function balanceOf(address account) external view returns (uint256) {
+        return _balances[account];
     }
 
-    function deposit(uint256 assets, address recipient) external onlyHub returns (uint256) {
-        uint256 shares = previewDeposit(assets);
-
-        sharesOf[receiver] = shares;
-        totalShares += shares;
-        emit Deposit(receiver, assets, shares);
-
-        // _mint(receiver, shares);
-        return shares;
+    function lastTimeRewardApplicable() public view returns (uint256) {
+        return Math.min(block.timestamp, periodFinish);
     }
 
-    function redeem(uint256 shares, address from) external onlyHub returns (uint256) {
-        uint256 assets = previewRedeem(shares);
-
-        sharesOf[from] -= shares;
-        totalShares -= shares;
-        emit Withdraw(from, assets, shares);
+    function rewardPerToken() public view returns (uint256) {
+        if (_totalSupply == 0) {
+            return rewardPerTokenStored;
+        }
+        return rewardPerTokenStored.add(
+            lastTimeRewardApplicable().sub(lastUpdateTime).mul(rewardRate).mul(1e18).div(_totalSupply)
+        );
+    }
 
-        // _burn(from, shares);
-        return assets;
+    function earned(address account) public view returns (uint256) {
+        return _balances[account].mul(rewardPerToken().sub(userRewardPerTokenPaid[account])).div(1e18).add(
+            rewards[account]
+        );
     }
 
-    function distributeReward(uint256 rewards) external onlyHub {
-        totalAssets += rewards;
+    function getRewardForDuration() external view returns (uint256) {
+        return rewardRate.mul(rewardsDuration);
     }
 
-    function convertToAssets(uint256 shares) public view returns (uint256) {
-        return _convertToAssets(shares, Math.Rounding.Down);
+    /* ========== MUTATIVE FUNCTIONS ========== */
+
+    function stake(uint256 amount, address account)
+        external
+        onlyRewardsDistribution
+        nonReentrant
+        updateReward(account)
+    {
+        require(amount > 0, "Cannot stake 0");
+        _totalSupply = _totalSupply.add(amount);
+        _balances[account] = _balances[account].add(amount);
+        emit Staked(msg.sender, amount);
     }
 
-    // assets = shares * (total_assets / total_shares) --- (== price_per_share * shares)
-    function _convertToAssets(uint256 shares, Math.Rounding rounding) internal view virtual returns (uint256) {
-        return shares.mulDiv(totalAssets() + 1, totalSupply() + 10 ** _decimalsOffset(), rounding);
+    function withdraw(uint256 amount, address account)
+        public
+        onlyRewardsDistribution
+        nonReentrant
+        updateReward(account)
+    {
+        require(amount > 0, "Cannot withdraw 0");
+        _totalSupply = _totalSupply.sub(amount);
+        _balances[account] = _balances[account].sub(amount);
+        emit Withdrawn(account, amount);
     }
 
-    function convertToShares(uint256 assets) public view returns (uint256) {
-        return _convertToShares(assets, Math.Rounding.Down);
+    function getReward(address account) public nonReentrant onlyRewardsDistribution updateReward(account) {
+        uint256 reward = rewards[account];
+        if (reward > 0) {
+            rewards[account] = 0;
+            (bool success,) = account.call.value(reward)("");
+            require(success, "Transfer failed");
+            emit RewardPaid(account, reward);
+        }
     }
 
-    // shares = amount * (total_shares / total_assets) --- (== amount / price_per_share)
-    function _convertToShares(uint256 assets, Math.Rounding rounding) internal view virtual returns (uint256) {
-        return assets.mulDiv(totalSupply() + 10 ** _decimalsOffset(), totalAssets() + 1, rounding);
+    function exit() external {
+        withdraw(_balances[msg.sender]);
+        getReward();
     }
 
-    function _decimalsOffset() internal view returns (uint8) {
-        return 0;
+    /* ========== RESTRICTED FUNCTIONS ========== */
+
+    function notifyRewardAmount() external payable onlyRewardsDistribution updateReward(address(0)) {
+        uint256 reward = msg.value;
+        if (block.timestamp >= periodFinish) {
+            rewardRate = reward.div(rewardsDuration);
+        } else {
+            uint256 remaining = periodFinish.sub(block.timestamp);
+            uint256 leftover = remaining.mul(rewardRate);
+            rewardRate = reward.add(leftover).div(rewardsDuration);
+        }
+
+        // Ensure the provided reward amount is not more than the balance in the contract.
+        // This keeps the reward rate in the right range, preventing overflows due to
+        // very high values of rewardRate in the earned and rewardsPerToken functions;
+        // Reward + leftover must be less than 2^256 / 10^18 to avoid overflow.
+        uint256 balance = address(this).balance;
+        require(rewardRate <= balance.div(rewardsDuration), "Provided reward too high");
+
+        lastUpdateTime = block.timestamp;
+        periodFinish = block.timestamp.add(rewardsDuration);
+        emit RewardAdded(reward);
     }
 
-    function decimals() public view returns (uint8) {
-        return 18 + _decimalsOffset();
+    /* ========== MODIFIERS ========== */
+
+    modifier updateReward(address account) {
+        rewardPerTokenStored = rewardPerToken();
+        lastUpdateTime = lastTimeRewardApplicable();
+        if (account != address(0)) {
+            rewards[account] = earned(account);
+            userRewardPerTokenPaid[account] = rewardPerTokenStored;
+        }
+        _;
     }
+
+    /* ========== EVENTS ========== */
+
+    event RewardAdded(uint256 reward);
+    event Staked(address indexed user, uint256 amount);
+    event Withdrawn(address indexed user, uint256 amount);
+    event RewardPaid(address indexed user, uint256 reward);
 }