perf: speed up valkeys generation (A-1361)

Author: alexghr

yarn-project/cli/src/cmds/validator_keys/shared.ts

@@ -1,4 +1,5 @@
 import { prettyPrintJSON } from '@aztec/cli/utils';
+import { times, timesParallel } from '@aztec/foundation/collection';
 import { deriveBlsPrivateKey } from '@aztec/foundation/crypto/bls';
 import { createBn254Keystore } from '@aztec/foundation/crypto/bls/bn254_keystore';
 import { computeBn254G1PublicKeyCompressed } from '@aztec/foundation/crypto/bn254';
@@ -91,62 +92,57 @@ export async function buildValidatorEntries(input: BuildValidatorsInput) {
     remoteSigner,
   } = input;
 
-  const summaries: ValidatorSummary[] = [];
-
-  const validators = await Promise.all(
-    Array.from({ length: validatorCount }, async (_unused, i) => {
-      const addressIndex = baseAddressIndex + i;
-      const basePath = blsPath ?? defaultBlsPath;
-      const perValidatorPath = withValidatorIndex(basePath, accountIndex, addressIndex);
-
-      const blsPrivKey = ikm || mnemonic ? deriveBlsPrivateKey(mnemonic, ikm, perValidatorPath) : undefined;
-      const blsPubCompressed = blsPrivKey ? await computeBlsPublicKeyCompressed(blsPrivKey) : undefined;
-
-      const ethAttester = deriveEthAttester(mnemonic, accountIndex, addressIndex, remoteSigner);
-      const attester = blsPrivKey ? { eth: ethAttester, bls: blsPrivKey } : ethAttester;
-
-      let publisherField: EthAccount | EthPrivateKey | (EthAccount | EthPrivateKey)[] | undefined;
-      const publisherAddresses: string[] = [];
-      if (publishers && publishers.length > 0) {
-        publisherAddresses.push(...publishers);
-        publisherField = publishers.length === 1 ? (publishers[0] as EthPrivateKey) : (publishers as EthPrivateKey[]);
-      } else if (publisherCount > 0) {
-        const publishersBaseIndex = baseAddressIndex + validatorCount + i * publisherCount;
-        const publisherAccounts = Array.from({ length: publisherCount }, (_unused2, j) => {
-          const publisherAddressIndex = publishersBaseIndex + j;
-          const pubAcct = mnemonicToAccount(mnemonic, {
-            accountIndex,
-            addressIndex: publisherAddressIndex,
-          });
-          publisherAddresses.push(pubAcct.address as unknown as string);
-          return remoteSigner
-            ? ({ address: pubAcct.address as unknown as EthAddress, remoteSignerUrl: remoteSigner } as EthAccount)
-            : (('0x' + Buffer.from(pubAcct.getHdKey().privateKey!).toString('hex')) as EthPrivateKey);
+  const entries = await timesParallel(validatorCount, async i => {
+    const addressIndex = baseAddressIndex + i;
+    const basePath = blsPath ?? defaultBlsPath;
+    const perValidatorPath = withValidatorIndex(basePath, accountIndex, addressIndex);
+
+    const blsPrivKey = ikm || mnemonic ? deriveBlsPrivateKey(mnemonic, ikm, perValidatorPath) : undefined;
+    const blsPubCompressed = blsPrivKey ? await computeBlsPublicKeyCompressed(blsPrivKey) : undefined;
+    const attesterAccount = mnemonicToAccount(mnemonic, { accountIndex, addressIndex });
+    const attesterEthAddress = attesterAccount.address as unknown as string;
+    const ethAttester = remoteSigner
+      ? ({ address: attesterAccount.address as unknown as EthAddress, remoteSignerUrl: remoteSigner } as EthAccount)
+      : (('0x' + Buffer.from(attesterAccount.getHdKey().privateKey!).toString('hex')) as EthPrivateKey);
+    const attester = blsPrivKey ? { eth: ethAttester, bls: blsPrivKey } : ethAttester;
+
+    let publisherField: EthAccount | EthPrivateKey | (EthAccount | EthPrivateKey)[] | undefined;
+    const publisherAddresses: string[] = [];
+    if (publishers && publishers.length > 0) {
+      publisherAddresses.push(...publishers);
+      publisherField = publishers.length === 1 ? (publishers[0] as EthPrivateKey) : (publishers as EthPrivateKey[]);
+    } else if (publisherCount > 0) {
+      const publishersBaseIndex = baseAddressIndex + validatorCount + i * publisherCount;
+      const publisherAccounts = times(publisherCount, j => {
+        const publisherAddressIndex = publishersBaseIndex + j;
+        const pubAcct = mnemonicToAccount(mnemonic, {
+          accountIndex,
+          addressIndex: publisherAddressIndex,
         });
-        publisherField = publisherCount === 1 ? publisherAccounts[0] : publisherAccounts;
-      }
-
-      const acct = mnemonicToAccount(mnemonic, {
-        accountIndex,
-        addressIndex,
-      });
-      const attesterEthAddress = acct.address as unknown as string;
-      summaries.push({
-        attesterEth: attesterEthAddress,
-        attesterBls: blsPubCompressed,
-        publisherEth: publisherAddresses.length > 0 ? publisherAddresses : undefined,
+        publisherAddresses.push(pubAcct.address as unknown as string);
+        return remoteSigner
+          ? ({ address: pubAcct.address as unknown as EthAddress, remoteSignerUrl: remoteSigner } as EthAccount)
+          : (('0x' + Buffer.from(pubAcct.getHdKey().privateKey!).toString('hex')) as EthPrivateKey);
       });
+      publisherField = publisherCount === 1 ? publisherAccounts[0] : publisherAccounts;
+    }
 
-      return {
+    return {
+      validator: {
         attester,
         ...(publisherField !== undefined ? { publisher: publisherField } : {}),
         feeRecipient,
         coinbase: coinbase ?? attesterEthAddress,
-      } as ValidatorKeyStore;
-    }),
-  );
+      } as ValidatorKeyStore,
+      summary: {
+        attesterEth: attesterEthAddress,
+        attesterBls: blsPubCompressed,
+        publisherEth: publisherAddresses.length > 0 ? publisherAddresses : undefined,
+      } as ValidatorSummary,
+    };
+  });
 
-  return { validators, summaries };
+  return { validators: entries.map(entry => entry.validator), summaries: entries.map(entry => entry.summary) };
 }
 
 export async function resolveKeystoreOutputPath(dataDir?: string, file?: string) {
@@ -228,7 +224,7 @@ export async function writeBn254BlsKeystore(
 ): Promise<string> {
   mkdirSync(outDir, { recursive: true });
 
-  const keystore = createBn254Keystore(password, privateKeyHex, pubkeyHex, derivationPath);
+  const keystore = await createBn254Keystore(password, privateKeyHex, pubkeyHex, derivationPath);
 
   const safeBase = fileNameBase.replace(/[^a-zA-Z0-9_-]/g, '_');
   const outPath = join(outDir, `keystore-${safeBase}.json`);
@@ -241,28 +237,29 @@ export async function writeBlsBn254ToFile(
   validators: ValidatorKeyStore[],
   options: { outDir: string; password: string; blsPath?: string },
 ): Promise<void> {
-  for (let i = 0; i < validators.length; i++) {
-    const v = validators[i];
-    if (!v || typeof v !== 'object' || !('attester' in v)) {
-      continue;
-    }
-    const att = (v as any).attester;
+  await Promise.all(
+    validators.map(async (v, i) => {
+      if (!v || typeof v !== 'object' || !('attester' in v)) {
+        return;
+      }
+      const att = (v as any).attester;
 
-    // Shapes: { bls: <hex> } or { eth: <ethAccount>, bls?: <hex> } or plain EthAccount
-    const blsKey: string | undefined = typeof att === 'object' && 'bls' in att ? (att as any).bls : undefined;
-    if (!blsKey || typeof blsKey !== 'string') {
-      continue;
-    }
+      // Shapes: { bls: <hex> } or { eth: <ethAccount>, bls?: <hex> } or plain EthAccount
+      const blsKey: string | undefined = typeof att === 'object' && 'bls' in att ? (att as any).bls : undefined;
+      if (!blsKey || typeof blsKey !== 'string') {
+        return;
+      }
 
-    const pub = await computeBlsPublicKeyCompressed(blsKey);
-    const path = options.blsPath ?? defaultBlsPath;
-    const fileBase = `${String(i + 1)}_${pub.slice(2, 18)}`;
-    const keystorePath = await writeBn254BlsKeystore(options.outDir, fileBase, options.password, blsKey, pub, path);
+      const pub = await computeBlsPublicKeyCompressed(blsKey);
+      const path = options.blsPath ?? defaultBlsPath;
+      const fileBase = `${String(i + 1)}_${pub.slice(2, 18)}`;
+      const keystorePath = await writeBn254BlsKeystore(options.outDir, fileBase, options.password, blsKey, pub, path);
 
-    if (typeof att === 'object') {
-      (att as any).bls = { path: keystorePath, password: options.password };
-    }
-  }
+      if (typeof att === 'object') {
+        (att as any).bls = { path: keystorePath, password: options.password };
+      }
+    }),
+  );
 }
 
 /** Writes an Ethereum JSON V3 keystore using ethers, returns absolute path */
@@ -295,32 +292,31 @@ export async function writeEthJsonV3ToFile(
     return account;
   };
 
-  for (let i = 0; i < validators.length; i++) {
-    const v = validators[i];
-    if (!v || typeof v !== 'object') {
-      continue;
-    }
+  await Promise.all(
+    validators.map(async (v, i) => {
+      if (!v || typeof v !== 'object') {
+        return;
+      }
 
-    // attester may be string (eth), object with eth, or remote signer
-    const att = (v as any).attester;
-    if (typeof att === 'string') {
-      (v as any).attester = await maybeEncryptEth(att, `attester_${i + 1}`);
-    } else if (att && typeof att === 'object' && 'eth' in att) {
-      (att as any).eth = await maybeEncryptEth((att as any).eth, `attester_${i + 1}`);
-    }
+      // attester may be string (eth), object with eth, or remote signer
+      const att = (v as any).attester;
+      if (typeof att === 'string') {
+        (v as any).attester = await maybeEncryptEth(att, `attester_${i + 1}`);
+      } else if (att && typeof att === 'object' && 'eth' in att) {
+        (att as any).eth = await maybeEncryptEth((att as any).eth, `attester_${i + 1}`);
+      }
 
-    // publisher can be single or array
-    if ('publisher' in v) {
-      const pub = (v as any).publisher;
-      if (Array.isArray(pub)) {
-        const out: any[] = [];
-        for (let j = 0; j < pub.length; j++) {
-          out.push(await maybeEncryptEth(pub[j], `publisher_${i + 1}_${j + 1}`));
+      // publisher can be single or array
+      if ('publisher' in v) {
+        const pub = (v as any).publisher;
+        if (Array.isArray(pub)) {
+          (v as any).publisher = await Promise.all(
+            pub.map((account, j) => maybeEncryptEth(account, `publisher_${i + 1}_${j + 1}`)),
+          );
+        } else if (pub !== undefined) {
+          (v as any).publisher = await maybeEncryptEth(pub, `publisher_${i + 1}`);
         }
-        (v as any).publisher = out;
-      } else if (pub !== undefined) {
-        (v as any).publisher = await maybeEncryptEth(pub, `publisher_${i + 1}`);
       }
-    }
-  }
+    }),
+  );
 }

yarn-project/cli/src/cmds/validator_keys/valkeys.test.ts

@@ -181,6 +181,21 @@ describe('validator keys utilities', () => {
       expect(summaries[0].publisherEth!.length).toBe(3);
     });
 
+    it('preserves summary order for multiple validators', async () => {
+      const { summaries } = await buildValidatorEntries({
+        validatorCount: 3,
+        publisherCount: 0,
+        accountIndex: 0,
+        baseAddressIndex: 0,
+        mnemonic: TEST_MNEMONIC,
+        feeRecipient: feeRecipient,
+      });
+
+      expect(summaries.map(summary => summary.attesterEth)).toEqual(
+        [0, 1, 2].map(addressIndex => mnemonicToAccount(TEST_MNEMONIC, { accountIndex: 0, addressIndex }).address),
+      );
+    });
+
     it('derives different BLS and ETH keys when account index changes', async () => {
       // Build with account index 0
       const { validators: v1, summaries: s1 } = await buildValidatorEntries({