refactor(wallets): narrow stub-class cache to just the class id

Callers only ever consume the `id` field from `getContractClassFromArtifact`,
so cache `Map<AccountType, Promise<Fr>>` instead of the full
`ContractClassWithId & ContractClassIdPreimage`. Method renamed
`#getStubClass` → `#getStubClassId`, map renamed `#stubClasses` →
`#stubClassIds`, callers drop the destructure.

Drop the verbose Promise-as-value rationale comment — the idiom is
established in this codebase (see ContractSyncService) and the PXE
precedents document the cache contract briefly without explaining
the dedupe pattern.

Author: dbanks12

yarn-project/cli-wallet/src/utils/wallet.ts

@@ -17,7 +17,6 @@ import type { PXEConfig } from '@aztec/pxe/config';
 import type { PXE } from '@aztec/pxe/server';
 import { createPXE, getPXEConfig } from '@aztec/pxe/server';
 import { AztecAddress } from '@aztec/stdlib/aztec-address';
-import type { ContractClassIdPreimage, ContractClassWithId } from '@aztec/stdlib/contract';
 import { deriveSigningKey } from '@aztec/stdlib/keys';
 import { NoteDao } from '@aztec/stdlib/note';
 import type { SimulationOverrides, TxExecutionRequest, TxProvingResult } from '@aztec/stdlib/tx';
