feat: contractOverrides for private-call simulation

Plumbs the existing PXE-side SimulationOverrides.contracts mechanism through
wallet.simulateTx and aztec.js .simulate() options as 'contractOverrides'.
This covers private-function dispatch during simulation against an upgraded
class - the PXE's ACIR simulator uses the override artifact for the given
address instead of its registered one.

Together with the existing node-side stateOverrides, .simulate() now handles
both private and public call flavors when simulating against an upgraded
contract.

fastForwardContractUpdate now takes a ContractArtifact and returns both
override layers in one blob so callers can spread it across .simulate():

    const overrides = await fastForwardContractUpdate({
      instanceAddress, newArtifact, node,
    });
    await updatedContract.methods.X().simulate({ ...overrides });

Adds e2e coverage for the private-call path (set_private_value, which only
exists on the upgraded class) to demonstrate the full flow.

Author: dbanks12

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

@@ -67,9 +67,12 @@ 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
+## Simulating with state and contract 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.
+`.simulate()` accepts two override options that are scoped to that single simulation; real chain state is untouched.
+
+- `stateOverrides`: state-tree overrides (e.g. `publicStorage` writes).
+- `contractOverrides`: contract-DB overrides keyed by address - per-address `instance` (always applied) and an optional `artifact` (for redirecting private-function dispatch).
 
 Override a public-storage slot:
 
@@ -81,26 +84,30 @@ const result = await contract.methods.read_balance(account).simulate({
 });
 ```
 
-Use this to:
+Use these 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
+- Simulate a contract instance as if it had been upgraded
 - Test branches that depend on rare values without orchestrating the contract calls that produce them
 
 ### Fast-forwarding a contract update
 
-`fastForwardContractUpdate` builds the full set of overrides needed to simulate a deployed instance as if it had already been upgraded to a new contract class. The new class must already be registered on chain. The cheat mirrors a real `pxe.updateContract` followed by waiting out the upgrade delay: the instance's `currentContractClassId` is bumped, and the `ContractInstanceRegistry`'s delayed-public-mutable storage is rewritten to look like the upgrade was scheduled in the past.
+`fastForwardContractUpdate` builds the override blobs needed to simulate a deployed instance as if it had already been upgraded to a new contract class. The new class must already be registered on chain. Mirrors a real `pxe.updateContract` followed by waiting out the upgrade delay.
+
+It returns both `stateOverrides` (node-side, public dispatch) and `contractOverrides` (PXE-side, private dispatch) so a single spread covers any mix of private and public function calls on the upgraded contract.
 
 ```typescript
 import { fastForwardContractUpdate } from '@aztec/aztec.js';
 
