feat: public storage overrides for simulation via stateOverrides

Adds a `stateOverrides?: StateOverrides` parameter to `simulatePublicCalls`
and threads it through aztec.js → wallet-sdk → PXE → node. The new shape is
modeled after Foundry / Geth `eth_call` state overrides — a single bundled
argument that can be extended with new override flavors (e.g. contract
class/instance shadows, nullifier overrides) without changing the RPC.

This first cut populates only `publicStorage`:

```ts
type StateOverrides = {
  publicStorage?: PublicStorageOverride[];
};
type PublicStorageOverride = {
  contract: AztecAddress;
  slot: Fr;       // raw storage slot, not leafSlot — derived internally
  value: Fr;
};
```

The node opens an ephemeral world-state fork per simulation and writes each
override via `applyPublicDataOverrides(fork, publicStorage)` (computes
`computePublicDataTreeLeafSlot(contract, slot)` then `sequentialInsert`).
Writes never reach committed state.

Includes the e2e wiring fix in TestWallet.simulateViaEntrypoint to forward
the option through to PXE so e2e tests can exercise the full chain.

Author: dbanks12

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

@@ -67,6 +67,26 @@ 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 state overrides
+
+`.simulate()` accepts a `stateOverrides` option that injects values into the simulator's ephemeral world-state fork before the call runs. The override is scoped to that single simulation; the real chain state is untouched.
+
+Override a public-storage slot:
+
+```typescript
+const result = await contract.methods.read_balance(account).simulate({
+  stateOverrides: {
+    publicStorage: [{ contract: contract.address, slot: BALANCE_SLOT, value: new Fr(1_000_000n) }],
+  },
+});
+```
+
+Use this to:
+
+- Set up state preconditions without running a full setup transaction
+- Reproduce a bug from production by pinning storage to the values seen at a specific block
+- Test branches that depend on rare values 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

@@ -9,6 +9,20 @@ Aztec is in active development. Each version may introduce breaking changes that
 
 ## TBD
 
+### [Aztec.js] `simulate` accepts `stateOverrides` for fork-style testing
+
+`Contract.methods.foo(...).simulate(...)` now accepts a `stateOverrides` option that injects values into the simulator's ephemeral world-state fork before the call runs. The first override flavor is `publicStorage`, which writes a `(contract, slot, value)` triple into the public-data tree as if a previous tx had set it. Overrides are scoped to the simulation; the real chain state is untouched.
+
+```typescript
+const result = await contract.methods.read_balance(account).simulate({
+  stateOverrides: {
+    publicStorage: [{ contract: contract.address, slot: BALANCE_SLOT, value: new Fr(1_000_000n) }],
+  },
+});
+```
+
+The same option flows through `wallet.simulateTx`, `pxe.simulateTx`, and the `simulatePublicCalls` RPC on `AztecNode`. The bundled `StateOverrides` shape is forward-compatible: future override flavors (contract class/instance shadows, etc.) will be added as additional optional fields without changing the RPC.
+
 ### [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,85 @@
+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 { jest } from '@jest/globals';
+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.
+ *
+ * The fork is ephemeral — these writes never reach the committed world state.
+ */
+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

@@ -80,6 +80,7 @@ import type {
   CheckpointResponse,
   GetContractClassLogsResponse,
   GetPublicLogsResponse,
+  StateOverrides,
 } from '@aztec/stdlib/interfaces/client';
 import { AztecNodeAdminConfigSchema } from '@aztec/stdlib/interfaces/client';
 import {
@@ -143,6 +144,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.
@@ -1481,11 +1483,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 stateOverrides - Optional state overrides applied to the ephemeral fork and contract DB before simulation.
    **/
   @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,
+    stateOverrides?: StateOverrides,
+  ): Promise<PublicSimulationOutput> {
     // Check total gas limit for simulation
     const gasSettings = tx.data.constants.txContext.gasSettings;
     const txGasLimit = gasSettings.gasLimits.l2Gas;
@@ -1530,6 +1538,7 @@ export class AztecNodeService implements AztecNode, AztecNodeAdmin, AztecNodeDeb
     await this.worldStateSynchronizer.syncImmediate(latestBlockNumber);
     const merkleTreeFork = await this.worldStateSynchronizer.fork();
     try {
+      await applyPublicDataOverrides(merkleTreeFork, stateOverrides?.publicStorage);
       const config = PublicSimulatorConfig.from({
         skipFeeEnforcement,
         collectDebugLogs: true,

yarn-project/aztec.js/README.md

@@ -42,3 +42,5 @@ const contract = await Contract.at(contractAddress, MyContractArtifact, wallet);
 const balance = await contract.methods.get_balance(wallet.getAddress()).simulate();
 console.log(`Account balance is ${balance}`);
 ```
+
+For fork-style testing, `simulate` accepts an optional `stateOverrides` to pre-populate public storage or shadow contract classes/instances for the duration of the call. See [Simulate with state overrides](https://docs.aztec.network/developers/aztec-js/how_to_simulate_with_overrides) for the reference, including the `spoof*` cheatcode helpers.

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

@@ -78,4 +78,11 @@ export { AccountManager } from '../wallet/account_manager.js';
 
 export { TxSimulationResultWithAppOffset } from '../wallet/tx_simulation_result_with_app_offset.js';
 
+export {
+  type PublicStorageOverride,
+  PublicStorageOverrideSchema,
+  type StateOverrides,
+  StateOverridesSchema,
+} 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 { StateOverrides } 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;
+  /** Pre-simulation state overrides applied to the ephemeral fork and contract DB. */
+  stateOverrides?: StateOverrides;
 };
 
 /**

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 { StateOverridesSchema } 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)),
+  stateOverrides: optional(StateOverridesSchema),
 });
 
 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]);
@@ -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 publicStorage: PublicStorageOverride[] = [
+        { contract: avmContract.address, slot: SINGLE_SLOT, value: overrideValue },
+      ];
+
+      const simResult = await avmContract.methods
+        .read_storage_single()
+        .simulate({ from: defaultAccountAddress, stateOverrides: { 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, stateOverrides: { 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;