feat: fastForwardContractUpdate cheatcode for simulating contract updates

Adds a high-level helper that returns a `SimulationOverrides` blob simulating a deployed instance as if it had already been upgraded to a new contract class:

- `overrides.publicStorage` rewrites the `ContractInstanceRegistry`'s delayed-public-mutable storage so the AVM's `UpdateCheck` resolves to the new class id.

- `overrides.contracts` swaps the deployed instance for one whose `currentContractClassId` is bumped to the new class. Drives both AVM-side public dispatch and PXE-side ACIR private dispatch.

Both pieces are required: a storage-only override would not redirect the AVM's class dispatch (which reads `currentContractClassId` from the contract DB); an instance-only override would cause the witgen `UpdateCheck` to throw on inconsistency. The new class must already be registered on chain.

Author: dbanks12

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

@@ -83,6 +83,24 @@ const result = await contract.methods.read_balance(account).simulate({
 
 Use this to set up state preconditions, reproduce production bugs against pinned storage, or exercise rare value branches without orchestrating the contract calls that produce them.
 
+### Fast-forwarding a contract update
+
+`fastForwardContractUpdate` returns a `SimulationOverrides` object that simulates 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

@@ -185,6 +185,19 @@ Direct callers of the `SimulationOverrides` constructor must switch from a posit
 + new SimulationOverrides({ contracts });
 ```
 
+`overrides.contracts` 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 onchain upgrade flow, use the `fastForwardContractUpdate` helper which returns a `SimulationOverrides` covering both registry storage rewrites and the upgraded instance entry:
+
+```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,
@@ -1508,7 +1508,11 @@ 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());
+      if (overrides?.contracts) {
+        contractsDB.addContracts(Object.values(overrides.contracts).map(({ instance }) => instance));
+      }
+      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

@@ -76,6 +76,17 @@ describe('Contract Class', () => {
     expect(result).toBe(42n);
   });
 
+  it('throws when overrides are passed to a utility function simulation', async () => {
+    const fooContract = Contract.at(contractAddress, testContractArtifact, wallet);
+    await expect(
+      fooContract.methods.qux(123n).simulate({
+        from: account.getAddress(),
+        overrides: { 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.overrides?.publicStorage?.length || options.overrides?.contracts) {
+        throw new Error('overrides 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,91 @@
+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 overrides = await fastForwardContractUpdate({ instanceAddress, newClassId, node });
+
+    const upgraded = overrides.contracts?.[instanceAddress.toString()];
+    expect(upgraded).toBeDefined();
+    expect(upgraded!.instance.address).toEqual(instanceAddress);
+    expect(upgraded!.instance.currentContractClassId).toEqual(newClassId);
+    expect(upgraded!.instance.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(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(
+      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,71 @@
+import { Fr } from '@aztec/foundation/curves/bn254';
+import { ProtocolContractAddress } from '@aztec/protocol-contracts';
+import type { AztecAddress } from '@aztec/stdlib/aztec-address';
+import {
+  DelayedPublicMutableValuesWithHash,
+  ScheduledDelayChange,
+  ScheduledValueChange,
+} from '@aztec/stdlib/delayed-public-mutable';
+import type { AztecNode } from '@aztec/stdlib/interfaces/client';
+import { SimulationOverrides } from '@aztec/stdlib/tx';
+
+/**
+ * Builds `SimulationOverrides` that simulate a deployed instance as if it had already been upgraded to a
+ * new contract class. Mirrors a real on-chain upgrade (`pxe.updateContract` followed by waiting out the delay):
+ *
+ * - `publicStorage` rewrites the `ContractInstanceRegistry`'s delayed-public-mutable storage so the AVM's
+ *   `UpdateCheck` resolves to the new class id.
+ * - `contracts` swaps the deployed instance for one whose `currentContractClassId` is bumped to the new class.
+ *
+ * The new class must already be registered on chain.
+ *
+ * @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: {
+  /** Address of the deployed instance to upgrade. */
+  instanceAddress: AztecAddress;
+  /** ID of the (already-registered) class to upgrade to. */
+  newClassId: Fr;
+  /** Node used to fetch the existing instance and validate the class is registered. */
+  node: AztecNode;
+}): Promise<SimulationOverrides> {
+  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.
+  const svc = new ScheduledValueChange(/*previous=*/ [new Fr(0)], /*post=*/ [newClassId], /*timestampOfChange=*/ 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 = { ...instance, currentContractClassId: newClassId };
+
+  return new SimulationOverrides({
+    publicStorage,
+    contracts: { [instanceAddress.toString()]: { instance: upgradedInstance } },
+  });
+}

yarn-project/end-to-end/src/e2e_avm_simulator.test.ts

@@ -3,8 +3,7 @@ import { BatchCall, type ContractInstanceWithAddress } from '@aztec/aztec.js/con
 import { Fr } from '@aztec/aztec.js/fields';
 import type { AztecNode } from '@aztec/aztec.js/node';
 import { TxExecutionResult } from '@aztec/aztec.js/tx';
-import type { PublicStorageOverride } from '@aztec/aztec.js/wallet';
-import type { Wallet } from '@aztec/aztec.js/wallet';
+import type { PublicStorageOverride, Wallet } from '@aztec/aztec.js/wallet';
 import { AvmInitializerTestContract } from '@aztec/noir-test-contracts.js/AvmInitializerTest';
 import { AvmTestContract } from '@aztec/noir-test-contracts.js/AvmTest';
 

yarn-project/end-to-end/src/e2e_contract_updates.test.ts

@@ -1,5 +1,5 @@
 import { getSchnorrAccountContractAddress } from '@aztec/accounts/schnorr';
-import { getContractClassFromArtifact } from '@aztec/aztec.js/contracts';
+import { fastForwardContractUpdate, getContractClassFromArtifact } from '@aztec/aztec.js/contracts';
 import { publishContractClass } from '@aztec/aztec.js/deployment';
 import { Fr } from '@aztec/aztec.js/fields';
 import type { AztecNode } from '@aztec/aztec.js/node';
@@ -177,4 +177,59 @@ describe('e2e_contract_updates', () => {
       'Could not update contract to a class different from the current one',
     );
   });
+
+  // UpdatableContract's `set_public_value(Field)` and UpdatedContract's `set_public_value()`
+  // have different function selectors. Without an upgrade, only the deployed Updatable's
+  // (Field) selector exists; with a fastForwardContractUpdate override, the AVM dispatches
+  // against UpdatedContract's bytecode and the no-args selector resolves.
+  it('fastForwardContractUpdate enables simulation of post-upgrade public calls', async () => {
+    // Local construction with the new artifact - no PXE/wallet side effect, no chain mutation.
+    const updatedContract = UpdatedContract.at(contract.address, wallet);
+
+    // Without overrides, UpdatedContract's no-args selector doesn't match the deployed class.
+    await expect(
+      updatedContract.methods.set_public_value().simulate({ from: defaultAccountAddress }),
+    ).rejects.toThrow();
+
+    // With the fastForwardContractUpdate overrides, the AVM dispatches against UpdatedContract's
+    // bytecode and the call simulates successfully.
+    const overrides = await fastForwardContractUpdate({
+      instanceAddress: contract.address,
+      newClassId: updatedContractClassId,
+      node: aztecNode,
+    });
+    await expect(
+      updatedContract.methods.set_public_value().simulate({ from: defaultAccountAddress, overrides }),
+    ).resolves.toBeDefined();
+
+    // Chain state is untouched: the original Updatable's set_public_value(Field) still simulates fine.
+    await expect(
+      contract.methods.set_public_value(5678n).simulate({ from: defaultAccountAddress }),
+    ).resolves.toBeDefined();
+  });
+
+  // UpdatedContract.set_private_value is a private function that doesn't exist on UpdatableContract.
+  // For PXE-side ACIR dispatch to find it, the artifact must be registered locally first via
+  // wallet.registerContractClass; the helper itself only takes the class id.
+  it('fastForwardContractUpdate enables simulation of post-upgrade private calls', async () => {
+    const updatedContract = UpdatedContract.at(contract.address, wallet);
+
+    // Without overrides (and without local artifact registration), the new private function isn't
+    // available on the deployed class.
+    await expect(
+      updatedContract.methods.set_private_value().simulate({ from: defaultAccountAddress }),
+    ).rejects.toThrow();
+
+    // Register the new artifact in the local PXE so the ACIR simulator can find its private functions.
+    await wallet.registerContractClass(UpdatedContract.artifact);
+
+    const overrides = await fastForwardContractUpdate({
+      instanceAddress: contract.address,
+      newClassId: updatedContractClassId,
+      node: aztecNode,
+    });
+    await expect(
+      updatedContract.methods.set_private_value().simulate({ from: defaultAccountAddress, overrides }),
+    ).resolves.toBeDefined();
+  });
 });