refactor(wallets): consolidate AccountType + stub dispatch in @aztec/accounts/stub

Lift AccountType (cli-wallet's 4-value superset, including the SSH variant) and
the schnorr-vs-ecdsa stub partition into @aztec/accounts/stub, with eager
(./stub) and lazy (./stub/lazy) entrypoints exposing
getStubArtifactForAccountType + createStubAccountForAccountType.

cli-wallet, embedded wallet, and test_wallet now share the same dispatch.
cli-wallet drops its ad-hoc 'schnorr' | 'ecdsa' flavor union and keys its
stub-class cache on AccountType, matching embedded; the resulting redundant
hash + registerContractClass round-trips are bounded by the (small)
AccountType cardinality and only cost on first-miss per type.

Author: dbanks12

yarn-project/accounts/package.json

@@ -10,6 +10,8 @@
     "./ecdsa/lazy": "./dest/ecdsa/lazy.js",
     "./schnorr": "./dest/schnorr/index.js",
     "./schnorr/lazy": "./dest/schnorr/lazy.js",
+    "./stub": "./dest/stub/index.js",
+    "./stub/lazy": "./dest/stub/lazy.js",
     "./stub/schnorr": "./dest/stub/schnorr/index.js",
     "./stub/schnorr/lazy": "./dest/stub/schnorr/lazy.js",
     "./stub/ecdsa": "./dest/stub/ecdsa/index.js",

yarn-project/accounts/src/stub/index.ts

@@ -0,0 +1,29 @@
+import type { Account } from '@aztec/aztec.js/account';
+import type { CompleteAddress } from '@aztec/aztec.js/addresses';
+import type { ContractArtifact } from '@aztec/stdlib/abi';
+
+import { StubEcdsaAccountContractArtifact, createStubEcdsaAccount } from './ecdsa/index.js';
+import { StubSchnorrAccountContractArtifact, createStubSchnorrAccount } from './schnorr/index.js';
+
+/** Account-contract families recognized by the wallets in this monorepo. */
+export const AccountTypes = ['schnorr', 'ecdsasecp256r1', 'ecdsasecp256r1ssh', 'ecdsasecp256k1'] as const;
+
+/** One of the account-contract families in {@link AccountTypes}. */
+export type AccountType = (typeof AccountTypes)[number];
+
+/**
+ * Returns the stub contract artifact whose constructor signature matches the given account type.
+ * All ECDSA variants share one stub because their real account contracts share a constructor
+ * signature (`[u8; 32], [u8; 32]`); schnorr has its own (`Field, Field`).
+ */
+export function getStubArtifactForAccountType(type: AccountType): ContractArtifact {
+  return type === 'schnorr' ? StubSchnorrAccountContractArtifact : StubEcdsaAccountContractArtifact;
+}
+
+/**
+ * Creates a stub account that impersonates the given address, picking the stub whose constructor
+ * signature matches the real account's so deployment-simulation selector lookups succeed.
+ */
+export function createStubAccountForAccountType(type: AccountType, address: CompleteAddress): Account {
+  return type === 'schnorr' ? createStubSchnorrAccount(address) : createStubEcdsaAccount(address);
+}

yarn-project/accounts/src/stub/lazy.ts

@@ -0,0 +1,33 @@
+import type { Account } from '@aztec/aztec.js/account';
+import type { CompleteAddress } from '@aztec/aztec.js/addresses';
+import type { ContractArtifact } from '@aztec/stdlib/abi';
+
+import type { AccountType } from './index.js';
+
+export { AccountTypes, type AccountType } from './index.js';
+
+/**
+ * Lazily resolves the stub contract artifact whose constructor signature matches the given
+ * account type, dynamically importing the artifact JSON to keep it out of the initial bundle.
+ */
+export async function getStubArtifactForAccountType(type: AccountType): Promise<ContractArtifact> {
+  if (type === 'schnorr') {
+    const { getStubSchnorrAccountContractArtifact } = await import('./schnorr/lazy.js');
+    return getStubSchnorrAccountContractArtifact();
+  }
+  const { getStubEcdsaAccountContractArtifact } = await import('./ecdsa/lazy.js');
+  return getStubEcdsaAccountContractArtifact();
+}
+
+/**
+ * Lazily creates a stub account that impersonates the given address, dynamically importing the
+ * matching stub artifact module.
+ */
+export async function createStubAccountForAccountType(type: AccountType, address: CompleteAddress): Promise<Account> {
+  if (type === 'schnorr') {
+    const { createStubSchnorrAccount } = await import('./schnorr/lazy.js');
+    return createStubSchnorrAccount(address);
+  }
+  const { createStubEcdsaAccount } = await import('./ecdsa/lazy.js');
+  return createStubEcdsaAccount(address);
+}

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

@@ -1,4 +1,3 @@
-export const MIN_FEE_PADDING = 0.5;
+export { AccountTypes, type AccountType } from '@aztec/accounts/stub';
 
-export const AccountTypes = ['schnorr', 'ecdsasecp256r1', 'ecdsasecp256r1ssh', 'ecdsasecp256k1'] as const;
-export type AccountType = (typeof AccountTypes)[number];
+export const MIN_FEE_PADDING = 0.5;

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

@@ -1,7 +1,6 @@
 import { EcdsaRAccountContract, EcdsaRSSHAccountContract } from '@aztec/accounts/ecdsa';
 import { SchnorrAccountContract } from '@aztec/accounts/schnorr';
-import { StubEcdsaAccountContractArtifact, createStubEcdsaAccount } from '@aztec/accounts/stub/ecdsa';
-import { StubSchnorrAccountContractArtifact, createStubSchnorrAccount } from '@aztec/accounts/stub/schnorr';
+import { createStubAccountForAccountType, getStubArtifactForAccountType } from '@aztec/accounts/stub';
 import { getIdentities } from '@aztec/accounts/utils';
 import { type Account, type AccountContract, NO_FROM } from '@aztec/aztec.js/account';
 import { type InteractionFeeOptions, getContractClassFromArtifact, getGasLimits } from '@aztec/aztec.js/contracts';
@@ -32,10 +31,10 @@ import { printGasEstimates } from './options/fees.js';
 export class CLIWallet extends BaseWallet {
   private accountCache = new Map<string, Account>();
   /**
-   * Per-stub-flavor cache of the stub class id and preimage. The Promise is stored (not the resolved
+   * 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<'schnorr' | 'ecdsa', Promise<ContractClassWithId & ContractClassIdPreimage>>();
+  #stubClasses = new Map<AccountType, Promise<ContractClassWithId & ContractClassIdPreimage>>();
 
   constructor(
     pxe: PXE,
@@ -204,27 +203,26 @@ 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 stubAccount = isSchnorr ? createStubSchnorrAccount(originalAddress) : createStubEcdsaAccount(originalAddress);
-    const { id: stubClassId } = await this.#getStubClass(isSchnorr ? 'schnorr' : 'ecdsa');
+    const stubAccount = createStubAccountForAccountType(type, originalAddress);
+    const { id: stubClassId } = await this.#getStubClass(type);
     const instance = { ...contractInstance, currentContractClassId: stubClassId };
     return { account: stubAccount, instance };
   }
 
   /**
-   * Lazily hashes and registers the stub class for the given flavor, caching the result so
+   * 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(flavor: 'schnorr' | 'ecdsa'): Promise<ContractClassWithId & ContractClassIdPreimage> {
-    let cached = this.#stubClasses.get(flavor);
+  #getStubClass(type: AccountType): Promise<ContractClassWithId & ContractClassIdPreimage> {
+    let cached = this.#stubClasses.get(type);
     if (!cached) {
       cached = (async () => {
-        const artifact = flavor === 'schnorr' ? StubSchnorrAccountContractArtifact : StubEcdsaAccountContractArtifact;
+        const artifact = getStubArtifactForAccountType(type);
         const stubClass = await getContractClassFromArtifact(artifact);
         await this.pxe.registerContractClass(artifact);
         return stubClass;
       })();
-      this.#stubClasses.set(flavor, cached);
+      this.#stubClasses.set(type, cached);
     }
     return cached;
   }

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

@@ -1,7 +1,6 @@
 import { EcdsaKAccountContract, EcdsaRAccountContract } from '@aztec/accounts/ecdsa';
 import { SchnorrAccountContract } from '@aztec/accounts/schnorr';
-import { StubEcdsaAccountContractArtifact, createStubEcdsaAccount } from '@aztec/accounts/stub/ecdsa';
-import { StubSchnorrAccountContractArtifact, createStubSchnorrAccount } from '@aztec/accounts/stub/schnorr';
+import { type AccountType, createStubAccountForAccountType, getStubArtifactForAccountType } from '@aztec/accounts/stub';
 import { type Account, type AccountContract, NO_FROM } from '@aztec/aztec.js/account';
 import type { CompleteAddress } from '@aztec/aztec.js/addresses';
 import {
@@ -42,7 +41,6 @@ import {
 } from '@aztec/stdlib/tx';
 import { ExecutionPayload, mergeExecutionPayloads } from '@aztec/stdlib/tx';
 import { BaseWallet, type SimulateViaEntrypointOptions } from '@aztec/wallet-sdk/base-wallet';
-import type { AccountType } from '@aztec/wallets/embedded';
 
 import { DEFAULT_MIN_FEE_PADDING } from '../fixtures/fixtures.js';
 import { AztecNodeProxy, ProvenTx } from './utils.js';
@@ -130,7 +128,7 @@ export class TestWallet extends BaseWallet {
     let cached = this.#stubClasses.get(type);
     if (!cached) {
       cached = (async () => {
-        const stubArtifact = type === 'schnorr' ? StubSchnorrAccountContractArtifact : StubEcdsaAccountContractArtifact;
+        const stubArtifact = getStubArtifactForAccountType(type);
         const stubClass = await getContractClassFromArtifact(stubArtifact);
         await this.pxe.registerContractClass(stubArtifact);
         return stubClass;
@@ -177,9 +175,7 @@ export class TestWallet extends BaseWallet {
   }
 
   private getStubAccountFor(address: AztecAddress, completeAddress: CompleteAddress) {
-    return this.getTypeFor(address) === 'schnorr'
-      ? createStubSchnorrAccount(completeAddress)
-      : createStubEcdsaAccount(completeAddress);
+    return createStubAccountForAccountType(this.getTypeFor(address), completeAddress);
   }
 
   /**

yarn-project/pxe/tsconfig.json

@@ -40,10 +40,10 @@
       "path": "../stdlib"
     },
     {
-      "path": "../world-state"
+      "path": "../noir-test-contracts.js"
     },
     {
-      "path": "../noir-test-contracts.js"
+      "path": "../world-state"
     }
   ],
   "include": ["src"]

yarn-project/wallets/src/embedded/account-contract-providers/bundle.ts

@@ -1,7 +1,6 @@
 import { EcdsaKAccountContract, EcdsaRAccountContract } from '@aztec/accounts/ecdsa';
 import { SchnorrAccountContract } from '@aztec/accounts/schnorr';
-import { StubEcdsaAccountContractArtifact, createStubEcdsaAccount } from '@aztec/accounts/stub/ecdsa';
-import { StubSchnorrAccountContractArtifact, createStubSchnorrAccount } from '@aztec/accounts/stub/schnorr';
+import { createStubAccountForAccountType, getStubArtifactForAccountType } from '@aztec/accounts/stub';
 import type { Account, AccountContract } from '@aztec/aztec.js/account';
 import type { Fq } from '@aztec/foundation/curves/bn254';
 import { getCanonicalMultiCallEntrypoint } from '@aztec/protocol-contracts/multi-call-entrypoint';
@@ -29,11 +28,11 @@ export class BundleAccountContractsProvider implements AccountContractsProvider
   }
 
   getStubAccountContractArtifact(type: AccountType): Promise<ContractArtifact> {
-    return Promise.resolve(type === 'schnorr' ? StubSchnorrAccountContractArtifact : StubEcdsaAccountContractArtifact);
+    return Promise.resolve(getStubArtifactForAccountType(type));
   }
 
   createStubAccount(address: CompleteAddress, type: AccountType): Promise<Account> {
-    return Promise.resolve(type === 'schnorr' ? createStubSchnorrAccount(address) : createStubEcdsaAccount(address));
+    return Promise.resolve(createStubAccountForAccountType(type, address));
   }
 
   getMulticallContract(): Promise<{ instance: ContractInstanceWithAddress; artifact: ContractArtifact }> {

yarn-project/wallets/src/embedded/account-contract-providers/lazy.ts

@@ -1,3 +1,4 @@
+import { createStubAccountForAccountType, getStubArtifactForAccountType } from '@aztec/accounts/stub/lazy';
 import type { Account, AccountContract } from '@aztec/aztec.js/account';
 import type { Fq } from '@aztec/foundation/curves/bn254';
 import { getCanonicalMultiCallEntrypoint } from '@aztec/protocol-contracts/multi-call-entrypoint/lazy';
@@ -27,24 +28,12 @@ export class LazyAccountContractsProvider implements AccountContractsProvider {
     return new EcdsaKAccountContract(signingKey);
   }
 
-  async getStubAccountContractArtifact(type: AccountType): Promise<ContractArtifact> {
-    if (type === 'schnorr') {
-      const { getStubSchnorrAccountContractArtifact } = await import('@aztec/accounts/stub/schnorr/lazy');
-      return getStubSchnorrAccountContractArtifact();
-    } else {
-      const { getStubEcdsaAccountContractArtifact } = await import('@aztec/accounts/stub/ecdsa/lazy');
-      return getStubEcdsaAccountContractArtifact();
-    }
+  getStubAccountContractArtifact(type: AccountType): Promise<ContractArtifact> {
+    return getStubArtifactForAccountType(type);
   }
 
-  async createStubAccount(address: CompleteAddress, type: AccountType): Promise<Account> {
-    if (type === 'schnorr') {
-      const { createStubSchnorrAccount } = await import('@aztec/accounts/stub/schnorr/lazy');
-      return createStubSchnorrAccount(address);
-    } else {
-      const { createStubEcdsaAccount } = await import('@aztec/accounts/stub/ecdsa/lazy');
-      return createStubEcdsaAccount(address);
-    }
+  createStubAccount(address: CompleteAddress, type: AccountType): Promise<Account> {
+    return createStubAccountForAccountType(type, address);
   }
 
   getMulticallContract(): Promise<{ instance: ContractInstanceWithAddress; artifact: ContractArtifact }> {

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

@@ -1,11 +1,11 @@
+import { type AccountType, AccountTypes } from '@aztec/accounts/stub';
 import type { Aliased } from '@aztec/aztec.js/wallet';
 import { Fq, Fr } from '@aztec/foundation/curves/bn254';
 import type { LogFn } from '@aztec/foundation/log';
 import type { AztecAsyncKVStore, AztecAsyncMap } from '@aztec/kv-store';
 import { AztecAddress } from '@aztec/stdlib/aztec-address';
 
-export const AccountTypes = ['schnorr', 'ecdsasecp256r1', 'ecdsasecp256k1'] as const;
-export type AccountType = (typeof AccountTypes)[number];
+export { AccountTypes, type AccountType };
 
 function accountKey(field: string, address: AztecAddress | string): string {
   return `${field}:${address.toString()}`;