refactor: drop artifact field from SimulationOverrides

docs/docs-developers/docs/resources/migration_notes.md

@@ -151,6 +151,19 @@ If you relied on a bundled bare-name binary for general use:
 
 If you set `Noir: Nargo Path` in the VS Code Noir extension to `$HOME/.aztec/current/bin/nargo`, change it to `$HOME/.aztec/current/bin/aztec-nargo` (the symlink is a drop-in for `nargo`). See the [Noir VSCode Extension guide](../aztec-nr/installation.md) for details.
 
+### [Stdlib] `SimulationOverrides.contracts` entries no longer carry an artifact
+
+`ContractOverrides` entries are now `{ instance }` only. To override a contract's artifact, pre-register the target class via `pxe.registerContractClass(artifact)` and set the override instance's `currentContractClassId` to that class id:
+
+```diff
+- const instance = await getContractInstanceFromInstantiationParams(stubArtifact, { salt: Fr.random() });
++ const instance = await pxe.getContractInstance(addr);
++ await pxe.registerContractClass(stubArtifact);
++ const stubClassId = (await getContractClassFromArtifact(stubArtifact)).id;
+- overrides = { contracts: { [addr.toString()]: { instance, artifact: stubArtifact } } };
++ overrides = { contracts: { [addr.toString()]: { instance: { ...instance, currentContractClassId: stubClassId } } } };
+```
+
 ### [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/cli-wallet/src/utils/wallet.ts

@@ -4,11 +4,7 @@ import { StubEcdsaAccountContractArtifact, createStubEcdsaAccount } from '@aztec
 import { StubSchnorrAccountContractArtifact, createStubSchnorrAccount } from '@aztec/accounts/stub/schnorr';
 import { getIdentities } from '@aztec/accounts/utils';
 import { type Account, type AccountContract, NO_FROM } from '@aztec/aztec.js/account';
-import {
-  type InteractionFeeOptions,
-  getContractInstanceFromInstantiationParams,
-  getGasLimits,
-} from '@aztec/aztec.js/contracts';
+import { type InteractionFeeOptions, getContractClassFromArtifact, getGasLimits } from '@aztec/aztec.js/contracts';
 import type { AztecNode } from '@aztec/aztec.js/node';
 import { AccountManager, type Aliased, type SimulateOptions } from '@aztec/aztec.js/wallet';
 import { TxSimulationResultWithAppOffset } from '@aztec/aztec.js/wallet';
