feat: add "IMulticall3" interface and "getTokenBalances" method (#271)

PaulRBerg · mds1 · web-flow · commit 066ff16c5c03 · 2023-01-17T16:34:37.000-08:00
* feat: add "IMulticall3" interface chore: fix typos in comments feat: add "getTokenBalances" util refactor: add header separators in "StdUtils" refactor: order functions alphabetically in "StdUtils" * fix: set pragma to ">=0.6.2 <0.9.0" in IMulticall3 fix: add experimental ABIEncoderV2 pragma style: apply "forge fmt" fixes * test: add header separators in "StdUtils.t.sol" test: organize tests alphabetically in "StdUtils.t.sol" * test: add revert test for "getTokenBalances" * Revert "test: add revert test for "getTokenBalances"" This reverts commit c3bedee. * fix/test: fix multicall fork test * refactor: disallow EOAs in "getTokenBalances" test: write "getTokenBalances" test when address is EOA test: test USDC and SHIB holders with "getTokenBalances" * test: add empty array test for "getTokenBalances" chore: delete "console2" log import in "StdUtils.sol" * chore: fix parentheses direction * fix: backwards compatibility with older solc versions * refactor: pull fork call into setUp * fix: codesize check * chore: add comment Co-authored-by: Matt Solomon <matt@mattsolomon.dev>
diff --git a/src/StdUtils.sol b/src/StdUtils.sol
@@ -1,18 +1,29 @@
 // SPDX-License-Identifier: MIT
 pragma solidity >=0.6.2 <0.9.0;
 
+pragma experimental ABIEncoderV2;
+
+import {IMulticall3} from "./interfaces/IMulticall3.sol";
 // TODO Remove import.
 import {VmSafe} from "./Vm.sol";
 
 abstract contract StdUtils {
+    /*//////////////////////////////////////////////////////////////////////////
+                                     CONSTANTS
+    //////////////////////////////////////////////////////////////////////////*/
+
+    IMulticall3 private constant multicall = IMulticall3(0xcA11bde05977b3631167028862bE2a173976CA11);
     VmSafe private constant vm = VmSafe(address(uint160(uint256(keccak256("hevm cheat code")))));
     address private constant CONSOLE2_ADDRESS = 0x000000000000000000636F6e736F6c652e6c6f67;
-
     uint256 private constant INT256_MIN_ABS =
         57896044618658097711785492504343953926634992332820282019728792003956564819968;
     uint256 private constant UINT256_MAX =
         115792089237316195423570985008687907853269984665640564039457584007913129639935;
 
+    /*//////////////////////////////////////////////////////////////////////////
+                                 INTERNAL FUNCTIONS
+    //////////////////////////////////////////////////////////////////////////*/
+
     function _bound(uint256 x, uint256 min, uint256 max) internal pure virtual returns (uint256 result) {
         require(min <= max, "StdUtils bound(uint256,uint256,uint256): Max is less than min.");
         // If x is between min and max, return x directly. This is to ensure that dictionary values
@@ -66,8 +77,13 @@ abstract contract StdUtils {
         console2_log("Bound result", vm.toString(result));
     }
 
+    function bytesToUint(bytes memory b) internal pure virtual returns (uint256) {
+        require(b.length <= 32, "StdUtils bytesToUint(bytes): Bytes length exceeds 32.");
+        return abi.decode(abi.encodePacked(new bytes(32 - b.length), b), (uint256));
+    }
+
     /// @dev Compute the address a contract will be deployed at for a given deployer address and nonce
-    /// @notice adapated from Solmate implementation (https://github.com/Rari-Capital/solmate/blob/main/src/utils/LibRLP.sol)
+    /// @notice adapted from Solmate implementation (https://github.com/Rari-Capital/solmate/blob/main/src/utils/LibRLP.sol)
     function computeCreateAddress(address deployer, uint256 nonce) internal pure virtual returns (address) {
         // forgefmt: disable-start
         // The integer zero is treated as an empty byte string, and as a result it only has a length prefix, 0x80, computed via 0x80 + 0.
@@ -100,11 +116,40 @@ abstract contract StdUtils {
         return addressFromLast20Bytes(keccak256(abi.encodePacked(bytes1(0xff), deployer, salt, initcodeHash)));
     }
 
-    function bytesToUint(bytes memory b) internal pure virtual returns (uint256) {
-        require(b.length <= 32, "StdUtils bytesToUint(bytes): Bytes length exceeds 32.");
-        return abi.decode(abi.encodePacked(new bytes(32 - b.length), b), (uint256));
+    // Performs a single call with Multicall3 to query the ERC-20 token balances of the given addresses.
+    function getTokenBalances(address token, address[] memory addresses)
+        internal
+        virtual
+        returns (uint256[] memory balances)
+    {
+        uint256 tokenCodeSize;
+        assembly {
+            tokenCodeSize := extcodesize(token)
+        }
+        require(tokenCodeSize > 0, "StdUtils getTokenBalances(address,address[]): Token address is not a contract.");
+
+        // ABI encode the aggregate call to Multicall3.
+        uint256 length = addresses.length;
+        IMulticall3.Call[] memory calls = new IMulticall3.Call[](length);
+        for (uint256 i = 0; i < length; ++i) {
+            // 0x70a08231 = bytes4("balanceOf(address)"))
+            calls[i] = IMulticall3.Call({target: token, callData: abi.encodeWithSelector(0x70a08231, (addresses[i]))});
+        }
+
+        // Make the aggregate call.
+        (, bytes[] memory returnData) = multicall.aggregate(calls);
+
+        // ABI decode the return data and return the balances.
+        balances = new uint256[](length);
+        for (uint256 i = 0; i < length; ++i) {
+            balances[i] = abi.decode(returnData[i], (uint256));
+        }
     }
 
+    /*//////////////////////////////////////////////////////////////////////////
+                                 PRIVATE FUNCTIONS
+    //////////////////////////////////////////////////////////////////////////*/
+
     function addressFromLast20Bytes(bytes32 bytesValue) private pure returns (address) {
         return address(uint160(uint256(bytesValue)));
     }
diff --git a/src/interfaces/IMulticall3.sol b/src/interfaces/IMulticall3.sol
@@ -0,0 +1,73 @@
+// SPDX-License-Identifier: MIT
+pragma solidity >=0.6.2 <0.9.0;
+
+pragma experimental ABIEncoderV2;
+
+interface IMulticall3 {
+    struct Call {
+        address target;
+        bytes callData;
+    }
+
+    struct Call3 {
+        address target;
+        bool allowFailure;
+        bytes callData;
+    }
+
+    struct Call3Value {
+        address target;
+        bool allowFailure;
+        uint256 value;
+        bytes callData;
+    }
+
+    struct Result {
+        bool success;
+        bytes returnData;
+    }
+
+    function aggregate(Call[] calldata calls)
+        external
+        payable
+        returns (uint256 blockNumber, bytes[] memory returnData);
+
+    function aggregate3(Call3[] calldata calls) external payable returns (Result[] memory returnData);
+
+    function aggregate3Value(Call3Value[] calldata calls) external payable returns (Result[] memory returnData);
+
+    function blockAndAggregate(Call[] calldata calls)
+        external
+        payable
+        returns (uint256 blockNumber, bytes32 blockHash, Result[] memory returnData);
+
+    function getBasefee() external view returns (uint256 basefee);
+
+    function getBlockHash(uint256 blockNumber) external view returns (bytes32 blockHash);
+
+    function getBlockNumber() external view returns (uint256 blockNumber);
+
+    function getChainId() external view returns (uint256 chainid);
+
+    function getCurrentBlockCoinbase() external view returns (address coinbase);
+
+    function getCurrentBlockDifficulty() external view returns (uint256 difficulty);
+
+    function getCurrentBlockGasLimit() external view returns (uint256 gaslimit);
+
+    function getCurrentBlockTimestamp() external view returns (uint256 timestamp);
+
+    function getEthBalance(address addr) external view returns (uint256 balance);
+
+    function getLastBlockHash() external view returns (bytes32 blockHash);
+
+    function tryAggregate(bool requireSuccess, Call[] calldata calls)
+        external
+        payable
+        returns (Result[] memory returnData);
+
+    function tryBlockAndAggregate(bool requireSuccess, Call[] calldata calls)
+        external
+        payable
+        returns (uint256 blockNumber, bytes32 blockHash, Result[] memory returnData);
+}
diff --git a/test/StdUtils.t.sol b/test/StdUtils.t.sol
@@ -3,7 +3,21 @@ pragma solidity >=0.7.0 <0.9.0;
 
 import "../src/Test.sol";
 
+contract StdUtilsMock is StdUtils {
+    // We deploy a mock version so we can properly test expected reverts.
+    function getTokenBalances_(address token, address[] memory addresses)
+        external
+        returns (uint256[] memory balances)
+    {
+        return getTokenBalances(token, addresses);
+    }
+}
+
 contract StdUtilsTest is Test {
+    /*//////////////////////////////////////////////////////////////////////////
+                                     BOUND UINT
+    //////////////////////////////////////////////////////////////////////////*/
+
     function testBound() public {
         assertEq(bound(uint256(5), 0, 4), 0);
         assertEq(bound(uint256(0), 69, 69), 69);
@@ -74,6 +88,10 @@ contract StdUtilsTest is Test {
         bound(num, min, max);
     }
 
+    /*//////////////////////////////////////////////////////////////////////////
+                                     BOUND INT
+    //////////////////////////////////////////////////////////////////////////*/
+
     function testBoundInt() public {
         assertEq(bound(-3, 0, 4), 2);
         assertEq(bound(0, -69, -69), -69);
@@ -158,34 +176,114 @@ contract StdUtilsTest is Test {
         bound(num, min, max);
     }
 
-    function testGenerateCreateAddress() external {
+    /*//////////////////////////////////////////////////////////////////////////
+                                   BYTES TO UINT
+    //////////////////////////////////////////////////////////////////////////*/
+
+    function testBytesToUint() external {
+        bytes memory maxUint = hex"ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff";
+        bytes memory two = hex"02";
+        bytes memory millionEther = hex"d3c21bcecceda1000000";
+
+        assertEq(bytesToUint(maxUint), type(uint256).max);
+        assertEq(bytesToUint(two), 2);
+        assertEq(bytesToUint(millionEther), 1_000_000 ether);
+    }
+
+    function testCannotConvertGT32Bytes() external {
+        bytes memory thirty3Bytes = hex"ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff";
+        vm.expectRevert("StdUtils bytesToUint(bytes): Bytes length exceeds 32.");
+        bytesToUint(thirty3Bytes);
+    }
+
+    /*//////////////////////////////////////////////////////////////////////////
+                               COMPUTE CREATE ADDRESS
+    //////////////////////////////////////////////////////////////////////////*/
+
+    function testComputeCreateAddress() external {
         address deployer = 0x6C9FC64A53c1b71FB3f9Af64d1ae3A4931A5f4E9;
         uint256 nonce = 14;
         address createAddress = computeCreateAddress(deployer, nonce);
         assertEq(createAddress, 0x68b3465833fb72A70ecDF485E0e4C7bD8665Fc45);
     }
 
-    function testGenerateCreate2Address() external {
+    /*//////////////////////////////////////////////////////////////////////////
+                              COMPUTE CREATE2 ADDRESS
+    //////////////////////////////////////////////////////////////////////////*/
+
+    function testComputeCreate2Address() external {
         bytes32 salt = bytes32(uint256(31415));
         bytes32 initcodeHash = keccak256(abi.encode(0x6080));
         address deployer = 0x6C9FC64A53c1b71FB3f9Af64d1ae3A4931A5f4E9;
         address create2Address = computeCreate2Address(salt, initcodeHash, deployer);
         assertEq(create2Address, 0xB147a5d25748fda14b463EB04B111027C290f4d3);
     }
+}
 
-    function testBytesToUint() external {
-        bytes memory maxUint = hex"ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff";
-        bytes memory two = hex"02";
-        bytes memory millionEther = hex"d3c21bcecceda1000000";
+contract StdUtilsForkTest is Test {
+    /*//////////////////////////////////////////////////////////////////////////
+                                  GET TOKEN BALANCES
+    //////////////////////////////////////////////////////////////////////////*/
 
-        assertEq(bytesToUint(maxUint), type(uint256).max);
-        assertEq(bytesToUint(two), 2);
-        assertEq(bytesToUint(millionEther), 1_000_000 ether);
+    address internal SHIB = 0x95aD61b0a150d79219dCF64E1E6Cc01f0B64C4cE;
+    address internal SHIB_HOLDER_0 = 0x855F5981e831D83e6A4b4EBFCAdAa68D92333170;
+    address internal SHIB_HOLDER_1 = 0x8F509A90c2e47779cA408Fe00d7A72e359229AdA;
+    address internal SHIB_HOLDER_2 = 0x0e3bbc0D04fF62211F71f3e4C45d82ad76224385;
+
+    address internal USDC = 0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48;
+    address internal USDC_HOLDER_0 = 0xDa9CE944a37d218c3302F6B82a094844C6ECEb17;
+    address internal USDC_HOLDER_1 = 0x3e67F4721E6d1c41a015f645eFa37BEd854fcf52;
+
+    function setUp() public {
+        // All tests of the `getTokenBalances` method are fork tests using live contracts.
+        vm.createSelectFork({urlOrAlias: "mainnet", blockNumber: 16_428_900});
     }
 
-    function testCannotConvertGT32Bytes() external {
-        bytes memory thirty3Bytes = hex"ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff";
-        vm.expectRevert("StdUtils bytesToUint(bytes): Bytes length exceeds 32.");
-        bytesToUint(thirty3Bytes);
+    function testCannotGetTokenBalances_NonTokenContract() external {
+        // We deploy a mock version so we can properly test the revert.
+        StdUtilsMock stdUtils = new StdUtilsMock();
+
+        // The UniswapV2Factory contract has neither a `balanceOf` function nor a fallback function,
+        // so the `balanceOf` call should revert.
+        address token = address(0x5C69bEe701ef814a2B6a3EDD4B1652CB9cc5aA6f);
+        address[] memory addresses = new address[](1);
+        addresses[0] = USDC_HOLDER_0;
+
+        vm.expectRevert("Multicall3: call failed");
+        stdUtils.getTokenBalances_(token, addresses);
+    }
+
+    function testCannotGetTokenBalances_EOA() external {
+        address eoa = vm.addr({privateKey: 1});
+        address[] memory addresses = new address[](1);
+        addresses[0] = USDC_HOLDER_0;
+        vm.expectRevert("StdUtils getTokenBalances(address,address[]): Token address is not a contract.");
+        getTokenBalances(eoa, addresses);
+    }
+
+    function testGetTokenBalances_Empty() external {
+        address[] memory addresses = new address[](0);
+        uint256[] memory balances = getTokenBalances(USDC, addresses);
+        assertEq(balances.length, 0);
+    }
+
+    function testGetTokenBalances_USDC() external {
+        address[] memory addresses = new address[](2);
+        addresses[0] = USDC_HOLDER_0;
+        addresses[1] = USDC_HOLDER_1;
+        uint256[] memory balances = getTokenBalances(USDC, addresses);
+        assertEq(balances[0], 159_000_000_000_000);
+        assertEq(balances[1], 131_350_000_000_000);
+    }
+
+    function testGetTokenBalances_SHIB() external {
+        address[] memory addresses = new address[](3);
+        addresses[0] = SHIB_HOLDER_0;
+        addresses[1] = SHIB_HOLDER_1;
+        addresses[2] = SHIB_HOLDER_2;
+        uint256[] memory balances = getTokenBalances(SHIB, addresses);
+        assertEq(balances[0], 3_323_256_285_484.42e18);
+        assertEq(balances[1], 1_271_702_771_149.99999928e18);
+        assertEq(balances[2], 606_357_106_247e18);
     }
 }
