feat: contractOverrides for private-call simulation

Extends `contractOverrides` (introduced downstack) to also drive PXE-side ACIR dispatch, so private function calls on an overridden instance simulate against the override's class instead of the deployed one.

- Translates aztec.js's `contractOverrides: ContractInstanceWithAddress[]` into PXE's existing `SimulationOverrides.contracts` record at the wallet boundary.

- Relaxes the PXE-internal `ContractOverrides` type so `artifact` is optional; `proxied_contract_data_source` falls through to the regular `ContractStore` when the override has no artifact.

- Exposes `wallet.registerContractClass(artifact)` so users can register the new class artifact locally before the override is used.

Author: dbanks12

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

@@ -67,9 +67,12 @@ Test that invalid operations revert as expected:
 
 Use `.simulate()` to test reverts without spending gas. The simulation will throw if the transaction would fail onchain.
 
-## Simulating with state overrides
+## Simulating with state and contract overrides
 
-`.simulate()` accepts a `stateOverrides` option that injects values into the simulator's ephemeral world-state fork before the call runs. The override is scoped to that single simulation; the real chain state is untouched.
+`.simulate()` accepts two override options that are scoped to that single simulation; real chain state is untouched.
+
+- `stateOverrides`: state-tree overrides (e.g. `publicStorage` writes).
+- `contractOverrides`: an array of `ContractInstanceWithAddress` to override deployed contract instances. Register the new class artifact locally first via `wallet.registerContractClass(artifact)`.
 
 Override a public-storage slot:
 
