refactor(pxe): prefetch updated class id hints per unique contract (#23130)

Author: nchamo

yarn-project/pxe/src/private_kernel/private_kernel_execution_prover.test.ts

@@ -54,16 +54,20 @@ describe('Private Kernel Sequencer', () => {
     {
       publicInputs,
       childPublicInputs = [],
+      address = contractAddress,
+      nestedResults,
     }: {
       publicInputs?: PrivateCircuitPublicInputs;
       childPublicInputs?: PrivateCircuitPublicInputs[];
+      address?: AztecAddress;
+      nestedResults?: PrivateCallExecutionResult[];
     } = {},
   ): PrivateCallExecutionResult => {
     if (!publicInputs) {
       publicInputs = PrivateCircuitPublicInputs.empty();
     }
     publicInputs.callContext.functionSelector = new FunctionSelector(fnName.charCodeAt(0));
-    publicInputs.callContext.contractAddress = contractAddress;
+    publicInputs.callContext.contractAddress = address;
 
     return new PrivateCallExecutionResult(
       Buffer.alloc(0),
@@ -75,9 +79,10 @@ describe('Private Kernel Sequencer', () => {
       [],
       [],
       [],
-      (dependencies[fnName] || []).map((name, i) =>
-        createCallExecutionResult(name, { publicInputs: childPublicInputs[i] }),
-      ),
+      nestedResults ??
+        (dependencies[fnName] || []).map((name, i) =>
+          createCallExecutionResult(name, { publicInputs: childPublicInputs[i] }),
+        ),
       [],
     );
   };
@@ -360,4 +365,22 @@ describe('Private Kernel Sequencer', () => {
     expect(proofCreator.simulateReset).toHaveBeenCalledTimes(3);
     expect(proofCreator.simulateTail).toHaveBeenCalledTimes(1);
   });
+
+  it('fetches updated class id hints once per unique contract address', async () => {
+    const contractAddressB = AztecAddress.fromBigInt(111111n);
+
+    // a { b {} c {} }
+    // a and c use contractAddress, b uses contractAddressB → 2 unique contracts, 3 executions.
+    dependencies = {};
+    const bExec = createCallExecutionResult('b', { address: contractAddressB });
+    const cExec = createCallExecutionResult('c');
+    const aExec = createCallExecutionResult('a', { nestedResults: [bExec, cExec] });
+
+    const executionResult = new PrivateExecutionResult(aExec, Fr.zero(), []);
+    await prove(executionResult);
+
+    expect(oracle.getUpdatedClassIdHints).toHaveBeenCalledTimes(2);
+    expect(oracle.getUpdatedClassIdHints).toHaveBeenCalledWith(contractAddress);
+    expect(oracle.getUpdatedClassIdHints).toHaveBeenCalledWith(contractAddressB);
+  });
 });

yarn-project/pxe/src/private_kernel/private_kernel_execution_prover.ts

@@ -1,3 +1,4 @@
+import { uniqueBy } from '@aztec/foundation/collection';
 import { vkAsFieldsMegaHonk } from '@aztec/foundation/crypto/keys';
 import { Fr } from '@aztec/foundation/curves/bn254';
 import { type Logger, type LoggerBindings, createLogger } from '@aztec/foundation/log';
@@ -22,12 +23,14 @@ import {
   PrivateKernelTailCircuitPrivateInputs,
   type PrivateKernelTailCircuitPublicInputs,
   PrivateVerificationKeyHints,
+  type UpdatedClassIdHints,
 } from '@aztec/stdlib/kernel';
 import { ChonkProof, ChonkProofWithPublicInputs } from '@aztec/stdlib/proofs';
 import {
   type PrivateCallExecutionResult,
   type PrivateExecutionResult,
   TxRequest,
+  collectNested,
   collectNoteHashNullifierCounterMap,
   getFinalMinRevertibleSideEffectCounter,
 } from '@aztec/stdlib/tx';
@@ -108,6 +111,8 @@ export class PrivateKernelExecutionProver {
     const minRevertibleSideEffectCounter = getFinalMinRevertibleSideEffectCounter(executionResult);
     const splitCounter = isPrivateOnlyTx ? 0 : minRevertibleSideEffectCounter;
 
+    const updatedClassIdHintsMap = await this.prefetchUpdatedClassIdHints(executionResult);
+
     while (executionStack.length) {
       if (!firstIteration) {
         let resetBuilder = new PrivateKernelResetPrivateInputsBuilder(
@@ -161,7 +166,7 @@ export class PrivateKernelExecutionProver {
         },
       });
 
-      const privateCallData = await this.createPrivateCallData(currentExecution);
+      const privateCallData = await this.createPrivateCallData(currentExecution, updatedClassIdHintsMap);
 
       if (firstIteration) {
         const witgenTimer = new Timer();
@@ -401,7 +406,28 @@ export class PrivateKernelExecutionProver {
     );
   }
 
-  private async createPrivateCallData({ publicInputs, vk: vkAsBuffer }: PrivateCallExecutionResult) {
+  /** Prefetches updated class id hints for all unique contracts in the execution tree in parallel. */
+  private async prefetchUpdatedClassIdHints(
+    executionResult: PrivateExecutionResult,
+  ): Promise<Map<string, UpdatedClassIdHints>> {
+    const allAddresses = collectNested([executionResult.entrypoint], exec => [
+      exec.publicInputs.callContext.contractAddress,
+    ]);
+    const uniqueAddresses = uniqueBy(allAddresses, a => a.toString());
+    return new Map<string, UpdatedClassIdHints>(
+      await Promise.all(
+        uniqueAddresses.map(
+          async addr =>
+            [addr.toString(), await this.oracle.getUpdatedClassIdHints(addr)] as [string, UpdatedClassIdHints],
+        ),
+      ),
+    );
+  }
+
+  private async createPrivateCallData(
+    { publicInputs, vk: vkAsBuffer }: PrivateCallExecutionResult,
+    updatedClassIdHintsMap: Map<string, UpdatedClassIdHints>,
+  ) {
     const { contractAddress, functionSelector } = publicInputs.callContext;
 
     const vkAsFields = await vkAsFieldsMegaHonk(vkAsBuffer);
@@ -417,7 +443,7 @@ export class PrivateKernelExecutionProver {
     const { artifactHash: contractClassArtifactHash, publicBytecodeCommitment: contractClassPublicBytecodeCommitment } =
       await this.oracle.getContractClassIdPreimage(currentContractClassId);
 
-    const updatedClassIdHints = await this.oracle.getUpdatedClassIdHints(contractAddress);
+    const updatedClassIdHints = updatedClassIdHintsMap.get(contractAddress.toString())!;
 
     return PrivateCallData.from({
       publicInputs,