feat: public data tree overrides for simulation (plumbing)

Thread publicDataOverrides through SimulateInteractionOptions,
SimulateOptions, SimulateOptionsSchema, SimulateViaEntrypointOptions,
PXE.simulateTx, and wallet-sdk's simulateViaNode path so overrides
set in ContractFunctionInteraction.simulate flow down to the node.

Re-export PublicDataTreeOverride and PublicDataTreeOverrideSchema from
@aztec/aztec.js. Add two e2e tests in e2e_avm_simulator that verify:
- overriding an unwritten slot returns the override value during simulation
  while real storage remains untouched
- overriding a previously-written slot returns the override during simulation
  while the committed value is preserved

Author: dbanks12

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

@@ -78,4 +78,6 @@ export { AccountManager } from '../wallet/account_manager.js';
 
 export { TxSimulationResultWithAppOffset } from '../wallet/tx_simulation_result_with_app_offset.js';
 
+export { type PublicDataTreeOverride, PublicDataTreeOverrideSchema } from '@aztec/stdlib/interfaces/client';
+
 export { type DeployAccountOptions, DeployAccountMethod } from '../wallet/deploy_account_method.js';

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

@@ -3,6 +3,7 @@ import type { FieldsOf } from '@aztec/foundation/types';
 import type { AuthWitness } from '@aztec/stdlib/auth-witness';
 import { AztecAddress } from '@aztec/stdlib/aztec-address';
 import type { GasSettings, ManaUsageEstimate } from '@aztec/stdlib/gas';
+import type { PublicDataTreeOverride } from '@aztec/stdlib/interfaces/client';
 import {
   type Capsule,
   OFFCHAIN_MESSAGE_IDENTIFIER,
@@ -157,6 +158,8 @@ export type SimulateInteractionOptions = Omit<SendInteractionOptions, 'fee'> & {
   /** Whether to include metadata such as performance statistics (e.g. timing information of the different circuits and oracles) and gas estimation
    * in the simulation result, in addition to the return value and offchain effects */
   includeMetadata?: boolean;
+  /** Public state overrides injected into the ephemeral world-state fork before simulation. */
+  publicDataOverrides?: PublicDataTreeOverride[];
 };
 
 /**

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

@@ -13,6 +13,7 @@ import { AuthWitness } from '@aztec/stdlib/auth-witness';
 import type { AztecAddress } from '@aztec/stdlib/aztec-address';
 import { type ContractInstanceWithAddress, ContractInstanceWithAddressSchema } from '@aztec/stdlib/contract';
 import { Gas, ManaUsageEstimate } from '@aztec/stdlib/gas';
+import { type PublicDataTreeOverride, PublicDataTreeOverrideSchema } from '@aztec/stdlib/interfaces/client';
 import { LogId } from '@aztec/stdlib/logs';
 import { AbiDecodedSchema, type ApiSchemaFor, optional, schemas, zodFor } from '@aztec/stdlib/schemas';
 import type { ExecutionPayload, InTx } from '@aztec/stdlib/tx';
@@ -335,6 +336,7 @@ export const SimulateOptionsSchema = z.object({
   skipFeeEnforcement: optional(z.boolean()),
   includeMetadata: optional(z.boolean()),
   additionalScopes: optional(z.array(schemas.AztecAddress)),
+  publicDataOverrides: optional(z.array(PublicDataTreeOverrideSchema)),
 });
 
 export const ProfileOptionsSchema = SimulateOptionsSchema.extend({

yarn-project/end-to-end/src/e2e_avm_simulator.test.ts

@@ -1,7 +1,9 @@
 import { AztecAddress } from '@aztec/aztec.js/addresses';
 import { BatchCall, type ContractInstanceWithAddress } from '@aztec/aztec.js/contracts';
 import { Fr } from '@aztec/aztec.js/fields';
+import type { AztecNode } from '@aztec/aztec.js/node';
 import { TxExecutionResult } from '@aztec/aztec.js/tx';
+import type { PublicDataTreeOverride } from '@aztec/aztec.js/wallet';
 import type { Wallet } from '@aztec/aztec.js/wallet';
 import { AvmInitializerTestContract } from '@aztec/noir-test-contracts.js/AvmInitializerTest';
 import { AvmTestContract } from '@aztec/noir-test-contracts.js/AvmTest';
@@ -16,13 +18,15 @@ describe('e2e_avm_simulator', () => {
   jest.setTimeout(TIMEOUT);
 
   let wallet: Wallet;
+  let aztecNode: AztecNode;
   let defaultAccountAddress: AztecAddress;
   let teardown: () => Promise<void>;
 
   beforeAll(async () => {
     ({
       teardown,
       wallet,
+      aztecNode,
       accounts: [defaultAccountAddress],
     } = await setup(1));
     await ensureAccountContractsPublished(wallet, [defaultAccountAddress]);
@@ -249,6 +253,51 @@ describe('e2e_avm_simulator', () => {
     });
   });
 
+  describe('publicDataOverrides', () => {
+    // AvmTestContract: `single` is the first storage variable and lives at raw slot 1.
+    const SINGLE_SLOT = new Fr(1n);
+    let avmContract: AvmTestContract;
+
+    beforeEach(async () => {
+      ({ contract: avmContract } = await AvmTestContract.deploy(wallet).send({ from: defaultAccountAddress }));
+    });
+
+    it('simulated read of an unwritten slot returns the override; real storage is untouched', async () => {
+      const overrideValue = new Fr(0xdeadbeefn);
+      const publicDataOverrides: PublicDataTreeOverride[] = [
+        { contract: avmContract.address, slot: SINGLE_SLOT, value: overrideValue },
+      ];
+
+      const simResult = await avmContract.methods
+        .read_storage_single()
+        .simulate({ from: defaultAccountAddress, publicDataOverrides });
+      expect(simResult.result).toEqual(overrideValue.toBigInt());
+
+      // Real state is untouched — the slot was never written.
+      const realValue = await aztecNode.getPublicStorageAt('latest', avmContract.address, SINGLE_SLOT);
+      expect(realValue.toBigInt()).toEqual(0n);
+    });
+
+    it('simulated read returns the override when a slot was previously written by a real tx', async () => {
+      const realValue = new Fr(100n);
+      await avmContract.methods.set_storage_single(realValue).send({ from: defaultAccountAddress });
+
+      const overrideValue = new Fr(999n);
+      const publicDataOverrides: PublicDataTreeOverride[] = [
+        { contract: avmContract.address, slot: SINGLE_SLOT, value: overrideValue },
+      ];
+
+      const simResult = await avmContract.methods
+        .read_storage_single()
+        .simulate({ from: defaultAccountAddress, publicDataOverrides });
+      expect(simResult.result).toEqual(overrideValue.toBigInt());
+
+      // Real storage still holds the original written value.
+      const storedValue = await aztecNode.getPublicStorageAt('latest', avmContract.address, SINGLE_SLOT);
+      expect(storedValue.toBigInt()).toEqual(realValue.toBigInt());
+    });
+  });
+
   describe('AvmInitializerTestContract', () => {
     let avmContract: AvmInitializerTestContract;
 

yarn-project/pxe/src/pxe.ts

@@ -27,7 +27,7 @@ import {
   getContractClassFromArtifact,
 } from '@aztec/stdlib/contract';
 import { SimulationError } from '@aztec/stdlib/errors';
-import type { AztecNode, PrivateKernelProver } from '@aztec/stdlib/interfaces/client';
+import type { AztecNode, PrivateKernelProver, PublicDataTreeOverride } from '@aztec/stdlib/interfaces/client';
 import type {
   PrivateExecutionStep,
   PrivateKernelExecutionProofOutput,
@@ -125,6 +125,8 @@ export type SimulateTxOpts = {
   scopes: AztecAddress[];
   /** Sender address used to derive discovery tags for private messages (notes, events, logs) this tx emits. */
   senderForTags?: AztecAddress;
+  /** Public state overrides injected into the ephemeral world-state fork before simulation. */
+  publicDataOverrides?: PublicDataTreeOverride[];
 };
 
 /** Options for PXE.executeUtility. */
