Added snapshot of exchange rate + search + updated power calculation

contracts/interfaces/IStakedTokenV3.sol

@@ -12,6 +12,8 @@ interface IStakedTokenV3 is IStakedToken {
 
   function slash(address destination, uint256 amount) external;
 
+  function donate(uint256 amount) external;
+
   function getMaxSlashablePercentage() external view returns (uint256);
 
   function setMaxSlashablePercentage(uint256 percentage) external;

contracts/stake/StakedTokenV3.sol

@@ -44,6 +44,9 @@ contract StakedTokenV3 is StakedTokenV2, IStakedTokenV3, RoleManager {
 
   address internal _claimHelper;
 
+  mapping(uint256 => Snapshot) internal _exchangeRateSnapshots;
+  uint256 internal _countExchangeRateSnapshots;
+
   modifier onlyAdmin {
     require(msg.sender == getAdmin(MAIN_ADMIN_ROLE), 'CALLER_NOT_MAIN_ADMIN');
     _;
@@ -69,6 +72,8 @@ contract StakedTokenV3 is StakedTokenV2, IStakedTokenV3, RoleManager {
   event CooldownPauseChanged(bool pause);
   event MaxSlashablePercentageChanged(uint256 newPercentage);
   event Slashed(address indexed destination, uint256 amount);
+  event Donated(address indexed sender, uint256 amount);
+  event EchangeRateSnapshotted(uint128 blockNumber, uint128 exchangeRate);
   event CooldownPauseAdminChanged(address indexed newAdmin);
   event SlashingAdminChanged(address indexed newAdmin);
   event ClaimHelperChanged(address indexed newClaimHelper);
@@ -301,21 +306,6 @@ contract StakedTokenV3 is StakedTokenV2, IStakedTokenV3, RoleManager {
     _redeem(from, to, redeemAmount);
   }
 
-  /**
-   * @dev Calculates the exchange rate between the amount of STAKED_TOKEN and the the StakeToken total supply.
-   * Slashing will reduce the exchange rate. Supplying STAKED_TOKEN to the stake contract
-   * can replenish the slashed STAKED_TOKEN and bring the exchange rate back to 1
-   **/
-  function exchangeRate() public view override returns (uint256) {
-    uint256 currentSupply = totalSupply();
-
-    if (currentSupply == 0) {
-      return 1e18; //initial exchange rate is 1:1
-    }
-
-    return STAKED_TOKEN.balanceOf(address(this)).mul(1e18).div(currentSupply);
-  }
-
   /**
    * @dev Executes a slashing of the underlying of a certain amount, transferring the seized funds
    * to destination. Decreasing the amount of underlying will automatically adjust the exchange rate
@@ -330,10 +320,18 @@ contract StakedTokenV3 is StakedTokenV2, IStakedTokenV3, RoleManager {
     require(amount <= maxSlashable, 'INVALID_SLASHING_AMOUNT');
 
     STAKED_TOKEN.safeTransfer(destination, amount);
+    _snapshotExchangeRate();
 
     emit Slashed(destination, amount);
   }
 
+  function donate(uint256 amount) external override {
+    REWARD_TOKEN.safeTransferFrom(msg.sender, address(this), amount);
+    _snapshotExchangeRate();
+
+    emit Donated(msg.sender, amount);
+  }
+
   /**
    * @dev Set the address of the contract with priviledge, the ClaimHelper contract
    * It speicifically enables to claim from several contracts at once
@@ -364,6 +362,21 @@ contract StakedTokenV3 is StakedTokenV2, IStakedTokenV3, RoleManager {
     emit MaxSlashablePercentageChanged(percentage);
   }
 
+  /**
+   * @dev Calculates the exchange rate between the amount of STAKED_TOKEN and the the StakeToken total supply.
+   * Slashing will reduce the exchange rate. Supplying STAKED_TOKEN to the stake contract
+   * can replenish the slashed STAKED_TOKEN and bring the exchange rate back to 1
+   **/
+  function exchangeRate() public view override returns (uint256) {
+    uint256 currentSupply = totalSupply();
+
+    if (currentSupply == 0) {
+      return 1e18; //initial exchange rate is 1:1
+    }
+
+    return STAKED_TOKEN.balanceOf(address(this)).mul(1e18).div(currentSupply);
+  }
+
   /**
    * @dev returns the current address of the claimHelper Contract, contract with priviledge
    * It speicifically enables to claim from several contracts at once
@@ -394,6 +407,43 @@ contract StakedTokenV3 is StakedTokenV2, IStakedTokenV3, RoleManager {
     return REVISION();
   }
 
+  function getPowerAtBlock(
+    address user,
+    uint256 blockNumber,
+    DelegationType delegationType
+  ) external view override returns (uint256) {
+    (
+      mapping(address => mapping(uint256 => Snapshot)) storage snapshots,
+      mapping(address => uint256) storage snapshotsCounts,
+
+    ) = _getDelegationDataByType(delegationType);
+
+    return (
+      _searchByBlockNumber(snapshots, snapshotsCounts, user, blockNumber)
+        .mul(_searchExchangeRateByBlockNumber(blockNumber))
+        .div(1e18)
+    );
+  }
+
+  function getPowerCurrent(address user, DelegationType delegationType)
+    external
+    view
+    override
+    returns (uint256)
+  {
+    (
+      mapping(address => mapping(uint256 => Snapshot)) storage snapshots,
+      mapping(address => uint256) storage snapshotsCounts,
+
+    ) = _getDelegationDataByType(delegationType);
+
+    return (
+      _searchByBlockNumber(snapshots, snapshotsCounts, user, block.number).mul(exchangeRate()).div(
+        1e18
+      )
+    );
+  }
+
   function _claimRewards(
     address from,
     address to,
@@ -404,7 +454,7 @@ contract StakedTokenV3 is StakedTokenV2, IStakedTokenV3, RoleManager {
     uint256 amountToClaim = (amount == type(uint256).max) ? newTotalRewards : amount;
 
     stakerRewardsToClaim[from] = newTotalRewards.sub(amountToClaim, 'INVALID_AMOUNT');
-    REWARD_TOKEN.safeTransferFrom(REWARDS_VAULT, to, amountToClaim);
+    STAKED_TOKEN.safeTransferFrom(REWARDS_VAULT, to, amountToClaim);
     emit RewardsClaimed(from, to, amountToClaim);
     return (amountToClaim);
   }
@@ -479,4 +529,58 @@ contract StakedTokenV3 is StakedTokenV2, IStakedTokenV3, RoleManager {
 
     emit Redeem(from, to, amountToRedeem, underlyingToRedeem);
   }
+
+  function _snapshotExchangeRate() internal {
+    uint128 currentBlock = uint128(block.number);
+    uint128 newExchangeRate = uint128(exchangeRate());
+
+    // Doing multiple operations in the same block
+    if (
+      _countExchangeRateSnapshots != 0 &&
+      _exchangeRateSnapshots[_countExchangeRateSnapshots - 1].blockNumber == currentBlock
+    ) {
+      _exchangeRateSnapshots[_countExchangeRateSnapshots - 1].value = newExchangeRate;
+    } else {
+      _exchangeRateSnapshots[_countExchangeRateSnapshots] = Snapshot(currentBlock, newExchangeRate);
+      _countExchangeRateSnapshots = _countExchangeRateSnapshots + 1;
+    }
+    emit EchangeRateSnapshotted(currentBlock, newExchangeRate);
+  }
+
+  /**
+   * @dev searches a exchange Rate by block number. Uses binary search.
+   * @param blockNumber the block number being searched
+   **/
+  function _searchExchangeRateByBlockNumber(uint256 blockNumber) internal view returns (uint256) {
+    require(blockNumber <= block.number, 'INVALID_BLOCK_NUMBER');
+
+    if (_countExchangeRateSnapshots == 0) {
+      return exchangeRate();
+    }
+
+    // First check most recent balance
+    if (_exchangeRateSnapshots[_countExchangeRateSnapshots - 1].blockNumber <= blockNumber) {
+      return _exchangeRateSnapshots[_countExchangeRateSnapshots - 1].value;
+    }
+
+    // Next check implicit zero balance
+    if (_exchangeRateSnapshots[0].blockNumber > blockNumber) {
+      return 1e18; //initial exchange rate is 1:1
+    }
+
+    uint256 lower = 0;
+    uint256 upper = _countExchangeRateSnapshots - 1;
+    while (upper > lower) {
+      uint256 center = upper - (upper - lower) / 2; // ceil, avoiding overflow
+      Snapshot memory snapshot = _exchangeRateSnapshots[center];
+      if (snapshot.blockNumber == blockNumber) {
+        return snapshot.value;
+      } else if (snapshot.blockNumber < blockNumber) {
+        lower = center;
+      } else {
+        upper = center - 1;
+      }
+    }
+    return _exchangeRateSnapshots[lower].value;
+  }
 }