refactor(p2p): merge checkpoint attestation methods into tryAddCheckpointAttestation

Replace separate hasCheckpointAttestation, canAddCheckpointAttestation, and
hasReachedCheckpointAttestationCap methods with a single tryAddCheckpointAttestation
method that handles all validation and addition in one atomic operation.

This simplifies the API and ensures consistent behavior between checking
and adding attestations.

Co-Authored-By: Claude Opus 4.5 <[email protected]>

Author: spalladino

yarn-project/p2p/src/mem_pools/attestation_pool/attestation_pool.ts

@@ -1,13 +1,13 @@
 import type { SlotNumber } from '@aztec/foundation/branded-types';
 import type { BlockProposal, CheckpointAttestation, CheckpointProposalCore } from '@aztec/stdlib/p2p';
 
-/** Result of trying to add a proposal (block or checkpoint) */
-export type TryAddProposalResult = {
-  /** Whether proposal was added */
+/** Result of trying to add an item (proposal or attestation) to the pool */
+export type TryAddResult = {
+  /** Whether the item was added */
   added: boolean;
-  /** Whether exact proposal already existed */
+  /** Whether the exact item already existed */
   alreadyExists: boolean;
-  /** Total proposals for this position - used for duplicate detection */
+  /** Total items for this position - used for duplicate detection */
   totalForPosition: number;
 };
 
@@ -29,7 +29,7 @@ export interface AttestationPool {
    * @param blockProposal - The block proposal to add
    * @returns Result indicating whether the proposal was added and duplicate detection info
    */
-  tryAddBlockProposal(blockProposal: BlockProposal): Promise<TryAddProposalResult>;
+  tryAddBlockProposal(blockProposal: BlockProposal): Promise<TryAddResult>;
 
   /**
    * Get block proposal by its ID.
@@ -55,7 +55,7 @@ export interface AttestationPool {
    * @param proposal - The checkpoint proposal core to add
    * @returns Result indicating whether the proposal was added and duplicate detection info
    */
-  tryAddCheckpointProposal(proposal: CheckpointProposalCore): Promise<TryAddProposalResult>;
+  tryAddCheckpointProposal(proposal: CheckpointProposalCore): Promise<TryAddResult>;
 
   /**
    * Get checkpoint proposal by its ID.
@@ -75,6 +75,20 @@ export interface AttestationPool {
    */
   addCheckpointAttestations(attestations: CheckpointAttestation[]): Promise<void>;
 
+  /**
+   * Attempts to add a checkpoint attestation to the pool.
+   *
+   * This method performs validation and addition in a single call:
+   * - Checks if the attestation already exists (returns alreadyExists: true if so)
+   * - Checks if the (slot, proposalId) has reached the attestation cap (returns added: false if so)
+   * - Adds the attestation if validation passes
+   *
+   * @param attestation - The checkpoint attestation to add
+   * @param committeeSize - Committee size for the attestation's slot
+   * @returns Result indicating whether the attestation was added and existence info
+   */
+  tryAddCheckpointAttestation(attestation: CheckpointAttestation, committeeSize: number): Promise<TryAddResult>;
+
   /**
    * Delete all pool data (attestations, proposals) older than the given slot
    *
@@ -99,33 +113,6 @@ export interface AttestationPool {
    */
   getCheckpointAttestationsForSlotAndProposal(slot: SlotNumber, proposalId: string): Promise<CheckpointAttestation[]>;
 
-  /**
-   * Check if a specific checkpoint attestation exists in the pool
-   *
-   * @param attestation - The attestation to check
-   * @return True if the attestation exists, false otherwise
-   */
-  hasCheckpointAttestation(attestation: CheckpointAttestation): Promise<boolean>;
-
-  /**
-   * Returns whether a checkpoint attestation would be accepted for (slot, proposalId).
-   *
-   * @param attestation - The attestation to check
-   * @param committeeSize - Committee size for the attestation's slot
-   * @returns True if the attestation can be added, false otherwise.
-   */
-  canAddCheckpointAttestation(attestation: CheckpointAttestation, committeeSize: number): Promise<boolean>;
-
-  /**
-   * Returns whether the checkpoint attestation cap for the given slot and proposal has been reached.
-   *
-   * @param slot - The slot to check
-   * @param proposalId - The proposal to check
-   * @param committeeSize - Committee size for the slot
-   * @returns True if the cap has been reached, false otherwise.
-   */
-  hasReachedCheckpointAttestationCap(slot: SlotNumber, proposalId: string, committeeSize: number): Promise<boolean>;
-
   /** Returns whether the pool is empty. */
   isEmpty(): Promise<boolean>;
 }

yarn-project/p2p/src/mem_pools/attestation_pool/attestation_pool_test_suite.ts

@@ -67,11 +67,6 @@ export function describeAttestationPool(getAttestationPool: () => AttestationPoo
       expect(retrievedAttestations.length).toBe(attestations.length);
       compareCheckpointAttestations(retrievedAttestations, attestations);
 
-      // Check hasCheckpointAttestation for added attestations
-      for (const attestation of attestations) {
-        expect(await ap.hasCheckpointAttestation(attestation)).toBe(true);
-      }
-
       const retrievedAttestationsForSlot = await ap.getCheckpointAttestationsForSlot(SlotNumber(slotNumber));
       expect(retrievedAttestationsForSlot.length).toBe(attestations.length);
       compareCheckpointAttestations(retrievedAttestationsForSlot, attestations);
@@ -85,7 +80,6 @@ export function describeAttestationPool(getAttestationPool: () => AttestationPoo
       );
       expect(retrievedAttestationsAfterAdd.length).toBe(attestations.length + 1);
       compareCheckpointAttestations(retrievedAttestationsAfterAdd, [...attestations, newAttestation]);
-      expect(await ap.hasCheckpointAttestation(newAttestation)).toBe(true);
       const retrievedAttestationsForSlotAfterAdd = await ap.getCheckpointAttestationsForSlot(SlotNumber(slotNumber));
       expect(retrievedAttestationsForSlotAfterAdd.length).toBe(attestations.length + 1);
       compareCheckpointAttestations(retrievedAttestationsForSlotAfterAdd, [...attestations, newAttestation]);
@@ -98,11 +92,6 @@ export function describeAttestationPool(getAttestationPool: () => AttestationPoo
         archive.toString(),
       );
       expect(retreivedAttestationsAfterDelete.length).toBe(0);
-      // Check hasCheckpointAttestation after deletion
-      for (const attestation of attestations) {
-        expect(await ap.hasCheckpointAttestation(attestation)).toBe(false);
-      }
-      expect(await ap.hasCheckpointAttestation(newAttestation)).toBe(false);
     });
 
     it('should handle duplicate proposals in a slot', async () => {

yarn-project/p2p/src/mem_pools/attestation_pool/kv_attestation_pool.test.ts

@@ -51,23 +51,25 @@ describe('KV Attestation Pool', () => {
       // Create 'limit' distinct checkpoint attestations for the same (slot, proposalId)
       const signers = Array.from({ length: limit }, () => Secp256k1Signer.random());
       const attestations = signers.map(s => mockCheckpointAttestation(s, slotNumber, archive));
-      await kvAttestationPool.addCheckpointAttestations(attestations);
 
-      // We should now be at cap
-      expect(
-        await kvAttestationPool.hasReachedCheckpointAttestationCap(
-          SlotNumber(slotNumber),
-          archive.toString(),
-          committeeSize,
-        ),
-      ).toBe(true);
+      // Add each attestation using tryAddCheckpointAttestation
+      for (let i = 0; i < attestations.length; i++) {
+        const result = await kvAttestationPool.tryAddCheckpointAttestation(attestations[i], committeeSize);
+        expect(result.added).toBe(true);
+        expect(result.totalForPosition).toBe(i + 1);
+      }
 
-      // A new attestation from a new signer should not be accepted (per validation helper semantics)
+      // A new attestation from a new signer should not be added (cap reached)
       const extra = mockCheckpointAttestation(Secp256k1Signer.random(), slotNumber, archive);
-      expect(await kvAttestationPool.canAddCheckpointAttestation(extra, committeeSize)).toBe(false);
+      const extraResult = await kvAttestationPool.tryAddCheckpointAttestation(extra, committeeSize);
+      expect(extraResult.added).toBe(false);
+      expect(extraResult.alreadyExists).toBe(false);
+      expect(extraResult.totalForPosition).toBe(limit);
 
-      // Re-adding an existing attestation should be allowed
-      expect(await kvAttestationPool.canAddCheckpointAttestation(attestations[0], committeeSize)).toBe(true);
+      // Re-adding an existing attestation should return alreadyExists
+      const existingResult = await kvAttestationPool.tryAddCheckpointAttestation(attestations[0], committeeSize);
+      expect(existingResult.added).toBe(false);
+      expect(existingResult.alreadyExists).toBe(true);
     });
   });
 });

