feat: fastForwardContractUpdate cheatcode for simulating contract updates

Adds a high-level helper that simulates a deployed instance as if it had been upgraded to a new contract class. The helper returns a coherent set of overrides covering both:

- `stateOverrides` — public-storage writes to the ContractInstanceRegistry's delayed-public-mutable storage, so the AVM's UpdateCheck resolves to the new class id.

- `contractOverrides` — an instance entry whose `currentContractClassId` is bumped to the new class.

Plumbs `contractOverrides` as a top-level option on `SimulateInteractionOptions` through wallet → PXE → `AztecNode.simulatePublicCalls`. The new class must already be registered on chain.

Author: dbanks12

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

@@ -87,6 +87,24 @@ Use this to:
 - 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
 
+### 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.
+
+```typescript
+import { fastForwardContractUpdate } from '@aztec/aztec.js';
+
+const overrides = await fastForwardContractUpdate({
+  instanceAddress: contract.address,
+  newClassId: upgradedClass.id,
+  node,
+});
+
+const result = await contract.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.
+
 ## Further reading
 
 - [How to read contract data](./how_to_read_data.md)

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

@@ -23,6 +23,19 @@ const result = await contract.methods.read_balance(account).simulate({
 
 The same option flows through `wallet.simulateTx` and eventually to `simulatePublicCalls` RPC on `AztecNode`.
 
+A separate top-level option, `contractOverrides`, swaps 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 returns both `stateOverrides` and `contractOverrides`:
+
+```typescript
+import { fastForwardContractUpdate } from '@aztec/aztec.js';
+
+const overrides = await fastForwardContractUpdate({
+  instanceAddress: contract.address,
+  newClassId: upgradedClass.id,
+  node,
+});
+const result = await contract.methods.upgraded_method().simulate({ ...overrides });
+```
+
 ### [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/server.ts

@@ -39,7 +39,7 @@ import {
   SequencerClient,
   type SequencerPublisher,
 } from '@aztec/sequencer-client';
-import { PublicProcessorFactory } from '@aztec/simulator/server';
+import { PublicContractsDB, PublicProcessorFactory } from '@aztec/simulator/server';
 import {
   AttestationsBlockWatcher,
   EpochPruneWatcher,
@@ -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 applied to the ephemeral fork before simulation.
+   * @param contractOverrides - Optional contract instance overrides applied to the contract DB.
    **/
   @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;
@@ -1549,7 +1551,9 @@ export class AztecNodeService implements AztecNode, AztecNodeAdmin, AztecNodeDeb
           maxDebugLogMemoryReads: this.config.rpcSimulatePublicMaxDebugLogMemoryReads,
         }),
       });
-      const processor = publicProcessorFactory.create(merkleTreeFork, newGlobalVariables, config);
+      const contractsDB = new PublicContractsDB(this.contractDataSource, this.log.getBindings());
+      contractsDB.addContracts(contractOverrides);
+      const processor = publicProcessorFactory.create(merkleTreeFork, newGlobalVariables, config, contractsDB);
 
       // REFACTOR: Consider merging ProcessReturnValues into ProcessedTx
       const [processedTxs, failedTxs, _usedTxs, returns, debugLogs] = await processor.process([tx]);

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

@@ -74,6 +74,7 @@ export {
 } from '../contract/deploy_method.js';
 export { waitForProven, type WaitForProvenOpts, DefaultWaitForProvenOpts } from '../contract/wait_for_proven.js';
 export { getGasLimits } from '../contract/get_gas_limits.js';
+export { fastForwardContractUpdate } from '../contract/fastforward_contract_update.js';
 
 export {
   type PartialAddress,

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

@@ -231,6 +231,17 @@ describe('Contract Class', () => {
     expect(result).toBe(42n);
   });
 
+  it('throws when stateOverrides or 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(),
+        stateOverrides: { publicStorage: [{ contract: contractAddress, slot: new Fr(1), value: new Fr(42) }] },
+      }),
+    ).rejects.toThrow(/not supported for utility/);
+    expect(wallet.executeUtility).not.toHaveBeenCalled();
+  });
+
   it('should extract offchain messages with anchor block timestamp on simulate', async () => {
     const recipient = await AztecAddress.random();
     const msgPayload = [Fr.random(), Fr.random()];

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

@@ -129,6 +129,9 @@ export class ContractFunctionInteraction extends BaseContractInteraction {
   ): Promise<SimulationResult> {
     // docs:end:simulate
     if (this.functionDao.functionType == FunctionType.UTILITY) {
+      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 ?? [])];
       const utilityResult = await this.wallet.executeUtility(call, {

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

@@ -0,0 +1,94 @@
+import { Fr } from '@aztec/foundation/curves/bn254';
+import { ProtocolContractAddress } from '@aztec/protocol-contracts';
+import { AztecAddress } from '@aztec/stdlib/aztec-address';
+import { SerializableContractInstance } 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 { type MockProxy, mock } from 'jest-mock-extended';
+
+import { fastForwardContractUpdate } from './fastforward_contract_update.js';
+
+describe('fastForwardContractUpdate', () => {
+  let node: MockProxy<AztecNode>;
+  let instanceAddress: AztecAddress;
+  let originalClassId: Fr;
+  let newClassId: Fr;
+
+  beforeEach(async () => {
+    node = mock<AztecNode>();
+    instanceAddress = await AztecAddress.random();
+    originalClassId = Fr.random();
+    newClassId = Fr.random();
+
+    const instance = (
+      await SerializableContractInstance.random({
+        currentContractClassId: originalClassId,
+        originalContractClassId: originalClassId,
+      })
+    ).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);
+  });
+
+  it('produces overrides with bumped currentContractClassId and registry storage writes', async () => {
+    const { stateOverrides, 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);
+
+    const expectedSlots = await DelayedPublicMutableValuesWithHash.getContractUpdateSlots(instanceAddress);
+    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(stateOverrides.publicStorage![0].slot).toEqual(baseSlot);
+    expect(stateOverrides.publicStorage![stateOverrides.publicStorage!.length - 1].slot).toEqual(
+      expectedSlots.delayedPublicMutableHashSlot,
+    );
+  });
+
+  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 () => {
+    node.getContractClass.mockResolvedValue(undefined);
+    await expect(fastForwardContractUpdate({ instanceAddress, newClassId, node })).rejects.toThrow(/not registered/);
+  });
+
+  it('throws when the instance is already on the target class', async () => {
+    const sameClassInstance = (
+      await SerializableContractInstance.random({
+        currentContractClassId: newClassId,
+        originalContractClassId: originalClassId,
+      })
+    ).withAddress(instanceAddress);
+    node.getContract.mockResolvedValue(sameClassInstance);
+
+    await expect(fastForwardContractUpdate({ instanceAddress, newClassId, node })).rejects.toThrow(/already on class/);
+  });
+});

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

@@ -0,0 +1,72 @@
+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,
+  ScheduledValueChange,
+} from '@aztec/stdlib/delayed-public-mutable';
+import type { AztecNode, StateOverrides } from '@aztec/stdlib/interfaces/client';
+
+/**
+ * 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). Returns:
+ *
+ * - `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.
+ *
+ * The new class must already be registered on chain.
+ *
+ * @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, 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: StateOverrides; contractOverrides: ContractInstanceWithAddress[] }> {
+  const { instanceAddress, newClassId, node } = args;
+
+  const instance = await node.getContract(instanceAddress);
+  if (!instance) {
+    throw new Error(`Instance not deployed at ${instanceAddress}; deploy it before fast-forwarding an update`);
+  }
+
+  const klass = await node.getContractClass(newClassId);
+  if (!klass) {
+    throw new Error(`Contract class ${newClassId} is not registered on chain; publish it before fast-forwarding to it`);
+  }
+
+  if (instance.currentContractClassId.equals(newClassId)) {
+    throw new Error(`Instance ${instanceAddress} is already on class ${newClassId}; nothing to fast-forward`);
+  }
+
+  // Build the SVC the same way `ContractInstanceRegistry::update` would have, but with a timestamp_of_change
+  // safely in the past so the AVM's UpdateCheck resolves to the post-upgrade class id at any sim timestamp.
+  // `pre = 0` is the canonical "first upgrade" form: the C++ check falls back to `originalContractClassId`.
+  const svc = new ScheduledValueChange([new Fr(0)], [newClassId], 1n);
+  const sdc = ScheduledDelayChange.empty();
+  const dpmv = new DelayedPublicMutableValuesWithHash(svc, sdc);
+
+  const { delayedPublicMutableSlot } = await DelayedPublicMutableValuesWithHash.getContractUpdateSlots(instanceAddress);
+  const fields = await dpmv.toFields();
+
+  const publicStorage = fields.map((value, i) => ({
+    contract: ProtocolContractAddress.ContractInstanceRegistry,
+    slot: delayedPublicMutableSlot.add(new Fr(i)),
+    value,
+  }));
+
+  const upgradedInstance: ContractInstanceWithAddress = { ...instance, currentContractClassId: newClassId };
+
+  return {
+    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,13 @@ 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.
+   */
+  contractOverrides?: ContractInstanceWithAddress[];
 };
 
 /**

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

@@ -337,6 +337,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({