refactor: encapsulate userCallOffset logic in TxSimulationResult

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

@@ -24,6 +24,8 @@ function mockTxSimResult(overrides: { anchorBlockTimestamp?: bigint; offchainEff
       },
     },
   });
+  // Batch calls always go through an account contract (not DefaultEntrypoint), with no FPC by default.
+  Object.defineProperty(txSimResult, 'userCallOffset', { value: 0, writable: true });
   return txSimResult;
 }
 
@@ -134,9 +136,7 @@ describe('BatchCall', () => {
       const publicReturnValues = [Fr.random()];
 
       const txSimResult = mockTxSimResult();
-      txSimResult.getPrivateReturnValues.mockReturnValue({
-        nested: [{ values: privateReturnValues }],
-      } as any);
+      txSimResult.getUserPrivateReturnValues.mockReturnValue({ values: privateReturnValues } as any);
       txSimResult.getPublicReturnValues.mockReturnValue([{ values: publicReturnValues }] as any);
 
       // Mock wallet.batch to return both utility results and simulateTx result
@@ -298,9 +298,9 @@ describe('BatchCall', () => {
           { data: txRawEffectData, contractAddress: emitterContract },
         ],
       });
-      txSimResult.getPrivateReturnValues.mockReturnValue({
-        nested: [{ values: [Fr.random()] }, { values: [Fr.random()] }],
-      } as any);
+      txSimResult.getUserPrivateReturnValues.mockImplementation((callIndex: number) => ({
+        values: [Fr.random()],
+      }));
 
       wallet.batch.mockResolvedValue([
         { name: 'executeUtility', result: utilityResult },
@@ -342,9 +342,7 @@ describe('BatchCall', () => {
       const publicReturnValues = [Fr.random()];
 
       const txSimResult = mockTxSimResult();
-      txSimResult.getPrivateReturnValues.mockReturnValue({
-        nested: [{ values: privateReturnValues }],
-      } as any);
+      txSimResult.getUserPrivateReturnValues.mockReturnValue({ values: privateReturnValues } as any);
       txSimResult.getPublicReturnValues.mockReturnValue([{ values: publicReturnValues }] as any);
 
       wallet.batch.mockResolvedValue([{ name: 'simulateTx', result: txSimResult }] as any);

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

@@ -129,12 +129,10 @@ export class BatchCall extends BaseContractInteraction {
         simulatedTx = txResultWrapper.result as TxSimulationResult;
         indexedExecutionPayloads.forEach(([request, callIndex, resultIndex]) => {
           const call = request.calls[0];
-          // As account entrypoints are private, for private functions we retrieve the return values from the first nested call
-          // since we're interested in the first set of values AFTER the account entrypoint
-          // For public functions we retrieve the first values directly from the public output.
+          // For public functions we retrieve the values directly from the public output.
           const rawReturnValues =
             call.type == FunctionType.PRIVATE
-              ? simulatedTx!.getPrivateReturnValues()?.nested?.[resultIndex].values
+              ? simulatedTx!.getUserPrivateReturnValues(resultIndex)?.values
               : simulatedTx!.getPublicReturnValues()?.[resultIndex].values;
 
           results[callIndex] = {

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

@@ -233,6 +233,8 @@ describe('Contract Class', () => {
 
     const txSimResult = mock<TxSimulationResult>();
     txSimResult.getPrivateReturnValues.mockReturnValue({ nested: [{ values: [] }] } as any);
+    // Called via account contract with no FPC → user fn is at nested[0].
+    Object.defineProperty(txSimResult, 'userCallOffset', { value: 0, writable: true });
     Object.defineProperty(txSimResult, 'offchainEffects', {
       value: [
         {

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

@@ -154,14 +154,7 @@ export class ContractFunctionInteraction extends BaseContractInteraction {
 
     let rawReturnValues;
     if (this.functionDao.functionType == FunctionType.PRIVATE) {
-      if (simulatedTx.getPrivateReturnValues().nested.length > 0) {
-        // The function invoked is private and it was called via an account contract
-        // TODO(#10631): There is a bug here: this branch might be triggered when there is no-account contract as well
-        rawReturnValues = simulatedTx.getPrivateReturnValues().nested[0].values;
-      } else {
-        // The function invoked is private and it was called directly (without account contract)
-        rawReturnValues = simulatedTx.getPrivateReturnValues().values;
-      }
+      rawReturnValues = simulatedTx.getUserPrivateReturnValues(0)?.values;
     } else {
       // For public functions we retrieve the first values directly from the public output.
       rawReturnValues = simulatedTx.getPublicReturnValues()?.[0]?.values;

yarn-project/end-to-end/src/test-wallet/test_wallet.ts

@@ -269,7 +269,7 @@ export class TestWallet extends BaseWallet {
 
   async proveTx(exec: ExecutionPayload, opts: Omit<SendOptions, 'wait'>): Promise<ProvenTx> {
     const fee = await this.completeFeeOptions(opts.from, exec.feePayer, opts.fee?.gasSettings);
-    const txRequest = await this.createTxExecutionRequestFromPayloadAndFee(exec, opts.from, fee);
+    const { txRequest } = await this.createTxExecutionRequestFromPayloadAndFee(exec, opts.from, fee);
     const txProvingResult = await this.pxe.proveTx(txRequest, this.scopesFrom(opts.from, opts.additionalScopes));
     return new ProvenTx(
       this.aztecNode,

yarn-project/stdlib/src/tx/simulated_tx.ts

@@ -79,6 +79,13 @@ export class PrivateSimulationResult {
 }
 
 export class TxSimulationResult {
+  /**
+   * Index into the private return values nested array where user function calls start. Set by the wallet after
+   * simulation based on how many fee payment calls precede user calls. Undefined when using DefaultEntrypoint
+   * (NO_FROM) — in that case the user fn is the root call.
+   */
+  public userCallOffset?: number;
+
   constructor(
     public privateExecutionResult: PrivateExecutionResult,
     public publicInputs: PrivateKernelTailCircuitPublicInputs,
@@ -147,6 +154,19 @@ export class TxSimulationResult {
     return new PrivateSimulationResult(this.privateExecutionResult, this.publicInputs).getPrivateReturnValues();
   }
 
+  /**
+   * Returns the private return values for the user call at the given index, accounting for any fee payment calls
+   * that precede user calls in the nested array. When userCallOffset is undefined (DefaultEntrypoint / NO_FROM),
+   * the user fn is the root call and callIndex is ignored.
+   */
+  getUserPrivateReturnValues(callIndex: number = 0): NestedProcessReturnValues | undefined {
+    const all = this.getPrivateReturnValues();
+    if (this.userCallOffset === undefined) {
+      return all;
+    }
+    return all?.nested?.[callIndex + this.userCallOffset];
+  }
+
   toSimulatedTx(): Promise<Tx> {
     return new PrivateSimulationResult(this.privateExecutionResult, this.publicInputs).toSimulatedTx();
   }

yarn-project/wallet-sdk/src/base-wallet/base_wallet.ts

@@ -141,16 +141,20 @@ export abstract class BaseWallet implements Wallet {
     executionPayload: ExecutionPayload,
     from: AztecAddress | NoFrom,
     feeOptions: FeeOptions,
-  ): Promise<TxExecutionRequest> {
+  ): Promise<{ txRequest: TxExecutionRequest; userCallOffset: number | undefined }> {
     const feeExecutionPayload = await feeOptions.walletFeePaymentMethod?.getExecutionPayload();
+    const feeCallCount = feeExecutionPayload?.calls.length ?? 0;
     const finalExecutionPayload = feeExecutionPayload
       ? mergeExecutionPayloads([feeExecutionPayload, executionPayload])
       : executionPayload;
     const chainInfo = await this.getChainInfo();
 
     if (from === NO_FROM) {
       const entrypoint = new DefaultEntrypoint();
-      return entrypoint.createTxExecutionRequest(finalExecutionPayload, feeOptions.gasSettings, chainInfo);
+      return {
+        txRequest: await entrypoint.createTxExecutionRequest(finalExecutionPayload, feeOptions.gasSettings, chainInfo),
+        userCallOffset: undefined,
+      };
     } else {
       const fromAccount = await this.getAccountFromAddress(from);
       const executionOptions: DefaultAccountEntrypointOptions = {
@@ -159,12 +163,15 @@ export abstract class BaseWallet implements Wallet {
         // If from is an address, feeOptions include the way the account contract should handle the fee payment
         feePaymentMethodOptions: feeOptions.accountFeePaymentMethodOptions!,
       };
-      return fromAccount.createTxExecutionRequest(
-        finalExecutionPayload,
-        feeOptions.gasSettings,
-        chainInfo,
-        executionOptions,
-      );
+      return {
+        txRequest: await fromAccount.createTxExecutionRequest(
+          finalExecutionPayload,
+          feeOptions.gasSettings,
+          chainInfo,
+          executionOptions,
+        ),
+        userCallOffset: feeCallCount,
+      };
     }
   }
 
@@ -330,17 +337,19 @@ export abstract class BaseWallet implements Wallet {
    * @param opts - Simulation options.
    */
   protected async simulateViaEntrypoint(executionPayload: ExecutionPayload, opts: SimulateViaEntrypointOptions) {
-    const txRequest = await this.createTxExecutionRequestFromPayloadAndFee(
+    const { txRequest, userCallOffset } = await this.createTxExecutionRequestFromPayloadAndFee(
       executionPayload,
       opts.from,
       opts.feeOptions,
     );
-    return this.pxe.simulateTx(txRequest, {
+    const result = await this.pxe.simulateTx(txRequest, {
       simulatePublic: true,
       skipTxValidation: opts.skipTxValidation,
       skipFeeEnforcement: opts.skipFeeEnforcement,
       scopes: opts.scopes,
     });
+    result.userCallOffset = userCallOffset;
+    return result;
   }
 
   /**
@@ -393,12 +402,14 @@ export abstract class BaseWallet implements Wallet {
         : Promise.resolve(null),
     ]);
 
-    return buildMergedSimulationResult(optimizedResults, normalResult);
+    const mergedResult = buildMergedSimulationResult(optimizedResults, normalResult);
+    mergedResult.userCallOffset = normalResult?.userCallOffset;
+    return mergedResult;
   }
 
   async profileTx(executionPayload: ExecutionPayload, opts: ProfileOptions): Promise<TxProfileResult> {
     const feeOptions = await this.completeFeeOptions(opts.from, executionPayload.feePayer, opts.fee?.gasSettings);
-    const txRequest = await this.createTxExecutionRequestFromPayloadAndFee(executionPayload, opts.from, feeOptions);
+    const { txRequest } = await this.createTxExecutionRequestFromPayloadAndFee(executionPayload, opts.from, feeOptions);
     return this.pxe.profileTx(txRequest, {
       profileMode: opts.profileMode,
       skipProofGeneration: opts.skipProofGeneration ?? true,
@@ -411,7 +422,7 @@ export abstract class BaseWallet implements Wallet {
     opts: SendOptions<W>,
   ): Promise<SendReturn<W>> {
     const feeOptions = await this.completeFeeOptions(opts.from, executionPayload.feePayer, opts.fee?.gasSettings);
-    const txRequest = await this.createTxExecutionRequestFromPayloadAndFee(executionPayload, opts.from, feeOptions);
+    const { txRequest } = await this.createTxExecutionRequestFromPayloadAndFee(executionPayload, opts.from, feeOptions);
     const provenTx = await this.pxe.proveTx(txRequest, this.scopesFrom(opts.from, opts.additionalScopes));
     const offchainOutput = extractOffchainOutput(
       provenTx.getOffchainEffects(),