@@ -34,6 +30,9 @@ import { printGasEstimates } from './options/fees.js';
 
 export class CLIWallet extends BaseWallet {
   private accountCache = new Map<string, Account>();
+  // Stub class ids, populated on wallet startup
+  // to avoid redundant work per simulation
+  private stubClassIds = new Map<AccountType, Fr>();
 
   constructor(
     pxe: PXE,
@@ -53,7 +52,27 @@ export class CLIWallet extends BaseWallet {
   ): Promise<CLIWallet> {
     const pxeConfig = Object.assign(getPXEConfig(), overridePXEConfig);
     const pxe = await createPXE(node, pxeConfig);
-    return new CLIWallet(pxe, node, log, db);
+    const wallet = new CLIWallet(pxe, node, log, db);
+    await wallet.initStubClasses();
+    return wallet;
+  }
+
+  /**
+   * Hashes and registers the stub class for every supported account type with PXE, populating
+   * stubClassIds. Called on wallet initialization.
+   */
+  private async initStubClasses(): Promise<void> {
+    const { id: schnorrClassId } = await getContractClassFromArtifact(StubSchnorrAccountContractArtifact);
+    await this.pxe.registerContractClass(StubSchnorrAccountContractArtifact);
+
+    // ecdsa stubs share the same class id
+    const { id: ecdsaClassId } = await getContractClassFromArtifact(StubEcdsaAccountContractArtifact);
+    await this.pxe.registerContractClass(StubEcdsaAccountContractArtifact);
+
+    this.stubClassIds.set('schnorr', schnorrClassId);
+    this.stubClassIds.set('ecdsasecp256k1', ecdsaClassId);
+    this.stubClassIds.set('ecdsasecp256r1', ecdsaClassId);
+    this.stubClassIds.set('ecdsasecp256r1ssh', ecdsaClassId);
   }
 
   override async getAccounts(): Promise<Aliased<AztecAddress>[]> {
@@ -202,11 +221,16 @@ export class CLIWallet extends BaseWallet {
       throw new Error(`No contract instance found for address: ${originalAddress.address}`);
     }
     const { type } = await this.db!.retrieveAccount(address);
-    const isSchnorr = type === 'schnorr';
-    const artifact = isSchnorr ? StubSchnorrAccountContractArtifact : StubEcdsaAccountContractArtifact;
-    const stubAccount = isSchnorr ? createStubSchnorrAccount(originalAddress) : createStubEcdsaAccount(originalAddress);
-    const instance = await getContractInstanceFromInstantiationParams(artifact, { salt: Fr.random() });
-    return { account: stubAccount, instance, artifact };
+    const stubAccount =
+      type === 'schnorr' ? createStubSchnorrAccount(originalAddress) : createStubEcdsaAccount(originalAddress);
+    const stubClassId = this.stubClassIds.get(type);
+    if (!stubClassId) {
+      throw new Error(
+        `Stub class for account type '${type}' was not registered at wallet init. This is a bug — initStubClasses should cover every supported AccountType.`,
+      );
+    }
+    const instance = { ...contractInstance, currentContractClassId: stubClassId };
+    return { account: stubAccount, instance };
   }
 
   override async simulateTx(
@@ -249,9 +273,9 @@ export class CLIWallet extends BaseWallet {
       const entrypoint = new DefaultEntrypoint();
       txRequest = await entrypoint.createTxExecutionRequest(finalExecutionPayload, feeOptions.gasSettings, chainInfo);
     } else {
-      const { account, instance, artifact } = await this.getFakeAccountDataFor(from);
+      const { account, instance } = await this.getFakeAccountDataFor(from);
       overrides = {
-        contracts: { [from.toString()]: { instance, artifact } },
+        contracts: { [from.toString()]: { instance } },
       };
       const executionOptions: DefaultAccountEntrypointOptions = {
         txNonce: Fr.random(),

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

@@ -25,7 +25,7 @@ import { type PXEConfig, getPXEConfig } from '@aztec/pxe/config';
 import { PXE, type PXECreationOptions, createPXE } from '@aztec/pxe/server';
 import { AuthWitness } from '@aztec/stdlib/auth-witness';
 import { AztecAddress } from '@aztec/stdlib/aztec-address';
-import { getContractInstanceFromInstantiationParams } from '@aztec/stdlib/contract';
+import { getContractClassFromArtifact } from '@aztec/stdlib/contract';
 import { deriveSigningKey } from '@aztec/stdlib/keys';
 import type { NoteDao } from '@aztec/stdlib/note';
 import {
@@ -78,7 +78,9 @@ export class TestWallet extends BaseWallet {
       ...overridePXEConfig,
     });
     const pxe = await createPXE(nodeRef, pxeConfig, options);
-    return new TestWallet(pxe, nodeRef);
+    const wallet = new TestWallet(pxe, nodeRef);
+    await wallet.initStubClasses();
+    return wallet;
   }
 
   /**
@@ -112,6 +114,27 @@ export class TestWallet extends BaseWallet {
     });
   }
 
+  // Stub class ids, populated on wallet startup
+  // to avoid redundant work per simulation
+  private stubClassIds = new Map<AccountType, Fr>();
+
+  /**
+   * Hashes and registers the stub class for every supported account type with PXE, populating
+   * stubClassIds. Called on wallet initialization.
+   */
+  private async initStubClasses(): Promise<void> {
+    const { id: schnorrClassId } = await getContractClassFromArtifact(StubSchnorrAccountContractArtifact);
+    await this.pxe.registerContractClass(StubSchnorrAccountContractArtifact);
+
+    // ecdsa stubs share the same class id
+    const { id: ecdsaClassId } = await getContractClassFromArtifact(StubEcdsaAccountContractArtifact);
+    await this.pxe.registerContractClass(StubEcdsaAccountContractArtifact);
+
+    this.stubClassIds.set('schnorr', schnorrClassId);
+    this.stubClassIds.set('ecdsasecp256k1', ecdsaClassId);
+    this.stubClassIds.set('ecdsasecp256r1', ecdsaClassId);
+  }
+
   /**
    * Builds contract overrides for all provided addresses by replacing their account contracts with stub implementations.
    */
@@ -132,17 +155,16 @@ export class TestWallet extends BaseWallet {
         );
       }
 
-      const stubArtifact = this.getStubArtifactFor(address);
-      const stubConstructorArgs =
-        this.getTypeFor(address) === 'schnorr' ? [Fr.ZERO, Fr.ZERO] : [Buffer.alloc(32), Buffer.alloc(32)];
-      const stubInstance = await getContractInstanceFromInstantiationParams(stubArtifact, {
-        salt: Fr.random(),
-        constructorArgs: stubConstructorArgs,
-      });
+      const type = this.getTypeFor(address);
+      const stubClassId = this.stubClassIds.get(type);
+      if (!stubClassId) {
+        throw new Error(
+          `Stub class for account type '${type}' was not registered at wallet init. This is a bug — initStubClasses should cover every supported AccountType.`,
+        );
+      }
 
       contracts[address.toString()] = {
-        instance: stubInstance,
-        artifact: stubArtifact,
+        instance: { ...contractInstance, currentContractClassId: stubClassId },
       };
     }
 
@@ -155,12 +177,6 @@ export class TestWallet extends BaseWallet {
     return this.accounts.get(address.toString())?.type ?? 'schnorr';
   }
 
-  private getStubArtifactFor(address: AztecAddress) {
-    return this.getTypeFor(address) === 'schnorr'
-      ? StubSchnorrAccountContractArtifact
-      : StubEcdsaAccountContractArtifact;
-  }
-
   private getStubAccountFor(address: AztecAddress, completeAddress: CompleteAddress) {
     return this.getTypeFor(address) === 'schnorr'
       ? createStubSchnorrAccount(completeAddress)

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

@@ -1,89 +1,69 @@
 import { FunctionSelector } from '@aztec/stdlib/abi';
-import { AztecAddress } from '@aztec/stdlib/aztec-address';
+import type { AztecAddress } from '@aztec/stdlib/aztec-address';
 import type { ContractOverrides } from '@aztec/stdlib/tx';
 
 import type { ContractStore } from '../storage/contract_store/contract_store.js';
 
 /*
- * Proxy generator for a ContractStore that allows overriding contract instances and artifacts, so
- * the contract function simulator can execute different bytecode on certain addresses. An example use case
- * would be overriding your own account contract so that valid signatures don't have to be provided while simulating.
+ * Proxy generator for a ContractStore that allows overriding contract instances at given addresses,
+ * so the contract function simulator can execute different bytecode on certain addresses. An example
+ * use case would be overriding your own account contract so that valid signatures don't have to be
+ * provided while simulating.
+ *
+ * Function artifact lookups for an overridden address are routed via the override-instance's
+ * `currentContractClassId` rather than the underlying ContractStore's address→class mapping —
+ * `ContractStore.getFunctionArtifact` calls `this.getContractInstance` internally, which the proxy
+ * cannot intercept (the binding sees the raw target, not the proxy). The target class must be
+ * registered ahead of time via `pxe.registerContractClass(...)`.
  */
 export class ProxiedContractStoreFactory {
   static create(contractStore: ContractStore, overrides?: ContractOverrides) {
     if (!overrides) {
       return contractStore;
     }
 
+    const findFunctionInArtifact = async (
+      artifact: { name: string; functions: { name: string; parameters: any[] }[] },
+      selector: FunctionSelector,
+      contractAddress: AztecAddress,
+    ) => {
+      for (const fn of artifact.functions) {
+        const fnSelector = await FunctionSelector.fromNameAndParameters(fn.name, fn.parameters);
+        if (fnSelector.equals(selector)) {
+          return { ...fn, contractName: artifact.name } as any;
+        }
+      }
+      throw new Error(
+        `Function with selector ${selector} not found in stub artifact for overridden contract at ${contractAddress}.`,
+      );
+    };
+
     return new Proxy(contractStore, {
       get(target, prop: keyof ContractStore) {
-        switch (prop) {
-          case 'getContractInstance': {
-            return async (address: AztecAddress) => {
-              if (overrides[address.toString()]) {
-                const { instance } = overrides[address.toString()]!;
-                instance.address = address;
-                const realInstance = await target.getContractInstance(address);
-                if (!realInstance) {
-                  throw new Error(`Contract instance not found for address: ${address}`);
-                }
-                instance.currentContractClassId = realInstance.currentContractClassId;
-                instance.originalContractClassId = realInstance.originalContractClassId;
-                instance.initializationHash = realInstance.initializationHash;
-                return instance;
-              } else {
-                return target.getContractInstance(address);
-              }
-            };
-          }
-          case 'getFunctionArtifact': {
-            return async (contractAddress: AztecAddress, selector: FunctionSelector) => {
-              if (overrides[contractAddress.toString()]) {
-                const { artifact } = overrides[contractAddress.toString()]!;
-                const functions = artifact.functions;
-                for (let i = 0; i < functions.length; i++) {
-                  const fn = functions[i];
-                  const fnSelector = await FunctionSelector.fromNameAndParameters(fn.name, fn.parameters);
-                  if (fnSelector.equals(selector)) {
-                    return fn;
-                  }
-                }
-                throw new Error(
-                  `Function with selector ${selector} not found in stub artifact for overridden contract at ${contractAddress}. The stub does not implement this function.`,
-                );
-              } else {
-                return target.getFunctionArtifact(contractAddress, selector);
-              }
-            };
-          }
-          case 'getFunctionArtifactWithDebugMetadata': {
-            return async (contractAddress: AztecAddress, selector: FunctionSelector) => {
-              if (overrides[contractAddress.toString()]) {
-                const { artifact } = overrides[contractAddress.toString()]!;
-                const functions = artifact.functions;
-                for (let i = 0; i < functions.length; i++) {
-                  const fn = functions[i];
-                  const fnSelector = await FunctionSelector.fromNameAndParameters(fn.name, fn.parameters);
-                  if (fnSelector.equals(selector)) {
-                    return fn;
-                  }
-                }
-                throw new Error(
-                  `Function with selector ${selector} not found in stub artifact for overridden contract at ${contractAddress}. The stub does not implement this function.`,
-                );
-              } else {
-                return target.getFunctionArtifactWithDebugMetadata(contractAddress, selector);
-              }
-            };
-          }
-          default: {
-            const value = Reflect.get(target, prop);
-            if (typeof value === 'function') {
-              return value.bind(target);
+        if (prop === 'getContractInstance') {
+          return (address: AztecAddress) => {
+            const override = overrides[address.toString()];
+            return override ? Promise.resolve(override.instance) : target.getContractInstance(address);
+          };
+        }
+        if (prop === 'getFunctionArtifact' || prop === 'getFunctionArtifactWithDebugMetadata') {
+          return async (contractAddress: AztecAddress, selector: FunctionSelector) => {
+            const override = overrides[contractAddress.toString()];
+            if (!override) {
+              return (target[prop] as Function).call(target, contractAddress, selector);
+            }
+            const artifact = await target.getContractArtifact(override.instance.currentContractClassId);
+            if (!artifact) {
+              throw new Error(
+                `No artifact registered for override class ${override.instance.currentContractClassId} ` +
+                  `at ${contractAddress}. Register it via pxe.registerContractClass(...) before simulating.`,
+              );
             }
-            return value;
-          }
+            return findFunctionInArtifact(artifact, selector, contractAddress);
+          };
         }
+        const value = Reflect.get(target, prop);
+        return typeof value === 'function' ? value.bind(target) : value;
       },
     }) satisfies ContractStore;
   }

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

@@ -3,7 +3,6 @@ import type { FieldsOf } from '@aztec/foundation/types';
 
 import { z } from 'zod';
 
-import { type ContractArtifact, ContractArtifactSchema } from '../abi/abi.js';
 import {
   type ContractInstanceWithAddress,
   ContractInstanceWithAddressSchema,
@@ -24,18 +23,17 @@ 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.
+ * If passed during the execution of a user circuit, the contract function simulator will replace
+ * the contract instance at that address with the one provided. An example use case would be
+ * overriding your own account contract so that valid signatures don't have to be provided while
+ * simulating. The override's `currentContractClassId` resolves through PXE's locally registered
+ * classes, so pre-register the target artifact via `pxe.registerContractClass(...)`.
  */
-export type ContractOverrides = Record<
-  string /* AztecAddress as string */,
-  { instance: ContractInstanceWithAddress; artifact: ContractArtifact }
->;
+export type ContractOverrides = Record<string /* AztecAddress as string */, { instance: ContractInstanceWithAddress }>;
 
 /*
  * 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
+ * set, it *must* be run without the kernel circuits, or validations will fail.
  */
 export class SimulationOverrides {
   public contracts?: ContractOverrides;
@@ -47,12 +45,7 @@ export class SimulationOverrides {
   static get schema() {
     return z
       .object({
-        contracts: optional(
-          z.record(
-            z.string(),
-            z.object({ instance: ContractInstanceWithAddressSchema, artifact: ContractArtifactSchema }),
-          ),
-        ),
+        contracts: optional(z.record(z.string(), z.object({ instance: ContractInstanceWithAddressSchema }))),
       })
       .transform(({ contracts }) => {
         return new SimulationOverrides(contracts);