feat: fastForwardContractUpdate cheatcode for simulating contract updates

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:

noir-projects/noir-contracts/contracts/account/simulated_ecdsa_account_contract/Nargo.toml

@@ -6,3 +6,4 @@ type = "contract"
 
 [dependencies]
 aztec = { path = "../../../../aztec-nr/aztec" }
+ecdsa_public_key_note = { path = "../../libs/ecdsa_public_key_note" }

noir-projects/noir-contracts/contracts/account/simulated_ecdsa_account_contract/src/main.nr

@@ -3,17 +3,26 @@ use aztec::macros::aztec;
 // Stub account contract for ECDSA accounts (both secp256k1 and secp256r1) used during simulation.
 // Matches the constructor signature of EcdsaKAccount / EcdsaRAccount so that deployment
 // simulations using this stub as an override do not fail on selector lookup.
-// See simulated_account_contract for the base stub without a constructor.
+// Mirrors the EcdsaKAccount/EcdsaRAccount storage layout so simulation sync can decode
+// the real account's signing_public_key note.
 #[aztec]
 pub contract SimulatedEcdsaAccount {
     use aztec::{
         authwit::{account::AccountActions, auth::IS_VALID_SELECTOR, entrypoint::app::AppPayload},
         context::PrivateContext,
-        macros::functions::{allow_phase_change, external, view},
+        macros::{functions::{allow_phase_change, external, view}, storage::storage},
         messages::encoding::MESSAGE_CIPHERTEXT_LEN,
         oracle::random::random,
+        state_vars::SinglePrivateImmutable,
     };
 
+    use ecdsa_public_key_note::EcdsaPublicKeyNote;
+
+    #[storage]
+    struct Storage<Context> {
+        signing_public_key: SinglePrivateImmutable<EcdsaPublicKeyNote, Context>,
+    }
+
     // Stub constructor matching the EcdsaKAccount / EcdsaRAccount constructor signature.
     // Does NOT use #[initializer] so that the macro does not inject
     // assert_initialization_matches_address_preimage_private, which would fail during kernelless
@@ -67,9 +76,4 @@ pub contract SimulatedEcdsaAccount {
     fn is_valid_impl(_context: &mut PrivateContext, _outer_hash: Field) -> bool {
         true
     }
-
-    #[external("utility")]
-    unconstrained fn sync_state() {
-        assert(false, "BUG ALERT: sync_state on a simulated account contract should never be triggered.");
-    }
 }

noir-projects/noir-contracts/contracts/account/simulated_schnorr_account_contract/src/main.nr

@@ -1,19 +1,30 @@
+mod public_key_note;
+
 use aztec::macros::aztec;
 
 // Stub account contract for Schnorr accounts used during simulation.
 // Matches the constructor signature of SchnorrAccount so that deployment
 // simulations using this stub as an override do not fail on selector lookup.
-// See simulated_account_contract for the base stub without a constructor.
+// Mirrors the SchnorrAccount storage layout so simulation sync can decode
+// the real account's signing_public_key note.
 #[aztec]
 pub contract SimulatedSchnorrAccount {
     use aztec::{
         authwit::{account::AccountActions, auth::IS_VALID_SELECTOR, entrypoint::app::AppPayload},
         context::PrivateContext,
-        macros::functions::{allow_phase_change, external, view},
+        macros::{functions::{allow_phase_change, external, view}, storage::storage},
         messages::encoding::MESSAGE_CIPHERTEXT_LEN,
         oracle::random::random,
+        state_vars::SinglePrivateImmutable,
     };
 
+    use crate::public_key_note::PublicKeyNote;
+
+    #[storage]
+    struct Storage<Context> {
+        signing_public_key: SinglePrivateImmutable<PublicKeyNote, Context>,
+    }
+
     // Stub constructor matching the SchnorrAccount constructor signature.
     // Does NOT use #[initializer] so that the macro does not inject
     // assert_initialization_matches_address_preimage_private, which would fail during kernelless
@@ -67,9 +78,4 @@ pub contract SimulatedSchnorrAccount {
     fn is_valid_impl(_context: &mut PrivateContext, _outer_hash: Field) -> bool {
         true
     }
-
-    #[external("utility")]
-    unconstrained fn sync_state() {
-        assert(false, "BUG ALERT: sync_state on a simulated account contract should never be triggered.");
-    }
 }

noir-projects/noir-contracts/contracts/account/simulated_schnorr_account_contract/src/public_key_note.nr

@@ -0,0 +1,10 @@
+use aztec::{macros::notes::note, protocol::traits::Packable};
+
+// Mirrors PublicKeyNote in schnorr_account_contract so the stub's auto-generated _compute_note_hash
+// can decode the real account's signing_public_key note during simulation sync.
+#[derive(Eq, Packable)]
+#[note]
+pub struct PublicKeyNote {
+    pub x: Field,
+    pub y: Field,
+}

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';