diff --git a/yarn-project/aztec.js/src/contract/batch_call.test.ts b/yarn-project/aztec.js/src/contract/batch_call.test.ts index 1b87a7d27d7d..a0e60fd3124a 100644 --- a/yarn-project/aztec.js/src/contract/batch_call.test.ts +++ b/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(); + 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(); + 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(); diff --git a/yarn-project/aztec.js/src/contract/batch_call.ts b/yarn-project/aztec.js/src/contract/batch_call.ts index 8a304c4d4dc7..7b7a7241aaaf 100644 --- a/yarn-project/aztec.js/src/contract/batch_call.ts +++ b/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 { + 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) : [], diff --git a/yarn-project/aztec.js/src/contract/contract_function_interaction.ts b/yarn-project/aztec.js/src/contract/contract_function_interaction.ts index 61ab4a6e4dd7..df4c7903bb24 100644 --- a/yarn-project/aztec.js/src/contract/contract_function_interaction.ts +++ b/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) : []; diff --git a/yarn-project/end-to-end/src/automine/phase_check.parallel.test.ts b/yarn-project/end-to-end/src/automine/phase_check.parallel.test.ts index d33eda729a71..93d97a6f2cdc 100644 --- a/yarn-project/end-to-end/src/automine/phase_check.parallel.test.ts +++ b/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'); }); });