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. For private-call simulation, 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.
+
+For private-call simulation against the new class, register the artifact in your local PXE first via `wallet.registerContractClass(artifact)` so PXE can find its ACIR. Public-only simulations don't need this step.
 
 ```typescript
 import { fastForwardContractUpdate } from '@aztec/aztec.js';
+import { getContractClassFromArtifact } from '@aztec/aztec.js/contracts';
+
+// One-time, only needed if simulating private calls on the new class
+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. For private dispatch to find the new class's ACIR, register the artifact locally first via `wallet.registerContractClass(artifact)` (newly exposed for this purpose).
+
+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). For private-call simulation, register the new artifact locally via `wallet.registerContractClass(artifact)` first so PXE can find its ACIR:
 
 ```typescript
-import { fastForwardContractUpdate } from '@aztec/aztec.js';
+import { fastForwardContractUpdate, getContractClassFromArtifact } from '@aztec/aztec.js/contracts';
+
+await wallet.registerContractClass(UpdatedContract.artifact); // only needed for private-call sim
 
-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,30 @@ 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. For private-call simulation against the new class, the artifact
+ * must also be registered in the local PXE — call `wallet.registerContractClass(artifact)` once before invoking this
+ * helper. Public-only simulations don't need the local artifact registration.
  *
  * @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 +43,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 +66,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);
   }