feat: batch matching stateOverrides in call

.changeset/curvy-owls-sparkle.md

@@ -0,0 +1,5 @@
+---
+"viem": patch
+---
+
+Added support for batching calls that share matching `stateOverride` values.

src/actions/public/__snapshots__/call.test.ts.snap

@@ -26,6 +26,20 @@ exports[`batch call (deployless: { deployless: false }) > args: blockNumber 1`]
 ]
 `;
 
+exports[`batch call (deployless: { deployless: false }) > args: stateOverride 1`] = `
+[
+  {
+    "data": "0x000000000000000000000000000000000000000000000000000000000000002000000000000000000000000000000000000000000000000000000000000000057761676d69000000000000000000000000000000000000000000000000000000",
+  },
+  {
+    "data": "0x0000000000000000000000000000000000000000000000000000000000000020000000000000000000000000000000000000000000000000000000000000000855534420436f696e000000000000000000000000000000000000000000000000",
+  },
+  {
+    "data": "0x00000000000000000000000000000000000000000000000000000000000000200000000000000000000000000000000000000000000000000000000000000011426f7265644170655961636874436c7562000000000000000000000000000000",
+  },
+]
+`;
+
 exports[`batch call (deployless: { deployless: false }) > chain not configured with multicall 1`] = `
 [
   {
@@ -147,6 +161,20 @@ exports[`batch call (deployless: { deployless: true }) > args: blockNumber 1`] =
 ]
 `;
 
+exports[`batch call (deployless: { deployless: true }) > args: stateOverride 1`] = `
+[
+  {
+    "data": "0x000000000000000000000000000000000000000000000000000000000000002000000000000000000000000000000000000000000000000000000000000000057761676d69000000000000000000000000000000000000000000000000000000",
+  },
+  {
+    "data": "0x0000000000000000000000000000000000000000000000000000000000000020000000000000000000000000000000000000000000000000000000000000000855534420436f696e000000000000000000000000000000000000000000000000",
+  },
+  {
+    "data": "0x00000000000000000000000000000000000000000000000000000000000000200000000000000000000000000000000000000000000000000000000000000011426f7265644170655961636874436c7562000000000000000000000000000000",
+  },
+]
+`;
+
 exports[`batch call (deployless: { deployless: true }) > client config 1`] = `
 [
   {

src/actions/public/call.test.ts

@@ -1028,6 +1028,110 @@ describe.each([{ deployless: true }, { deployless: false }])(
       expect(results).toMatchSnapshot()
     })
 
+    test('args: stateOverride', async () => {
+      const client_2 = anvilMainnet.getClient({
+        batch: {
+          multicall: {
+            deployless,
+          },
+        },
+      })
+
+      const spy = vi.spyOn(client_2, 'request')
+
+      const stateOverride = [
+        {
+          address: wagmiContractAddress,
+          balance: parseEther('100'),
+        },
+      ] as const
+
+      const p = []
+      // Two calls with the same stateOverride should batch together
+      p.push(
+        call(client_2, {
+          data: name4bytes,
+          to: wagmiContractAddress,
+          stateOverride: [...stateOverride],
+        }),
+      )
+      p.push(
+        call(client_2, {
+          data: name4bytes,
+          to: usdcContractConfig.address,
+          stateOverride: [...stateOverride],
+        }),
+      )
+      // A call with a different stateOverride should go to a separate batch
+      p.push(
+        call(client_2, {
+          data: name4bytes,
+          to: baycContractConfig.address,
+          stateOverride: [
+            {
+              address: wagmiContractAddress,
+              balance: parseEther('200'),
+            },
+          ],
+        }),
+      )
+
+      const results = await Promise.all(p)
+
+      // 2 batched eth_call requests: one for the shared stateOverride, one for the different stateOverride
+      expect(spy).toBeCalledTimes(2)
+      expect(results).toMatchSnapshot()
+    })
+
+    test.runIf(deployless === false)(
+      'args: stateOverride for multicall address',
+      async () => {
+        const client_2 = anvilMainnet.getClient({
+          batch: {
+            multicall: {
+              deployless,
+            },
+          },
+        })
+
+        const multicallAddress = client_2.chain!.contracts!.multicall3!.address
+        const spy = vi
+          .spyOn(client_2, 'request')
+          .mockImplementation(async ({ params }) => {
+            const [{ to }] = params as [{ to?: Hex }]
+            if (to?.toLowerCase() === multicallAddress.toLowerCase())
+              throw new Error('multicall should not be used')
+            return '0x'
+          })
+
+        const results = await Promise.all([
+          call(client_2, {
+            data: name4bytes,
+            to: wagmiContractAddress,
+            stateOverride: [
+              {
+                address: multicallAddress,
+                code: '0x',
+              },
+            ],
+          }),
+          call(client_2, {
+            data: name4bytes,
+            to: usdcContractConfig.address,
+            stateOverride: [
+              {
+                address: multicallAddress,
+                code: '0x',
+              },
+            ],
+          }),
+        ])
+
+        expect(spy).toBeCalledTimes(2)
+        expect(results).toEqual([{ data: undefined }, { data: undefined }])
+      },
+    )
+
     test.runIf(deployless === false)(
       'chain not configured with multicall',
       async () => {

src/actions/public/call.ts

@@ -34,7 +34,10 @@ import type { BlockTag } from '../../types/block.js'
 import type { Chain } from '../../types/chain.js'
 import type { EIP1193RequestOptions } from '../../types/eip1193.js'
 import type { Hex } from '../../types/misc.js'
-import type { RpcTransactionRequest } from '../../types/rpc.js'
+import type {
+  RpcStateOverride,
+  RpcTransactionRequest,
+} from '../../types/rpc.js'
 import type { StateOverride } from '../../types/stateOverride.js'
 import type { TransactionRequest } from '../../types/transaction.js'
 import type { ExactPartial, UnionOmit } from '../../types/utils.js'
@@ -50,6 +53,7 @@ import {
   type EncodeFunctionDataErrorType,
   encodeFunctionData,
 } from '../../utils/abi/encodeFunctionData.js'
+import { isAddressEqual } from '../../utils/address/isAddressEqual.js'
 import type { RequestErrorType } from '../../utils/buildRequest.js'
 import {
   type GetChainContractAddressErrorType,
@@ -255,19 +259,29 @@ export async function call<chain extends Chain | undefined>(
       'call',
     ) as TransactionRequest
 
-    if (
-      batch &&
-      shouldPerformMulticall({ request }) &&
-      !rpcStateOverride &&
-      !rpcBlockOverrides
-    ) {
+    if (batch && shouldPerformMulticall({ request }) && !rpcBlockOverrides) {
       try {
-        return await scheduleMulticall(client, {
-          ...request,
+        const { deployless = false } =
+          typeof client.batch?.multicall === 'object'
+            ? client.batch.multicall
+            : {}
+        const multicallAddress = getMulticallAddress(client, {
           blockNumber,
-          blockTag,
-          requestOptions,
-        } as unknown as ScheduleMulticallParameters<chain>)
+          deployless,
+        })
+
+        if (
+          !multicallAddress ||
+          !hasStateOverrideForAddress(rpcStateOverride, multicallAddress)
+        )
+          return await scheduleMulticall(client, {
+            ...request,
+            blockNumber,
+            blockTag,
+            multicallAddress,
+            requestOptions,
+            rpcStateOverride,
+          } as unknown as ScheduleMulticallParameters<chain>)
       } catch (err) {
         if (
           !(err instanceof ClientChainNotConfiguredError) &&
@@ -365,9 +379,10 @@ type ScheduleMulticallParameters<chain extends Chain | undefined> = Pick<
   'blockNumber' | 'blockTag'
 > & {
   data: Hex
-  multicallAddress?: Address | undefined
+  multicallAddress?: Address | null | undefined
   requestOptions?: EIP1193RequestOptions | undefined
   to: Address
+  rpcStateOverride?: RpcStateOverride | undefined
 }
 
 type ScheduleMulticallErrorType =
@@ -392,29 +407,30 @@ async function scheduleMulticall<chain extends Chain | undefined>(
     blockNumber,
     blockTag = client.experimental_blockTag ?? 'latest',
     data,
+    multicallAddress: multicallAddress_,
     requestOptions,
+    rpcStateOverride,
     to,
   } = args
 
-  const multicallAddress = (() => {
-    if (deployless) return null
-    if (args.multicallAddress) return args.multicallAddress
-    if (client.chain) {
-      return getChainContractAddress({
-        blockNumber,
-        chain: client.chain,
-        contract: 'multicall3',
-      })
-    }
-    throw new ClientChainNotConfiguredError()
-  })()
+  const multicallAddress =
+    multicallAddress_ !== undefined
+      ? multicallAddress_
+      : getMulticallAddress(client, {
+          blockNumber,
+          deployless,
+        })
 
   const blockNumberHex =
     typeof blockNumber === 'bigint' ? numberToHex(blockNumber) : undefined
   const block = blockNumberHex || blockTag
 
+  const stateOverrideKey = rpcStateOverride
+    ? `.${JSON.stringify(rpcStateOverride)}`
+    : ''
+
   const { schedule } = createBatchScheduler({
-    id: `${client.uid}.${block}.${getRequestOptionsId(requestOptions)}`,
+    id: `${client.uid}.${block}.${getRequestOptionsId(requestOptions)}${stateOverrideKey}`,
     wait,
     shouldSplitBatch(args) {
       const size = args.reduce((size, { data }) => size + (data.length - 2), 0)
@@ -438,22 +454,22 @@ async function scheduleMulticall<chain extends Chain | undefined>(
         functionName: 'aggregate3',
       })
 
+      const multicallRequest = {
+        ...(multicallAddress === null
+          ? {
+              data: toDeploylessCallViaBytecodeData({
+                code: multicall3Bytecode,
+                data: calldata,
+              }),
+            }
+          : { to: multicallAddress, data: calldata }),
+      }
       const data = await client.request(
         {
           method: 'eth_call',
-          params: [
-            {
-              ...(multicallAddress === null
-                ? {
-                    data: toDeploylessCallViaBytecodeData({
-                      code: multicall3Bytecode,
-                      data: calldata,
-                    }),
-                  }
-                : { to: multicallAddress, data: calldata }),
-            },
-            block,
-          ],
+          params: rpcStateOverride
+            ? [multicallRequest, block, rpcStateOverride]
+            : [multicallRequest, block],
         },
         requestOptions,
       )
@@ -474,6 +490,34 @@ async function scheduleMulticall<chain extends Chain | undefined>(
   return { data: returnData }
 }
 
+function getMulticallAddress(
+  client: Client<Transport>,
+  parameters: {
+    blockNumber?: bigint | undefined
+    deployless?: boolean | undefined
+  },
+): Address | null {
+  const { blockNumber, deployless } = parameters
+  if (deployless) return null
+  if (client.chain)
+    return getChainContractAddress({
+      blockNumber,
+      chain: client.chain,
+      contract: 'multicall3',
+    })
+  throw new ClientChainNotConfiguredError()
+}
+
+function hasStateOverrideForAddress(
+  rpcStateOverride: RpcStateOverride | undefined,
+  address: Address,
+) {
+  if (!rpcStateOverride) return false
+  return Object.keys(rpcStateOverride).some((stateOverrideAddress) =>
+    isAddressEqual(stateOverrideAddress as Address, address),
+  )
+}
+
 type ToDeploylessCallViaBytecodeDataErrorType =
   | EncodeDeployDataErrorType
   | ErrorType