yarn-project/p2p/src/mem_pools/attestation_pool/kv_attestation_pool.ts

@@ -12,7 +12,7 @@ import {
 import { type TelemetryClient, getTelemetryClient } from '@aztec/telemetry-client';
 
 import { PoolInstrumentation, PoolName, type PoolStatsCallback } from '../instrumentation.js';
-import type { AttestationPool, TryAddProposalResult } from './attestation_pool.js';
+import type { AttestationPool, TryAddResult } from './attestation_pool.js';
 
 export const MAX_PROPOSALS_PER_SLOT = 5;
 export const MAX_PROPOSALS_PER_POSITION = 3;
@@ -117,7 +117,7 @@ export class KvAttestationPool implements AttestationPool {
     return (slot << KvAttestationPool.INDEX_BITS) | indexWithinCheckpoint;
   }
 
-  public async tryAddBlockProposal(blockProposal: BlockProposal): Promise<TryAddProposalResult> {
+  public async tryAddBlockProposal(blockProposal: BlockProposal): Promise<TryAddResult> {
     const proposalId = blockProposal.archive.toString();
 
     // Check if already exists
@@ -180,7 +180,7 @@ export class KvAttestationPool implements AttestationPool {
     return Promise.resolve(undefined);
   }
 
-  public async tryAddCheckpointProposal(proposal: CheckpointProposalCore): Promise<TryAddProposalResult> {
+  public async tryAddCheckpointProposal(proposal: CheckpointProposalCore): Promise<TryAddResult> {
     const proposalId = proposal.archive.toString();
 
     // Check if already exists
@@ -325,50 +325,52 @@ export class KvAttestationPool implements AttestationPool {
     });
   }
 
-  public async hasCheckpointAttestation(attestation: CheckpointAttestation): Promise<boolean> {
+  public async tryAddCheckpointAttestation(
+    attestation: CheckpointAttestation,
+    committeeSize: number,
+  ): Promise<TryAddResult> {
     const slotNumber = attestation.payload.header.slotNumber;
-    const proposalId = attestation.archive;
+    const proposalId = attestation.archive.toString();
     const sender = attestation.getSender();
 
-    // Attestations with invalid signatures are never in the pool
     if (!sender) {
-      return false;
+      return { added: false, alreadyExists: false, totalForPosition: 0 };
     }
 
-    const address = sender.toString();
-    const key = this.getAttestationKey(slotNumber, proposalId, address);
+    const key = this.getAttestationKey(slotNumber, proposalId, sender.toString());
+    const alreadyExists = await this.checkpointAttestations.hasAsync(key);
 
-    return await this.checkpointAttestations.hasAsync(key);
-  }
+    if (alreadyExists) {
+      const total = await this.getAttestationCount(slotNumber, proposalId);
+      return { added: false, alreadyExists: true, totalForPosition: total };
+    }
 
-  public async canAddCheckpointAttestation(
-    attestation: CheckpointAttestation,
-    committeeSize: number,
-  ): Promise<boolean> {
-    return (
-      (await this.hasCheckpointAttestation(attestation)) ||
-      !(await this.hasReachedCheckpointAttestationCap(
-        attestation.payload.header.slotNumber,
-        attestation.archive.toString(),
-        committeeSize,
-      ))
-    );
+    const limit = committeeSize + ATTESTATION_CAP_BUFFER;
+    const currentCount = await this.getAttestationCount(slotNumber, proposalId);
+
+    if (currentCount >= limit) {
+      return { added: false, alreadyExists: false, totalForPosition: currentCount };
+    }
+
+    await this.checkpointAttestations.set(key, attestation.toBuffer());
+
+    this.log.verbose(`Added checkpoint attestation for slot ${slotNumber} from ${sender.toString()}`, {
+      signature: attestation.signature.toString(),
+      slotNumber,
+      address: sender.toString(),
+      proposalId,
+    });
+
+    return { added: true, alreadyExists: false, totalForPosition: currentCount + 1 };
   }
 
-  public async hasReachedCheckpointAttestationCap(
-    slot: SlotNumber,
-    proposalId: string,
-    committeeSize: number,
-  ): Promise<boolean> {
-    const limit = committeeSize + ATTESTATION_CAP_BUFFER;
+  /** Gets the count of attestations for a given (slot, proposalId). */
+  private async getAttestationCount(slot: SlotNumber, proposalId: string): Promise<number> {
     const range = this.getAttestationKeyRangeForProposal(slot, proposalId);
     let count = 0;
     for await (const _ of this.checkpointAttestations.keysAsync(range)) {
       count++;
-      if (count >= limit) {
-        return true;
-      }
     }
-    return false;
+    return count;
   }
 }

yarn-project/p2p/src/mem_pools/attestation_pool/memory_attestation_pool.ts

@@ -4,7 +4,7 @@ import type { BlockProposal, CheckpointAttestation, CheckpointProposalCore } fro
 import { type TelemetryClient, getTelemetryClient } from '@aztec/telemetry-client';
 
 import { PoolInstrumentation, PoolName, type PoolStatsCallback } from '../instrumentation.js';
-import type { AttestationPool, TryAddProposalResult } from './attestation_pool.js';
+import type { AttestationPool, TryAddResult } from './attestation_pool.js';
 import { ATTESTATION_CAP_BUFFER, MAX_PROPOSALS_PER_POSITION, MAX_PROPOSALS_PER_SLOT } from './kv_attestation_pool.js';
 
 export class InMemoryAttestationPool implements AttestationPool {
@@ -52,7 +52,7 @@ export class InMemoryAttestationPool implements AttestationPool {
     return Promise.resolve(this.checkpointAttestations.size === 0 && this.proposals.size === 0);
   }
 
-  public tryAddBlockProposal(blockProposal: BlockProposal): Promise<TryAddProposalResult> {
+  public tryAddBlockProposal(blockProposal: BlockProposal): Promise<TryAddResult> {
     const proposalId = blockProposal.archive.toString();
     const slot = blockProposal.slotNumber;
     const index = blockProposal.indexWithinCheckpoint;
@@ -113,7 +113,7 @@ export class InMemoryAttestationPool implements AttestationPool {
 
   // Checkpoint attestation methods
 
-  public tryAddCheckpointProposal(proposal: CheckpointProposalCore): Promise<TryAddProposalResult> {
+  public tryAddCheckpointProposal(proposal: CheckpointProposalCore): Promise<TryAddResult> {
     const proposalId = proposal.archive.toString();
 
     // 1. Check if already exists
@@ -245,51 +245,43 @@ export class InMemoryAttestationPool implements AttestationPool {
     return Promise.resolve();
   }
 
-  public hasReachedCheckpointAttestationCap(
-    slot: SlotNumber,
-    proposalId: string,
-    committeeSize: number,
-  ): Promise<boolean> {
-    const limit = committeeSize + ATTESTATION_CAP_BUFFER;
-    const count = this.checkpointAttestations.get(slot)?.get(proposalId)?.size ?? 0;
-    return Promise.resolve(limit <= 0 || count >= limit);
-  }
-
-  public async canAddCheckpointAttestation(
-    attestation: CheckpointAttestation,
-    committeeSize: number,
-  ): Promise<boolean> {
-    const sender = attestation.getSender();
-    const slot = attestation.payload.header.slotNumber;
-    const pid = attestation.archive.toString();
-    return (
-      !!sender &&
-      ((this.checkpointAttestations.get(slot)?.get(pid)?.has(sender.toString()) ?? false) ||
-        !(await this.hasReachedCheckpointAttestationCap(slot, pid, committeeSize)))
-    );
-  }
-
-  public hasCheckpointAttestation(attestation: CheckpointAttestation): Promise<boolean> {
+  public tryAddCheckpointAttestation(attestation: CheckpointAttestation, committeeSize: number): Promise<TryAddResult> {
     const slotNumber = attestation.payload.header.slotNumber;
     const proposalId = attestation.archive.toString();
     const sender = attestation.getSender();
 
-    // Attestations with invalid signatures are never in the pool
     if (!sender) {
-      return Promise.resolve(false);
+      return Promise.resolve({ added: false, alreadyExists: false, totalForPosition: 0 });
     }
 
-    const slotAttestationMap = this.checkpointAttestations.get(slotNumber);
-    if (!slotAttestationMap) {
-      return Promise.resolve(false);
+    const address = sender.toString();
+    const slotAttestationMap = getCheckpointSlotOrDefault(this.checkpointAttestations, slotNumber);
+    const proposalAttestationMap = getCheckpointProposalOrDefault(slotAttestationMap, proposalId);
+
+    // Check if already exists
+    const alreadyExists = proposalAttestationMap.has(address);
+    if (alreadyExists) {
+      return Promise.resolve({ added: false, alreadyExists: true, totalForPosition: proposalAttestationMap.size });
     }
 
-    const proposalAttestationMap = slotAttestationMap.get(proposalId);
-    if (!proposalAttestationMap) {
-      return Promise.resolve(false);
+    // Check cap
+    const limit = committeeSize + ATTESTATION_CAP_BUFFER;
+    const currentCount = proposalAttestationMap.size;
+    if (currentCount >= limit) {
+      return Promise.resolve({ added: false, alreadyExists: false, totalForPosition: currentCount });
     }
 
-    return Promise.resolve(proposalAttestationMap.has(sender.toString()));
+    // Add the attestation
+    proposalAttestationMap.set(address, attestation);
+
+    this.log.verbose(`Added checkpoint attestation for slot ${slotNumber} from ${address}`, {
+      signature: attestation.signature.toString(),
+      slotNumber,
+      address,
+      proposalId,
+    });
+
+    return Promise.resolve({ added: true, alreadyExists: false, totalForPosition: currentCount + 1 });
   }
 }