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

@@ -89,18 +89,21 @@ Use this to:
 
 ### 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 two coherent override layers - `stateOverrides` (node-side, redirects public-call dispatch) and `contractOverrides` (PXE-side, redirects private-call 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,22 @@ 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:
+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.
+
+`.simulate(...)` also accepts `contractOverrides`, which replaces the (instance, artifact) pair the PXE-side ACIR simulator uses for a given address. Required to redirect private-function dispatch when simulating against an upgraded class.
+
+To simulate a complete on-chain upgrade flow, use the `fastForwardContractUpdate` helper. It returns both `stateOverrides` (node-side, public dispatch) and `contractOverrides` (PXE-side, private dispatch) 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.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,11 @@ 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 &&
+          (options.stateOverrides.publicStorage?.length || options.stateOverrides.contractInstances?.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,47 @@ 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 });
+  it('produces stateOverrides with bumped currentContractClassId and registry storage writes', async () => {
+    const { stateOverrides } = await fastForwardContractUpdate({ instanceAddress, newArtifact, 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);
+    expect(stateOverrides.contractInstances).toHaveLength(1);
+    expect(stateOverrides.contractInstances![0].address).toEqual(instanceAddress);
+    expect(stateOverrides.contractInstances![0].currentContractClassId).toEqual(newClassId);
+    expect(stateOverrides.contractInstances![0].originalContractClassId).toEqual(originalClassId);
 
     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 +88,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,17 @@ export async function fastForwardContractUpdate(args: {
 
   const upgradedInstance = { ...instance, currentContractClassId: newClassId };
 
-  return {
+  const stateOverrides: StateOverrides = {
     publicStorage,
     contractInstances: [upgradedInstance],
   };
+
+  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({