@@ -81,26 +84,37 @@ const result = await contract.methods.read_balance(account).simulate({
 });
 ```
 
-Use this to:
+Use these to:
 
 - Set up state preconditions without running a full setup transaction
 - Reproduce a bug from production by pinning storage to the values seen at a specific block
+- Simulate a contract instance as if it had been upgraded
 - 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.
+`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 `stateOverrides` (registry storage rewrites) and `contractOverrides` (instance with bumped class id). A single spread covers any mix of private and public function calls on the upgraded contract.
+
+Register the new class artifact in your local PXE first via `wallet.registerContractClass(artifact)`.
 
 ```typescript
 import { fastForwardContractUpdate } from '@aztec/aztec.js';
+import { getContractClassFromArtifact } from '@aztec/aztec.js/contracts';
+
+// One-time local PXE registration of the new class artifact
+await wallet.registerContractClass(UpdatedContract.artifact);
 
+const newClassId = (await getContractClassFromArtifact(UpdatedContract.artifact)).id;
 const overrides = await fastForwardContractUpdate({
   instanceAddress: contract.address,
-  newClassId: upgradedClass.id,
+  newClassId,
   node,
 });
 
-const result = await contract.methods.upgraded_method().simulate({ ...overrides });
+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,23 @@ 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`:
+`.simulate(...)` also accepts `contractOverrides` - an array of `ContractInstanceWithAddress` that override deployed instances at their addresses. Applied on both the AVM (public dispatch) and PXE (private dispatch via ACIR) sides. Register the new class artifact locally first via `wallet.registerContractClass(artifact)` (newly exposed for this purpose) so PXE has the ACIR.
+
+To simulate a complete on-chain upgrade flow, use the `fastForwardContractUpdate` helper. It returns both `stateOverrides` (registry storage rewrites) and `contractOverrides` (instance with bumped class id). Register the new artifact locally via `wallet.registerContractClass(artifact)` first:
 
 ```typescript
-import { fastForwardContractUpdate } from '@aztec/aztec.js';
+import { fastForwardContractUpdate, getContractClassFromArtifact } from '@aztec/aztec.js/contracts';
+
+await wallet.registerContractClass(UpdatedContract.artifact);
 
+const newClassId = (await getContractClassFromArtifact(UpdatedContract.artifact)).id;
 const overrides = await fastForwardContractUpdate({
   instanceAddress: contract.address,
-  newClassId: upgradedClass.id,
+  newClassId,
   node,
 });
-const result = await contract.methods.upgraded_method().simulate({ ...overrides });
+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/wallet/wallet.test.ts

@@ -155,6 +155,18 @@ describe('WalletSchema', () => {
     });
   });
 
+  it('registerContractClass', async () => {
+    const mockArtifact: ContractArtifact = {
+      name: 'TestContract',
+      functions: [],
+      nonDispatchPublicFunctions: [],
+      outputs: { structs: {}, globals: {} },
+      fileMap: {},
+      storageLayout: {},
+    };
+    await context.client.registerContractClass(mockArtifact);
+  });
+
   it('simulateTx', async () => {
     const exec: ExecutionPayload = {
       calls: [],
@@ -448,6 +460,8 @@ class MockWallet implements Wallet {
     };
   }
 
+  async registerContractClass(_artifact: any): Promise<void> {}
+
   async simulateTx(_exec: ExecutionPayload, _opts: SimulateOptions): Promise<TxSimulationResultWithAppOffset> {
     return TxSimulationResultWithAppOffset.fromResultAndOffset(await TxSimulationResult.random(), 0);
   }

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

@@ -272,6 +272,12 @@ export type Wallet = {
     artifact?: ContractArtifact,
     secretKey?: Fr,
   ): Promise<ContractInstanceWithAddress>;
+  /**
+   * Registers a contract class artifact in the local PXE without binding it to any instance.
+   * Useful for simulation flows that need the artifact available locally before any on-chain
+   * upgrade has taken effect. The artifact's class id must match its content; no chain check.
+   */
+  registerContractClass(artifact: ContractArtifact): Promise<void>;
   simulateTx(exec: ExecutionPayload, opts: SimulateOptions): Promise<TxSimulationResultWithAppOffset>;
   executeUtility(call: FunctionCall, opts: ExecuteUtilityOptions): Promise<UtilityExecutionResult>;
   profileTx(exec: ExecutionPayload, opts: ProfileOptions): Promise<TxProfileResult>;
@@ -560,6 +566,7 @@ const WalletMethodSchemas = {
     .function()
     .args(ContractInstanceWithAddressSchema, optional(ContractArtifactSchema), optional(schemas.Fr))
     .returns(ContractInstanceWithAddressSchema),
+  registerContractClass: z.function().args(ContractArtifactSchema).returns(z.void()),
   simulateTx: z
     .function()
     .args(ExecutionPayloadSchema, SimulateOptionsSchema)

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

@@ -208,4 +208,29 @@ describe('e2e_contract_updates', () => {
       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();
+  });
 });

yarn-project/end-to-end/src/test-wallet/worker_wallet.ts

@@ -166,6 +166,10 @@ export class WorkerWallet implements Wallet {
     return this.call('registerContract', instance, artifact, secretKey);
   }
 
+  registerContractClass(artifact: ContractArtifact): Promise<void> {
+    return this.call('registerContractClass', artifact);
+  }
+
   simulateTx(exec: ExecutionPayload, opts: SimulateOptions): Promise<TxSimulationResultWithAppOffset> {
     return this.call('simulateTx', exec, opts);
   }

yarn-project/pxe/src/contract_function_simulator/proxied_contract_data_source.ts

@@ -38,9 +38,9 @@ export class ProxiedContractStoreFactory {
           }
           case 'getFunctionArtifact': {
             return async (contractAddress: AztecAddress, selector: FunctionSelector) => {
-              if (overrides[contractAddress.toString()]) {
-                const { artifact } = overrides[contractAddress.toString()]!;
-                const functions = artifact.functions;
+              const override = overrides[contractAddress.toString()];
+              if (override?.artifact) {
+                const functions = override.artifact.functions;
                 for (let i = 0; i < functions.length; i++) {
                   const fn = functions[i];
                   const fnSelector = await FunctionSelector.fromNameAndParameters(fn.name, fn.parameters);
@@ -58,9 +58,9 @@ export class ProxiedContractStoreFactory {
           }
           case 'getFunctionArtifactWithDebugMetadata': {
             return async (contractAddress: AztecAddress, selector: FunctionSelector) => {
-              if (overrides[contractAddress.toString()]) {
-                const { artifact } = overrides[contractAddress.toString()]!;
-                const functions = artifact.functions;
+              const override = overrides[contractAddress.toString()];
+              if (override?.artifact) {
+                const functions = override.artifact.functions;
                 for (let i = 0; i < functions.length; i++) {
                   const fn = functions[i];
                   const fnSelector = await FunctionSelector.fromNameAndParameters(fn.name, fn.parameters);

yarn-project/stdlib/src/tx/simulated_tx.ts

@@ -24,15 +24,28 @@ import { NestedProcessReturnValues, PublicSimulationOutput } from './public_simu
 import { Tx } from './tx.js';
 
 /*
- * If passed during the execution of a user circuit, the contract function simulator will replace the instance and class
- * of the contract with the one provided in the overrides for that address. An example use case
- * would be overriding your own account contract so that valid signatures don't have to be provided while simulating.
+ * Per-address contract DB overrides applied during simulation. For each address, replaces the
+ * (instance, optional artifact) pair the simulator uses for that address:
+ *
+ * - `instance` is always applied. The AVM-side PublicContractsDB uses it for public-call dispatch;
+ *   the PXE-side ACIR simulator uses it to resolve address → class id.
+ * - `artifact` is optional. When present, the PXE-side simulator uses the override artifact for
+ *   ACIR/function lookups during private execution. When absent, PXE falls back to whatever artifact
+ *   is registered for the override's `instance.currentContractClassId`.
+ *
+ * Common use cases: simulating an upgraded class without scheduling a real upgrade; overriding your
+ * own account contract so signatures don't need to be valid during simulation.
  */
 export type ContractOverrides = Record<
   string /* AztecAddress as string */,
-  { instance: ContractInstanceWithAddress; artifact: ContractArtifact }
+  { instance: ContractInstanceWithAddress; artifact?: ContractArtifact }
 >;
 
+export const ContractOverridesSchema = z.record(
+  z.string(),
+  z.object({ instance: ContractInstanceWithAddressSchema, artifact: ContractArtifactSchema.optional() }),
+);
+
 /*
  * Optional values that can be overridden during simulation. In order to simulate a transaction with these
  * set, it *must* be run without the kernel circuits, or validations will fail
@@ -50,7 +63,7 @@ export class SimulationOverrides {
         contracts: optional(
           z.record(
             z.string(),
-            z.object({ instance: ContractInstanceWithAddressSchema, artifact: ContractArtifactSchema }),
+            z.object({ instance: ContractInstanceWithAddressSchema, artifact: ContractArtifactSchema.optional() }),
           ),
         ),
       })

yarn-project/wallet-sdk/src/base-wallet/base_wallet.ts

@@ -57,6 +57,7 @@ import type { AztecNode } from '@aztec/stdlib/interfaces/client';
 import {
   BlockHeader,
   ExecutionPayload,
+  SimulationOverrides,
   type TxExecutionRequest,
   type TxProfileResult,
   type UtilityExecutionResult,
@@ -354,6 +355,10 @@ export abstract class BaseWallet implements Wallet {
     return instance;
   }
 
+  registerContractClass(artifact: ContractArtifact): Promise<void> {
+    return this.pxe.registerContractClass(artifact);
+  }
+
   /**
    * Simulates calls through the standard PXE path (account entrypoint).
    * @param executionPayload - The execution payload to simulate.
@@ -373,6 +378,11 @@ export abstract class BaseWallet implements Wallet {
       senderForTags: this.senderForTagsFrom(opts.from, opts.sendMessagesAs),
       stateOverrides: opts.stateOverrides,
       contractOverrides: opts.contractOverrides,
+      overrides: opts.contractOverrides?.length
+        ? new SimulationOverrides(
+            Object.fromEntries(opts.contractOverrides.map(instance => [instance.address.toString(), { instance }])),
+          )
+        : undefined,
     });
     const appCallOffset = await this.computeAppCallOffset(opts.from, opts.feeOptions);
     return TxSimulationResultWithAppOffset.fromResultAndOffset(result, appCallOffset);