feat: public data tree overrides for simulation

docs/docs-developers/docs/aztec-js/how_to_test.md

@@ -67,6 +67,22 @@ Test that invalid operations revert as expected:
 
 Use `.simulate()` to test reverts without spending gas. The simulation will throw if the transaction would fail onchain.
 
+## Simulating with overrides
+
+`.simulate()` accepts an `overrides` option that injects values into the simulator's (ephemeral) world-state fork and contract DB before the call runs. The override is scoped to that single simulation and thrown away afterwards.
+
+Override a public-storage slot:
+
+```typescript
+const result = await contract.methods.read_balance(account).simulate({
+  overrides: {
+    publicStorage: [{ contract: contract.address, slot: BALANCE_SLOT, value: new Fr(1_000_000n) }],
+  },
+});
+```
+
+Use this to set up state preconditions, reproduce production bugs against pinned storage, or exercise rare value branches without orchestrating the contract calls that produce them.
+
 ## Further reading
 
 - [How to read contract data](./how_to_read_data.md)

docs/docs-developers/docs/resources/migration_notes.md

@@ -164,6 +164,27 @@ If you set `Noir: Nargo Path` in the VS Code Noir extension to `$HOME/.aztec/cur
 + overrides = { contracts: { [addr.toString()]: { instance: { ...instance, currentContractClassId: stubClassId } } } };
 ```
 
+### [Aztec.js] `simulate` accepts `overrides` for testing "what if storage value was X?"
+
+`Contract.methods.foo(...).simulate(...)` now accepts an `overrides` option that injects values into the simulator's (ephemeral) world-state fork and contract DB before the call runs. The supported field is `publicStorage`, which writes a `(contract, slot, value)` into the public-data tree as if a previous tx had set it. Overrides are thrown away after simulation completes.
+
+```typescript
+const result = await contract.methods.read_balance(account).simulate({
+  overrides: {
+    publicStorage: [{ contract: contract.address, slot: BALANCE_SLOT, value: new Fr(1_000_000n) }],
+  },
+});
+```
+
+The same option flows through `wallet.simulateTx` and eventually to `simulatePublicCalls` RPC on `AztecNode`.
+
+Direct callers of the `SimulationOverrides` constructor must switch from a positional `contracts` argument to an options bag:
+
+```diff
+- new SimulationOverrides(contracts);
++ new SimulationOverrides({ contracts });
+```
+
 ### [PXE] `proveTx` takes an options bag
 
 `PXE.proveTx` used to accept `scopes` as a positional argument; it now takes an options bag consistent with `simulateTx` and `profileTx`, and adds an optional `senderForTags` field. Update direct callers:

yarn-project/aztec-node/src/aztec-node/public_data_overrides.test.ts

@@ -0,0 +1,84 @@
+import { Fr } from '@aztec/foundation/curves/bn254';
+import { PublicDataWrite } from '@aztec/stdlib/avm';
+import { AztecAddress } from '@aztec/stdlib/aztec-address';
+import { computePublicDataTreeLeafSlot } from '@aztec/stdlib/hash';
+import type { MerkleTreeWriteOperations } from '@aztec/stdlib/interfaces/server';
+import { MerkleTreeId } from '@aztec/stdlib/trees';
+
+import { type MockProxy, mock } from 'jest-mock-extended';
+
+import { applyPublicDataOverrides } from './public_data_overrides.js';
+
+describe('applyPublicDataOverrides', () => {
+  let fork: MockProxy<MerkleTreeWriteOperations>;
+
+  beforeEach(() => {
+    fork = mock<MerkleTreeWriteOperations>();
+    fork.sequentialInsert.mockResolvedValue({ lowLeavesWitnesses: [], insertionWitnesses: [] } as any);
+  });
+
+  it('does nothing when overrides is undefined', async () => {
+    await applyPublicDataOverrides(fork, undefined);
+    expect(fork.sequentialInsert).not.toHaveBeenCalled();
+  });
+
+  it('does nothing when overrides is empty', async () => {
+    await applyPublicDataOverrides(fork, []);
+    expect(fork.sequentialInsert).not.toHaveBeenCalled();
+  });
+
+  it('inserts a single override at the correct leaf slot', async () => {
+    const contract = await AztecAddress.random();
+    const slot = Fr.random();
+    const value = Fr.random();
+
+    await applyPublicDataOverrides(fork, [{ contract, slot, value }]);
+
+    const expectedLeafSlot = await computePublicDataTreeLeafSlot(contract, slot);
+    const expectedWrite = new PublicDataWrite(expectedLeafSlot, value);
+
+    expect(fork.sequentialInsert).toHaveBeenCalledTimes(1);
+    expect(fork.sequentialInsert).toHaveBeenCalledWith(MerkleTreeId.PUBLIC_DATA_TREE, [expectedWrite.toBuffer()]);
+  });
+
+  it('inserts multiple overrides in a single batch call', async () => {
+    const contract = await AztecAddress.random();
+    const overrides = [
+      { contract, slot: Fr.random(), value: Fr.random() },
+      { contract, slot: Fr.random(), value: Fr.random() },
+    ];
+
+    await applyPublicDataOverrides(fork, overrides);
+
+    const expectedWrites = await Promise.all(
+      overrides.map(async o => {
+        const leafSlot = await computePublicDataTreeLeafSlot(o.contract, o.slot);
+        return new PublicDataWrite(leafSlot, o.value);
+      }),
+    );
+
+    expect(fork.sequentialInsert).toHaveBeenCalledTimes(1);
+    expect(fork.sequentialInsert).toHaveBeenCalledWith(
+      MerkleTreeId.PUBLIC_DATA_TREE,
+      expectedWrites.map(w => w.toBuffer()),
+    );
+  });
+
+  it('passes duplicate (contract, slot) writes through — last write wins via tree semantics', async () => {
+    const contract = await AztecAddress.random();
+    const slot = Fr.random();
+    const firstValue = Fr.random();
+    const secondValue = Fr.random();
+
+    await applyPublicDataOverrides(fork, [
+      { contract, slot, value: firstValue },
+      { contract, slot, value: secondValue },
+    ]);
+
+    // Both writes are passed to sequentialInsert; the tree handles last-wins ordering.
+    expect(fork.sequentialInsert).toHaveBeenCalledTimes(1);
+    const [treeId, leaves] = fork.sequentialInsert.mock.calls[0] as [MerkleTreeId, Buffer[]];
+    expect(treeId).toBe(MerkleTreeId.PUBLIC_DATA_TREE);
+    expect(leaves).toHaveLength(2);
+  });
+});

