Merge pull request #27 from rareprotocol/feature/further-gas-optimiza…

…tions

chore: lessons in gas savings

src/staking/registry/IRareStakingRegistry.sol

@@ -201,18 +201,6 @@ interface IRareStakingRegistry {
   /// @return uint256 Amount of rare being staked on the user.
   function getTotalAmountStakedOnUser(address _user) external view returns (uint256);
 
-  /// @notice Retrieves a list of all the ERC20 staking contracts.
-  /// @return list of contracts users can use to stake.
-  function getAllStakingContracts() external view returns (address[] memory);
-
-  /// @notice Retrieves a list of all the users participating in staking.
-  /// @return list of addresses of all the users who are currently staking.
-  function getAllStakers() external view returns (address[] memory);
-
-  /// @notice Retrieves a list of all the users being staked on.
-  /// @return list of addresses of all the users who are being staked on.
-  function getAllStakedOn() external view returns (address[] memory);
-
   /// @notice Query the users for the following staking addresseses.
   /// @param _stakingAddrs Addresses of staking contracts being queried.
   function getUsersForStakingAddresses(address[] calldata _stakingAddrs) external view returns (address[] memory);

src/staking/registry/RareStakingRegistry.sol

@@ -45,8 +45,17 @@ contract RareStakingRegistry is IRareStakingRegistry, AccessControlEnumerableUpg
   using strings for *;
   using SafeCast for uint256;
   using SafeCast for uint128;
-  using EnumerableMapUpgradeable for EnumerableMapUpgradeable.AddressToUintMap;
-  using EnumerableSetUpgradeable for EnumerableSetUpgradeable.AddressSet;
+
+  /*//////////////////////////////////////////////////////////////////////////
+                              Structs
+  //////////////////////////////////////////////////////////////////////////*/
+
+  /// @notice A struct holding the Rarity pool staking address and the reward accumulator address.
+  /// @dev Mainly for internal use since `Info` is exposed externally.
+  struct RarityPoolContractPair {
+    address stakingAddress;
+    address rewardAddress;
+  }
 
   /*//////////////////////////////////////////////////////////////////////////
                               Constants
@@ -75,19 +84,19 @@ contract RareStakingRegistry is IRareStakingRegistry, AccessControlEnumerableUpg
   //////////////////////////////////////////////////////////////////////////*/
 
   // Mapping of address to the User's staking info.
-  mapping(address => Info) private userToStakingInfo;
+  mapping(address => RarityPoolContractPair) private userToRarityPoolPair;
+
+  // Reverse map of staking pool address to the staking target.
+  mapping(address => address) private rarityPoolToUser;
 
   // Mapping of ERC20 token address to the ETH/ERC20 Uniswap pool.
   mapping(address => address) private swapPools;
 
-  // Enumerable set of staking contracts.
-  EnumerableSetUpgradeable.AddressSet private stakingContracts;
-
-  // Enumerable map of total RARE staked by a user.
-  EnumerableMapUpgradeable.AddressToUintMap private amountStakedByUser;
+  // Mapping of total RARE staked by a user.
+  mapping(address => uint256) private amountStakedByUser;
 
-  // Enumerable map of total RARE staked by on a target.
-  EnumerableMapUpgradeable.AddressToUintMap private amountStakedOnTarget;
+  // Mapping of total RARE staked by on a target.
+  mapping(address => uint256) private amountStakedOnTarget;
 
   // ENS reverse registrar
   ReverseRegistrar private reverseRegistrar;
@@ -185,41 +194,25 @@ contract RareStakingRegistry is IRareStakingRegistry, AccessControlEnumerableUpg
     if (_user == address(0)) revert ZeroAddressUnsupported();
     if (_stakingAddress == address(0)) revert ZeroAddressUnsupported();
     if (_rewardSwapAddress == address(0)) revert ZeroAddressUnsupported();
-    if (userToStakingInfo[_user].stakingAddress != address(0)) revert StakingContractAlreadyExists();
-    userToStakingInfo[_user] = Info("", "", _stakingAddress, _rewardSwapAddress);
-    stakingContracts.add(_stakingAddress);
+    if (userToRarityPoolPair[_user].stakingAddress != address(0)) revert StakingContractAlreadyExists();
+    userToRarityPoolPair[_user] = RarityPoolContractPair(_stakingAddress, _rewardSwapAddress);
+    rarityPoolToUser[_stakingAddress] = _user;
   }
 
   /// @inheritdoc IRareStakingRegistry
   /// @dev Requires the caller to have the {STAKING_STAT_SETTER_ROLE} access control role.
   function increaseAmountStaked(address _staker, address _stakedOn, uint256 _amount) external {
     if (!hasRole(STAKING_STAT_SETTER_ROLE, msg.sender)) revert Unauthorized();
-    (, uint256 amtStaked) = amountStakedByUser.tryGet(_staker);
-    amountStakedByUser.set(_staker, amtStaked + _amount);
-
-    (, uint256 amtStakedOn) = amountStakedOnTarget.tryGet(_stakedOn);
-    amountStakedOnTarget.set(_stakedOn, amtStakedOn + _amount);
+    amountStakedByUser[_staker] += _amount;
+    amountStakedOnTarget[_stakedOn] += _amount;
   }
 
   /// @inheritdoc IRareStakingRegistry
   /// @dev Requires the caller to have the {STAKING_STAT_SETTER_ROLE} access control role.
   function decreaseAmountStaked(address _staker, address _stakedOn, uint256 _amount) external {
     if (!hasRole(STAKING_STAT_SETTER_ROLE, msg.sender)) revert Unauthorized();
-    (, uint256 amtStaked) = amountStakedByUser.tryGet(_staker);
-
-    if (amtStaked - _amount == 0) {
-      amountStakedByUser.remove(_staker);
-    } else {
-      amountStakedByUser.set(_staker, amtStaked - _amount);
-    }
-
-    (, uint256 amtStakedOn) = amountStakedOnTarget.tryGet(_stakedOn);
-
-    if (amtStakedOn - _amount == 0) {
-      amountStakedOnTarget.remove(_stakedOn);
-    } else {
-      amountStakedOnTarget.set(_stakedOn, amtStakedOn - _amount);
-    }
+    amountStakedByUser[_staker] -= _amount;
+    amountStakedOnTarget[_stakedOn] -= _amount;
   }
 
   /// @inheritdoc IRareStakingRegistry