@@ -31,13 +30,9 @@ import { printGasEstimates } from './options/fees.js';
 
 export class CLIWallet extends BaseWallet {
   private accountCache = new Map<string, Account>();
-  /**
-   * Per-account-type cache of the stub class id and preimage. The Promise is stored (not the resolved
-   * value) so concurrent first-time callers dedupe on the same hashing + registration work. ECDSA
-   * variants all map to the same stub artifact but get their own cache slot, so we may hash + register
-   * the same artifact one extra time per variant on first miss; that cost is bounded and one-time.
-   */
-  #stubClasses = new Map<AccountType, Promise<ContractClassWithId & ContractClassIdPreimage>>();
+  // Stub class ids, populated on wallet startup
+  // to avoid redundant work per simulation
+  private stubClassIds = new Map<AccountType, Fr>();
 
   constructor(
     pxe: PXE,
@@ -57,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>[]> {
@@ -208,29 +223,16 @@ export class CLIWallet extends BaseWallet {
     const { type } = await this.db!.retrieveAccount(address);
     const stubAccount =
       type === 'schnorr' ? createStubSchnorrAccount(originalAddress) : createStubEcdsaAccount(originalAddress);
-    const { id: stubClassId } = await this.#getStubClass(type);
+    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 };
   }
 
-  /**
-   * Lazily hashes and registers the stub class for the given account type, caching the result so
-   * subsequent simulations skip the artifact-hashing + registration round-trip.
-   */
-  #getStubClass(type: AccountType): Promise<ContractClassWithId & ContractClassIdPreimage> {
-    let cached = this.#stubClasses.get(type);
-    if (!cached) {
-      cached = (async () => {
-        const artifact = type === 'schnorr' ? StubSchnorrAccountContractArtifact : StubEcdsaAccountContractArtifact;
-        const stubClass = await getContractClassFromArtifact(artifact);
-        await this.pxe.registerContractClass(artifact);
-        return stubClass;
-      })();
-      this.#stubClasses.set(type, cached);
-    }
-    return cached;
-  }
-
   override async simulateTx(
     executionPayload: ExecutionPayload,
     opts: SimulateOptions,

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

@@ -25,11 +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 {
-  type ContractClassIdPreimage,
-  type ContractClassWithId,
-  getContractClassFromArtifact,
-} from '@aztec/stdlib/contract';
+import { getContractClassFromArtifact } from '@aztec/stdlib/contract';
 import { deriveSigningKey } from '@aztec/stdlib/keys';
 import type { NoteDao } from '@aztec/stdlib/note';
 import {
@@ -82,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;
   }
 
   /**
@@ -116,28 +114,25 @@ export class TestWallet extends BaseWallet {
     });
   }
 
-  /**
-   * Per-account-type cache of the stub class id and preimage. The Promise is stored (not the resolved
-   * value) so concurrent first-time callers dedupe on the same hashing + registration work.
-   */
-  #stubClasses = new Map<AccountType, Promise<ContractClassWithId & ContractClassIdPreimage>>();
+  // Stub class ids, populated on wallet startup
+  // to avoid redundant work per simulation
+  private stubClassIds = new Map<AccountType, Fr>();
 
   /**
-   * Lazily hashes and registers the stub class for the given account type, caching the result so
-   * subsequent simulations skip the artifact-hashing + registration round-trip.
+   * Hashes and registers the stub class for every supported account type with PXE, populating
+   * stubClassIds. Called on wallet initialization.
    */
-  #getStubClass(type: AccountType): Promise<ContractClassWithId & ContractClassIdPreimage> {
-    let cached = this.#stubClasses.get(type);
-    if (!cached) {
-      cached = (async () => {
-        const stubArtifact = type === 'schnorr' ? StubSchnorrAccountContractArtifact : StubEcdsaAccountContractArtifact;
-        const stubClass = await getContractClassFromArtifact(stubArtifact);
-        await this.pxe.registerContractClass(stubArtifact);
-        return stubClass;
-      })();
-      this.#stubClasses.set(type, cached);
-    }
-    return cached;
+  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);
   }
 
   /**
@@ -160,7 +155,13 @@ export class TestWallet extends BaseWallet {
         );
       }
 
-      const { id: stubClassId } = await this.#getStubClass(this.getTypeFor(address));
+      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: { ...contractInstance, currentContractClassId: stubClassId },

yarn-project/wallets/src/embedded/embedded_wallet.ts

@@ -11,11 +11,7 @@ import type { AztecAsyncKVStore } from '@aztec/kv-store';
 import type { PXEConfig, PXECreationOptions } from '@aztec/pxe/client/lazy';
 import type { PXE } from '@aztec/pxe/server';
 import { AztecAddress } from '@aztec/stdlib/aztec-address';
-import {
-  type ContractClassIdPreimage,
-  type ContractClassWithId,
-  getContractClassFromArtifact,
-} from '@aztec/stdlib/contract';
+import { getContractClassFromArtifact } from '@aztec/stdlib/contract';
 import { GasSettings } from '@aztec/stdlib/gas';
 import type { AztecNode } from '@aztec/stdlib/interfaces/client';
 import { deriveSigningKey } from '@aztec/stdlib/keys';
@@ -80,11 +76,9 @@ const DEFAULT_ESTIMATED_GAS_PADDING = 0.1;
 export class EmbeddedWallet extends BaseWallet {
   protected estimatedGasPadding = DEFAULT_ESTIMATED_GAS_PADDING;
 
-  /**
-   * Per-account-type cache of the stub class id and preimage. The Promise is stored (not the resolved
-   * value) so concurrent first-time callers dedupe on the same hashing + registration work.
-   */
-  #stubClasses = new Map<AccountType, Promise<ContractClassWithId & ContractClassIdPreimage>>();
+  // Stub class ids, populated on wallet startup
+  // to avoid redundant work per simulation
+  protected stubClassIds = new Map<AccountType, Fr>();
 
   constructor(
     pxe: PXE,
@@ -203,21 +197,22 @@ export class EmbeddedWallet extends BaseWallet {
   }
 
   /**
-   * Lazily hashes and registers the stub class for the given account type, caching the result so
-   * subsequent simulations skip the artifact-hashing + registration round-trip.
+   * Hashes and registers the stub class for every supported account type with PXE, populating
+   * stubClassIds. Called on wallet initialization.
    */
-  #getStubClass(type: AccountType): Promise<ContractClassWithId & ContractClassIdPreimage> {
-    let cached = this.#stubClasses.get(type);
-    if (!cached) {
-      cached = (async () => {
-        const stubArtifact = await this.accountContracts.getStubAccountContractArtifact(type);
-        const stubClass = await getContractClassFromArtifact(stubArtifact);
-        await this.pxe.registerContractClass(stubArtifact);
-        return stubClass;
-      })();
-      this.#stubClasses.set(type, cached);
-    }
-    return cached;
+  async initStubClasses(): Promise<void> {
+    const schnorrArtifact = await this.accountContracts.getStubAccountContractArtifact('schnorr');
+    const { id: schnorrClassId } = await getContractClassFromArtifact(schnorrArtifact);
+    await this.pxe.registerContractClass(schnorrArtifact);
+
+    // ecdsa stubs share the same class id
+    const ecdsaArtifact = await this.accountContracts.getStubAccountContractArtifact('ecdsasecp256r1');
+    const { id: ecdsaClassId } = await getContractClassFromArtifact(ecdsaArtifact);
+    await this.pxe.registerContractClass(ecdsaArtifact);
+
+    this.stubClassIds.set('schnorr', schnorrClassId);
+    this.stubClassIds.set('ecdsasecp256k1', ecdsaClassId);
+    this.stubClassIds.set('ecdsasecp256r1', ecdsaClassId);
   }
 
   /**
@@ -233,7 +228,12 @@ export class EmbeddedWallet extends BaseWallet {
     for (const account of filtered) {
       const address = account.item;
       const { type } = await this.walletDB.retrieveAccount(address);
-      const { id: stubClassId } = await this.#getStubClass(type);
+      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 originalAccount = await this.getAccountFromAddress(address);
       const completeAddress = originalAccount.getCompleteAddress();

yarn-project/wallets/src/embedded/entrypoints/browser.ts

@@ -69,7 +69,9 @@ export class BrowserEmbeddedWallet extends EmbeddedWallet {
           ));
     const walletDB = new WalletDB(walletDBStore, rootLogger.createChild('wallet:db').info);
 
-    return new this(pxe, aztecNode, walletDB, new LazyAccountContractsProvider(), rootLogger) as T;
+    const wallet = new this(pxe, aztecNode, walletDB, new LazyAccountContractsProvider(), rootLogger) as T;
+    await wallet.initStubClasses();
+    return wallet;
   }
 }
 

yarn-project/wallets/src/embedded/entrypoints/node.ts

@@ -76,7 +76,9 @@ export class NodeEmbeddedWallet extends EmbeddedWallet {
           ));
     const walletDB = new WalletDB(walletDBStore, rootLogger.createChild('wallet:db').info);
 
-    return new this(pxe, aztecNode, walletDB, new BundleAccountContractsProvider(), rootLogger) as T;
+    const wallet = new this(pxe, aztecNode, walletDB, new BundleAccountContractsProvider(), rootLogger) as T;
+    await wallet.initStubClasses();
+    return wallet;
   }
 }