staking: use multicall on the Staking contract

Author: abarmat

contracts/staking/IStaking.sol

@@ -113,25 +113,10 @@ interface IStaking is IStakingData {
 
     function closeAllocation(address _allocationID, bytes32 _poi) external;
 
-    function closeAllocationMany(CloseAllocationRequest[] calldata _requests) external;
-
-    function closeAndAllocate(
-        address _oldAllocationID,
-        bytes32 _poi,
-        address _indexer,
-        bytes32 _subgraphDeploymentID,
-        uint256 _tokens,
-        address _allocationID,
-        bytes32 _metadata,
-        bytes calldata _proof
-    ) external;
-
     function collect(uint256 _tokens, address _allocationID) external;
 
     function claim(address _allocationID, bool _restake) external;
 
-    function claimMany(address[] calldata _allocationID, bool _restake) external;
-
     // -- Getters and calculations --
 
     function hasStake(address _indexer) external view returns (bool);

contracts/staking/Staking.sol

@@ -5,6 +5,7 @@ pragma experimental ABIEncoderV2;
 
 import "@openzeppelin/contracts/cryptography/ECDSA.sol";
 
+import "../base/Multicall.sol";
 import "../upgrades/GraphUpgradeable.sol";
 import "../utils/TokenUtils.sol";
 
@@ -20,7 +21,7 @@ import "./libs/Stakes.sol";
  * Allocations on a Subgraph. It also allows Delegators to Delegate towards an Indexer. The
  * contract also has the slashing functionality.
  */
-contract Staking is StakingV2Storage, GraphUpgradeable, IStaking {
+contract Staking is StakingV2Storage, GraphUpgradeable, IStaking, Multicall {
     using SafeMath for uint256;
     using Stakes for Stakes.Indexer;
     using Rebates for Rebates.Pool;
@@ -908,49 +909,6 @@ contract Staking is StakingV2Storage, GraphUpgradeable, IStaking {
         _closeAllocation(_allocationID, _poi);
     }
 
-    /**
-     * @dev Close multiple allocations and free the staked tokens.
-     * To be eligible for rewards a proof of indexing must be presented.
-     * Presenting a bad proof is subject to slashable condition.
-     * To opt out for rewards set _poi to 0x0
-     * @param _requests An array of CloseAllocationRequest
-     */
-    function closeAllocationMany(CloseAllocationRequest[] calldata _requests)
-        external
-        override
-        notPaused
-    {
-        for (uint256 i = 0; i < _requests.length; i++) {
-            _closeAllocation(_requests[i].allocationID, _requests[i].poi);
-        }
-    }
-
-    /**
-     * @dev Close and allocate. This will perform a close and then create a new Allocation
-     * atomically on the same transaction.
-     * @param _closingAllocationID The identifier of the allocation to be closed
-     * @param _poi Proof of indexing submitted for the allocated period
-     * @param _indexer Indexer address to allocate funds from.
-     * @param _subgraphDeploymentID ID of the SubgraphDeployment where tokens will be allocated
-     * @param _tokens Amount of tokens to allocate
-     * @param _allocationID The allocation identifier
-     * @param _metadata IPFS hash for additional information about the allocation
-     * @param _proof A 65-bytes Ethereum signed message of `keccak256(indexerAddress,allocationID)`
-     */
-    function closeAndAllocate(
-        address _closingAllocationID,
-        bytes32 _poi,
-        address _indexer,
-        bytes32 _subgraphDeploymentID,
-        uint256 _tokens,
-        address _allocationID,
-        bytes32 _metadata,
-        bytes calldata _proof
-    ) external override notPaused {
-        _closeAllocation(_closingAllocationID, _poi);
-        _allocate(_indexer, _subgraphDeploymentID, _tokens, _allocationID, _metadata, _proof);
-    }
-
     /**
      * @dev Collect query fees from state channels and assign them to an allocation.
      * Funds received are only accepted from a valid sender.
@@ -1036,21 +994,6 @@ contract Staking is StakingV2Storage, GraphUpgradeable, IStaking {
         _claim(_allocationID, _restake);
     }
 
-    /**
-     * @dev Claim tokens from the rebate pool for many allocations.
-     * @param _allocationID Array of allocations from where we are claiming tokens
-     * @param _restake True if restake fees instead of transfer to indexer
-     */
-    function claimMany(address[] calldata _allocationID, bool _restake)
-        external
-        override
-        notPaused
-    {
-        for (uint256 i = 0; i < _allocationID.length; i++) {
-            _claim(_allocationID[i], _restake);
-        }
-    }
-
     /**
      * @dev Stake tokens on the indexer.
      * This function does not check minimum indexer stake requirement to allow

test/staking/allocation.test.ts

@@ -1,5 +1,5 @@
 import { expect } from 'chai'
-import { constants, BigNumber } from 'ethers'
+import { constants, BigNumber, PopulatedTransaction } from 'ethers'
 
 import { Curation } from '../../build/types/Curation'
 import { EpochManager } from '../../build/types/EpochManager'
@@ -644,17 +644,21 @@ describe('Staking:Allocation', () => {
       await advanceToNextEpoch(epochManager)
 
       // Close multiple allocations in one tx
-      const requests = [
-        {
-          allocationID: allocationID,
-          poi: poi,
-        },
-        {
-          allocationID: allocationID2,
-          poi: poi,
-        },
-      ]
-      await staking.connect(indexer.signer).closeAllocationMany(requests)
+      const requests = await Promise.all(
+        [
+          {
+            allocationID: allocationID,
+            poi: poi,
+          },
+          {
+            allocationID: allocationID2,
+            poi: poi,
+          },
+        ].map(({ allocationID, poi }) =>
+          staking.connect(indexer.signer).populateTransaction.closeAllocation(allocationID, poi),
+        ),
+      ).then((e) => e.map((e: PopulatedTransaction) => e.data))
+      await staking.connect(indexer.signer).multicall(requests)
     })
   })
 
@@ -672,19 +676,21 @@ describe('Staking:Allocation', () => {
       // Close and allocate
       const newChannelKey = deriveChannelKey()
       const newAllocationID = newChannelKey.address
-      const tx = staking
-        .connect(indexer.signer)
-        .closeAndAllocate(
-          allocationID,
-          HashZero,
-          indexer.address,
-          subgraphDeploymentID,
-          tokensToAllocate,
-          newAllocationID,
-          metadata,
-          await newChannelKey.generateProof(indexer.address),
-        )
-      await tx
+
+      // Close multiple allocations in one tx
+      const requests = await Promise.all([
+        staking.connect(indexer.signer).populateTransaction.closeAllocation(allocationID, poi),
+        staking
+          .connect(indexer.signer)
+          .populateTransaction.allocate(
+            subgraphDeploymentID,
+            tokensToAllocate,
+            newAllocationID,
+            metadata,
+            await newChannelKey.generateProof(indexer.address),
+          ),
+      ]).then((e) => e.map((e: PopulatedTransaction) => e.data))
+      await staking.connect(indexer.signer).multicall(requests)
     })
   })
 
@@ -870,7 +876,10 @@ describe('Staking:Allocation', () => {
 
         // Claim with restake
         expect(await staking.getAllocationState(allocationID)).eq(AllocationState.Finalized)
-        await staking.connect(indexer.signer).claimMany([allocationID], true)
+        const tx = await staking
+          .connect(indexer.signer)
+          .populateTransaction.claim(allocationID, true)
+        await staking.connect(indexer.signer).multicall([tx.data])
 
         // Verify that the claimed tokens are restaked
         const afterIndexerStake = await staking.getIndexerStakedTokens(indexer.address)