yarn-project/aztec-node/src/aztec-node/public_data_overrides.ts

@@ -0,0 +1,35 @@
+import { PublicDataWrite } from '@aztec/stdlib/avm';
+import { computePublicDataTreeLeafSlot } from '@aztec/stdlib/hash';
+import type { PublicStorageOverride } from '@aztec/stdlib/interfaces/client';
+import type { MerkleTreeWriteOperations } from '@aztec/stdlib/interfaces/server';
+import { MerkleTreeId } from '@aztec/stdlib/trees';
+
+/**
+ * Injects public-state overrides into an (ephemeral) world-state fork before simulation.
+ *
+ * Each override is written via the same `sequentialInsert` path the public processor
+ * uses during real transaction execution, so low-leaf updates and root coherence are
+ * handled identically for both simulation and proof generation.
+ *
+ * Writes never reach committed world state — the fork is thrown away after simulation.
+ */
+export async function applyPublicDataOverrides(
+  fork: MerkleTreeWriteOperations,
+  publicStorage: PublicStorageOverride[] | undefined,
+): Promise<void> {
+  if (!publicStorage?.length) {
+    return;
+  }
+
+  const writes = await Promise.all(
+    publicStorage.map(async o => {
+      const leafSlot = await computePublicDataTreeLeafSlot(o.contract, o.slot);
+      return new PublicDataWrite(leafSlot, o.value);
+    }),
+  );
+
+  await fork.sequentialInsert(
+    MerkleTreeId.PUBLIC_DATA_TREE,
+    writes.map(w => w.toBuffer()),
+  );
+}

yarn-project/aztec-node/src/aztec-node/server.ts