@@ -467,11 +469,11 @@ export class PXE {
    * It can also be used for estimating gas in the future.
    * @param tx - The transaction to be simulated.
    */
-  async #simulatePublicCalls(tx: Tx, skipFeeEnforcement: boolean) {
+  async #simulatePublicCalls(tx: Tx, skipFeeEnforcement: boolean, publicDataOverrides?: PublicDataTreeOverride[]) {
     // Simulating public calls can throw if the TX fails in a phase that doesn't allow reverts (setup)
     // Or return as reverted if it fails in a phase that allows reverts (app logic, teardown)
     try {
-      const result = await this.node.simulatePublicCalls(tx, skipFeeEnforcement);
+      const result = await this.node.simulatePublicCalls(tx, skipFeeEnforcement, publicDataOverrides);
       if (result.revertReason) {
         throw result.revertReason;
       }
@@ -956,6 +958,7 @@ export class PXE {
       overrides,
       scopes,
       senderForTags,
+      publicDataOverrides,
     }: SimulateTxOpts,
   ): Promise<TxSimulationResult> {
     // We disable concurrent simulations since those might execute oracles which read and write to the PXE stores (e.g.
@@ -1037,7 +1040,7 @@ export class PXE {
         let publicOutput: PublicSimulationOutput | undefined;
         if (simulatePublic && publicInputs.forPublic) {
           const publicSimulationTimer = new Timer();
-          publicOutput = await this.#simulatePublicCalls(simulatedTx, skipFeeEnforcement);
+          publicOutput = await this.#simulatePublicCalls(simulatedTx, skipFeeEnforcement, publicDataOverrides);
           publicSimulationTime = publicSimulationTimer.ms();
           if (publicOutput?.debugLogs?.length) {
             await displayDebugLogs(publicOutput.debugLogs, addr => this.contractStore.getDebugContractName(addr));

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

@@ -85,7 +85,7 @@ export type FeeOptions = {
 /** Options for `simulateViaEntrypoint`. */
 export type SimulateViaEntrypointOptions = Pick<
   SimulateOptions,
-  'from' | 'additionalScopes' | 'skipTxValidation' | 'skipFeeEnforcement' | 'sendMessagesAs'
+  'from' | 'additionalScopes' | 'skipTxValidation' | 'skipFeeEnforcement' | 'sendMessagesAs' | 'publicDataOverrides'
 > & {
   /** Fee options for the entrypoint */
   feeOptions: FeeOptions;
@@ -365,6 +365,7 @@ export abstract class BaseWallet implements Wallet {
       skipFeeEnforcement: opts.skipFeeEnforcement,
       scopes: this.scopesFrom(opts.from, opts.additionalScopes),
       senderForTags: this.senderForTagsFrom(opts.from, opts.sendMessagesAs),
+      publicDataOverrides: opts.publicDataOverrides,
     });
     const appCallOffset = await this.computeAppCallOffset(opts.from, opts.feeOptions);
     return TxSimulationResultWithAppOffset.fromResultAndOffset(result, appCallOffset);
@@ -428,6 +429,7 @@ export abstract class BaseWallet implements Wallet {
             blockHeader,
             opts.skipFeeEnforcement ?? true,
             this.getContractName.bind(this),
+            opts.publicDataOverrides,
           )
         : Promise.resolve([]),
       remainingCalls.length > 0
@@ -438,6 +440,7 @@ export abstract class BaseWallet implements Wallet {
             skipTxValidation: opts.skipTxValidation,
             skipFeeEnforcement: opts.skipFeeEnforcement ?? true,
             sendMessagesAs: opts.sendMessagesAs,
+            publicDataOverrides: opts.publicDataOverrides,
           })
         : Promise.resolve(null),
     ]);

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

@@ -11,6 +11,7 @@ import { generateSimulatedProvingResult } from '@aztec/pxe/simulator';
 import { type FunctionCall, FunctionSelector } from '@aztec/stdlib/abi';
 import type { AztecAddress } from '@aztec/stdlib/aztec-address';
 import type { GasSettings } from '@aztec/stdlib/gas';
+import type { PublicDataTreeOverride } from '@aztec/stdlib/interfaces/client';
 import {
   ClaimedLengthArray,
   CountedPublicCallRequest,
@@ -65,6 +66,8 @@ export function extractOptimizablePublicStaticCalls(payload: ExecutionPayload):
  * @param gasSettings - Gas settings for the transaction.
  * @param blockHeader - Block header to use as anchor.
  * @param skipFeeEnforcement - Whether to skip fee enforcement during simulation.
+ * @param getContractName - Resolver for contract names (used for debug log display).
+ * @param publicDataOverrides - Optional public state overrides injected before simulation.
  * @returns TxSimulationResult with public return values.
  */
 async function simulateBatchViaNode(
@@ -76,6 +79,7 @@ async function simulateBatchViaNode(
   blockHeader: BlockHeader,
   skipFeeEnforcement: boolean,
   getContractName: ContractNameResolver,
+  publicDataOverrides?: PublicDataTreeOverride[],
 ): Promise<TxSimulationResult> {
   const txContext = new TxContext(chainInfo.chainId, chainInfo.version, gasSettings);
 
@@ -143,7 +147,7 @@ async function simulateBatchViaNode(
     publicFunctionCalldata: publicFunctionCalldata,
   });
 
-  const publicOutput = await node.simulatePublicCalls(tx, skipFeeEnforcement);
+  const publicOutput = await node.simulatePublicCalls(tx, skipFeeEnforcement, publicDataOverrides);
 
   if (publicOutput.revertReason) {
     throw publicOutput.revertReason;
@@ -166,6 +170,8 @@ async function simulateBatchViaNode(
  * @param gasSettings - Gas settings for the transaction.
  * @param blockHeader - Block header to use as anchor.
  * @param skipFeeEnforcement - Whether to skip fee enforcement during simulation.
+ * @param getContractName - Resolver for contract names (used for debug log display).
+ * @param publicDataOverrides - Optional public state overrides injected before simulation.
  * @returns Array of TxSimulationResult, one per batch.
  */
 export async function simulateViaNode(
@@ -177,6 +183,7 @@ export async function simulateViaNode(
   blockHeader: BlockHeader,
   skipFeeEnforcement: boolean = true,
   getContractName: ContractNameResolver,
+  publicDataOverrides?: PublicDataTreeOverride[],
 ): Promise<TxSimulationResult[]> {
   const batches: FunctionCall[][] = [];
 
@@ -196,6 +203,7 @@ export async function simulateViaNode(
       blockHeader,
       skipFeeEnforcement,
       getContractName,
+      publicDataOverrides,
     );
     results.push(result);
   }