-const stateOverrides = await fastForwardContractUpdate({
+const overrides = await fastForwardContractUpdate({
   instanceAddress: contract.address,
-  newClassId: upgradedClass.id,
+  newArtifact: UpdatedContract.artifact,
   node,
 });
 
-const result = await contract.methods.upgraded_method().simulate({ stateOverrides });
+const upgradedContract = UpdatedContract.at(contract.address, wallet);
+const result = await upgradedContract.methods.upgraded_method().simulate({ ...overrides });
 ```
 
 Use this to test code paths that only execute after an upgrade, without orchestrating the full delayed-mutable upgrade flow.

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

@@ -23,17 +23,20 @@ const result = await contract.methods.read_balance(account).simulate({
 
 The same option flows through `wallet.simulateTx` and eventually to `simulatePublicCalls` RPC on `AztecNode`.
 
-A second override flavor, `contractInstances`, shadows contract instances in the simulator's contract DB - useful for simulating a contract being on a different class than the one it was deployed with. To simulate a complete on-chain upgrade flow, use the `fastForwardContractUpdate` helper which produces a coherent set of `publicStorage` and `contractInstances` overrides:
+`.simulate(...)` also accepts `contractOverrides` - an address-keyed map of `{ instance, artifact? }` that replaces the contract DB entry the simulator uses for that address. The instance is applied on both the AVM (public dispatch) and PXE (private dispatch via ACIR) sides; the optional artifact controls private-call ACIR lookup.
+
+To simulate a complete on-chain upgrade flow, use the `fastForwardContractUpdate` helper. It returns both `stateOverrides` (state-tree writes - registry public storage) and `contractOverrides` (instance/artifact bumped to the new class) so a single spread covers any mix of call flavors:
 
 ```typescript
 import { fastForwardContractUpdate } from '@aztec/aztec.js';
 
-const stateOverrides = await fastForwardContractUpdate({
+const overrides = await fastForwardContractUpdate({
   instanceAddress: contract.address,
-  newClassId: upgradedClass.id,
+  newArtifact: UpdatedContract.artifact,
   node,
 });
-const result = await contract.methods.upgraded_method().simulate({ stateOverrides });
+const upgradedContract = UpdatedContract.at(contract.address, wallet);
+const result = await upgradedContract.methods.upgraded_method().simulate({ ...overrides });
 ```
 
 ### [PXE] `proveTx` takes an options bag

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

@@ -100,6 +100,7 @@ import type { NullifierLeafPreimage, PublicDataTreeLeafPreimage } from '@aztec/s
 import { MerkleTreeId, NullifierMembershipWitness, PublicDataWitness } from '@aztec/stdlib/trees';
 import {
   type BlockHeader,
+  type ContractOverrides,
   type FeeProvider,
   type GlobalVariableBuilder as GlobalVariableBuilderInterface,
   type IndexedTxEffect,
@@ -1484,7 +1485,8 @@ 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.
+   * @param stateOverrides - Optional state-tree overrides (publicStorage writes) applied to the ephemeral fork.
+   * @param contractOverrides - Optional contract-DB overrides; per-address `instance` is used for AVM dispatch.
    **/
   @trackSpan('AztecNodeService.simulatePublicCalls', (tx: Tx) => ({
     [Attributes.TX_HASH]: tx.getTxHash().toString(),
@@ -1493,6 +1495,7 @@ export class AztecNodeService implements AztecNode, AztecNodeAdmin, AztecNodeDeb
     tx: Tx,
     skipFeeEnforcement = false,
     stateOverrides?: StateOverrides,
+    contractOverrides?: ContractOverrides,
   ): Promise<PublicSimulationOutput> {
     // Check total gas limit for simulation
     const gasSettings = tx.data.constants.txContext.gasSettings;
@@ -1550,7 +1553,10 @@ export class AztecNodeService implements AztecNode, AztecNodeAdmin, AztecNodeDeb
         }),
       });
       const contractsDB = new PublicContractsDB(this.contractDataSource, this.log.getBindings());
-      contractsDB.addContracts(stateOverrides?.contractInstances);
+      const overrideInstances = contractOverrides
+        ? Object.values(contractOverrides).map(({ instance }) => instance)
+        : undefined;
+      contractsDB.addContracts(overrideInstances);
       const processor = publicProcessorFactory.create(merkleTreeFork, newGlobalVariables, config, contractsDB);
 
       // REFACTOR: Consider merging ProcessReturnValues into ProcessedTx

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

@@ -238,7 +238,20 @@ describe('Contract Class', () => {
         from: account.getAddress(),
         stateOverrides: { publicStorage: [{ contract: contractAddress, slot: new Fr(1), value: new Fr(42) }] },
       }),
-    ).rejects.toThrow(/stateOverrides are not supported for utility/);
+    ).rejects.toThrow(/not supported for utility function/);
+    expect(wallet.executeUtility).not.toHaveBeenCalled();
+  });
+
+  it('throws when contractOverrides are passed to a utility function simulation', async () => {
+    const fooContract = Contract.at(contractAddress, defaultArtifact, wallet);
+    await expect(
+      fooContract.methods.qux(123n).simulate({
+        from: account.getAddress(),
+        contractOverrides: {
+          [contractAddress.toString()]: { instance: contractInstance, artifact: defaultArtifact },
+        },
+      }),
+    ).rejects.toThrow(/not supported for utility function/);
     expect(wallet.executeUtility).not.toHaveBeenCalled();
   });
 

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

@@ -130,10 +130,10 @@ export class ContractFunctionInteraction extends BaseContractInteraction {
     // docs:end:simulate
     if (this.functionDao.functionType == FunctionType.UTILITY) {
       if (
-        options.stateOverrides &&
-        (options.stateOverrides.publicStorage?.length || options.stateOverrides.contractInstances?.length)
+        options.stateOverrides?.publicStorage?.length ||
+        (options.contractOverrides && Object.keys(options.contractOverrides).length > 0)
       ) {
-        throw new Error('stateOverrides are not supported for utility function simulation.');
+        throw new Error('stateOverrides and contractOverrides are not supported for utility function simulation.');
       }
       const call = await this.getFunctionCall();
       const scopes = [...(options.additionalScopes ?? [])];

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

@@ -1,12 +1,14 @@
 import { Fr } from '@aztec/foundation/curves/bn254';
 import { ProtocolContractAddress } from '@aztec/protocol-contracts';
+import type { ContractArtifact } from '@aztec/stdlib/abi';
 import { AztecAddress } from '@aztec/stdlib/aztec-address';
-import { SerializableContractInstance } from '@aztec/stdlib/contract';
+import { SerializableContractInstance, getContractClassFromArtifact } from '@aztec/stdlib/contract';
 import {
   DELAYED_PUBLIC_MUTABLE_VALUES_LEN,
   DelayedPublicMutableValuesWithHash,
 } from '@aztec/stdlib/delayed-public-mutable';
 import type { AztecNode } from '@aztec/stdlib/interfaces/client';
+import { getTestContractArtifact, getTokenContractArtifact } from '@aztec/stdlib/testing/fixtures';
 
 import { type MockProxy, mock } from 'jest-mock-extended';
 
@@ -16,13 +18,15 @@ describe('fastForwardContractUpdate', () => {
   let node: MockProxy<AztecNode>;
   let instanceAddress: AztecAddress;
   let originalClassId: Fr;
+  let newArtifact: ContractArtifact;
   let newClassId: Fr;
 
   beforeEach(async () => {
     node = mock<AztecNode>();
     instanceAddress = await AztecAddress.random();
-    originalClassId = Fr.random();
-    newClassId = Fr.random();
+    originalClassId = (await getContractClassFromArtifact(getTokenContractArtifact())).id;
+    newArtifact = getTestContractArtifact();
+    newClassId = (await getContractClassFromArtifact(newArtifact)).id;
 
     const instance = (
       await SerializableContractInstance.random({
@@ -32,48 +36,42 @@ describe('fastForwardContractUpdate', () => {
     ).withAddress(instanceAddress);
 
     node.getContract.mockResolvedValue(instance);
-    node.getContractClass.mockResolvedValue({
-      id: newClassId,
-      artifactHash: Fr.random(),
-      packedBytecodeCommitments: [],
-      privateFunctionsRoot: Fr.random(),
-      publicBytecodeCommitment: Fr.random(),
-      version: 1,
-      privateFunctions: [],
-      utilityFunctions: [],
-      publicFunctions: [],
-      packedBytecode: Buffer.alloc(0),
-    } as any);
+    node.getContractClass.mockResolvedValue({ id: newClassId } as any);
   });
 
-  it('produces overrides with bumped currentContractClassId and registry storage writes', async () => {
-    const overrides = await fastForwardContractUpdate({ instanceAddress, newClassId, node });
-
-    expect(overrides.contractInstances).toHaveLength(1);
-    expect(overrides.contractInstances![0].address).toEqual(instanceAddress);
-    expect(overrides.contractInstances![0].currentContractClassId).toEqual(newClassId);
-    expect(overrides.contractInstances![0].originalContractClassId).toEqual(originalClassId);
+  it('produces stateOverrides with registry storage writes', async () => {
+    const { stateOverrides } = await fastForwardContractUpdate({ instanceAddress, newArtifact, node });
 
     const expectedSlots = await DelayedPublicMutableValuesWithHash.getContractUpdateSlots(instanceAddress);
-    expect(overrides.publicStorage).toHaveLength(DELAYED_PUBLIC_MUTABLE_VALUES_LEN + 1);
-    for (const entry of overrides.publicStorage!) {
+    expect(stateOverrides.publicStorage).toHaveLength(DELAYED_PUBLIC_MUTABLE_VALUES_LEN + 1);
+    for (const entry of stateOverrides.publicStorage!) {
       expect(entry.contract).toEqual(ProtocolContractAddress.ContractInstanceRegistry);
     }
-    const baseSlot = expectedSlots.delayedPublicMutableSlot;
-    expect(overrides.publicStorage![0].slot).toEqual(baseSlot);
-    expect(overrides.publicStorage![overrides.publicStorage!.length - 1].slot).toEqual(
+    expect(stateOverrides.publicStorage![0].slot).toEqual(expectedSlots.delayedPublicMutableSlot);
+    expect(stateOverrides.publicStorage![stateOverrides.publicStorage!.length - 1].slot).toEqual(
       expectedSlots.delayedPublicMutableHashSlot,
     );
   });
 
+  it('produces contractOverrides keyed by the instance address with the new artifact', async () => {
+    const { contractOverrides } = await fastForwardContractUpdate({ instanceAddress, newArtifact, node });
+
+    const entry = contractOverrides[instanceAddress.toString()];
+    expect(entry).toBeDefined();
+    expect(entry.artifact).toBe(newArtifact);
+    expect(entry.instance.address).toEqual(instanceAddress);
+    expect(entry.instance.currentContractClassId).toEqual(newClassId);
+    expect(entry.instance.originalContractClassId).toEqual(originalClassId);
+  });
+
   it('throws when the instance is not deployed', async () => {
     node.getContract.mockResolvedValue(undefined);
-    await expect(fastForwardContractUpdate({ instanceAddress, newClassId, node })).rejects.toThrow(/not deployed/);
+    await expect(fastForwardContractUpdate({ instanceAddress, newArtifact, node })).rejects.toThrow(/not deployed/);
   });
 
   it('throws when the new class is not registered', async () => {
     node.getContractClass.mockResolvedValue(undefined);
-    await expect(fastForwardContractUpdate({ instanceAddress, newClassId, node })).rejects.toThrow(/not registered/);
+    await expect(fastForwardContractUpdate({ instanceAddress, newArtifact, node })).rejects.toThrow(/not registered/);
   });
 
   it('throws when the instance is already on the target class', async () => {
@@ -85,6 +83,6 @@ describe('fastForwardContractUpdate', () => {
     ).withAddress(instanceAddress);
     node.getContract.mockResolvedValue(sameClassInstance);
 
-    await expect(fastForwardContractUpdate({ instanceAddress, newClassId, node })).rejects.toThrow(/already on class/);
+    await expect(fastForwardContractUpdate({ instanceAddress, newArtifact, node })).rejects.toThrow(/already on class/);
   });
 });

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

@@ -1,35 +1,47 @@
 import { Fr } from '@aztec/foundation/curves/bn254';
 import { ProtocolContractAddress } from '@aztec/protocol-contracts';
+import type { ContractArtifact } from '@aztec/stdlib/abi';
 import type { AztecAddress } from '@aztec/stdlib/aztec-address';
+import { getContractClassFromArtifact } from '@aztec/stdlib/contract';
 import {
   DelayedPublicMutableValuesWithHash,
   ScheduledDelayChange,
   ScheduledValueChange,
 } from '@aztec/stdlib/delayed-public-mutable';
 import type { AztecNode, StateOverrides } from '@aztec/stdlib/interfaces/client';
+import type { ContractOverrides } from '@aztec/stdlib/tx';
 
 /**
- * Builds `StateOverrides` that simulate a contract instance having already been upgraded to a new contract class.
+ * Builds the override blobs that simulate a contract instance having already been upgraded to a new contract class.
  *
- * Mirrors a real on-chain upgrade flow (`pxe.updateContract` followed by waiting out the delay): the contract
- * instance's `currentContractClassId` is bumped to `newClassId`, and the `ContractInstanceRegistry`'s delayed
- * public mutable storage is rewritten to look like the upgrade was scheduled in the past.
+ * Mirrors a real on-chain upgrade flow (`pxe.updateContract` followed by waiting out the delay). Returns two
+ * coherent override layers:
  *
- * The new class must already be registered on chain. To upgrade to an unregistered class, register it first
- * (or use a higher-level helper that bundles class registration with the upgrade).
+ * - `stateOverrides` — applied node-side. Bumps `currentContractClassId` on the contract instance and rewrites
+ *   the `ContractInstanceRegistry`'s delayed-public-mutable storage to look like the upgrade was scheduled in
+ *   the past. Redirects public-call dispatch.
+ * - `contractOverrides` — applied PXE-side. Replaces the (instance, artifact) pair the ACIR simulator uses for
+ *   the upgraded address. Redirects private-call dispatch.
+ *
+ * Both layers are returned so a single `simulate({ ...overrides })` spread works for any mix of private and
+ * public function calls on the upgraded contract.
+ *
+ * The new class must already be registered on chain. To upgrade to an unregistered class, register it first.
  *
  * @param args.instanceAddress - Address of the deployed instance to upgrade.
- * @param args.newClassId - ID of the (already-registered) class to upgrade to.
+ * @param args.newArtifact - Artifact of the class to upgrade to. Must be registered on chain.
  * @param args.node - Node used to fetch the existing instance and validate the class is registered.
- * @returns `StateOverrides` to spread into a `simulate({ stateOverrides })` call.
+ * @returns `{ stateOverrides, contractOverrides }` to spread into a `simulate(...)` call.
  * @throws If the instance is not deployed, the class is not registered, or the instance is already on the target class.
  */
 export async function fastForwardContractUpdate(args: {
   instanceAddress: AztecAddress;
-  newClassId: Fr;
+  newArtifact: ContractArtifact;
   node: AztecNode;
-}): Promise<StateOverrides> {
-  const { instanceAddress, newClassId, node } = args;
+}): Promise<{ stateOverrides: StateOverrides; contractOverrides: ContractOverrides }> {
+  const { instanceAddress, newArtifact, node } = args;
+
+  const newClassId = (await getContractClassFromArtifact(newArtifact)).id;
 
   const instance = await node.getContract(instanceAddress);
   if (!instance) {
@@ -63,8 +75,14 @@ export async function fastForwardContractUpdate(args: {
 
   const upgradedInstance = { ...instance, currentContractClassId: newClassId };
 
-  return {
-    publicStorage,
-    contractInstances: [upgradedInstance],
+  const stateOverrides: StateOverrides = { publicStorage };
+
+  const contractOverrides: ContractOverrides = {
+    [instanceAddress.toString()]: {
+      instance: upgradedInstance,
+      artifact: newArtifact,
+    },
   };
+
+  return { stateOverrides, contractOverrides };
 }

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

@@ -6,6 +6,7 @@ import type { GasSettings, ManaUsageEstimate } from '@aztec/stdlib/gas';
 import type { StateOverrides } from '@aztec/stdlib/interfaces/client';
 import {
   type Capsule,
+  type ContractOverrides,
   OFFCHAIN_MESSAGE_IDENTIFIER,
   type OffchainEffect,
   type SimulationStats,
@@ -160,6 +161,12 @@ export type SimulateInteractionOptions = Omit<SendInteractionOptions, 'fee'> & {
   includeMetadata?: boolean;
   /** Pre-simulation state overrides applied to the ephemeral fork and contract DB. */
   stateOverrides?: StateOverrides;
+  /**
+   * Per-simulation contract overrides applied PXE-side. For each address, replaces the (instance, artifact)
+   * pair the ACIR simulator uses for the duration of the simulation. Required to redirect private-function
+   * dispatch when simulating against an upgraded class.
+   */
+  contractOverrides?: ContractOverrides;
 };
 
 /**

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

@@ -19,6 +19,7 @@ import { AbiDecodedSchema, type ApiSchemaFor, optional, schemas, zodFor } from '
 import type { ExecutionPayload, InTx } from '@aztec/stdlib/tx';
 import {
   Capsule,
+  ContractOverridesSchema,
   HashedValues,
   TxHash,
   TxProfileResult,
@@ -337,6 +338,7 @@ export const SimulateOptionsSchema = z.object({
   includeMetadata: optional(z.boolean()),
   additionalScopes: optional(z.array(schemas.AztecAddress)),
   stateOverrides: optional(StateOverridesSchema),
+  contractOverrides: optional(ContractOverridesSchema),
 });
 
 export const ProfileOptionsSchema = SimulateOptionsSchema.extend({