@@ -106,6 +106,7 @@ import {
   type GlobalVariableBuilder as GlobalVariableBuilderInterface,
   type IndexedTxEffect,
   PublicSimulationOutput,
+  type SimulationOverrides,
   Tx,
   type TxHash,
   TxReceipt,
@@ -146,6 +147,7 @@ import {
 } from './block_response_helpers.js';
 import { type AztecNodeConfig, createKeyStoreForValidator } from './config.js';
 import { NodeMetrics } from './node_metrics.js';
+import { applyPublicDataOverrides } from './public_data_overrides.js';
 
 /**
  * The aztec node.
@@ -1440,11 +1442,17 @@ export class AztecNodeService implements AztecNode, AztecNodeAdmin, AztecNodeDeb
   /**
    * Simulates the public part of a transaction with the current state.
    * @param tx - The transaction to simulate.
+   * @param skipFeeEnforcement - If true, fee enforcement is skipped.
+   * @param overrides - Optional pre-simulation overrides applied to the ephemeral fork and contract DB.
    **/
   @trackSpan('AztecNodeService.simulatePublicCalls', (tx: Tx) => ({
     [Attributes.TX_HASH]: tx.getTxHash().toString(),
   }))
-  public async simulatePublicCalls(tx: Tx, skipFeeEnforcement = false): Promise<PublicSimulationOutput> {
+  public async simulatePublicCalls(
+    tx: Tx,
+    skipFeeEnforcement = false,
+    overrides?: SimulationOverrides,
+  ): Promise<PublicSimulationOutput> {
     // Check total gas limit for simulation
     const gasSettings = tx.data.constants.txContext.gasSettings;
     const txGasLimit = gasSettings.gasLimits.l2Gas;
@@ -1489,6 +1497,7 @@ export class AztecNodeService implements AztecNode, AztecNodeAdmin, AztecNodeDeb
     await this.worldStateSynchronizer.syncImmediate(latestBlockNumber);
     const merkleTreeFork = await this.worldStateSynchronizer.fork();
     try {
+      await applyPublicDataOverrides(merkleTreeFork, overrides?.publicStorage);
       const config = PublicSimulatorConfig.from({
         skipFeeEnforcement,
         collectDebugLogs: true,

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

@@ -78,4 +78,7 @@ export { AccountManager } from '../wallet/account_manager.js';
 
 export { TxSimulationResultWithAppOffset } from '../wallet/tx_simulation_result_with_app_offset.js';
 
+export { type PublicStorageOverride, PublicStorageOverrideSchema } from '@aztec/stdlib/interfaces/client';
+export { SimulationOverrides } from '@aztec/stdlib/tx';
+
 export { type DeployAccountOptions, DeployAccountMethod } from '../wallet/deploy_account_method.js';

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

@@ -7,6 +7,7 @@ import {
   type Capsule,
   OFFCHAIN_MESSAGE_IDENTIFIER,
   type OffchainEffect,
+  type SimulationOverrides,
   type SimulationStats,
   type TxHash,
   type TxReceipt,
@@ -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;
+  /** Pre-simulation overrides applied to the ephemeral fork and contract DB (publicStorage writes, contract instance overrides). */
+  overrides?: SimulationOverrides;
 };
 
 /**

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

@@ -19,6 +19,7 @@ import type { ExecutionPayload, InTx } from '@aztec/stdlib/tx';
 import {
   Capsule,
   HashedValues,
+  SimulationOverrides,
   TxHash,
   TxProfileResult,
   TxReceipt,
@@ -335,6 +336,7 @@ export const SimulateOptionsSchema = z.object({
   skipFeeEnforcement: optional(z.boolean()),
   includeMetadata: optional(z.boolean()),
   additionalScopes: optional(z.array(schemas.AztecAddress)),
+  overrides: optional(SimulationOverrides.schema),
 });
 
 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 { PublicStorageOverride } 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]);
@@ -251,6 +255,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 publicStorage: PublicStorageOverride[] = [
+        { contract: avmContract.address, slot: SINGLE_SLOT, value: overrideValue },
+      ];
+
+      const simResult = await avmContract.methods
+        .read_storage_single()
+        .simulate({ from: defaultAccountAddress, overrides: { publicStorage } });
+      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 publicStorage: PublicStorageOverride[] = [
+        { contract: avmContract.address, slot: SINGLE_SLOT, value: overrideValue },
+      ];
+
+      const simResult = await avmContract.methods
+        .read_storage_single()
+        .simulate({ from: defaultAccountAddress, overrides: { publicStorage } });
+      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/end-to-end/src/test-wallet/test_wallet.ts

@@ -287,11 +287,14 @@ export class TestWallet extends BaseWallet {
       : executionPayload;
     const chainInfo = await this.getChainInfo();
 
-    let overrides: SimulationOverrides | undefined;
+    let overrides = opts.overrides;
     let txRequest: TxExecutionRequest;
     if (useOverride) {
       const accountOverrides = await this.buildAccountOverrides(scopes);
-      overrides = new SimulationOverrides(accountOverrides);
+      overrides = new SimulationOverrides({
+        publicStorage: overrides?.publicStorage,
+        contracts: { ...overrides?.contracts, ...accountOverrides },
+      });
     }
 
     if (from === NO_FROM) {

yarn-project/pxe/src/pxe.ts

@@ -120,7 +120,11 @@ export type SimulateTxOpts = {
   skipFeeEnforcement?: boolean;
   /** If true, kernel logic is emulated in TS for simulation */
   skipKernels?: boolean;
-  /** State overrides for the simulation, such as contract instances and artifacts. Requires skipKernels: true */
+  /**
+   * Pre-simulation overrides applied to the ephemeral fork and contract DB. Bundles publicStorage
+   * writes (no skipKernels required) and per-address (instance, artifact?) overrides used by both
+   * AVM-side public dispatch and PXE-side ACIR private dispatch (requires skipKernels: true).
+   */
   overrides?: SimulationOverrides;
   /** Addresses whose private state and keys are accessible during private execution */
   scopes: AztecAddress[];
@@ -484,11 +488,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, overrides?: SimulationOverrides) {
     // 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, overrides);
       if (result.revertReason) {
         throw result.revertReason;
       }
@@ -1054,7 +1058,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, overrides);
           publicSimulationTime = publicSimulationTimer.ms();
           if (publicOutput?.debugLogs?.length) {
             await displayDebugLogs(publicOutput.debugLogs, addr => this.contractStore.getDebugContractName(addr));