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`: an array of `ContractInstanceWithAddress` to override deployed contract instances. Register the new class artifact locally first via `wallet.registerContractClass(artifact)`.
 
 Override a public-storage slot:
 
@@ -81,26 +84,37 @@ 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 `stateOverrides` (registry storage rewrites) and `contractOverrides` (instance with bumped class id). A single spread covers any mix of private and public function calls on the upgraded contract.
+
+Register the new class artifact in your local PXE first via `wallet.registerContractClass(artifact)`.
 
 ```typescript
 import { fastForwardContractUpdate } from '@aztec/aztec.js';
+import { getContractClassFromArtifact } from '@aztec/aztec.js/contracts';
+
+// One-time local PXE registration of the new class artifact
+await wallet.registerContractClass(UpdatedContract.artifact);
 
-const stateOverrides = await fastForwardContractUpdate({
+const newClassId = (await getContractClassFromArtifact(UpdatedContract.artifact)).id;
+const overrides = await fastForwardContractUpdate({
   instanceAddress: contract.address,
-  newClassId: upgradedClass.id,
+  newClassId,
   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,23 @@ 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 array of `ContractInstanceWithAddress` that override deployed instances at their addresses. Applied on both the AVM (public dispatch) and PXE (private dispatch via ACIR) sides. Register the new class artifact locally first via `wallet.registerContractClass(artifact)` (newly exposed for this purpose) so PXE has the ACIR.
+
+To simulate a complete on-chain upgrade flow, use the `fastForwardContractUpdate` helper. It returns both `stateOverrides` (registry storage rewrites) and `contractOverrides` (instance with bumped class id). Register the new artifact locally via `wallet.registerContractClass(artifact)` first:
 
 ```typescript
-import { fastForwardContractUpdate } from '@aztec/aztec.js';
+import { fastForwardContractUpdate, getContractClassFromArtifact } from '@aztec/aztec.js/contracts';
+
+await wallet.registerContractClass(UpdatedContract.artifact);
 
-const stateOverrides = await fastForwardContractUpdate({
+const newClassId = (await getContractClassFromArtifact(UpdatedContract.artifact)).id;
+const overrides = await fastForwardContractUpdate({
   instanceAddress: contract.address,
-  newClassId: upgradedClass.id,
+  newClassId,
   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

@@ -1484,7 +1484,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 instance overrides applied to AVM dispatch.
    **/
   @trackSpan('AztecNodeService.simulatePublicCalls', (tx: Tx) => ({
     [Attributes.TX_HASH]: tx.getTxHash().toString(),
@@ -1493,6 +1494,7 @@ export class AztecNodeService implements AztecNode, AztecNodeAdmin, AztecNodeDeb
     tx: Tx,
     skipFeeEnforcement = false,
     stateOverrides?: StateOverrides,
+    contractOverrides?: ContractInstanceWithAddress[],
   ): Promise<PublicSimulationOutput> {
     // Check total gas limit for simulation
     const gasSettings = tx.data.constants.txContext.gasSettings;
@@ -1550,7 +1552,7 @@ export class AztecNodeService implements AztecNode, AztecNodeAdmin, AztecNodeDeb
         }),
       });
       const contractsDB = new PublicContractsDB(this.contractDataSource, this.log.getBindings());
-      contractsDB.addContracts(stateOverrides?.contractInstances);
+      contractsDB.addContracts(contractOverrides);
       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,18 @@ 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: [contractInstance],
+      }),
+    ).rejects.toThrow(/not supported for utility function/);
     expect(wallet.executeUtility).not.toHaveBeenCalled();
   });
 

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

@@ -129,11 +129,8 @@ export class ContractFunctionInteraction extends BaseContractInteraction {
   ): Promise<SimulationResult> {
     // docs:end:simulate
     if (this.functionDao.functionType == FunctionType.UTILITY) {
-      if (
-        options.stateOverrides &&
-        (options.stateOverrides.publicStorage?.length || options.stateOverrides.contractInstances?.length)
-      ) {
-        throw new Error('stateOverrides are not supported for utility function simulation.');
+      if (options.stateOverrides?.publicStorage?.length || options.contractOverrides?.length) {
+        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

@@ -32,46 +32,38 @@ 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, newClassId, 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 a contractOverrides entry with bumped currentContractClassId', async () => {
+    const { contractOverrides } = await fastForwardContractUpdate({ instanceAddress, newClassId, node });
+
+    expect(contractOverrides).toHaveLength(1);
+    expect(contractOverrides[0].address).toEqual(instanceAddress);
+    expect(contractOverrides[0].currentContractClassId).toEqual(newClassId);
+    expect(contractOverrides[0].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/);
   });
 
-  it('throws when the new class is not registered', async () => {
+  it('throws when the new class is not registered on chain', async () => {
     node.getContractClass.mockResolvedValue(undefined);
     await expect(fastForwardContractUpdate({ instanceAddress, newClassId, node })).rejects.toThrow(/not registered/);
   });

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

@@ -1,6 +1,7 @@
 import { Fr } from '@aztec/foundation/curves/bn254';
 import { ProtocolContractAddress } from '@aztec/protocol-contracts';
 import type { AztecAddress } from '@aztec/stdlib/aztec-address';
+import type { ContractInstanceWithAddress } from '@aztec/stdlib/contract';
 import {
   DelayedPublicMutableValuesWithHash,
   ScheduledDelayChange,
@@ -9,26 +10,29 @@ import {
 import type { AztecNode, StateOverrides } from '@aztec/stdlib/interfaces/client';
 
 /**
- * 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:
  *
- * 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` — node-side state-tree writes that rewrite the `ContractInstanceRegistry`'s delayed-public-mutable
+ *   storage to look like the upgrade was scheduled in the past.
+ * - `contractOverrides` — an instance entry whose `currentContractClassId` is bumped to the new class. Applied both
+ *   AVM-side (public dispatch) and PXE-side (private dispatch).
+ *
+ * The new class must already be registered on chain. The artifact must also be registered in the local PXE — call
+ * `wallet.registerContractClass(artifact)` once before invoking this helper.
  *
  * @param args.instanceAddress - Address of the deployed instance to upgrade.
  * @param args.newClassId - ID of the (already-registered) class to upgrade to.
  * @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.
- * @throws If the instance is not deployed, the class is not registered, or the instance is already on the target class.
+ * @returns `{ stateOverrides, contractOverrides }` to spread into a `simulate(...)` call.
+ * @throws If the instance is not deployed, the class is not registered on chain, or the instance is already on the target class.
  */
 export async function fastForwardContractUpdate(args: {
   instanceAddress: AztecAddress;
   newClassId: Fr;
   node: AztecNode;
-}): Promise<StateOverrides> {
+}): Promise<{ stateOverrides: StateOverrides; contractOverrides: ContractInstanceWithAddress[] }> {
   const { instanceAddress, newClassId, node } = args;
 
   const instance = await node.getContract(instanceAddress);
@@ -38,7 +42,7 @@ export async function fastForwardContractUpdate(args: {
 
   const klass = await node.getContractClass(newClassId);
   if (!klass) {
-    throw new Error(`Contract class ${newClassId} is not registered; register it before fast-forwarding to it`);
+    throw new Error(`Contract class ${newClassId} is not registered on chain; publish it before fast-forwarding to it`);
   }
 
   if (instance.currentContractClassId.equals(newClassId)) {
@@ -61,10 +65,10 @@ export async function fastForwardContractUpdate(args: {
     value,
   }));
 
-  const upgradedInstance = { ...instance, currentContractClassId: newClassId };
+  const upgradedInstance: ContractInstanceWithAddress = { ...instance, currentContractClassId: newClassId };
 
   return {
-    publicStorage,
-    contractInstances: [upgradedInstance],
+    stateOverrides: { publicStorage },
+    contractOverrides: [upgradedInstance],
   };
 }

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

@@ -2,6 +2,7 @@ import type { Fr } from '@aztec/foundation/curves/bn254';
 import type { FieldsOf } from '@aztec/foundation/types';
 import type { AuthWitness } from '@aztec/stdlib/auth-witness';
 import { AztecAddress } from '@aztec/stdlib/aztec-address';
+import type { ContractInstanceWithAddress } from '@aztec/stdlib/contract';
 import type { GasSettings, ManaUsageEstimate } from '@aztec/stdlib/gas';
 import type { StateOverrides } from '@aztec/stdlib/interfaces/client';
 import {
@@ -158,8 +159,15 @@ 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. */
+  /** Pre-simulation state-tree overrides (e.g. publicStorage writes) applied to the ephemeral fork. */
   stateOverrides?: StateOverrides;
+  /**
+   * Per-simulation contract instance overrides. The simulator uses each provided instance in place of the
+   * chain-registered one at its address — applied both AVM-side (public dispatch) and PXE-side (private
+   * dispatch). For private dispatch to find the right ACIR, register the new class first via
+   * `wallet.registerContractClass(artifact)`.
+   */
+  contractOverrides?: ContractInstanceWithAddress[];
 };
 
 /**

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

@@ -155,6 +155,18 @@ describe('WalletSchema', () => {
     });
   });
 
+  it('registerContractClass', async () => {
+    const mockArtifact: ContractArtifact = {
+      name: 'TestContract',
+      functions: [],
+      nonDispatchPublicFunctions: [],
+      outputs: { structs: {}, globals: {} },
+      fileMap: {},
+      storageLayout: {},
+    };
+    await context.client.registerContractClass(mockArtifact);
+  });
+
   it('simulateTx', async () => {
     const exec: ExecutionPayload = {
       calls: [],
@@ -448,6 +460,8 @@ class MockWallet implements Wallet {
     };
   }
 
+  async registerContractClass(_artifact: any): Promise<void> {}
+
   async simulateTx(_exec: ExecutionPayload, _opts: SimulateOptions): Promise<TxSimulationResultWithAppOffset> {
     return TxSimulationResultWithAppOffset.fromResultAndOffset(await TxSimulationResult.random(), 0);
   }

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

@@ -272,6 +272,12 @@ export type Wallet = {
     artifact?: ContractArtifact,
     secretKey?: Fr,
   ): Promise<ContractInstanceWithAddress>;
+  /**
+   * Registers a contract class artifact in the local PXE without binding it to any instance.
+   * Useful for simulation flows that need the artifact available locally before any on-chain
+   * upgrade has taken effect. The artifact's class id must match its content; no chain check.
+   */
+  registerContractClass(artifact: ContractArtifact): Promise<void>;
   simulateTx(exec: ExecutionPayload, opts: SimulateOptions): Promise<TxSimulationResultWithAppOffset>;
   executeUtility(call: FunctionCall, opts: ExecuteUtilityOptions): Promise<UtilityExecutionResult>;
   profileTx(exec: ExecutionPayload, opts: ProfileOptions): Promise<TxProfileResult>;
@@ -337,6 +343,7 @@ export const SimulateOptionsSchema = z.object({
   includeMetadata: optional(z.boolean()),
   additionalScopes: optional(z.array(schemas.AztecAddress)),
   stateOverrides: optional(StateOverridesSchema),
+  contractOverrides: optional(z.array(ContractInstanceWithAddressSchema)),
 });
 
 export const ProfileOptionsSchema = SimulateOptionsSchema.extend({
@@ -559,6 +566,7 @@ const WalletMethodSchemas = {
     .function()
     .args(ContractInstanceWithAddressSchema, optional(ContractArtifactSchema), optional(schemas.Fr))
     .returns(ContractInstanceWithAddressSchema),
+  registerContractClass: z.function().args(ContractArtifactSchema).returns(z.void()),
   simulateTx: z
     .function()
     .args(ExecutionPayloadSchema, SimulateOptionsSchema)