@@ -302,7 +295,7 @@ contract RareStakingRegistry is IRareStakingRegistry, AccessControlEnumerableUpg
     if (_rare.allowance(_from, address(this)) < _amount) {
       revert InsufficientRareAllowance();
     }
-    if (!stakingContracts.contains(msg.sender)) revert Unauthorized();
+    if (rarityPoolToUser[msg.sender] == address(0)) revert Unauthorized();
     SafeERC20Upgradeable.safeTransferFrom(_rare, _from, _to, _amount);
   }
 
@@ -347,8 +340,10 @@ contract RareStakingRegistry is IRareStakingRegistry, AccessControlEnumerableUpg
 
   /// @inheritdoc IRareStakingRegistry
   function getStakingInfoForUser(address _user) external view returns (Info memory) {
-    Info memory info = userToStakingInfo[_user];
     strings.slice memory name = resolver.name((reverseRegistrar.node(_user))).toSlice();
+    Info memory info;
+    info.stakingAddress = userToRarityPoolPair[_user].stakingAddress;
+    info.rewardAddress = userToRarityPoolPair[_user].rewardAddress;
     if (name.len() != 0) {
       name.rsplit(".".toSlice());
       info.name = ("Synthetic RARE | ".toSlice()).concat(name);
@@ -366,57 +361,22 @@ contract RareStakingRegistry is IRareStakingRegistry, AccessControlEnumerableUpg
 
   /// @inheritdoc IRareStakingRegistry
   function getStakingAddressForUser(address _user) external view returns (address) {
-    return userToStakingInfo[_user].stakingAddress;
+    return userToRarityPoolPair[_user].stakingAddress;
   }
 
   /// @inheritdoc IRareStakingRegistry
   function getRewardAccumulatorAddressForUser(address _user) external view returns (address) {
-    return userToStakingInfo[_user].rewardAddress;
+    return userToRarityPoolPair[_user].rewardAddress;
   }
 
   /// @inheritdoc IRareStakingRegistry
   function getTotalAmountStakedByUser(address _user) external view returns (uint256 amount) {
-    (, amount) = amountStakedByUser.tryGet(_user);
+    return amountStakedByUser[_user];
   }
 
   /// @inheritdoc IRareStakingRegistry
   function getTotalAmountStakedOnUser(address _user) external view returns (uint256 amount) {
-    (, amount) = amountStakedOnTarget.tryGet(_user);
-  }
-
-  /// @inheritdoc IRareStakingRegistry
-  function getAllStakingContracts() external view returns (address[] memory) {
-    return stakingContracts.values();
-  }
-
-  /// @inheritdoc IRareStakingRegistry
-  /// @dev This function is intended to be called off chain.
-  function getAllStakers() external view returns (address[] memory) {
-    uint256 length = amountStakedByUser.length();
-
-    address[] memory stakers = new address[](length);
-
-    for (uint256 i = 0; i < length; i++) {
-      (address staker, ) = amountStakedByUser.at(i);
-      stakers[i] = staker;
-    }
-
-    return stakers;
-  }
-
-  /// @inheritdoc IRareStakingRegistry
-  /// @dev This function is intended to be called off chain.
-  function getAllStakedOn() external view returns (address[] memory) {
-    uint256 length = amountStakedOnTarget.length();
-
-    address[] memory stakedOn = new address[](length);
-
-    for (uint256 i = 0; i < length; i++) {
-      (address staker, ) = amountStakedOnTarget.at(i);
-      stakedOn[i] = staker;
-    }
-
-    return stakedOn;
+    return amountStakedOnTarget[_user];
   }
 
   /// @inheritdoc IRareStakingRegistry

src/staking/token/IRarityPool.sol

@@ -132,10 +132,6 @@ interface IRarityPool is IERC20Upgradeable {
   /// @return Address of target being staked on;
   function getTargetBeingStakedOn() external view returns (address);
 
-  /// @notice Retrieves a list of all the users that have staked.
-  /// @return list of addresses of all the users who are being staked on.
-  function getAllStakers() external view returns (address[] memory);
-
   /// @notice Total rewards available for the supplied round.
   /// @return uint256 Amount of $RARE tokens allocated as rewards for round.
   function getRoundRewards(uint256 _round) external view returns (uint256);

src/staking/token/RarityPool.sol

@@ -21,7 +21,6 @@ import {IRareStakingRegistry} from "../registry/IRareStakingRegistry.sol";
 /// @dev It is one base user per contract. This is the implementation contract for a beacon proxy.
 contract RarityPool is IRarityPool, ERC20SnapshotUpgradeable, ReentrancyGuardUpgradeable {
   using EnumerableMapUpgradeable for EnumerableMapUpgradeable.AddressToUintMap;
-  using EnumerableSetUpgradeable for EnumerableSetUpgradeable.UintSet;
 
   using Address for address payable;
   using SafeCast for uint256;
@@ -36,7 +35,7 @@ contract RarityPool is IRarityPool, ERC20SnapshotUpgradeable, ReentrancyGuardUpg
   mapping(address => uint256) private lastRoundClaimedByUser;
 
   // Amount staked per user
-  EnumerableMapUpgradeable.AddressToUintMap private amountStakedByUser;
+  mapping(address => uint256) private amountStakedByUser;
 
   // The address of the target being staked on
   address private targetStakedTo;
@@ -80,7 +79,6 @@ contract RarityPool is IRarityPool, ERC20SnapshotUpgradeable, ReentrancyGuardUpg
     periodStart = block.timestamp;
     lastSnapshotTimestamp = 0;
     _mint(_creator, 1 ether);
-    amountStakedByUser.set(_creator, 0);
     _snapshot(); // we do this increment the counter so rounds and snapshot IDs are equal.
     takeSnapshot();
   }
@@ -157,13 +155,12 @@ contract RarityPool is IRarityPool, ERC20SnapshotUpgradeable, ReentrancyGuardUpg
     registry.transferRareFrom(msg.sender, address(this), _amount);
 
     // Update amount staked by user on pool and on registry
-    (, uint256 amtStaked) = amountStakedByUser.tryGet(msg.sender);
-    amountStakedByUser.set(msg.sender, amtStaked + _amount);
+    amountStakedByUser[msg.sender] += _amount;
     registry.increaseAmountStaked(msg.sender, targetStakedTo, _amount);
 
     // Mint new SRARE to staker
     _mint(msg.sender, amountSRare);
-    emit Stake(msg.sender, _amount, amtStaked + _amount, amountSRare);
+    emit Stake(msg.sender, _amount, amountStakedByUser[msg.sender], amountSRare);
   }
 
   /// @inheritdoc IRarityPool
@@ -180,16 +177,14 @@ contract RarityPool is IRarityPool, ERC20SnapshotUpgradeable, ReentrancyGuardUpg
     }
 
     // Calculate and check amount of staked RARE to return
-    (, uint256 amtStaked) = amountStakedByUser.tryGet(msg.sender);
+    uint256 amtStaked = amountStakedByUser[msg.sender];
     uint256 amountRareReturned = calculateSaleReturn(balanceOf(msg.sender), amtStaked, _amount);
     if (amountRareReturned > amtStaked) {
       revert InsufficientStakedRare();
     }
 
     // Clean up staked amount on pool and on registry
-    amtStaked - amountRareReturned == 0
-      ? amountStakedByUser.remove(msg.sender)
-      : amountStakedByUser.set(msg.sender, amtStaked - amountRareReturned);
+    amountStakedByUser[msg.sender] = amtStaked - amountRareReturned;
     stakingRegistry.decreaseAmountStaked(msg.sender, targetStakedTo, amountRareReturned);
 
     // Burn SRARE
@@ -250,8 +245,7 @@ contract RarityPool is IRarityPool, ERC20SnapshotUpgradeable, ReentrancyGuardUpg
 
   /// @inheritdoc IRarityPool
   function getAmountStakedByUser(address _user) external view returns (uint256) {
-    (, uint256 amtStaked) = amountStakedByUser.tryGet(_user);
-    return amtStaked;
+    return amountStakedByUser[_user];
   }
 
   /// @inheritdoc IRarityPool
@@ -271,20 +265,6 @@ contract RarityPool is IRarityPool, ERC20SnapshotUpgradeable, ReentrancyGuardUpg
     return targetStakedTo;
   }
 
-  /// @inheritdoc IRarityPool
-  function getAllStakers() external view returns (address[] memory) {
-    uint256 length = amountStakedByUser.length();
-
-    address[] memory stakers = new address[](length);
-
-    for (uint256 i = 0; i < length; i++) {
-      (address staker, ) = amountStakedByUser.at(i);
-      stakers[i] = staker;
-    }
-
-    return stakers;
-  }
-
   /// @inheritdoc IRarityPool
   function getRoundRewards(uint256 _round) public view returns (uint256) {
     return roundRewardAmount[_round];