refactor(optimization): significantly reduce sloads/sstores

missed some stuff
Layr-Labs · Sep 18, 2024 · 536a97d · 536a97d
1 parent 902ec11
commit 536a97d
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 27 deletions.
diff --git a/src/contracts/core/StakeRootCompendium.sol b/src/contracts/core/StakeRootCompendium.sol
@@ -269,23 +269,24 @@ contract StakeRootCompendium is StakeRootCompendiumStorage {
         uint96 _maxTotalCharge
     ) external onlyOwner {
         require(_maxTotalCharge >= totalChargeHistory.latest(), MaxTotalChargeMustBeGreaterThanTheCurrentTotalCharge());
-        maxTotalCharge = _maxTotalCharge;
+        stakerootCharges.maxTotalCharge = _maxTotalCharge;
     }
 
     /// @inheritdoc IStakeRootCompendium
     function setChargePerProof(uint96 _chargePerStrategy, uint96 _chargePerOperatorSet) external onlyOwner {
-        _updateTotalCharge();
-        chargePerStrategy = _chargePerStrategy;
-        chargePerOperatorSet = _chargePerOperatorSet;
-        _updateTotalCharge();
+        StakerootCharges storage charges = stakerootCharges;
+        _updateTotalCharge(charges);
+        charges.chargePerStrategy = _chargePerStrategy;
+        charges.chargePerOperatorSet = _chargePerOperatorSet;
+        _updateTotalCharge(charges);
     }
 
     /// @inheritdoc IStakeRootCompendium
     function setProofIntervalSeconds(
         uint32 proofIntervalSeconds
     ) external onlyOwner {
         StakerootCumulativeCharges storage cumulativeCharges = stakerootCumulativeCharges;
-        _updateTotalCharge();
+        _updateTotalCharge(stakerootCharges);
         // we must not interrupt pending proof calculations by rugging the outstanding calculationTimestamps
         require(
             stakeRootSubmissions[stakeRootSubmissions.length - 1].calculationTimestamp
@@ -328,13 +329,13 @@ contract StakeRootCompendium is StakeRootCompendiumStorage {
 
     function _updateTotalStrategies(uint256 _countStrategiesBefore, uint256 _countStrategiesAfter) internal {
         totalStrategies = totalStrategies - _countStrategiesBefore + _countStrategiesAfter;
-        _updateTotalCharge();
+        _updateTotalCharge(stakerootCharges);
     }
 
-    function _updateTotalCharge() internal {
+    function _updateTotalCharge(StakerootCharges memory charges) internal {
         // note if totalStrategies is 0, the charge per proof will be 0, and provers should not post a proof
-        uint256 totalCharge = operatorSets.length * chargePerOperatorSet + totalStrategies * chargePerStrategy;
-        require(totalCharge <= maxTotalCharge, ChargePerProofExceedsMaxTotalCharge());
+        uint256 totalCharge = operatorSets.length * charges.chargePerOperatorSet + totalStrategies * charges.chargePerStrategy;
+        require(totalCharge <= charges.maxTotalCharge, ChargePerProofExceedsMaxTotalCharge());
         totalChargeHistory.push(uint32(block.timestamp), uint224(totalCharge));
     }
 
@@ -458,7 +459,8 @@ contract StakeRootCompendium is StakeRootCompendiumStorage {
     function minDepositBalance(
         uint256 numStrategies
     ) public view returns (uint256) {
-        return (numStrategies * chargePerStrategy + chargePerOperatorSet) * MIN_PREPAID_PROOFS;
+        StakerootCharges memory charges = stakerootCharges;
+        return (numStrategies * charges.chargePerStrategy + charges.chargePerOperatorSet) * MIN_PREPAID_PROOFS;
     }
 
     /// @inheritdoc IStakeRootCompendium
@@ -517,21 +519,22 @@ contract StakeRootCompendium is StakeRootCompendiumStorage {
 
     /// @inheritdoc IStakeRootCompendium
     function getStakeRoot(
-        OperatorSet[] calldata operatorSetsInStakeTree,
+        address avs,
+        uint32[] calldata operatorSetIdsInStakeTree,
         bytes32[] calldata operatorSetRoots
     ) external view returns (bytes32) {
         // TODO: This fn should revert if mismatched parameters are passed in due to out of bounds
         // array access, see if these checks can be removed.
-        require(operatorSets.length == operatorSetsInStakeTree.length, InputArrayLengthMismatch());
-        require(operatorSetsInStakeTree.length == operatorSetRoots.length, InputArrayLengthMismatch());
-        for (uint256 i = 0; i < operatorSetsInStakeTree.length; i++) {
-            /// TODO: Should take in an array of operatorSetIds instead, and a constant avs.
-            require(operatorSets[i].avs == operatorSetsInStakeTree[i].avs, InputCorrelatedVariableMismatch());
+        require(operatorSets.length == operatorSetIdsInStakeTree.length, InputArrayLengthMismatch());
+        require(operatorSetIdsInStakeTree.length == operatorSetRoots.length, InputArrayLengthMismatch());
+
+        for (uint256 i = 0; i < operatorSetIdsInStakeTree.length; i++) {            
             require(
-                operatorSets[i].operatorSetId == operatorSetsInStakeTree[i].operatorSetId,
+                operatorSets[i].operatorSetId == operatorSetIdsInStakeTree[i],
                 InputCorrelatedVariableMismatch()
             );
         }
+
         return Merkle.merkleizeKeccak256(operatorSetRoots);
     }
 

diff --git a/src/contracts/core/StakeRootCompendiumStorage.sol b/src/contracts/core/StakeRootCompendiumStorage.sol
@@ -44,13 +44,6 @@ abstract contract StakeRootCompendiumStorage is IStakeRootCompendium, OwnableUpg
     /// @notice the address allowed to confirm roots
     address public rootConfirmer;
 
-    /// @notice the linear charge per proof in the number of strategies
-    uint96 public chargePerOperatorSet;
-    /// @notice the constant charge per proof
-    uint96 public chargePerStrategy;
-    /// @notice the max total charge for a stakeRoot proof. used to bound computation offchain
-    uint96 public maxTotalCharge;
-
     /// @notice list of operator sets that have been configured to be in the StakeTree
     OperatorSet[] public operatorSets;
     /// @notice the total charge for a proofs at a certain time depending on the number of strategies

diff --git a/src/contracts/interfaces/IStakeRootCompendium.sol b/src/contracts/interfaces/IStakeRootCompendium.sol
@@ -127,13 +127,14 @@ interface IStakeRootCompendium {
 
     /**
      * @notice called offchain with the operatorSet roots ordered by the operatorSet index at the timestamp to calculate the stake root
-     * @param operatorSetsInStakeTree the operatorSets that each of the operatorSetRoots correspond to. must be the same as operatorSets storage var at the time of call
+     * @param operatorSetIdsInStakeTree the operatorSets ids that each of the operatorSetRoots correspond to. must be the same as operatorSets storage var at the time of call
      * @param operatorSetRoots the ordered operatorSet roots (not verified)
      * @dev operatorSetsInStakeTree must be the same as operatorSets storage var at the time of call
      * @dev operatorSetRoots must be ordered by the operatorSet index at the time of call
      */
     function getStakeRoot(
-        OperatorSet[] calldata operatorSetsInStakeTree,
+        address avs,
+        uint32[] calldata operatorSetIdsInStakeTree,
         bytes32[] calldata operatorSetRoots
     ) external view returns (bytes32);