fix(node): fix index misalignment in findLeavesIndexes

yarn-project/aztec-node/src/aztec-node/server.test.ts

@@ -16,7 +16,7 @@ import { computeFeePayerBalanceLeafSlot } from '@aztec/protocol-contracts/fee-ju
 import type { GlobalVariableBuilder, SequencerClient } from '@aztec/sequencer-client';
 import type { SlasherClientInterface } from '@aztec/slasher';
 import { AztecAddress } from '@aztec/stdlib/aztec-address';
-import { L2Block, type L2BlockSource } from '@aztec/stdlib/block';
+import { BlockHash, L2Block, type L2BlockSource } from '@aztec/stdlib/block';
 import type { ContractDataSource } from '@aztec/stdlib/contract';
 import { EmptyL1RollupConstants } from '@aztec/stdlib/epoch-helpers';
 import { GasFees } from '@aztec/stdlib/gas';
@@ -397,6 +397,130 @@ describe('aztec node', () => {
         expect(l2BlockSource.getL2Block).toHaveBeenCalledWith(3);
       });
     });
+
+    describe('findLeavesIndexes', () => {
+      const blockHash1 = Fr.random();
+      const blockHash2 = Fr.random();
+
+      beforeEach(() => {
+        lastBlockNumber = BlockNumber(2);
+      });
+
+      it('returns results for all found leaves', async () => {
+        merkleTreeOps.findLeafIndices.mockResolvedValue([10n, 20n]);
+        merkleTreeOps.getBlockNumbersForLeafIndices.mockResolvedValue([BlockNumber(1), BlockNumber(2)]);
+        (merkleTreeOps as any).getLeafValue.mockImplementation((_treeId: any, index: bigint) => {
+          if (index === 1n) {
+            return Promise.resolve(blockHash1);
+          }
+          if (index === 2n) {
+            return Promise.resolve(blockHash2);
+          }
+          return Promise.resolve(undefined);
+        });
+
+        const result = await node.findLeavesIndexes('latest', MerkleTreeId.NOTE_HASH_TREE, [Fr.random(), Fr.random()]);
+
+        expect(result).toEqual([
+          { l2BlockNumber: BlockNumber(1), l2BlockHash: new BlockHash(blockHash1), data: 10n },
+          { l2BlockNumber: BlockNumber(2), l2BlockHash: new BlockHash(blockHash2), data: 20n },
+        ]);
+      });
+
+      it('returns undefined for leaves not found', async () => {
+        merkleTreeOps.findLeafIndices.mockResolvedValue([undefined, undefined]);
+
+        const result = await node.findLeavesIndexes('latest', MerkleTreeId.NOTE_HASH_TREE, [Fr.random(), Fr.random()]);
+
+        expect(result).toEqual([undefined, undefined]);
+      });
+
+      it('returns correct results when some leaves are not found', async () => {
+        merkleTreeOps.findLeafIndices.mockResolvedValue([undefined, 10n, 20n]);
+        merkleTreeOps.getBlockNumbersForLeafIndices.mockResolvedValue([BlockNumber(1), BlockNumber(2)]);
+        (merkleTreeOps as any).getLeafValue.mockImplementation((_treeId: any, index: bigint) => {
+          if (index === 1n) {
+            return Promise.resolve(blockHash1);
+          }
+          if (index === 2n) {
+            return Promise.resolve(blockHash2);
+          }
+          return Promise.resolve(undefined);
+        });
+
+        const result = await node.findLeavesIndexes('latest', MerkleTreeId.NOTE_HASH_TREE, [
+          Fr.random(),
+          Fr.random(),
+          Fr.random(),
+        ]);
+
+        expect(result).toEqual([
+          undefined,
+          { l2BlockNumber: BlockNumber(1), l2BlockHash: new BlockHash(blockHash1), data: 10n },
+          { l2BlockNumber: BlockNumber(2), l2BlockHash: new BlockHash(blockHash2), data: 20n },
+        ]);
+        // Only defined indices should be passed
+        expect(merkleTreeOps.getBlockNumbersForLeafIndices).toHaveBeenCalledWith(MerkleTreeId.NOTE_HASH_TREE, [
+          10n,
+          20n,
+        ]);
+      });
+
+      it('handles multiple leaves in the same block', async () => {
+        merkleTreeOps.findLeafIndices.mockResolvedValue([10n, 20n, 30n]);
+        merkleTreeOps.getBlockNumbersForLeafIndices.mockResolvedValue([BlockNumber(1), BlockNumber(1), BlockNumber(2)]);
+        (merkleTreeOps as any).getLeafValue.mockImplementation((_treeId: any, index: bigint) => {
+          if (index === 1n) {
+            return Promise.resolve(blockHash1);
+          }
+          if (index === 2n) {
+            return Promise.resolve(blockHash2);
+          }
+          return Promise.resolve(undefined);
+        });
+
+        const result = await node.findLeavesIndexes('latest', MerkleTreeId.NOTE_HASH_TREE, [
+          Fr.random(),
+          Fr.random(),
+          Fr.random(),
+        ]);
+
+        expect(result).toEqual([
+          { l2BlockNumber: BlockNumber(1), l2BlockHash: new BlockHash(blockHash1), data: 10n },
+          { l2BlockNumber: BlockNumber(1), l2BlockHash: new BlockHash(blockHash1), data: 20n },
+          { l2BlockNumber: BlockNumber(2), l2BlockHash: new BlockHash(blockHash2), data: 30n },
+        ]);
+        // getLeafValue should be called only for unique block numbers
+        expect(merkleTreeOps.getLeafValue).toHaveBeenCalledTimes(2);
+      });
+
+      it('returns empty array for empty input', async () => {
+        merkleTreeOps.findLeafIndices.mockResolvedValue([]);
+
+        const result = await node.findLeavesIndexes('latest', MerkleTreeId.NOTE_HASH_TREE, []);
+
+        expect(result).toEqual([]);
+      });
+
+      it('throws when block number is undefined for a found leaf', async () => {
+        merkleTreeOps.findLeafIndices.mockResolvedValue([10n]);
+        merkleTreeOps.getBlockNumbersForLeafIndices.mockResolvedValue([undefined]);
+
+        await expect(node.findLeavesIndexes('latest', MerkleTreeId.NOTE_HASH_TREE, [Fr.random()])).rejects.toThrow(
+          /Block number is undefined/,
+        );
+      });
+
+      it('throws when block hash is undefined for a found block number', async () => {
+        merkleTreeOps.findLeafIndices.mockResolvedValue([10n]);
+        merkleTreeOps.getBlockNumbersForLeafIndices.mockResolvedValue([BlockNumber(1)]);
+        merkleTreeOps.getLeafValue.mockResolvedValue(undefined);
+
+        await expect(node.findLeavesIndexes('latest', MerkleTreeId.NOTE_HASH_TREE, [Fr.random()])).rejects.toThrow(
+          /Block hash is undefined/,
+        );
+      });
+    });
   });
 
   describe('simulatePublicCalls', () => {

yarn-project/aztec-node/src/aztec-node/server.ts

@@ -971,53 +971,59 @@ export class AztecNodeService implements AztecNode, AztecNodeAdmin, Traceable {
       treeId,
       leafValues.map(x => x.toBuffer()),
     );
-    // We filter out undefined values
-    const indices = maybeIndices.filter(x => x !== undefined) as bigint[];
+    // Filter out undefined values to query block numbers only for found leaves
+    const definedIndices = maybeIndices.filter(x => x !== undefined);
 
-    // Now we find the block numbers for the indices
-    const blockNumbers = await committedDb.getBlockNumbersForLeafIndices(treeId, indices);
+    // Now we find the block numbers for the defined indices
+    const blockNumbers = await committedDb.getBlockNumbersForLeafIndices(treeId, definedIndices);
 
-    // If any of the block numbers are undefined, we throw an error.
-    for (let i = 0; i < indices.length; i++) {
-      if (blockNumbers[i] === undefined) {
-        throw new Error(`Block number is undefined for leaf index ${indices[i]} in tree ${MerkleTreeId[treeId]}`);
+    // Build a map from leaf index to block number
+    const indexToBlockNumber = new Map<bigint, BlockNumber>();
+    for (let i = 0; i < definedIndices.length; i++) {
+      const blockNumber = blockNumbers[i];
+      if (blockNumber === undefined) {
+        throw new Error(
+          `Block number is undefined for leaf index ${definedIndices[i]} in tree ${MerkleTreeId[treeId]}`,
+        );
       }
+      indexToBlockNumber.set(definedIndices[i], blockNumber);
     }
 
     // Get unique block numbers in order to optimize num calls to getLeafValue function.
-    const uniqueBlockNumbers = [...new Set(blockNumbers.filter(x => x !== undefined))];
+    const uniqueBlockNumbers = [...new Set(indexToBlockNumber.values())];
 
-    // Now we obtain the block hashes from the archive tree by calling await `committedDb.getLeafValue(treeId, index)`
-    // (note that block number corresponds to the leaf index in the archive tree).
+    // Now we obtain the block hashes from the archive tree (block number = leaf index in archive tree).
     const blockHashes = await Promise.all(
       uniqueBlockNumbers.map(blockNumber => {
         return committedDb.getLeafValue(MerkleTreeId.ARCHIVE, BigInt(blockNumber));
       }),
     );
 
-    // If any of the block hashes are undefined, we throw an error.
+    // Build a map from block number to block hash
+    const blockNumberToHash = new Map<BlockNumber, Fr>();
     for (let i = 0; i < uniqueBlockNumbers.length; i++) {
-      if (blockHashes[i] === undefined) {
+      const blockHash = blockHashes[i];
+      if (blockHash === undefined) {
         throw new Error(`Block hash is undefined for block number ${uniqueBlockNumbers[i]}`);
       }
+      blockNumberToHash.set(uniqueBlockNumbers[i], blockHash);
     }
 
     // Create DataInBlock objects by combining indices, blockNumbers and blockHashes and return them.
-    return maybeIndices.map((index, i) => {
+    return maybeIndices.map(index => {
       if (index === undefined) {
         return undefined;
       }
-      const blockNumber = blockNumbers[i];
+      const blockNumber = indexToBlockNumber.get(index);
       if (blockNumber === undefined) {
-        return undefined;
+        throw new Error(`Block number not found for leaf index ${index} in tree ${MerkleTreeId[treeId]}`);
       }
-      const blockHashIndex = uniqueBlockNumbers.indexOf(blockNumber);
-      const blockHash = blockHashes[blockHashIndex];
-      if (!blockHash) {
-        return undefined;
+      const blockHash = blockNumberToHash.get(blockNumber);
+      if (blockHash === undefined) {
+        throw new Error(`Block hash not found for block number ${blockNumber}`);
       }
       return {
-        l2BlockNumber: BlockNumber(Number(blockNumber)),
+        l2BlockNumber: blockNumber,
         l2BlockHash: new BlockHash(blockHash),
         data: index,
       };