UniStaker supports multiple operations via multicall

src/UniStaker.sol

@@ -7,6 +7,7 @@ import {IERC20Delegates} from "src/interfaces/IERC20Delegates.sol";
 import {IERC20} from "openzeppelin/token/ERC20/IERC20.sol";
 import {SafeERC20} from "openzeppelin/token/ERC20/utils/SafeERC20.sol";
 import {ReentrancyGuard} from "openzeppelin/utils/ReentrancyGuard.sol";
+import {Multicall} from "openzeppelin/utils/Multicall.sol";
 
 /// @title UniStaker
 /// @author ScopeLift
@@ -25,7 +26,7 @@ import {ReentrancyGuard} from "openzeppelin/utils/ReentrancyGuard.sol";
 /// received, the reward duration restarts, and the rate at which rewards are streamed is updated
 /// to include the newly received rewards along with any remaining rewards that have finished
 /// streaming since the last time a reward was received.
-contract UniStaker is INotifiableRewardReceiver, ReentrancyGuard {
+contract UniStaker is INotifiableRewardReceiver, ReentrancyGuard, Multicall {
   type DepositIdentifier is uint256;
 
   /// @notice Thrown when an account attempts a call for which it lacks appropriate permission.
@@ -205,7 +206,7 @@ contract UniStaker is INotifiableRewardReceiver, ReentrancyGuard {
   /// @dev Neither the delegatee nor the beneficiary may be the zero address. The deposit will be
   /// owned by the message sender.
   function stake(uint256 _amount, address _delegatee, address _beneficiary)
-    public
+    external
     nonReentrant
     returns (DepositIdentifier _depositId)
   {
@@ -239,7 +240,10 @@ contract UniStaker is INotifiableRewardReceiver, ReentrancyGuard {
   /// @param _newDelegatee Address of the new governance delegate.
   /// @dev The new delegatee may not be the zero address. The message sender must be the owner of
   /// the deposit.
-  function alterDelegatee(DepositIdentifier _depositId, address _newDelegatee) public nonReentrant {
+  function alterDelegatee(DepositIdentifier _depositId, address _newDelegatee)
+    external
+    nonReentrant
+  {
     _revertIfAddressZero(_newDelegatee);
     Deposit storage deposit = deposits[_depositId];
     _revertIfNotDepositOwner(deposit);
@@ -257,7 +261,7 @@ contract UniStaker is INotifiableRewardReceiver, ReentrancyGuard {
   /// @dev The new beneficiary may not be the zero address. The message sender must be the owner of
   /// the deposit.
   function alterBeneficiary(DepositIdentifier _depositId, address _newBeneficiary)
-    public
+    external
     nonReentrant
   {
     _revertIfAddressZero(_newBeneficiary);

test/UniStaker.t.sol

@@ -1938,3 +1938,131 @@ contract ClaimReward is UniStakerRewardsTest {
     assertEq(uniStaker.earned(_depositor), 0);
   }
 }
+
+contract Multicall is UniStakerRewardsTest {
+  function _encodeStake(address _delegatee, uint256 _stakeAmount)
+    internal
+    pure
+    returns (bytes memory)
+  {
+    return
+      abi.encodeWithSelector(bytes4(keccak256("stake(uint256,address)")), _stakeAmount, _delegatee);
+  }
+
+  function _encodeStake(address _delegatee, uint256 _stakeAmount, address _beneficiary)
+    internal
+    pure
+    returns (bytes memory)
+  {
+    return abi.encodeWithSelector(
+      bytes4(keccak256("stake(uint256,address,address)")), _stakeAmount, _delegatee, _beneficiary
+    );
+  }
+
+  function _encodeStakeMore(UniStaker.DepositIdentifier _depositId, uint256 _stakeAmount)
+    internal
+    pure
+    returns (bytes memory)
+  {
+    return abi.encodeWithSelector(
+      bytes4(keccak256("stakeMore(uint256,uint256)")), _depositId, _stakeAmount
+    );
+  }
+
+  function _encodeWithdraw(UniStaker.DepositIdentifier _depositId, uint256 _amount)
+    internal
+    pure
+    returns (bytes memory)
+  {
+    return
+      abi.encodeWithSelector(bytes4(keccak256("withdraw(uint256,uint256)")), _depositId, _amount);
+  }
+
+  function _encodeAlterBeneficiary(UniStaker.DepositIdentifier _depositId, address _beneficiary)
+    internal
+    pure
+    returns (bytes memory)
+  {
+    return abi.encodeWithSelector(
+      bytes4(keccak256("alterBeneficiary(uint256,address)")), _depositId, _beneficiary
+    );
+  }
+
+  function _encodeAlterDelegatee(UniStaker.DepositIdentifier _depositId, address _delegatee)
+    internal
+    pure
+    returns (bytes memory)
+  {
+    return abi.encodeWithSelector(
+      bytes4(keccak256("alterDelegatee(uint256,address)")), _depositId, _delegatee
+    );
+  }
+
+  function testFuzz_CanUseMulticallToStakeMultipleTimes(
+    address _depositor,
+    address _delegatee1,
+    address _delegatee2,
+    uint256 _stakeAmount1,
+    uint256 _stakeAmount2
+  ) public {
+    _stakeAmount1 = _boundToRealisticStake(_stakeAmount1);
+    _stakeAmount2 = _boundToRealisticStake(_stakeAmount2);
+    vm.assume(_delegatee1 != address(0) && _delegatee2 != address(0));
+    _mintGovToken(_depositor, _stakeAmount1 + _stakeAmount2);
+
+    vm.prank(_depositor);
+    govToken.approve(address(uniStaker), _stakeAmount1 + _stakeAmount2);
+
+    bytes[] memory _calls = new bytes[](2);
+    _calls[0] = _encodeStake(_delegatee1, _stakeAmount1);
+    _calls[1] = _encodeStake(_delegatee2, _stakeAmount2);
+    vm.prank(_depositor);
+    uniStaker.multicall(_calls);
+    assertEq(uniStaker.totalDeposits(_depositor), _stakeAmount1 + _stakeAmount2);
+  }
+
+  function testFuzz_CanUseMulticallToStakeAndAlterBeneficiaryAndDelegatee(
+    address _depositor,
+    address _delegatee0,
+    address _delegatee1,
+    address _beneficiary0,
+    address _beneficiary1,
+    uint256 _stakeAmount0,
+    uint256 _stakeAmount1,
+    uint256 _timeElapsed
+  ) public {
+    _stakeAmount0 = _boundToRealisticStake(_stakeAmount0);
+    _stakeAmount1 = _boundToRealisticStake(_stakeAmount1);
+
+    vm.assume(
+      _depositor != address(0) && _delegatee0 != address(0) && _delegatee1 != address(0)
+        && _beneficiary0 != address(0) && _beneficiary1 != address(0)
+    );
+    _mintGovToken(_depositor, _stakeAmount0 + _stakeAmount1);
+
+    vm.startPrank(_depositor);
+    govToken.approve(address(uniStaker), _stakeAmount0 + _stakeAmount1);
+
+    // first, do initial stake without multicall
+    UniStaker.DepositIdentifier _depositId =
+      uniStaker.stake(_stakeAmount0, _delegatee0, _beneficiary0);
+
+    // some time goes by...
+    vm.warp(_timeElapsed);
+
+    // now I want to stake more, and also change my delegatee and beneficiary
+    bytes[] memory _calls = new bytes[](3);
+    _calls[0] = _encodeStakeMore(_depositId, _stakeAmount1);
+    _calls[1] = _encodeAlterBeneficiary(_depositId, _beneficiary1);
+    _calls[2] = _encodeAlterDelegatee(_depositId, _delegatee1);
+    uniStaker.multicall(_calls);
+    vm.stopPrank();
+
+    (uint256 _amountResult,, address _delegateeResult, address _beneficiaryResult) =
+      uniStaker.deposits(_depositId);
+    assertEq(uniStaker.totalDeposits(_depositor), _stakeAmount0 + _stakeAmount1);
+    assertEq(_amountResult, _stakeAmount0 + _stakeAmount1);
+    assertEq(_delegateeResult, _delegatee1);
+    assertEq(_beneficiaryResult, _beneficiary1);
+  }
+}