feat: expose wallet.registerContractClass for simulating contract upgrades

Adds `wallet.registerContractClass(artifact)` (a thin pass-through to PXE) so callers can register a new class artifact locally before passing an instance override that targets it. Without this, PXE-side ACIR dispatch can't resolve private functions of the override's class.

With `fastForwardContractUpdate` (downstack) plus this method, a single `.simulate({ overrides })` covers both private and public function calls on an upgraded contract:

```ts
await wallet.registerContractClass(UpdatedContract.artifact);
const overrides = await fastForwardContractUpdate({ instanceAddress, newClassId, node });
await updatedContract.methods.set_private_value().simulate({ from, overrides });
```

Author: dbanks12

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>;
@@ -559,6 +565,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/wallet-sdk/src/base-wallet/base_wallet.ts

@@ -348,6 +348,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.