diff --git a/packages/assets-controller/CHANGELOG.md b/packages/assets-controller/CHANGELOG.md index 55924b951ee..355b120d5c0 100644 --- a/packages/assets-controller/CHANGELOG.md +++ b/packages/assets-controller/CHANGELOG.md @@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed +- `MulticallClient` memoizes `balanceOf` and `getEthBalance` call encodings per account address when building multicall batches, reducing redundant ABI encoding for wallets with many tokens ([#9425](https://github.com/MetaMask/core/pull/9425)) - Bump `@metamask/transaction-controller` from `^68.2.2` to `^68.3.0` ([#9421](https://github.com/MetaMask/core/pull/9421)) ## [10.1.0] diff --git a/packages/assets-controller/src/data-sources/evm-rpc-services/clients/MulticallClient.test.ts b/packages/assets-controller/src/data-sources/evm-rpc-services/clients/MulticallClient.test.ts index 9e532fb3c9c..01528dbdee1 100644 --- a/packages/assets-controller/src/data-sources/evm-rpc-services/clients/MulticallClient.test.ts +++ b/packages/assets-controller/src/data-sources/evm-rpc-services/clients/MulticallClient.test.ts @@ -1,4 +1,5 @@ import { defaultAbiCoder, Interface } from '@ethersproject/abi'; +import * as controllerUtils from '@metamask/controller-utils'; import type { Hex } from '@metamask/utils'; import type { Address, BalanceOfRequest, ChainId, Provider } from '../types'; @@ -212,6 +213,52 @@ describe('MulticallClient', () => { expect(result[1].balance).toBe('2000000000'); }); + it('encodes balance call data once per account address in a batch', async () => { + const encodeSpy = jest.spyOn(controllerUtils, 'encodeFunctionData'); + const countBalanceOf = (): number => + encodeSpy.mock.calls.filter(([, method]) => method === 'balanceOf') + .length; + const countGetEthBalance = (): number => + encodeSpy.mock.calls.filter( + ([, method]) => method === 'getEthBalance', + ).length; + const accountA: Address = + '0x1111111111111111111111111111111111111111' as Address; + const accountB: Address = + '0x2222222222222222222222222222222222222222' as Address; + + const requests: BalanceOfRequest[] = [ + { tokenAddress: TEST_TOKEN_1, accountAddress: accountA }, + { tokenAddress: TEST_TOKEN_2, accountAddress: accountA }, + { tokenAddress: TEST_TOKEN_1, accountAddress: accountB }, + { tokenAddress: ZERO_ADDRESS, accountAddress: accountA }, + ]; + + const mockResponse = buildMockAggregate3Response([ + { success: true, balance: '1000000000' }, + { success: true, balance: '2000000000' }, + { success: true, balance: '3000000000' }, + { success: true, balance: '1000000000000000000' }, + ]); + + mockProvider.call.mockResolvedValue(mockResponse); + + const balanceOfBefore = countBalanceOf(); + const getEthBalanceBefore = countGetEthBalance(); + + await client.batchBalanceOf(MAINNET_CHAIN_ID, requests); + + expect(countBalanceOf() - balanceOfBefore).toBe(2); + expect(countGetEthBalance() - getEthBalanceBefore).toBe(1); + + await client.batchBalanceOf(MAINNET_CHAIN_ID, requests); + + expect(countBalanceOf() - balanceOfBefore).toBe(2); + expect(countGetEthBalance() - getEthBalanceBefore).toBe(1); + + encodeSpy.mockRestore(); + }); + it('should fetch native token balance using getEthBalance', async () => { const requests: BalanceOfRequest[] = [ { tokenAddress: ZERO_ADDRESS, accountAddress: TEST_ACCOUNT }, diff --git a/packages/assets-controller/src/data-sources/evm-rpc-services/clients/MulticallClient.ts b/packages/assets-controller/src/data-sources/evm-rpc-services/clients/MulticallClient.ts index c7d866444f5..cefcf3f42a4 100644 --- a/packages/assets-controller/src/data-sources/evm-rpc-services/clients/MulticallClient.ts +++ b/packages/assets-controller/src/data-sources/evm-rpc-services/clients/MulticallClient.ts @@ -402,6 +402,70 @@ function encodeGetEthBalance(accountAddress: Address): Hex { ]); } +type BalanceCallDataCache = { + getErc20BalanceCallData: (accountAddress: Address) => Hex; + getNativeBalanceCallData: (accountAddress: Address) => Hex; +}; + +function createBalanceCallDataCache(): BalanceCallDataCache { + const erc20ByAccount = new Map(); + const nativeByAccount = new Map(); + + return { + getErc20BalanceCallData(accountAddress: Address): Hex { + const key = accountAddress.toLowerCase(); + const existing = erc20ByAccount.get(key); + if (existing !== undefined) { + return existing; + } + const callData = encodeBalanceOf(accountAddress); + erc20ByAccount.set(key, callData); + return callData; + }, + getNativeBalanceCallData(accountAddress: Address): Hex { + const key = accountAddress.toLowerCase(); + const existing = nativeByAccount.get(key); + if (existing !== undefined) { + return existing; + } + const callData = encodeGetEthBalance(accountAddress); + nativeByAccount.set(key, callData); + return callData; + }, + }; +} + +/** + * Cache balance call encodings per account address. + * ERC-20 `balanceOf` and native `getEthBalance` call data depend only on the + * account being queried, not the token contract (which is the multicall target). + */ +const balanceCallDataCache = createBalanceCallDataCache(); + +/** + * Build Multicall3 aggregate3 calls for a batch of balance requests. + * + * @param batch - Balance requests in the current batch. + * @param multicallAddress - Multicall3 contract address for native balance calls. + * @returns Aggregate3 call descriptors. + */ +function buildAggregate3BalanceCalls( + batch: BalanceOfRequest[], + multicallAddress: Hex, +): { target: Address; allowFailure: boolean; callData: Hex }[] { + return batch.map((req) => { + const isNative = req.tokenAddress === ZERO_ADDRESS; + const target = isNative ? multicallAddress : req.tokenAddress; + return { + target, + allowFailure: true, + callData: isNative + ? balanceCallDataCache.getNativeBalanceCallData(req.accountAddress) + : balanceCallDataCache.getErc20BalanceCallData(req.accountAddress), + }; + }); +} + /** * Encode a Multicall3 aggregate3 call. * @@ -568,23 +632,12 @@ export class MulticallClient { batchSize, initialResult: [], eachBatch: async (workingResult, batch) => { + const calls = buildAggregate3BalanceCalls(batch, multicallAddress); + try { await createServicePolicy({ maxRetries: MULTICALL_MAX_RETRIES, }).execute(async () => { - // Build aggregate3 calls - const calls = batch.map((req) => { - const isNative = req.tokenAddress === ZERO_ADDRESS; - const target = isNative ? multicallAddress : req.tokenAddress; - return { - target, - allowFailure: true, - callData: isNative - ? encodeGetEthBalance(req.accountAddress) - : encodeBalanceOf(req.accountAddress), - }; - }); - // Encode and send aggregate3 call const callData = encodeAggregate3(calls); const result = await provider.call({ @@ -710,7 +763,8 @@ export class MulticallClient { } // ERC-20 token - const callData = encodeBalanceOf(accountAddress); + const callData = + balanceCallDataCache.getErc20BalanceCallData(accountAddress); const result = await provider.call({ to: tokenAddress, data: callData,