Skip to content

Commit

Permalink
feat: add "IMulticall3" interface and "getTokenBalances" method (#271)
Browse files Browse the repository at this point in the history 
* 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 <[email protected]>
  • Loading branch information
@PaulRBerg @mds1
PaulRBerg and mds1 authored Jan 18, 2023
1 parent 7bb0abb commit 066ff16
Show file tree
Hide file tree
Showing 3 changed files with 234 additions and 18 deletions.
55 changes: 50 additions & 5 deletions src/StdUtils.sol
View file
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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)));
}
Expand Down
73 changes: 73 additions & 0 deletions src/interfaces/IMulticall3.sol
View file
Original file line number Diff line number Diff line change
@@ -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);
}
124 changes: 111 additions & 13 deletions test/StdUtils.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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);
}
}

0 comments on commit 066ff16

Please sign in to comment.