refactor(pxe): deduplicate tx hash lookups in MessageContextService (#23075)

Author: nchamo

yarn-project/pxe/src/messages/message_context_service.test.ts

@@ -3,7 +3,7 @@ import { Fr } from '@aztec/foundation/curves/bn254';
 import { BlockHash } from '@aztec/stdlib/block';
 import type { AztecNode } from '@aztec/stdlib/interfaces/server';
 import { MessageContext } from '@aztec/stdlib/logs';
-import { TxHash } from '@aztec/stdlib/tx';
+import { TxEffect, TxHash } from '@aztec/stdlib/tx';
 
 import { mock } from 'jest-mock-extended';
 
@@ -123,4 +123,26 @@ describe('MessageContextService', () => {
     // Zero hash should not trigger getTxEffect
     expect(aztecNode.getTxEffect).toHaveBeenCalledTimes(3);
   });
+
+  it('deduplicates repeated tx hashes so each is fetched only once', async () => {
+    const txEffect = TxEffect.from({
+      ...(await TxEffect.random({ numNoteHashes: 1, numNullifiers: 1 })),
+    });
+
+    aztecNode.getTxEffect.mockResolvedValue({
+      l2BlockNumber: BlockNumber(anchorBlockNumber - 1),
+      l2BlockHash: BlockHash.random(),
+      txIndexInBlock: 0,
+      data: txEffect,
+    });
+
+    const results = await service.getMessageContextsByTxHash(
+      [txEffect.txHash.hash, txEffect.txHash.hash, txEffect.txHash.hash],
+      anchorBlockNumber,
+    );
+
+    const expected = new MessageContext(txEffect.txHash, txEffect.noteHashes, txEffect.nullifiers[0]);
+    expect(results).toEqual([expected, expected, expected]);
+    expect(aztecNode.getTxEffect).toHaveBeenCalledTimes(1);
+  });
 });

yarn-project/pxe/src/messages/message_context_service.ts

@@ -1,7 +1,8 @@
+import { uniqueBy } from '@aztec/foundation/collection';
 import { Fr } from '@aztec/foundation/curves/bn254';
 import type { AztecNode } from '@aztec/stdlib/interfaces/server';
 import { MessageContext } from '@aztec/stdlib/logs';
-import { TxHash } from '@aztec/stdlib/tx';
+import { type IndexedTxEffect, TxHash } from '@aztec/stdlib/tx';
 
 /** Resolves transaction hashes into the context needed to process messages. */
 export class MessageContextService {
@@ -14,31 +15,31 @@ export class MessageContextService {
    * process messages that originated from that transaction. Returns `null` for tx hashes that are zero, not yet
    * available, or in blocks beyond the anchor block.
    */
-  getMessageContextsByTxHash(txHashes: Fr[], anchorBlockNumber: number): Promise<(MessageContext | null)[]> {
-    // TODO: optimize, we might be hitting the node to get the same txHash repeatedly
-    return Promise.all(
-      txHashes.map(async txHashField => {
-        // A zero tx hash indicates a tx-less offchain message (e.g. one not tied to any onchain transaction).
-        // These messages don't have a transaction context to resolve, so we return null.
-        if (txHashField.isZero()) {
-          return null;
-        }
+  async getMessageContextsByTxHash(txHashes: Fr[], anchorBlockNumber: number): Promise<(MessageContext | null)[]> {
+    const nonZeroTxHashes = txHashes.filter(h => !h.isZero()).map(h => TxHash.fromField(h));
+    const uniqueTxHashes = uniqueBy(nonZeroTxHashes, h => h.toString());
+    const fetched = await Promise.all(uniqueTxHashes.map(h => this.aztecNode.getTxEffect(h)));
+    const txEffects = new Map(
+      uniqueTxHashes
+        .map((h, i): [string, IndexedTxEffect | undefined] => [h.toString(), fetched[i]])
+        .filter((entry): entry is [string, IndexedTxEffect] => entry[1] !== undefined),
+    );
 
-        const txHash = TxHash.fromField(txHashField);
-        const txEffect = await this.aztecNode.getTxEffect(txHash);
-        if (!txEffect || txEffect.l2BlockNumber > anchorBlockNumber) {
-          return null;
-        }
+    return txHashes.map(txHashField => {
+      const txHash = TxHash.fromField(txHashField);
+      const txEffect = txEffects.get(txHash.toString());
+      if (!txEffect || txEffect.l2BlockNumber > anchorBlockNumber) {
+        return null;
+      }
 
-        // Every tx has at least one nullifier (the first nullifier derived from the tx hash). Hitting this condition
-        // would mean a buggy node, but since we need to access data.nullifiers[0], the defensive check does no harm.
-        const data = txEffect.data;
-        if (data.nullifiers.length === 0) {
-          throw new Error(`Tx effect for ${txHash} has no nullifiers`);
-        }
+      // Every tx has at least one nullifier (the first nullifier derived from the tx hash). Hitting this condition
+      // would mean a buggy node, but since we need to access data.nullifiers[0], the defensive check does no harm.
+      const data = txEffect.data;
+      if (data.nullifiers.length === 0) {
+        throw new Error(`Tx effect for ${txHash} has no nullifiers`);
+      }
 
-        return new MessageContext(data.txHash, data.noteHashes, data.nullifiers[0]);
-      }),
-    );
+      return new MessageContext(data.txHash, data.noteHashes, data.nullifiers[0]);
+    });
   }
 }