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 produces a coherent set
of state overrides covering both the contract instance (bumped
currentContractClassId) and the ContractInstanceRegistry's delayed-public-
mutable storage (rewritten so the AVM's UpdateCheck resolves to the new
class). The new class must already be registered on chain.

Also adds the minimum plumbing to the StateOverrides type and
PublicContractsDB to support contract-instance overrides; contract-class
overrides are introduced separately upstack.

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 stateOverrides = await fastForwardContractUpdate({
+  instanceAddress: contract.address,
+  newClassId: upgradedClass.id,
+  node,
+});
+
+const result = await contract.methods.upgraded_method().simulate({ stateOverrides });
+```
+
+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 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:
+
+```typescript
+import { fastForwardContractUpdate } from '@aztec/aztec.js';
+
+const stateOverrides = await fastForwardContractUpdate({
+  instanceAddress: contract.address,
+  newClassId: upgradedClass.id,
+  node,
+});
+const result = await contract.methods.upgraded_method().simulate({ stateOverrides });
+```
+
 ### [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,
@@ -1549,7 +1549,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(stateOverrides?.contractInstances);
+      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 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(/stateOverrides are 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,12 @@ 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.');
+      }
       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,90 @@
+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 });
+
+    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);
+
+    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,70 @@
+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, StateOverrides } from '@aztec/stdlib/interfaces/client';
+
+/**
+ * Builds `StateOverrides` 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.
+ *
+ * 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).
+ *
+ * @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.
+ */
+export async function fastForwardContractUpdate(args: {
+  instanceAddress: AztecAddress;
+  newClassId: Fr;
+  node: AztecNode;
+}): Promise<StateOverrides> {
+  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; register 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 = { ...instance, currentContractClassId: newClassId };
+
+  return {
+    publicStorage,
+    contractInstances: [upgradedInstance],
+  };
+}

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';
@@ -178,4 +178,34 @@ 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 override, the AVM dispatches against UpdatedContract's
+    // bytecode and the call simulates successfully.
+    const stateOverrides = await fastForwardContractUpdate({
+      instanceAddress: contract.address,
+      newClassId: updatedContractClassId,
+      node: aztecNode,
+    });
+    await expect(
+      updatedContract.methods.set_public_value().simulate({ from: defaultAccountAddress, stateOverrides }),
+    ).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();
+  });
 });

yarn-project/simulator/src/public/public_db_sources.ts

@@ -55,7 +55,8 @@ export class PublicContractsDB implements PublicContractsDBInterface {
     this.log = createLogger('simulator:contracts-data-source', bindings);
   }
 
-  public addContracts(contractDeploymentData: ContractDeploymentData): void {
+  /** Parses raw log data from the C++/NAPI bridge and inserts the resulting contracts into the current checkpoint. */
+  public addContractsFromLogs(contractDeploymentData: ContractDeploymentData): void {
     const currentState = this.getCurrentState();
 
     this.addContractClassesFromEvents(
@@ -71,8 +72,16 @@ export class PublicContractsDB implements PublicContractsDBInterface {
 
   public addNewContracts(tx: Tx): void {
     const contractDeploymentData = AllContractDeploymentData.fromTx(tx);
-    this.addContracts(contractDeploymentData.getNonRevertibleContractDeploymentData());
-    this.addContracts(contractDeploymentData.getRevertibleContractDeploymentData());
+    this.addContractsFromLogs(contractDeploymentData.getNonRevertibleContractDeploymentData());
+    this.addContractsFromLogs(contractDeploymentData.getRevertibleContractDeploymentData());
+  }
+
+  /** Inserts typed contract instances directly into the current checkpoint. */
+  public addContracts(contractInstances?: ContractInstanceWithAddress[]): void {
+    const currentState = this.getCurrentState();
+    for (const instance of contractInstances ?? []) {
+      currentState.addInstance(instance.address, instance);
+    }
   }
 
   /**