Add multicall

uniswapfoundation · Jan 26, 2024 · 7a367ca · 7a367ca
1 parent 12eb7ec
commit 7a367ca
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 2 deletions.
diff --git a/src/UniStaker.sol b/src/UniStaker.sol
@@ -7,8 +7,9 @@ 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";
 
-contract UniStaker is INotifiableRewardReceiver, ReentrancyGuard {
+contract UniStaker is INotifiableRewardReceiver, ReentrancyGuard, Multicall {
   type DepositIdentifier is uint256;
 
   error UniStaker__Unauthorized(bytes32 reason, address caller);
@@ -73,7 +74,7 @@ contract UniStaker is INotifiableRewardReceiver, ReentrancyGuard {
   }
 
   function stake(uint256 _amount, address _delegatee)
-    external
+    public
     nonReentrant
     returns (DepositIdentifier _depositId)
   {

diff --git a/test/UniStaker.t.sol b/test/UniStaker.t.sol
@@ -1802,3 +1802,34 @@ contract ClaimReward is UniStakerRewardsTest {
     assertEq(uniStaker.earned(_depositor), 0);
   }
 }
+
+contract Multicall is UniStakerRewardsTest {
+  function testFuzz_CanCallMultipleMethodsInASingleTransaction(
+    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);
+    abi.encodeWithSelector(bytes4(keccak256("stake(uint,address)")), _stakeAmount1, _delegatee1);
+    abi.encodeWithSelector(bytes4(keccak256("stake(uint,address)")), _stakeAmount2, _delegatee2);
+
+    bytes[] memory _calls = new bytes[](2);
+    _calls[0] = abi.encodeWithSelector(
+      bytes4(keccak256("stake(uint256,address)")), _stakeAmount1, _delegatee1
+    );
+    _calls[1] = abi.encodeWithSelector(
+      bytes4(keccak256("stake(uint256,address)")), _stakeAmount2, _delegatee2
+    );
+    vm.prank(_depositor);
+    uniStaker.multicall(_calls);
+    assertEq(uniStaker.totalDeposits(_depositor), _stakeAmount1 + _stakeAmount2);
+  }
+}