refactor: parallelizing kernel hints building

yarn-project/end-to-end/src/bench/client_flows/amm.test.ts

@@ -22,7 +22,7 @@ const AMOUNT_PER_NOTE = 1_000_000;
 const MINIMUM_NOTES_FOR_RECURSION_LEVEL = [0, 2, 10];
 
 // Set to true to print out the round trip information to the console.
-const DEBUG_ROUND_TRIPS = true;
+const DEBUG_ROUND_TRIPS = false;
 
 // Expected number of node round trips per account contract and payment method.
 const EXPECTED_ROUND_TRIPS: Record<string, number> = {

yarn-project/pxe/src/private_kernel/hints/private_kernel_reset_private_inputs_builder.ts

@@ -157,16 +157,6 @@ export class PrivateKernelResetPrivateInputsBuilder {
       allowRemainder,
     );
 
-    const previousVkMembershipWitness = await oracle.getVkMembershipWitness(
-      this.previousKernelOutput.verificationKey.keyAsFields,
-    );
-    const vkData = new VkData(
-      this.previousKernelOutput.verificationKey,
-      Number(previousVkMembershipWitness.leafIndex),
-      previousVkMembershipWitness.siblingPath,
-    );
-    const previousKernelData = new PrivateKernelData(this.previousKernelOutput.publicInputs, vkData);
-
     this.reduceReadRequestActions(
       this.noteHashResetActions,
       dimensions.NOTE_HASH_PENDING_READ,
@@ -178,29 +168,51 @@ export class PrivateKernelResetPrivateInputsBuilder {
       dimensions.NULLIFIER_SETTLED_READ,
     );
 
-    // TODO: Enable padding when we have a better idea what are the final amounts we should pad to.
-    const paddedSideEffects = PaddedSideEffects.empty();
-
-    return new PrivateKernelResetCircuitPrivateInputs(
-      previousKernelData,
-      paddedSideEffects,
-      new PrivateKernelResetHints(
-        await buildNoteHashReadRequestHintsFromResetActions(
+    // Execute all the expensive node querying operations in parallel.
+    const [previousVkMembershipWitness, noteHashReadRequestHints, nullifierReadRequestHints, keyValidationHints] =
+      await Promise.all([
+        oracle.getVkMembershipWitness(this.previousKernelOutput.verificationKey.keyAsFields),
+        buildNoteHashReadRequestHintsFromResetActions<
+          typeof MAX_NOTE_HASH_READ_REQUESTS_PER_TX,
+          typeof MAX_NOTE_HASH_READ_REQUESTS_PER_TX
+        >(
           oracle,
           this.previousKernel.validationRequests.noteHashReadRequests,
           this.previousKernel.end.noteHashes,
           this.noteHashResetActions,
         ),
-        await buildNullifierReadRequestHintsFromResetActions(
+        buildNullifierReadRequestHintsFromResetActions<
+          typeof MAX_NULLIFIER_READ_REQUESTS_PER_TX,
+          typeof MAX_NULLIFIER_READ_REQUESTS_PER_TX
+        >(
           { getNullifierMembershipWitness: getNullifierMembershipWitnessResolver(oracle) },
           this.previousKernel.validationRequests.nullifierReadRequests,
           this.nullifierResetActions,
         ),
-        await getMasterSecretKeysAndAppKeyGenerators(
+        getMasterSecretKeysAndAppKeyGenerators(
           this.previousKernel.validationRequests.scopedKeyValidationRequestsAndGenerators,
           dimensions.KEY_VALIDATION,
           oracle,
         ),
+      ]);
+
+    const vkData = new VkData(
+      this.previousKernelOutput.verificationKey,
+      Number(previousVkMembershipWitness.leafIndex),
+      previousVkMembershipWitness.siblingPath,
+    );
+    const previousKernelData = new PrivateKernelData(this.previousKernelOutput.publicInputs, vkData);
+
+    // TODO: Enable padding when we have a better idea what are the final amounts we should pad to.
+    const paddedSideEffects = PaddedSideEffects.empty();
+
+    return new PrivateKernelResetCircuitPrivateInputs(
+      previousKernelData,
+      paddedSideEffects,
+      new PrivateKernelResetHints(
+        noteHashReadRequestHints,
+        nullifierReadRequestHints,
+        keyValidationHints,
         this.transientDataSquashingHints,
       ),
       dimensions,

yarn-project/stdlib/src/kernel/hints/build_note_hash_read_request_hints.ts

@@ -76,17 +76,28 @@ export async function buildNoteHashReadRequestHintsFromResetActions<PENDING exte
     builder.addPendingReadRequest(hint.readRequestIndex, hint.pendingValueIndex);
   });
 
+  // Collect all settled read requests
+  const settledRequests: { index: number; readRequest: ScopedReadRequest }[] = [];
   for (let i = 0; i < resetActions.actions.length; i++) {
     if (resetActions.actions[i] === ReadRequestActionEnum.READ_AS_SETTLED) {
-      const readRequest = noteHashReadRequests.array[i];
-      const membershipWitness = await oracle.getNoteHashMembershipWitness(readRequest.value);
-      if (!membershipWitness) {
-        throw new Error('Read request is reading an unknown note hash.');
-      }
-      builder.addSettledReadRequest(i, membershipWitness, readRequest.value);
+      settledRequests.push({ index: i, readRequest: noteHashReadRequests.array[i] });
     }
   }
 
+  // Fetch all membership witnesses in parallel
+  const membershipWitnesses = await Promise.all(
+    settledRequests.map(({ readRequest }) => oracle.getNoteHashMembershipWitness(readRequest.value)),
+  );
+
+  // Add settled read requests to builder
+  for (let i = 0; i < settledRequests.length; i++) {
+    const membershipWitness = membershipWitnesses[i];
+    if (!membershipWitness) {
+      throw new Error('Read request is reading an unknown note hash.');
+    }
+    builder.addSettledReadRequest(settledRequests[i].index, membershipWitness, settledRequests[i].readRequest.value);
+  }
+
   const noteHashMap: Map<bigint, { noteHash: ScopedNoteHash; index: number }[]> = new Map();
   noteHashes.getActiveItems().forEach((noteHash, index) => {
     const value = noteHash.value.toBigInt();

yarn-project/stdlib/src/kernel/hints/build_nullifier_read_request_hints.ts

@@ -82,21 +82,33 @@ export async function buildNullifierReadRequestHintsFromResetActions<PENDING ext
     builder.addPendingReadRequest(hint.readRequestIndex, hint.pendingValueIndex);
   });
 
+  // Collect all settled read requests
+  const settledRequests: { index: number; readRequest: ScopedReadRequest }[] = [];
   for (let i = 0; i < resetActions.actions.length; i++) {
     if (resetActions.actions[i] === ReadRequestActionEnum.READ_AS_SETTLED) {
-      const readRequest = nullifierReadRequests.array[i];
-      const siloedValue = siloed
-        ? readRequest.value
-        : await siloNullifier(readRequest.contractAddress, readRequest.value);
-      const membershipWitnessWithPreimage = await oracle.getNullifierMembershipWitness(siloedValue);
-      builder.addSettledReadRequest(
-        i,
-        membershipWitnessWithPreimage.membershipWitness,
-        membershipWitnessWithPreimage.leafPreimage,
-      );
+      settledRequests.push({ index: i, readRequest: nullifierReadRequests.array[i] });
     }
   }
 
+  // Compute siloed values in parallel (if not already siloed)
+  const siloedValues = siloed
+    ? settledRequests.map(({ readRequest }) => readRequest.value)
+    : await Promise.all(
+        settledRequests.map(({ readRequest }) => siloNullifier(readRequest.contractAddress, readRequest.value)),
+      );
+
+  // Fetch all membership witnesses in parallel
+  const membershipWitnesses = await Promise.all(siloedValues.map(value => oracle.getNullifierMembershipWitness(value)));
+
+  // Add settled read requests to builder
+  for (let i = 0; i < settledRequests.length; i++) {
+    builder.addSettledReadRequest(
+      settledRequests[i].index,
+      membershipWitnesses[i].membershipWitness,
+      membershipWitnesses[i].leafPreimage,
+    );
+  }
+
   return builder.toHints();
 }