Assets 3365 perf optimise multicall client to calculate address specific data once

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]

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 },

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<string, Hex>();
+  const nativeByAccount = new Map<string, Hex>();
+
+  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,