Implement reward accounting and test the earn calculation

This commit implements the basic reward calculation mechanics associated
with staking. The mechanics are largely modeled off the Synthetix StakingRewards
contract. This commit also adds numerous tests, and related testing
infrastructure for assesing whether the calucations for earned rewards is being
done correctly.

src/UniStaker.sol

@@ -11,6 +11,8 @@ contract UniStaker is ReentrancyGuard {
   type DepositIdentifier is uint256;
 
   error UniStaker__Unauthorized(bytes32 reason, address caller);
+  error UniStaker__InvalidRewardRate();
+  error UniStaker__InsufficientRewardBalance();
 
   struct Deposit {
     uint256 balance;
@@ -21,6 +23,7 @@ contract UniStaker is ReentrancyGuard {
 
   IERC20 public immutable REWARDS_TOKEN;
   IERC20Delegates public immutable STAKE_TOKEN;
+  uint256 private SCALE_FACTOR = 1e24;
 
   DepositIdentifier private nextDepositId;
 
@@ -34,11 +37,37 @@ contract UniStaker is ReentrancyGuard {
 
   mapping(address delegatee => DelegationSurrogate surrogate) public surrogates;
 
+  uint256 public rewardDuration = 7 days;
+  uint256 public finishAt;
+  uint256 public updatedAt;
+  uint256 public rewardRate;
+  uint256 public rewardPerTokenStored;
+  mapping(address account => uint256) public userRewardPerTokenPaid;
+  mapping(address account => uint256 amount) public rewards;
+
   constructor(IERC20 _rewardsToken, IERC20Delegates _stakeToken) {
     REWARDS_TOKEN = _rewardsToken;
     STAKE_TOKEN = _stakeToken;
   }
 
+  function lastTimeRewardApplicable() public view returns (uint256) {
+    if (finishAt <= block.timestamp) return finishAt;
+    else return block.timestamp;
+  }
+
+  function rewardPerToken() public view returns (uint256) {
+    if (totalSupply == 0) return rewardPerTokenStored;
+
+    return rewardPerTokenStored
+      + (rewardRate * (lastTimeRewardApplicable() - updatedAt) * SCALE_FACTOR) / totalSupply;
+  }
+
+  function earned(address _account) public view returns (uint256) {
+    return rewards[_account]
+      + (earningPower[_account] * (rewardPerToken() - userRewardPerTokenPaid[_account]))
+        / SCALE_FACTOR;
+  }
+
   function stake(uint256 _amount, address _delegatee)
     external
     nonReentrant
@@ -59,13 +88,37 @@ contract UniStaker is ReentrancyGuard {
     Deposit storage deposit = deposits[_depositId];
     if (msg.sender != deposit.owner) revert UniStaker__Unauthorized("not owner", msg.sender);
 
+    _updateReward(deposit.beneficiary);
+
     deposit.balance -= _amount; // overflow prevents withdrawing more than balance
     totalSupply -= _amount;
     totalDeposits[msg.sender] -= _amount;
     earningPower[deposit.beneficiary] -= _amount;
     _stakeTokenSafeTransferFrom(address(surrogates[deposit.delegatee]), deposit.owner, _amount);
   }
 
+  // TODO: this needs to be a restricted method
+  function notifyRewardsAmount(uint256 _amount) external {
+    _updateReward(address(0));
+
+    if (block.timestamp >= finishAt) {
+      // TODO: Can we move the scale factor into the rewardRate? This should reduce rounding errors
+      // introduced here when truncating on this division.
+      rewardRate = _amount / rewardDuration;
+    } else {
+      uint256 remainingRewards = rewardRate * (finishAt - block.timestamp);
+      rewardRate = (remainingRewards + _amount) / rewardDuration;
+    }
+
+    if (rewardRate == 0) revert UniStaker__InvalidRewardRate();
+    if ((rewardRate * rewardDuration) > REWARDS_TOKEN.balanceOf(address(this))) {
+      revert UniStaker__InsufficientRewardBalance();
+    }
+
+    finishAt = block.timestamp + rewardDuration;
+    updatedAt = block.timestamp;
+  }
+
   function _fetchOrDeploySurrogate(address _delegatee)
     internal
     returns (DelegationSurrogate _surrogate)
@@ -91,6 +144,8 @@ contract UniStaker is ReentrancyGuard {
     internal
     returns (DepositIdentifier _depositId)
   {
+    _updateReward(_beneficiary);
+
     DelegationSurrogate _surrogate = _fetchOrDeploySurrogate(_delegatee);
     _stakeTokenSafeTransferFrom(msg.sender, address(_surrogate), _amount);
     _depositId = _useDepositId();
@@ -105,4 +160,16 @@ contract UniStaker is ReentrancyGuard {
       beneficiary: _beneficiary
     });
   }
+
+  // TODO: rename snapshotReward?
+  // Extract into two methods global + user
+  function _updateReward(address _account) internal {
+    rewardPerTokenStored = rewardPerToken();
+    updatedAt = lastTimeRewardApplicable();
+
+    if (_account == address(0)) return;
+
+    rewards[_account] = earned(_account);
+    userRewardPerTokenPaid[_account] = rewardPerTokenStored;
+  }
 }