fix: passing in user call info from wallet (backport #21937)

yarn-project/aztec.js/src/api/wallet.ts

@@ -76,4 +76,6 @@ export {
 
 export { AccountManager } from '../wallet/account_manager.js';
 
+export { TxSimulationResultWithAppOffset } from '../wallet/tx_simulation_result_with_app_offset.js';
+
 export { type DeployAccountOptions, DeployAccountMethod } from '../wallet/deploy_account_method.js';

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

@@ -3,19 +3,20 @@ import { FunctionCall, FunctionSelector, FunctionType } from '@aztec/stdlib/abi'
 import { AztecAddress } from '@aztec/stdlib/aztec-address';
 import {
   ExecutionPayload,
+  NestedProcessReturnValues,
   OFFCHAIN_MESSAGE_IDENTIFIER,
   type OffchainEffect,
-  TxSimulationResult,
   UtilityExecutionResult,
 } from '@aztec/stdlib/tx';
 
 import { type MockProxy, mock } from 'jest-mock-extended';
 
+import { TxSimulationResultWithAppOffset } from '../wallet/tx_simulation_result_with_app_offset.js';
 import type { Wallet } from '../wallet/wallet.js';
 import { BatchCall } from './batch_call.js';
 
 function mockTxSimResult(overrides: { anchorBlockTimestamp?: bigint; offchainEffects?: OffchainEffect[] } = {}) {
-  const txSimResult = mock<TxSimulationResult>();
+  const txSimResult = mock<TxSimulationResultWithAppOffset>();
   Object.defineProperty(txSimResult, 'offchainEffects', { value: overrides.offchainEffects ?? [] });
   Object.defineProperty(txSimResult, 'publicInputs', {
     value: {
@@ -134,9 +135,7 @@ describe('BatchCall', () => {
       const publicReturnValues = [Fr.random()];
 
       const txSimResult = mockTxSimResult();
-      txSimResult.getPrivateReturnValues.mockReturnValue({
-        nested: [{ values: privateReturnValues }],
-      } as any);
+      txSimResult.getPrivateReturnValuesOfAppCall.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 +297,9 @@ describe('BatchCall', () => {
           { data: txRawEffectData, contractAddress: emitterContract },
         ],
       });
-      txSimResult.getPrivateReturnValues.mockReturnValue({
-        nested: [{ values: [Fr.random()] }, { values: [Fr.random()] }],
-      } as any);
+      txSimResult.getPrivateReturnValuesOfAppCall.mockImplementation(
+        () => new NestedProcessReturnValues([Fr.random()]),
+      );
 
       wallet.batch.mockResolvedValue([
         { name: 'executeUtility', result: utilityResult },
@@ -342,9 +341,7 @@ describe('BatchCall', () => {
       const publicReturnValues = [Fr.random()];
 
       const txSimResult = mockTxSimResult();
-      txSimResult.getPrivateReturnValues.mockReturnValue({
-        nested: [{ values: privateReturnValues }],
-      } as any);
+      txSimResult.getPrivateReturnValuesOfAppCall.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

@@ -1,6 +1,7 @@
 import { type FunctionCall, FunctionType, decodeFromAbi } from '@aztec/stdlib/abi';
-import { ExecutionPayload, TxSimulationResult, UtilityExecutionResult, mergeExecutionPayloads } from '@aztec/stdlib/tx';
+import { ExecutionPayload, UtilityExecutionResult, mergeExecutionPayloads } from '@aztec/stdlib/tx';
 
+import type { TxSimulationResultWithAppOffset } from '../wallet/tx_simulation_result_with_app_offset.js';
 import type { BatchedMethod, Wallet } from '../wallet/wallet.js';
 import { BaseContractInteraction } from './base_contract_interaction.js';
 import {
@@ -120,25 +121,24 @@ export class BatchCall extends BaseContractInteraction {
     }
 
     // Process tx simulation result (it comes last if present)
+    let simulatedTx: TxSimulationResultWithAppOffset | undefined;
     if (indexedExecutionPayloads.length > 0) {
       const txResultWrapper = batchResults[utility.length];
       if (txResultWrapper.name === 'simulateTx') {
-        const simulatedTx = txResultWrapper.result as TxSimulationResult;
+        simulatedTx = txResultWrapper.result as TxSimulationResultWithAppOffset;
         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.getPublicReturnValues()?.[resultIndex].values;
+              ? simulatedTx!.getPrivateReturnValuesOfAppCall(resultIndex)?.values
+              : simulatedTx!.getPublicReturnValues()?.[resultIndex].values;
 
           results[callIndex] = {
             result: rawReturnValues ? decodeFromAbi(call.returnTypes, rawReturnValues) : [],
             ...extractOffchainOutput(
-              simulatedTx.offchainEffects,
-              simulatedTx.publicInputs.constants.anchorBlockHeader.globalVariables.timestamp,
+              simulatedTx!.offchainEffects,
+              simulatedTx!.publicInputs.constants.anchorBlockHeader.globalVariables.timestamp,
             ),
           };
         });

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

@@ -6,12 +6,13 @@ import {
   type ContractInstanceWithAddress,
   getContractClassFromArtifact,
 } from '@aztec/stdlib/contract';
-import type { TxExecutionRequest, TxReceipt, TxSimulationResult, UtilityExecutionResult } from '@aztec/stdlib/tx';
+import type { TxExecutionRequest, TxReceipt, UtilityExecutionResult } from '@aztec/stdlib/tx';
 import { OFFCHAIN_MESSAGE_IDENTIFIER } from '@aztec/stdlib/tx';
 
 import { type MockProxy, mock } from 'jest-mock-extended';
 
 import type { Account } from '../account/account.js';
+import type { TxSimulationResultWithAppOffset } from '../wallet/tx_simulation_result_with_app_offset.js';
 import type { Wallet } from '../wallet/wallet.js';
 import { Contract } from './contract.js';
 
@@ -24,7 +25,10 @@ describe('Contract Class', () => {
 
   const mockTxRequest = { type: 'TxRequest' } as any as TxExecutionRequest;
   const mockTxReceipt = { type: 'TxReceipt' } as any as TxReceipt;
-  const mockTxSimulationResult = { type: 'TxSimulationResult', result: 1n } as any as TxSimulationResult;
+  const mockTxSimulationResultWithAppOffset = {
+    type: 'TxSimulationResultWithAppOffset',
+    result: 1n,
+  } as any as TxSimulationResultWithAppOffset;
   const mockUtilityResultValue = {
     result: [new Fr(42)],
     offchainEffects: [],
@@ -198,7 +202,7 @@ describe('Contract Class', () => {
     } as ContractInstanceWithAddress;
 
     wallet = mock<Wallet>();
-    wallet.simulateTx.mockResolvedValue(mockTxSimulationResult);
+    wallet.simulateTx.mockResolvedValue(mockTxSimulationResultWithAppOffset);
     account.createTxExecutionRequest.mockResolvedValue(mockTxRequest);
     wallet.registerContract.mockResolvedValue(contractInstance);
     wallet.sendTx.mockResolvedValue({ receipt: mockTxReceipt, offchainEffects: [], offchainMessages: [] });
@@ -231,8 +235,8 @@ describe('Contract Class', () => {
     const msgPayload = [Fr.random(), Fr.random()];
     const anchorBlockTimestamp = 9999n;
 
-    const txSimResult = mock<TxSimulationResult>();
-    txSimResult.getPrivateReturnValues.mockReturnValue({ nested: [{ values: [] }] } as any);
+    const txSimResult = mock<TxSimulationResultWithAppOffset>();
+    txSimResult.getPrivateReturnValuesOfAppCall.mockReturnValue({ values: [] } as any);
     Object.defineProperty(txSimResult, 'offchainEffects', {
       value: [
         {

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

@@ -155,14 +155,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.getPrivateReturnValuesOfAppCall(0)?.values;
     } else {
       // For public functions we retrieve the first values directly from the public output.
       rawReturnValues = simulatedTx.getPublicReturnValues()?.[0]?.values;

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

@@ -4,10 +4,11 @@ import { AztecAddress } from '@aztec/stdlib/aztec-address';
 import type { ContractInstanceWithAddress } from '@aztec/stdlib/contract';
 import { Gas } from '@aztec/stdlib/gas';
 import { PublicKeys } from '@aztec/stdlib/keys';
-import { OFFCHAIN_MESSAGE_IDENTIFIER, type OffchainEffect, type TxSimulationResult } from '@aztec/stdlib/tx';
+import { OFFCHAIN_MESSAGE_IDENTIFIER, type OffchainEffect } from '@aztec/stdlib/tx';
 
 import { type MockProxy, mock } from 'jest-mock-extended';
 
+import type { TxSimulationResultWithAppOffset } from '../wallet/tx_simulation_result_with_app_offset.js';
 import type { Wallet } from '../wallet/wallet.js';
 import type { ContractBase } from './contract_base.js';
 import { DeployMethod } from './deploy_method.js';
@@ -70,7 +71,7 @@ describe('DeployMethod', () => {
       },
     ];
 
-    const txSimResult = mock<TxSimulationResult>();
+    const txSimResult = mock<TxSimulationResultWithAppOffset>();
     Object.defineProperty(txSimResult, 'offchainEffects', { value: offchainEffects });
     Object.defineProperty(txSimResult, 'publicInputs', {
       value: {

yarn-project/aztec.js/src/wallet/index.ts

@@ -1,3 +1,4 @@
 export * from './wallet.js';
 export * from './account_manager.js';
 export * from './capabilities.js';
+export * from './tx_simulation_result_with_app_offset.js';

yarn-project/aztec.js/src/wallet/tx_simulation_result_with_app_offset.test.ts

@@ -0,0 +1,76 @@
+import { Fr } from '@aztec/foundation/curves/bn254';
+import { PrivateKernelTailCircuitPublicInputs } from '@aztec/stdlib/kernel';
+import { NestedProcessReturnValues, PrivateExecutionResult } from '@aztec/stdlib/tx';
+
+import { TxSimulationResultWithAppOffset } from './tx_simulation_result_with_app_offset.js';
+
+/**
+ * Builds a NestedProcessReturnValues tree that mimics what the private kernel produces.
+ *
+ * The structure reflects the flattened call order:
+ *   index 0 = entrypoint (root), tag Fr(0)
+ *   index 1..feeCallCount = fee calls, tags Fr(100), Fr(101), ...
+ *   index feeCallCount+1.. = app calls, tags Fr(200), Fr(201), ...
+ */
+function buildReturnValues(feeCallCount: number, appCallCount: number): NestedProcessReturnValues {
+  const makeLeaf = (tag: number) => new NestedProcessReturnValues([new Fr(tag)]);
+  const nested = [
+    ...Array.from({ length: feeCallCount }, (_, i) => makeLeaf(100 + i)),
+    ...Array.from({ length: appCallCount }, (_, i) => makeLeaf(200 + i)),
+  ];
+  return new NestedProcessReturnValues([new Fr(0)], nested);
+}
+
+/** Subclass that injects a controlled return values tree, using real helpers for all other fields. */
+class TestResult extends TxSimulationResultWithAppOffset {
+  private constructor(
+    privateExecutionResult: PrivateExecutionResult,
+    appCallOffset: number | undefined,
+    private returnValues: NestedProcessReturnValues,
+  ) {
+    super(privateExecutionResult, PrivateKernelTailCircuitPublicInputs.empty(), undefined, undefined, appCallOffset);
+  }
+
+  static async create(appCallOffset: number | undefined, returnValues: NestedProcessReturnValues) {
+    const executionResult = await PrivateExecutionResult.random();
+    return new TestResult(executionResult, appCallOffset, returnValues);
+  }
+
+  override getPrivateReturnValues() {
+    return this.returnValues;
+  }
+}
+
+describe('TxSimulationResultWithAppOffset.getPrivateReturnValuesOfAppCall', () => {
+  describe('with appCallOffset defined', () => {
+    it('offset=0 returns root values for appCallIndex=0', async () => {
+      const result = await TestResult.create(0, buildReturnValues(0, 1));
+      expect(result.getPrivateReturnValuesOfAppCall(0)?.values?.[0]).toEqual(new Fr(0));
+    });
+
+    it('offset=1 (entrypoint, no fee calls) returns correct app calls', async () => {
+      const result = await TestResult.create(1, buildReturnValues(0, 2));
+      expect(result.getPrivateReturnValuesOfAppCall(0)?.values?.[0]).toEqual(new Fr(200));
+      expect(result.getPrivateReturnValuesOfAppCall(1)?.values?.[0]).toEqual(new Fr(201));
+    });
+
+    it('offset=2 (entrypoint + one fee call) skips fee call and returns app calls', async () => {
+      const result = await TestResult.create(2, buildReturnValues(1, 2));
+      expect(result.getPrivateReturnValuesOfAppCall(0)?.values?.[0]).toEqual(new Fr(200));
+      expect(result.getPrivateReturnValuesOfAppCall(1)?.values?.[0]).toEqual(new Fr(201));
+    });
+  });
+
+  describe('with appCallOffset undefined (heuristic fallback)', () => {
+    it('returns nested[appCallIndex] when nested calls exist (app wrapped in entrypoint)', async () => {
+      const result = await TestResult.create(undefined, buildReturnValues(0, 2));
+      expect(result.getPrivateReturnValuesOfAppCall(0)?.values?.[0]).toEqual(new Fr(200));
+      expect(result.getPrivateReturnValuesOfAppCall(1)?.values?.[0]).toEqual(new Fr(201));
+    });
+
+    it('returns root values when there are no nested calls (direct/NO_FROM call)', async () => {
+      const result = await TestResult.create(undefined, new NestedProcessReturnValues([new Fr(42)]));
+      expect(result.getPrivateReturnValuesOfAppCall(0)?.values?.[0]).toEqual(new Fr(42));
+    });
+  });
+});

yarn-project/aztec.js/src/wallet/tx_simulation_result_with_app_offset.ts

@@ -0,0 +1,101 @@
+import type { ZodFor } from '@aztec/foundation/schemas';
+import { optional } from '@aztec/foundation/schemas';
+import { PrivateKernelTailCircuitPublicInputs } from '@aztec/stdlib/kernel';
+import {
+  type NestedProcessReturnValues,
+  PrivateExecutionResult,
+  PublicSimulationOutput,
+  type SimulationStats,
+  SimulationStatsSchema,
+  TxSimulationResult,
+} from '@aztec/stdlib/tx';
+
+import { z } from 'zod';
+
+/**
+ * Extends TxSimulationResult with the app call offset, which tracks where the app's calls begin in the flattened
+ * array of calls. Tracking of app call offset is a wallet-level concern: the wallet may wrap the app payload in an
+ * entrypoint or may prepend calls (this is typically done for fee payments).
+ */
+export class TxSimulationResultWithAppOffset extends TxSimulationResult {
+  constructor(
+    privateExecutionResult: PrivateExecutionResult,
+    publicInputs: PrivateKernelTailCircuitPublicInputs,
+    publicOutput?: PublicSimulationOutput,
+    stats?: SimulationStats,
+    /**
+     * Index of the app's first call in a flattened array of calls.
+     *   0 = app call is the root execution itself (DefaultEntrypoint / NO_FROM).
+     *   1..N = wallet prepended calls before the app call.
+     *   undefined = wallet did not send the field; use heuristic fallback.
+     */
+    public readonly appCallOffset: number | undefined = undefined,
+  ) {
+    super(privateExecutionResult, publicInputs, publicOutput, stats);
+  }
+
+  /**
+   * Returns the private return values that correspond to the provided app call.
+   * @param appCallIndex - Index of the app call within the app calls.
+   */
+  getPrivateReturnValuesOfAppCall(appCallIndex: number = 0): NestedProcessReturnValues | undefined {
+    const all = this.getPrivateReturnValues();
+
+    if (this.appCallOffset === undefined) {
+      // appCallOffset was not sent. Apply the pre-offset heuristic for backwards compatibility.
+      // If there are nested calls, the app was wrapped in an account contract entrypoint and its return values
+      // live in nested[appCallIndex]. Otherwise the app call is the root execution.
+      if ((all?.nested?.length ?? 0) > 0) {
+        return all?.nested?.[appCallIndex];
+      }
+      if (appCallIndex !== 0) {
+        throw new Error('App call index cannot be defined when there is single root app call');
+      }
+      return all;
+    }
+
+    // appCallOffset is defined on the flattened array of calls where entrypoint occupies index 0 and the app's
+    // first call occupies index appCallOffset. appCallIndex 0 with appCallOffset 0 means the root itself.
+    if (this.appCallOffset === 0 && appCallIndex === 0) {
+      return all;
+    }
+    // For all other cases, index into nested: subtract 1 because nested[0] corresponds to flat index 1 (first nested).
+    return all?.nested?.[appCallIndex + this.appCallOffset - 1];
+  }
+
+  /**
+   * Creates a TxSimulationResultWithAppOffset from an existing TxSimulationResult, attaching the app call offset
+   * computed by the wallet (i.e. how many calls precede the first app call in the flattened execution tree).
+   * @param result - The simulation result to wrap.
+   * @param appCallOffset - The index of the app's first call in the flattened execution tree.
+   */
+  static fromResultAndOffset(result: TxSimulationResult, appCallOffset: number): TxSimulationResultWithAppOffset {
+    return new TxSimulationResultWithAppOffset(
+      result.privateExecutionResult,
+      result.publicInputs,
+      result.publicOutput,
+      result.stats,
+      appCallOffset,
+    );
+  }
+
+  static override async random() {
+    const base = await TxSimulationResult.random();
+    return TxSimulationResultWithAppOffset.fromResultAndOffset(base, 0);
+  }
+
+  static override get schema(): ZodFor<TxSimulationResultWithAppOffset> {
+    return z
+      .object({
+        privateExecutionResult: PrivateExecutionResult.schema,
+        publicInputs: PrivateKernelTailCircuitPublicInputs.schema,
+        publicOutput: PublicSimulationOutput.schema.optional(),
+        stats: optional(SimulationStatsSchema),
+        appCallOffset: optional(z.number()),
+      })
+      .transform(
+        ({ privateExecutionResult, publicInputs, publicOutput, stats, appCallOffset }) =>
+          new TxSimulationResultWithAppOffset(privateExecutionResult, publicInputs, publicOutput, stats, appCallOffset),
+      );
+  }
+}