fix(aztec.js): honor fee payment method in BatchCall.simulate

yarn-project/aztec.js/src/contract/batch_call.test.ts

@@ -14,6 +14,7 @@ import {
 
 import { type MockProxy, mock } from 'jest-mock-extended';
 
+import type { FeePaymentMethod } from '../fee/fee_payment_method.js';
 import { TxSimulationResultWithAppOffset } from '../wallet/tx_simulation_result_with_app_offset.js';
 import type { Wallet } from '../wallet/wallet.js';
 import { BatchCall } from './batch_call.js';
@@ -383,6 +384,101 @@ describe('BatchCall', () => {
     });
   });
 
+  describe('simulate with fee payment method', () => {
+    it('offsets return-value indices by the calls the fee payment method prepends', async () => {
+      const appContract = await AztecAddress.random();
+      const feeContract = await AztecAddress.random();
+
+      const appPrivatePayload = createPrivateExecutionPayload('appPrivate', [Fr.random()], appContract, 1);
+      const appPublicPayload = createPublicExecutionPayload('appPublic', [Fr.random()], appContract);
+
+      batchCall = new BatchCall(wallet, [appPrivatePayload, appPublicPayload]);
+
+      // The fee method contributes one private and one public call, prepended ahead of the batch.
+      const feePrivateCall = createPrivateExecutionPayload('feePrivate', [], feeContract).calls[0];
+      const feePublicCall = createPublicExecutionPayload('feePublic', [], feeContract).calls[0];
+      const feePayload = new ExecutionPayload([feePrivateCall, feePublicCall], [], [], [], await AztecAddress.random());
+
+      const paymentMethod = mock<FeePaymentMethod>();
+      paymentMethod.getExecutionPayload.mockResolvedValue(feePayload);
+
+      const appPrivateReturnValues = [Fr.random()];
+      const appPublicReturnValues = [Fr.random()];
+
+      const txSimResult = mockTxSimResult();
+      // Private nested index 0 is the fee call, index 1 is the app call. Returning distinct values lets us detect
+      // an off-by-one that would decode the fee call's return values as the app call's.
+      txSimResult.getPrivateReturnValuesOfAppCall.mockImplementation(
+        (idx?: number) => (idx === 1 ? { values: appPrivateReturnValues } : { values: [Fr.random()] }) as any,
+      );
+      // Public index 0 is the fee call, index 1 is the app call.
+      txSimResult.getPublicReturnValues.mockReturnValue([
+        { values: [Fr.random()] },
+        { values: appPublicReturnValues },
+      ] as any);
+
+      wallet.batch.mockResolvedValue([{ name: 'simulateTx', result: txSimResult }] as any);
+
+      const { result: results } = await batchCall.simulate({
+        from: await AztecAddress.random(),
+        fee: { paymentMethod },
+      });
+
+      expect(txSimResult.getPrivateReturnValuesOfAppCall).toHaveBeenCalledWith(1);
+      expect(results).toHaveLength(2);
+      expect(results[0].result).toEqual(appPrivateReturnValues[0].toBigInt());
+      expect(results[1].result).toEqual(appPublicReturnValues[0].toBigInt());
+    });
+
+    it('merges the fee payment method payload and preserves its fee payer', async () => {
+      const appContract = await AztecAddress.random();
+      const feeContract = await AztecAddress.random();
+      const feePayer = await AztecAddress.random();
+
+      const appPayload = createPrivateExecutionPayload('app', [Fr.random()], appContract, 1);
+      batchCall = new BatchCall(wallet, [appPayload]);
+
+      const feeCall = createPrivateExecutionPayload('payFee', [], feeContract).calls[0];
+      const feePayload = new ExecutionPayload([feeCall], [], [], [], feePayer);
+
+      const paymentMethod = mock<FeePaymentMethod>();
+      paymentMethod.getExecutionPayload.mockResolvedValue(feePayload);
+
+      const txSimResult = mockTxSimResult();
+      txSimResult.getPrivateReturnValuesOfAppCall.mockReturnValue({ values: [Fr.random()] } as any);
+      wallet.batch.mockResolvedValue([{ name: 'simulateTx', result: txSimResult }] as any);
+
+      await batchCall.simulate({ from: await AztecAddress.random(), fee: { paymentMethod } });
+
+      const methods = wallet.batch.mock.calls[0][0] as any[];
+      const { args } = methods.find(m => m.name === 'simulateTx')!;
+      const [executionPayload] = args;
+      expect(executionPayload.calls).toHaveLength(2);
+      expect(executionPayload.calls[0]).toEqual(feeCall);
+      expect(executionPayload.calls[1]).toEqual(appPayload.calls[0]);
+      expect(executionPayload.feePayer).toEqual(feePayer);
+    });
+
+    it('preserves a fee payer carried by a batched execution payload', async () => {
+      const appContract = await AztecAddress.random();
+      const feePayer = await AztecAddress.random();
+
+      const appCall = createPrivateExecutionPayload('app', [Fr.random()], appContract, 1).calls[0];
+      const payloadWithFeePayer = new ExecutionPayload([appCall], [], [], [], feePayer);
+      batchCall = new BatchCall(wallet, [payloadWithFeePayer]);
+
+      const txSimResult = mockTxSimResult();
+      txSimResult.getPrivateReturnValuesOfAppCall.mockReturnValue({ values: [Fr.random()] } as any);
+      wallet.batch.mockResolvedValue([{ name: 'simulateTx', result: txSimResult }] as any);
+
+      await batchCall.simulate({ from: await AztecAddress.random() });
+
+      const methods = wallet.batch.mock.calls[0][0] as any[];
+      const { args } = methods.find(m => m.name === 'simulateTx')!;
+      expect(args[0].feePayer).toEqual(feePayer);
+    });
+  });
+
   describe('request', () => {
     it('should include fee payment method if provided', async () => {
       const contractAddress = await AztecAddress.random();

yarn-project/aztec.js/src/contract/batch_call.ts

@@ -59,6 +59,16 @@ export class BatchCall extends BaseContractInteraction {
    * @returns The results of all the interactions that make up the batch
    */
   public async simulate(options: SimulateInteractionOptions): Promise<SimulationResult> {
+    const feeExecutionPayload = options.fee?.paymentMethod
+      ? await options.fee.paymentMethod.getExecutionPayload()
+      : undefined;
+    // A call-contributing fee payment method prepends its calls in front of the batch (see the merge below). Those
+    // calls shift the return-value indices, and the wallet-side app-call offset does not absorb them: it only counts
+    // the wallet's own fee payment method, not one supplied through options.fee. We offset each app call's result
+    // index by the number of prepended fee calls of the matching type.
+    const feePrivateCallCount = feeExecutionPayload?.calls.filter(c => c.type === FunctionType.PRIVATE).length ?? 0;
+    const feePublicCallCount = feeExecutionPayload?.calls.filter(c => c.type === FunctionType.PUBLIC).length ?? 0;
+
     const { indexedExecutionPayloads, utility } = (await this.getExecutionPayloads()).reduce<{
       /** Keep track of the number of private calls to retrieve the return values */
       privateIndex: 0;
@@ -98,12 +108,15 @@ export class BatchCall extends BaseContractInteraction {
     // Add tx simulation to batch if there are any private/public calls
     if (indexedExecutionPayloads.length > 0) {
       const payloads = indexedExecutionPayloads.map(([request]) => request);
-      const combinedPayload = mergeExecutionPayloads(payloads);
+      const combinedPayload = mergeExecutionPayloads(
+        feeExecutionPayload ? [feeExecutionPayload, ...payloads] : payloads,
+      );
       const executionPayload = new ExecutionPayload(
         combinedPayload.calls,
         combinedPayload.authWitnesses.concat(options.authWitnesses ?? []),
         combinedPayload.capsules.concat(options.capsules ?? []),
         combinedPayload.extraHashedArgs,
+        combinedPayload.feePayer,
       );
 
       batchRequests.push({
@@ -139,11 +152,12 @@ export class BatchCall extends BaseContractInteraction {
         simulatedTx = txResultWrapper.result as TxSimulationResultWithAppOffset;
         indexedExecutionPayloads.forEach(([request, callIndex, resultIndex]) => {
           const call = request.calls[0];
-          // For public functions we retrieve the values directly from the public output.
+          // For public functions we retrieve the values directly from the public output. Both indices are offset by
+          // the fee payment method's own calls, which are prepended ahead of the batch.
           const rawReturnValues =
             call.type == FunctionType.PRIVATE
-              ? simulatedTx!.getPrivateReturnValuesOfAppCall(resultIndex)?.values
-              : simulatedTx!.getPublicReturnValues()?.[resultIndex].values;
+              ? simulatedTx!.getPrivateReturnValuesOfAppCall(feePrivateCallCount + resultIndex)?.values
+              : simulatedTx!.getPublicReturnValues()?.[feePublicCallCount + resultIndex].values;
 
           results[callIndex] = {
             result: rawReturnValues ? decodeFromAbi(call.returnTypes, rawReturnValues) : [],

yarn-project/aztec.js/src/contract/contract_function_interaction.ts

@@ -157,10 +157,16 @@ export class ContractFunctionInteraction extends BaseContractInteraction {
 
     let rawReturnValues;
     if (this.functionDao.functionType == FunctionType.PRIVATE) {
-      rawReturnValues = simulatedTx.getPrivateReturnValuesOfAppCall(0)?.values;
+      // request() prepends the fee payment method's calls (if any) before this interaction's single call, so the app
+      // call is the last call of its type in the payload. Its position among the private return values is the number
+      // of private calls that precede it.
+      const appCallIndex = executionPayload.calls.filter(c => c.type === FunctionType.PRIVATE).length - 1;
+      rawReturnValues = simulatedTx.getPrivateReturnValuesOfAppCall(appCallIndex)?.values;
     } else {
-      // For public functions we retrieve the first values directly from the public output.
-      rawReturnValues = simulatedTx.getPublicReturnValues()?.[0]?.values;
+      // For public functions we retrieve the values directly from the public output, offset by any public fee calls
+      // that request() prepended ahead of the app call.
+      const appCallIndex = executionPayload.calls.filter(c => c.type === FunctionType.PUBLIC).length - 1;
+      rawReturnValues = simulatedTx.getPublicReturnValues()?.[appCallIndex]?.values;
     }
 
     const returnValue = rawReturnValues ? decodeFromAbi(this.functionDao.returnTypes, rawReturnValues) : [];

yarn-project/end-to-end/src/automine/phase_check.parallel.test.ts

@@ -93,19 +93,16 @@ describe('automine/phase_check', () => {
   });
 
   it('should fail when the fee payer is elected after the setup phase has ended', async () => {
-    // BatchCall.simulate ignores the fee payment method, which would make the wallet fall back to
-    // PREEXISTING_FEE_JUICE and end setup in the account entrypoint before any app call runs. Build the payload via
-    // request() instead, which merges the payment method's payload (and thus its fee payer), and simulate it directly.
-    const lateElection = await new BatchCall(wallet, [
-      contract.methods.call_function_that_ends_setup_without_phase_check(),
-      sponsoredFPC.methods.sponsor_unconditionally(),
-    ]).request({
-      fee: {
-        paymentMethod: new DeferredSponsoredFeePaymentMethod(sponsoredFPC.address),
-      },
-    });
-    await expect(wallet.simulateTx(lateElection, { from: defaultAccountAddress })).rejects.toThrow(
-      'fee payer must be elected during the setup phase',
-    );
+    await expect(
+      new BatchCall(wallet, [
+        contract.methods.call_function_that_ends_setup_without_phase_check(),
+        sponsoredFPC.methods.sponsor_unconditionally(),
+      ]).simulate({
+        from: defaultAccountAddress,
+        fee: {
+          paymentMethod: new DeferredSponsoredFeePaymentMethod(sponsoredFPC.address),
+        },
+      }),
+    ).rejects.toThrow('fee payer must be elected during the setup phase');
   });
 });