refactor(p2p): (A-588) maintain sorted array in tx pool instead of sorting on read

yarn-project/p2p/src/mem_pools/tx_pool_v2/tx_metadata.ts

@@ -158,13 +158,13 @@ export function txHashFromBigInt(value: bigint): string {
 }
 
 /** Minimal fields required for priority comparison. */
-type PriorityComparable = Pick<TxMetaData, 'txHashBigInt' | 'priorityFee'>;
+export type PriorityComparable = Pick<TxMetaData, 'txHash' | 'txHashBigInt' | 'priorityFee'>;
 
 /**
  * Compares two priority fees in ascending order.
  * Returns negative if a < b, positive if a > b, 0 if equal.
  */
-export function compareFee(a: bigint, b: bigint): number {
+export function compareFee(a: bigint, b: bigint): -1 | 0 | 1 {
   return a < b ? -1 : a > b ? 1 : 0;
 }
 
@@ -173,7 +173,7 @@ export function compareFee(a: bigint, b: bigint): number {
  * Uses field element comparison for deterministic ordering.
  * Returns negative if a < b, positive if a > b, 0 if equal.
  */
-export function compareTxHash(a: bigint, b: bigint): number {
+export function compareTxHash(a: bigint, b: bigint): -1 | 0 | 1 {
   return Fr.cmpAsBigInt(a, b);
 }
 
@@ -182,7 +182,7 @@ export function compareTxHash(a: bigint, b: bigint): number {
  * Returns negative if a < b, positive if a > b, 0 if equal.
  * Use with sort() for ascending order, or negate/reverse for descending.
  */
-export function comparePriority(a: PriorityComparable, b: PriorityComparable): number {
+export function comparePriority(a: PriorityComparable, b: PriorityComparable): -1 | 0 | 1 {
   const feeComparison = compareFee(a.priorityFee, b.priorityFee);
   if (feeComparison !== 0) {
     return feeComparison;

yarn-project/p2p/src/mem_pools/tx_pool_v2/tx_pool_indices.test.ts

@@ -0,0 +1,184 @@
+import { Fr } from '@aztec/foundation/curves/bn254';
+
+import { stubTxMetaData, txHashFromBigInt } from './tx_metadata.js';
+import { TxPoolIndices } from './tx_pool_indices.js';
+
+describe('TxPoolIndices', () => {
+  let indices: TxPoolIndices;
+
+  const makeMeta = (seed: number, priorityFee: bigint) =>
+    stubTxMetaData(new Fr(seed).toString(), { priorityFee, nullifiers: [`nullifier-${seed}`] });
+
+  beforeEach(() => {
+    indices = new TxPoolIndices();
+  });
+
+  describe('sorted pending order', () => {
+    it('iterates descending by fee then hash', () => {
+      const low = makeMeta(1, 10n);
+      const mid = makeMeta(2, 50n);
+      const high = makeMeta(3, 100n);
+
+      indices.addPending(low);
+      indices.addPending(high);
+      indices.addPending(mid);
+
+      const desc = [...indices.iteratePendingByPriority('desc')];
+      expect(desc).toEqual([high.txHash, mid.txHash, low.txHash]);
+    });
+
+    it('iterates ascending by fee then hash', () => {
+      const low = makeMeta(1, 10n);
+      const mid = makeMeta(2, 50n);
+      const high = makeMeta(3, 100n);
+
+      indices.addPending(high);
+      indices.addPending(low);
+      indices.addPending(mid);
+
+      const asc = [...indices.iteratePendingByPriority('asc')];
+      expect(asc).toEqual([low.txHash, mid.txHash, high.txHash]);
+    });
+
+    it('uses txHash as tiebreaker for equal fees', () => {
+      const a = makeMeta(10, 50n);
+      const b = makeMeta(20, 50n);
+      const c = makeMeta(30, 50n);
+
+      indices.addPending(c);
+      indices.addPending(a);
+      indices.addPending(b);
+
+      const asc = [...indices.iteratePendingByPriority('asc')];
+      expect(asc).toHaveLength(3);
+
+      const hashes = [a, b, c].map(m => m.txHashBigInt);
+      hashes.sort((x, y) => (x < y ? -1 : x > y ? 1 : 0));
+      const expectedAsc = hashes.map(h => txHashFromBigInt(h));
+      expect(asc).toEqual(expectedAsc);
+    });
+  });
+
+  describe('remove', () => {
+    it('maintains order after removal', () => {
+      const a = makeMeta(1, 10n);
+      const b = makeMeta(2, 50n);
+      const c = makeMeta(3, 100n);
+
+      indices.addPending(a);
+      indices.addPending(b);
+      indices.addPending(c);
+
+      indices.remove(b.txHash);
+
+      const desc = [...indices.iteratePendingByPriority('desc')];
+      expect(desc).toEqual([c.txHash, a.txHash]);
+    });
+
+    it('handles removing non-existent tx gracefully', () => {
+      const a = makeMeta(1, 10n);
+      indices.addPending(a);
+
+      indices.remove('0xdeadbeef');
+      expect(indices.getPendingTxCount()).toBe(1);
+    });
+  });
+
+  describe('count', () => {
+    it('returns correct count after adds and removes', () => {
+      expect(indices.getPendingTxCount()).toBe(0);
+
+      const a = makeMeta(1, 10n);
+      const b = makeMeta(2, 20n);
+      indices.addPending(a);
+      indices.addPending(b);
+      expect(indices.getPendingTxCount()).toBe(2);
+
+      indices.remove(a.txHash);
+      expect(indices.getPendingTxCount()).toBe(1);
+
+      indices.remove(b.txHash);
+      expect(indices.getPendingTxCount()).toBe(0);
+    });
+  });
+
+  describe('getLowestPriorityPendingTx', () => {
+    it('returns the lowest priority tx', () => {
+      const low = makeMeta(1, 5n);
+      const high = makeMeta(2, 100n);
+
+      indices.addPending(high);
+      indices.addPending(low);
+
+      expect(indices.getLowestPriorityPendingTx()?.txHash).toBe(low.txHash);
+    });
+
+    it('returns undefined for empty pool', () => {
+      expect(indices.getLowestPriorityPendingTx()).toBeUndefined();
+    });
+  });
+
+  describe('filter', () => {
+    it('applies filter during iteration', () => {
+      const a = makeMeta(1, 10n);
+      const b = makeMeta(2, 50n);
+      const c = makeMeta(3, 100n);
+
+      indices.addPending(a);
+      indices.addPending(b);
+      indices.addPending(c);
+
+      const filtered = [...indices.iteratePendingByPriority('desc', hash => hash !== b.txHash)];
+      expect(filtered).toEqual([c.txHash, a.txHash]);
+    });
+  });
+
+  describe('eligible pending', () => {
+    it('filters by receivedAt', () => {
+      const old = makeMeta(1, 10n);
+      old.receivedAt = 100;
+      const recent = makeMeta(2, 50n);
+      recent.receivedAt = 500;
+
+      indices.addPending(old);
+      indices.addPending(recent);
+
+      const eligible = [...indices.iterateEligiblePendingByPriority('desc', 200)];
+      expect(eligible).toEqual([old.txHash]);
+    });
+  });
+
+  describe('edge cases', () => {
+    it('iterates empty pool without error', () => {
+      expect([...indices.iteratePendingByPriority('desc')]).toEqual([]);
+      expect([...indices.iteratePendingByPriority('asc')]).toEqual([]);
+    });
+
+    it('handles single element', () => {
+      const a = makeMeta(1, 10n);
+      indices.addPending(a);
+
+      expect([...indices.iteratePendingByPriority('desc')]).toEqual([a.txHash]);
+      expect([...indices.iteratePendingByPriority('asc')]).toEqual([a.txHash]);
+    });
+
+    it('does not add duplicates', () => {
+      const a = makeMeta(1, 10n);
+      indices.addPending(a);
+      indices.addPending(a);
+
+      expect(indices.getPendingTxCount()).toBe(1);
+    });
+
+    it('add-remove-add cycle works', () => {
+      const a = makeMeta(1, 10n);
+      indices.addPending(a);
+      indices.remove(a.txHash);
+      expect(indices.getPendingTxCount()).toBe(0);
+
+      indices.addPending(a);
+      expect(indices.getPendingTxCount()).toBe(1);
+      expect([...indices.iteratePendingByPriority('desc')]).toEqual([a.txHash]);
+    });
+  });
+});

yarn-project/p2p/src/mem_pools/tx_pool_v2/tx_pool_indices.ts

@@ -1,7 +1,8 @@
+import { insertIntoSortedArray, removeFromSortedArray } from '@aztec/foundation/array';
 import { SlotNumber } from '@aztec/foundation/branded-types';
 import type { L2BlockId } from '@aztec/stdlib/block';
 
-import { type TxMetaData, type TxState, compareFee, compareTxHash, txHashFromBigInt } from './tx_metadata.js';
+import { type PriorityComparable, type TxMetaData, type TxState, comparePriority } from './tx_metadata.js';
 
 /**
  * Manages in-memory indices for the transaction pool.
@@ -22,8 +23,8 @@ export class TxPoolIndices {
   #nullifierToTxHash: Map<string, string> = new Map();
   /** Fee payer to txHashes index (pending txs only) */
   #feePayerToTxHashes: Map<string, Set<string>> = new Map();
-  /** Pending txHash bigints grouped by priority fee */
-  #pendingByPriority: Map<bigint, Set<bigint>> = new Map();
+  /** Pending transactions sorted ascending by priority fee, ties broken by txHash */
+  #pendingByPriority: PriorityComparable[] = [];
   /** Protected transactions: txHash -> slotNumber */
   #protectedTransactions: Map<string, SlotNumber> = new Map();
 
@@ -73,20 +74,14 @@ export class TxPoolIndices {
    * @param order - 'desc' for highest priority first, 'asc' for lowest priority first
    */
   *iteratePendingByPriority(order: 'asc' | 'desc', filter?: (hash: string) => boolean): Generator<string> {
-    const feeCompareFn = order === 'desc' ? (a: bigint, b: bigint) => compareFee(b, a) : compareFee;
-    const hashCompareFn =
-      order === 'desc' ? (a: bigint, b: bigint) => compareTxHash(b, a) : (a: bigint, b: bigint) => compareTxHash(a, b);
-
-    const sortedFees = [...this.#pendingByPriority.keys()].sort(feeCompareFn);
-
-    for (const fee of sortedFees) {
-      const hashesAtFee = this.#pendingByPriority.get(fee)!;
-      const sortedHashes = [...hashesAtFee].sort(hashCompareFn);
-      for (const hashBigInt of sortedHashes) {
-        const hash = txHashFromBigInt(hashBigInt);
-        if (filter === undefined || filter(hash)) {
-          yield hash;
-        }
+    const arr = this.#pendingByPriority;
+    const start = order === 'asc' ? 0 : arr.length - 1;
+    const step = order === 'asc' ? 1 : -1;
+    const inBounds = order === 'asc' ? (i: number) => i < arr.length : (i: number) => i >= 0;
+
+    for (let i = start; inBounds(i); i += step) {
+      if (filter === undefined || filter(arr[i].txHash)) {
+        yield arr[i].txHash;
       }
     }
   }
@@ -227,11 +222,7 @@ export class TxPoolIndices {
 
   /** Gets the count of pending transactions */
   getPendingTxCount(): number {
-    let count = 0;
-    for (const hashes of this.#pendingByPriority.values()) {
-      count += hashes.size;
-    }
-    return count;
+    return this.#pendingByPriority.length;
   }
 
   /** Gets the lowest priority pending transaction hashes (up to limit) */
@@ -264,12 +255,10 @@ export class TxPoolIndices {
   /** Gets all pending transactions */
   getPendingTxs(): TxMetaData[] {
     const result: TxMetaData[] = [];
-    for (const hashSet of this.#pendingByPriority.values()) {
-      for (const txHashBigInt of hashSet) {
-        const meta = this.#metadata.get(txHashFromBigInt(txHashBigInt));
-        if (meta) {
-          result.push(meta);
-        }
+    for (const entry of this.#pendingByPriority) {
+      const meta = this.#metadata.get(entry.txHash);
+      if (meta) {
+        result.push(meta);
       }
     }
     return result;
@@ -408,13 +397,12 @@ export class TxPoolIndices {
     }
     feePayerSet.add(meta.txHash);
 
-    // Add to priority bucket
-    let prioritySet = this.#pendingByPriority.get(meta.priorityFee);
-    if (!prioritySet) {
-      prioritySet = new Set();
-      this.#pendingByPriority.set(meta.priorityFee, prioritySet);
-    }
-    prioritySet.add(meta.txHashBigInt);
+    insertIntoSortedArray(
+      this.#pendingByPriority,
+      { txHash: meta.txHash, priorityFee: meta.priorityFee, txHashBigInt: meta.txHashBigInt },
+      comparePriority,
+      false,
+    );
   }
 
   #removeFromPendingIndices(meta: TxMetaData): void {
@@ -432,13 +420,11 @@ export class TxPoolIndices {
       }
     }
 
-    // Remove from priority map
-    const hashSet = this.#pendingByPriority.get(meta.priorityFee);
-    if (hashSet) {
-      hashSet.delete(meta.txHashBigInt);
-      if (hashSet.size === 0) {
-        this.#pendingByPriority.delete(meta.priorityFee);
-      }
-    }
+    // Remove from priority array
+    removeFromSortedArray(
+      this.#pendingByPriority,
+      { txHash: meta.txHash, priorityFee: meta.priorityFee, txHashBigInt: meta.txHashBigInt },
+      comparePriority,
+    );
   }
 }